feat(hub): add pretrained_revision to pin Hub model versions

- Add pretrained_revision field to PreTrainedConfig (policies) and
RewardModelConfig (reward models), and thread it through make_policy(),
make_pre_post_processors(), and make_reward_model() so that weights and
processor configs can be loaded from a specific Hub commit, branch, or
tag. Defaults to None (latest version, preserving current behavior).
Dataset and env hub loading already supported revision pinning.

.pre-commit-config.yaml

@@ -65,6 +65,9 @@ repos:
         name: Format Markdown with Prettier
         types_or: [markdown, mdx]
         args: [--prose-wrap=preserve]
+        # Jinja2 model-card templates use a .md extension but contain {% ... %} /
+        # {{ ... }} tags that prettier's Markdown formatter mangles (e.g. table loops).
+        exclude: ^src/lerobot/templates/.*\.md$
 
   ##### Security #####
   - repo: https://github.com/gitleaks/gitleaks

Makefile

@@ -178,3 +178,9 @@ test-smolvla-ete-eval:
 		--env.episode_length=5 \
 		--eval.n_episodes=1 \
 		--eval.batch_size=1
+
+# E2E annotation pipeline smoke test against a tiny in-memory fixture
+# dataset. Opt-in (not part of `make test-end-to-end`) and uses a stub VLM
+# backend, so it does not require a real model checkpoint or GPU.
+annotation-e2e:
+	uv run python -m tests.annotations.run_e2e_smoke

README.md

@@ -58,7 +58,7 @@ action = model.select_action(obs)
 robot.send_action(action)
 ```
 
-**Supported Hardware:** SO100, LeKiwi, Koch, HopeJR, OMX, EarthRover, Reachy2, Gamepads, Keyboards, Phones, OpenARM, Unitree G1.
+**Supported Hardware:** SO100, LeKiwi, Koch, HopeJR, OMX, EarthRover, Reachy2, Gamepads, Keyboards, Phones, OpenARM, Unitree G1, reBot B601.
 
 While these devices are natively integrated into the LeRobot codebase, the library is designed to be extensible. You can easily implement the Robot interface to utilize LeRobot's data collection, training, and visualization tools for your own custom robot.
 
@@ -101,11 +101,13 @@ lerobot-train \
   --dataset.repo_id=lerobot/aloha_mobile_cabinet
 ```
 
-| Category                   | Models                                                                                                                                                                                                                  |
-| -------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| **Imitation Learning**     | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md) |
-| **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon)                                                                                                             |
-| **VLAs Models**            | [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx)            |
+| Category                   | Models                                                                                                                                                                                                                                                                                                                                                     |
+| -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| **Imitation Learning**     | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md)                                                                                                                                    |
+| **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon)                                                                                                                                                                                                                                                |
+| **VLAs Models**            | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.5](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx) |
+| **World Models**           | [VLA-JEPA](./docs/source/vla_jepa.mdx) (more coming soon)                                                                                                                                                                                                                                                                                                  |
+| **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
 
@@ -133,6 +135,7 @@ Learn how to implement your own simulation environment or benchmark and distribu
 - **[Discord](https://discord.gg/q8Dzzpym3f):** Join the `LeRobot` server to discuss with the community.
 - **[X](https://x.com/LeRobotHF):** Follow us on X to stay up-to-date with the latest developments.
 - **[Robot Learning Tutorial](https://huggingface.co/spaces/lerobot/robot-learning-tutorial):** A free, hands-on course to learn robot learning using LeRobot.
+- **[T-Shirt Folding Experiment](https://huggingface.co/spaces/lerobot/robot-folding):** An end-to-end demonstration of folding t-shirts with LeRobot.
 
 ## Citation
 
@@ -140,7 +143,7 @@ If you use LeRobot in your project, please cite the GitHub repository to acknowl
 
 ```bibtex
 @misc{cadene2024lerobot,
-    author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Pascal, Caroline and Choghari, Jade and Moss, Jess and Wolf, Thomas},
+    author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Pascal, Caroline and Choghari, Jade and Meftah, Khalil and Ellerbach, Maxime and Moss, Jess and Wolf, Thomas},
     title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch},
     howpublished = "\url{https://github.com/huggingface/lerobot}",
     year = {2024}

docs/source/_toctree.yml

@@ -45,6 +45,8 @@
     title: Language Columns and Recipes
   - local: tools
     title: Tools
+  - local: annotation_pipeline
+    title: Annotation Pipeline
   - local: video_encoding_parameters
     title: Video encoding parameters
   - local: streaming_video_encoding
@@ -68,7 +70,7 @@
   - local: eo1
     title: EO-1
   - local: groot
-    title: NVIDIA GR00T
+    title: NVIDIA GR00T N1.5
   - local: xvla
     title: X-VLA
   - local: multi_task_dit

docs/source/annotation_pipeline.mdx

@@ -0,0 +1,291 @@
+# Annotation Pipeline
+
+`lerobot-annotate` watches each episode's video with a vision-language
+model (VLM) and writes natural-language annotations back into your
+dataset. It fills the two language columns from the
+[Language Columns and Recipes](./language_and_recipes) page —
+`language_persistent` and `language_events` — straight into
+`data/chunk-*/file-*.parquet`.
+
+In short: point it at a LeRobot dataset, and it adds subtasks, plans,
+memory, interjections, speech, and visual Q&A that a policy can be
+trained on.
+
+## How it fits together
+
+```text
+  your dataset                  lerobot-annotate
+  (LeRobot v3.1)
+        │
+        ▼
+  ┌─────────────────────────────────────────────────────┐
+  │                    read episodes                     │
+  └──────────────────────────┬──────────────────────────┘
+                             │
+        ┌────────────────────┼────────────────────┐
+        ▼                    ▼                     ▼
+  ┌──────────┐      ┌───────────────┐        ┌──────────┐       one shared Qwen-VL
+  │   plan   │      │ interjections │        │   vqa    │  ◀──   server (vLLM, OpenAI
+  └────┬─────┘      └───────┬───────┘        └────┬─────┘        API) drives all three
+       └────────────────────┼─────────────────────┘
+                            │   each module stages raw JSONL
+                            ▼   into .annotate_staging/
+                  ┌─────────────────┐
+                  │    validator    │  ◀──  checks everything
+                  └────────┬────────┘
+                           ▼
+                  ┌─────────────────┐
+                  │     writer      │
+                  └────────┬────────┘
+                           ▼
+              data/chunk-*/file-*.parquet
+              (+ meta/info.json tools)
+```
+
+Three modules (`plan`, `interjections`, `vqa`) all talk to **one** shared
+VLM. Each module stages its output to disk, a validator checks it, and a
+single writer rewrites the dataset shards in place.
+
+## What the pipeline produces
+
+Each module emits a few kinds of annotation ("styles"), routed to one of
+the two language columns:
+
+| Style / atom                                | Column                | Module          |
+| ------------------------------------------- | --------------------- | --------------- |
+| `subtask` (Pi0.7-style "how, not what")     | `language_persistent` | `plan`          |
+| `plan` (initial + refresh on interjection)  | `language_persistent` | `plan`          |
+| `memory` (MEM-style compression)            | `language_persistent` | `plan`          |
+| `task_aug` (rephrasings of the task)        | `language_persistent` | `plan`          |
+| `interjection`                              | `language_events`     | `interjections` |
+| speech tool-call atom (`style=null`, `say`) | `language_events`     | `interjections` |
+| `vqa` (user / assistant pair)               | `language_events`     | `vqa`           |
+
+### How subtasks are generated
+
+The `plan` module doesn't ask the VLM for subtasks in one shot. Instead
+it uses a two-step **describe → segment** flow:
+
+1. **Describe** — the VLM narrates only what it actually sees in the
+   chosen camera (no guessing about the task).
+2. **Segment** — that description is fed back in, and the VLM splits the
+   episode into consecutive atomic subtasks.
+
+Both passes see the episode as **timestamped contact sheets** — frames
+sampled at `frames_per_second` (0.5s by default) and packed into JPEG
+grids with each frame's time burned into its corner, so the VLM cites
+exact boundary times directly. This is far cheaper in vision tokens than
+one image per frame, so the sampling can stay dense; episodes longer than
+`max_frames_per_prompt` are split into windows at the same density and
+merged. Both prompts also carry a causal **event-boundary** definition (a
+new event starts when an object becomes held / is released / reaches a new
+location / a lid changes state / contents move) to sharpen where cuts land.
+
+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,
+auto-windowed subtask generation).
+
+### Tools
+
+The writer does **not** add a `tools` column to the parquet. The tool
+catalog lives in `meta/info.json["tools"]` instead (see [Tools](./tools)).
+After every run, the pipeline makes sure the canonical `say` schema is in
+that list, keeping any tools you declared beforehand.
+
+Want to add your own tool? Edit `meta/info.json["tools"]` directly — the
+pipeline preserves whatever is already there. That makes the tool visible
+to the chat template, so the model can learn to _generate_ the call. The
+runtime layer that actually _executes_ a generated call (the `Tool`
+protocol / `TOOL_REGISTRY` under `src/lerobot/tools/`) is not part of
+this PR — the [Tools](./tools) doc marks those pieces as
+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:
+
+```bash
+HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
+```
+
+[`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:
+
+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`,
+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).
+
+## Key options
+
+These are the flags you'll reach for most often. Run
+`lerobot-annotate --help` for everything else; the defaults are tuned for
+short manipulation episodes.
+
+### Dataset in / out
+
+| Flag              | Default | What it does                                                            |
+| ----------------- | ------- | ----------------------------------------------------------------------- |
+| `--repo_id`       | —       | Hub dataset to annotate (downloaded if `--root` unset).                 |
+| `--root`          | —       | Annotate a local dataset directory instead.                             |
+| `--new_repo_id`   | —       | Push the result to a new repo (leaves the source repo untouched).       |
+| `--push_to_hub`   | `false` | Upload after annotating (to `--new_repo_id`, else back to `--repo_id`). |
+| `--only_episodes` | all     | Annotate just these episode indices (handy for a test run).             |
+| `--seed`          | `1729`  | Seeds the RNGs that pick interjection timestamps + VQA question types.  |
+
+### Which modules run
+
+Every module is on by default and can be toggled independently (set to
+`false` to skip it, e.g. to iterate on one module at a time):
+
+| Flag                      | Default | Turns off                           |
+| ------------------------- | ------- | ----------------------------------- |
+| `--plan.enabled`          | `true`  | subtasks + plan + memory + task_aug |
+| `--interjections.enabled` | `true`  | interjections + speech atoms        |
+| `--vqa.enabled`           | `true`  | the VQA pairs                       |
+
+### The VLM (`--vlm.*`)
+
+| Flag                       | Default            | What it does                                                                        |
+| -------------------------- | ------------------ | ----------------------------------------------------------------------------------- |
+| `--vlm.model_id`           | `Qwen/Qwen3.6-27B` | The model to serve and prompt.                                                      |
+| `--vlm.camera_key`         | first `images.*`   | Which camera every prompt is grounded on.                                           |
+| `--vlm.serve_command`      | auto               | The exact `vllm serve …` command (set TP size, GPU memory, `--max-model-len` here). |
+| `--vlm.parallel_servers`   | `1`                | Independent servers for round-robin routing (one per GPU).                          |
+| `--vlm.num_gpus`           | `0`                | GPUs per server (`0` = one each).                                                   |
+| `--vlm.client_concurrency` | `16`               | In-flight requests across all servers.                                              |
+| `--vlm.max_new_tokens`     | `512`              | Generation cap per call.                                                            |
+| `--vlm.temperature`        | `0.2`              | Sampling temperature.                                                               |
+
+### Subtasks / plan / memory (`--plan.*`)
+
+| Flag                            | Default    | What it does                                                                                                              |
+| ------------------------------- | ---------- | ------------------------------------------------------------------------------------------------------------------------- |
+| `--plan.frames_per_second`      | `2.0`      | Frame sampling rate for the contact sheets (`2.0` = one frame every 0.5s).                                                |
+| `--plan.max_frames_per_prompt`  | `60`       | Frame budget per VLM call. Episodes whose sampling exceeds this are auto-windowed at the same density, then stitched.     |
+| `--plan.contact_sheet_columns`  | `5`        | Columns per contact-sheet grid (`contact_sheet_frames_per_sheet` tiles, time row-major).                                  |
+| `--plan.plan_max_steps`         | `8`        | Upper bound on subtasks per episode.                                                                                      |
+| `--plan.subtask_describe_first` | `true`     | Run the describe→segment grounding pass (best subtask quality; +1 call/episode).                                          |
+| `--plan.emit_plan`              | `true`     | Emit the numbered `plan` rows (`false` = subtasks + memory only).                                                         |
+| `--plan.emit_memory`            | `true`     | Emit the `memory` rows (`false` = subtasks + plan only); symmetric to `emit_plan`.                                        |
+| `--plan.n_task_rephrasings`     | `10`       | How many `task_aug` rephrasings to emit (`0` disables).                                                                   |
+| `--plan.derive_task_from_video` | `if_short` | Use the dataset task as-is (`off`), only when it's missing/short (`if_short`), or always re-derive from video (`always`). |
+
+### Interjections + VQA
+
+| Flag                                            | Default | What it does                                               |
+| ----------------------------------------------- | ------- | ---------------------------------------------------------- |
+| `--interjections.max_interjections_per_episode` | `3`     | Cap on interjection/speech pairs per episode.              |
+| `--vqa.vqa_emission_hz`                         | `1.0`   | How often VQA pairs are emitted.                           |
+| `--vqa.restrict_to_default_camera`              | `false` | Ground VQA only on `--vlm.camera_key` (else every camera). |
+| `--executor.episode_parallelism`                | `16`    | Episodes processed concurrently within each phase.         |
+
+## Contributing new modules
+
+The pipeline is built to grow, and **contributions are very welcome** —
+a brand-new module (say, trajectory traces or affordances), a new prompt
+template, a smarter grounding flow, or quality fixes to the existing
+`plan` / `interjections` / `vqa` modules.
+
+Every module lives under
+`src/lerobot/annotations/steerable_pipeline/modules/`, shares the VLM
+client and the keyframe cache, writes its raw output to the staging
+tree, and plugs into the executor as its own phase. Got an idea? Open an
+issue or PR on [the repo](https://github.com/huggingface/lerobot).
+
+## How recipes consume the output
+
+The annotations are meant to be read by recipes (see
+[Language Columns and Recipes](./language_and_recipes)). Typically:
+
+- low-level / high-level / memory-update branches read
+  `subtask` / `plan` / `memory` from `language_persistent`.
+- an interjection-response branch reads `interjection` events plus the
+  paired speech atom (merged into one assistant turn via `tool_calls_from`)
+  and the matching `plan` refresh at the same timestamp.
+- a VQA branch reads the `(vqa, user)` and `(vqa, assistant)` pairs from
+  `language_events`.
+
+## Why state and events are split
+
+Two ideas shape the design:
+
+1. **Persistent state vs. exact events.** Persistent rows (`subtask`,
+   `plan`, `memory`) apply to the whole episode and answer "what's true
+   right now?". Event rows (`interjection`, `vqa`, speech) appear only on
+   the one frame whose timestamp matches. Timestamps are copied straight
+   from the source parquet — never recomputed in floating point.
+2. **One VLM pass.** All three modules share a single VLM client (the
+   OpenAI-compatible client talking to the job's vLLM server), so you pay
+   for one model load per dataset, not three.
+
+## Re-running a single module
+
+Each module stages its raw output to
+`<root>/.annotate_staging/episode_{N:06d}/<module>.jsonl`. This makes
+prompt iteration cheap: re-running one module overwrites only its own
+JSONL, then the writer recomposes the final parquet. Disable modules you
+don't want with `--plan.enabled=false` (and likewise
+`--interjections.enabled` / `--vqa.enabled`) to test one at a time.
+
+## What the validator checks
+
+Before the writer runs, `StagingValidator` confirms:
+
+- every event row lands exactly on a real frame timestamp;
+- no speech / interjection pairs are left orphaned;
+- `plan` is refreshed at every interjection timestamp;
+- `memory` rows fall on subtask boundaries (a warning, not an error);
+- each VQA assistant `content` is valid JSON in one of the
+  bbox / keypoint / count / attribute / spatial shapes;
+- every row goes to the column chosen by `column_for_style(style)`.
+
+Any error aborts the writer. Pass `--skip_validation=true` to override
+while debugging.
+
+## Where each module's ideas come from
+
+- **`plan` — subtasks.** Hi Robot ([Shi 2025](https://arxiv.org/abs/2502.19417))
+  for atom granularity ("pick up one piece of lettuce", "place bowl to
+  box"); Pi0.7 ([Physical Intelligence 2025](https://pi.website/pi07))
+  for "how, not what" detail.
+- **`plan` — memory.** MEM ([Torne 2026](https://arxiv.org/abs/2603.03596)):
+  keep only the minimal relevant information — preserve outcomes, drop
+  specific attributes.
+- **`interjections`.** Hi Robot's scenario taxonomy: negative task,
+  situated correction, specific constraint, preference. Speech is a
+  tool-call-only atom
+  (`tool_calls=[{type:function, function:{name:"say", arguments:{text:...}}}]`).
+- **`vqa`.** ECoT ([Zawalski 2024](https://arxiv.org/abs/2407.08693)) for
+  grounded features (pixel bounding boxes `[x_min, y_min, x_max, y_max]`,
+  keypoints) and Steerable VLA Policies
+  ([Zhao 2025](https://arxiv.org/abs/2509.07626)) for multi-abstraction
+  grounding. Pi0.7 also grounds answers across abstraction levels.
+
+When improving a module, tweak its prompt template in
+`src/lerobot/annotations/steerable_pipeline/prompts/` rather than
+rewriting from scratch.
+
+## Roughly how much it costs
+
+Per episode, the pipeline makes about `max_steps` plan calls,
+`max_interjections_per_episode` interjection calls, and
+`vqa_emission_hz × episode_seconds` VQA calls. With the defaults (8
+subtasks, 1 interjection, 1 Hz × 3 pairs) on a 30-second episode, that's
+~50 VLM calls.
+
+Storage stays small: `language_persistent` is at most tens of KB per
+episode (parquet dictionary-encodes the one entry that repeats across
+frames), and `language_events` is empty on most frames — its size scales
+with the number of emissions, not `num_frames × num_emissions`.

docs/source/envhub_isaaclab_arena.mdx

@@ -193,7 +193,7 @@ To learn more about training policies with LeRobot, please refer to the training
 
 - [SmolVLA](./smolvla)
 - [Pi0.5](./pi05)
-- [GR00T N1.7](./groot)
+- [GR00T N1.5](./groot)
 
 Sample IsaacLab Arena datasets are available on HuggingFace Hub for experimentation:
 

docs/source/groot.mdx

@@ -1,19 +1,16 @@
-# GR00T Policy
+# GR00T N1.5 Policy
 
-GR00T is an NVIDIA foundation model family for generalized humanoid robot reasoning and skills. It is a cross-embodiment policy that accepts multimodal input, including language, images, and proprioception, to perform manipulation tasks in diverse environments.
+GR00T N1.5 is an open foundation model from NVIDIA designed for generalized humanoid robot reasoning and skills. It is a cross-embodiment model that accepts multimodal input, including language and images, to perform manipulation tasks in diverse environments.
 
-LeRobot integrates GR00T N1.7 through the `groot` policy type.
-
-> [!WARNING]
-> **Breaking change:** GR00T N1.5 support was removed from LeRobot, and current releases support GR00T N1.7 only. N1.5 checkpoints and configs are rejected with a migration note. To keep using an N1.5 checkpoint, pin the last release that supports it: `pip install 'lerobot==0.5.1'`. To use the current release, migrate to GR00T N1.7 (base model [`nvidia/GR00T-N1.7-3B`](https://huggingface.co/nvidia/GR00T-N1.7-3B)).
+This document outlines the specifics of its integration and usage within the LeRobot framework.
 
 ## Model Overview
 
-GR00T N1.7 uses a Cosmos-Reason2/Qwen3-VL backbone and provides checkpoints for SimplerEnv, DROID, and LIBERO.
+NVIDIA Isaac GR00T N1.5 is an upgraded version of the GR00T N1 foundation model. It is built to improve generalization and language-following abilities for humanoid robots.
 
-Developers and researchers can post-train GR00T with their own real or synthetic data to adapt it for specific humanoid robots or tasks.
+Developers and researchers can post-train GR00T N1.5 with their own real or synthetic data to adapt it for specific humanoid robots or tasks.
 
-GR00T uses pre-trained vision and language encoders with a flow matching action transformer to model a chunk of actions conditioned on vision, language, and proprioception.
+GR00T N1.5 (specifically the GR00T-N1.5-3B model) is built using pre-trained vision and language encoders. It utilizes a flow matching action transformer to model a chunk of actions, conditioned on vision, language, and proprioception.
 
 <img
   src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-groot-paper1%20(1).png"
@@ -31,43 +28,33 @@ This approach allows the model to be highly adaptable through post-training for
 
 ## Installation Requirements
 
-GR00T is intended for NVIDIA GPU-accelerated systems. Install LeRobot with the GR00T extra:
+As of today, GR00T N1.5 requires flash attention for it's internal working.
+
+We are working on making this optional, but in the meantime that means that we require an extra installation step and it can only be used in CUDA enabled devices.
+
+1. Following the Environment Setup of our [Installation Guide](./installation). **Attention** don't install `lerobot` in this step.
+2. Install [Flash Attention](https://github.com/Dao-AILab/flash-attention) by running:
 
 ```bash
-pip install "lerobot[groot]"
-```
-
-For a source checkout:
-
-```bash
-pip install -e ".[groot]"
-```
-
-### Optional: Flash Attention acceleration
-
-Flash Attention is a purely optional performance optimization. **LeRobot neither installs nor requires it**, and setting it up is up to the user as it has environment-specific build requirements (a matching PyTorch/CUDA toolchain). To enable it:
-
-1. Install a `flash-attn` build matching your PyTorch/CUDA environment (see the [Flash Attention project](https://github.com/Dao-AILab/flash-attention)):
-
-```bash
-# Check https://pytorch.org/get-started/locally/ for the right CUDA wheel index for your system.
-pip install "torch>=2.7,<2.12.0" "torchvision>=0.22.0,<0.27.0" \
-  --index-url https://download.pytorch.org/whl/cu128
-pip install "ninja>=1.11.1,<2.0.0" "packaging>=24.2,<26.0"
+# Check https://pytorch.org/get-started/locally/ for your system
+pip install "torch>=2.2.1,<2.8.0" "torchvision>=0.21.0,<0.23.0" # --index-url https://download.pytorch.org/whl/cu1XX
+pip install ninja "packaging>=24.2,<26.0" # flash attention dependencies
 pip install "flash-attn>=2.5.9,<3.0.0" --no-build-isolation
 python -c "import flash_attn; print(f'Flash Attention {flash_attn.__version__} imported successfully')"
 ```
 
-2. Install lerobot with the groot extra.
+3. Install LeRobot by running:
 
-3. Opt in by passing `--policy.use_flash_attention=true` when training/evaluating GR00T. If the kernel is missing or fails to import, the backbone transparently falls back to SDPA.
+```bash
+pip install lerobot[groot]
+```
 
 ## Usage
 
-To use GR00T N1.7:
+To use GR00T in your LeRobot configuration, specify the policy type as:
 
-```bash
---policy.type=groot
+```python
+policy.type=groot
 ```
 
 ## Training
@@ -100,53 +87,21 @@ accelerate launch \
 
 ## Performance Results
 
-### LIBERO Benchmark Results
+### Libero Benchmark Results
 
 > [!NOTE]
-> Follow the [LIBERO](./libero) setup instructions before running `lerobot-eval`.
+> Follow our instructions for Libero usage: [Libero](./libero)
 
-GR00T N1.7 has demonstrated strong performance on the LIBERO benchmark suite. To reproduce LeRobot results, follow the instructions in the [LIBERO](./libero) section.
+GR00T has demonstrated strong performance on the Libero benchmark suite. To compare and test its LeRobot implementation, we finetuned the GR00T N1.5 model for 30k steps on the Libero dataset and compared the results to the GR00T reference results.
 
-### GR00T N1.7 LIBERO Checkpoints
+| Benchmark          | LeRobot Implementation | GR00T Reference |
+| ------------------ | ---------------------- | --------------- |
+| **Libero Spatial** | 82.0%                  | 92.0%           |
+| **Libero Object**  | 99.0%                  | 92.0%           |
+| **Libero Long**    | 82.0%                  | 76.0%           |
+| **Average**        | 87.0%                  | 87.0%           |
 
-NVIDIA publishes GR00T N1.7 LIBERO checkpoints at [`nvidia/GR00T-N1.7-LIBERO`](https://huggingface.co/nvidia/GR00T-N1.7-LIBERO), with one subdirectory per LIBERO suite:
-
-| Suite          | Checkpoint subdirectory |
-| -------------- | ----------------------- |
-| LIBERO Spatial | `libero_spatial`        |
-| LIBERO Object  | `libero_object`         |
-| LIBERO Goal    | `libero_goal`           |
-| LIBERO 10      | `libero_10`             |
-
-Preliminary LeRobot integration results:
-
-| Suite          | Status | Success rate | n_episodes |
-| -------------- | ------ | -----------: | ---------: |
-| LIBERO Spatial | ✓      |         ~95% |         XX |
-| LIBERO Object  | ✓      |          XX% |         XX |
-| LIBERO Goal    | ✓      |          XX% |         XX |
-| LIBERO 10      | ✓      |          XX% |         XX |
-| **Average**    | ✓      |      **XX%** |     **XX** |
-
-Replace the `XX` placeholders with final eval artifacts before merge.
-
-Download the suite checkpoint locally, then point `--policy.base_model_path` at the downloaded subdirectory. `--policy.path` is reserved for LeRobot checkpoints that contain a LeRobot `config.json` with a `type` field.
-
-```bash
-hf download nvidia/GR00T-N1.7-LIBERO \
-  --include "libero_spatial/*" \
-  --local-dir ./GR00T-N1.7-LIBERO
-
-lerobot-eval \
-  --policy.type=groot \
-  --policy.base_model_path=./GR00T-N1.7-LIBERO/libero_spatial \
-  --policy.embodiment_tag=libero_sim \
-  --env.type=libero \
-  --env.task=libero_spatial \
-  --eval.n_episodes=50
-```
-
-Use `eval.n_episodes >= 50` per suite when reporting success rates.
+These results demonstrate GR00T's strong generalization capabilities across diverse robotic manipulation tasks. To reproduce these results, you can follow the instructions in the [Libero](https://huggingface.co/docs/lerobot/libero) section.
 
 ### Evaluate in your hardware setup
 
@@ -176,4 +131,4 @@ lerobot-rollout\
 
 ## License
 
-GR00T N1.7 is released under the [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/).
+This model follows NVIDIA's proprietary license, consistent with the original [GR00T repository](https://github.com/NVIDIA/Isaac-GR00T). Future versions (starting from N1.7) will follow **Apache 2.0 License**.

docs/source/hil_data_collection.mdx

@@ -57,11 +57,11 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
 
 **Compatible teleoperators:**
 
-- `openarm_mini` - OpenArm Mini
+- `bi_openarm_mini` - Bimanual OpenArm Mini
 - `so_leader` - SO100 / SO101 leader arm
 
 > [!IMPORTANT]
-> The provided commands default to `bi_openarm_follower` + `openarm_mini`.
+> The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.
 > `so_follower` + `so_leader` configs are also registered and can be used via CLI flags.
 
 ---
@@ -104,9 +104,9 @@ lerobot-rollout --strategy.type=dagger \
     --robot.right_arm_config.port=can0 \
     --robot.right_arm_config.side=right \
     --robot.cameras='{left_wrist: {type: opencv, index_or_path: "/dev/video0", width: 1280, height: 720, fps: 30}, right_wrist: {type: opencv, index_or_path: "/dev/video4", width: 1280, height: 720, fps: 30}, base: {type: opencv, index_or_path: "/dev/video2", width: 640, height: 480, fps: 30}}' \
-    --teleop.type=openarm_mini \
-    --teleop.port_left=/dev/ttyACM0 \
-    --teleop.port_right=/dev/ttyACM1 \
+    --teleop.type=bi_openarm_mini \
+    --teleop.left_arm_config.port=/dev/ttyACM0 \
+    --teleop.right_arm_config.port=/dev/ttyACM1 \
     --policy.path=outputs/pretrain/checkpoints/last/pretrained_model \
     --dataset.repo_id=your-username/rollout_hil_dataset \
     --dataset.single_task="Fold the T-shirt properly" \
@@ -131,9 +131,9 @@ lerobot-rollout --strategy.type=dagger \
     --robot.right_arm_config.port=can0 \
     --robot.right_arm_config.side=right \
     --robot.cameras='{left_wrist: {type: opencv, index_or_path: "/dev/video0", width: 1280, height: 720, fps: 30}, right_wrist: {type: opencv, index_or_path: "/dev/video4", width: 1280, height: 720, fps: 30}, base: {type: opencv, index_or_path: "/dev/video2", width: 640, height: 480, fps: 30}}' \
-    --teleop.type=openarm_mini \
-    --teleop.port_left=/dev/ttyACM0 \
-    --teleop.port_right=/dev/ttyACM1 \
+    --teleop.type=bi_openarm_mini \
+    --teleop.left_arm_config.port=/dev/ttyACM0 \
+    --teleop.right_arm_config.port=/dev/ttyACM1 \
     --policy.path=outputs/pretrain/checkpoints/last/pretrained_model \
     --dataset.repo_id=your-username/rollout_hil_rtc_dataset \
     --dataset.single_task="Fold the T-shirt properly" \

docs/source/il_robots.mdx

@@ -647,5 +647,6 @@ The `--strategy.type` flag selects the execution mode:
 - `sentry`: Continuous recording with auto-upload (useful for large-scale evaluation)
 - `highlight`: Ring buffer recording with keystroke save (useful for capturing interesting events)
 - `dagger`: Human-in-the-loop data collection (see [HIL Data Collection](./hil_data_collection))
+- `episodic`: Episode-oriented policy recording with reset phases between episodes
 
 All strategies support `--inference.type=rtc` for smooth execution with slow VLA models (Pi0, Pi0.5, SmolVLA).

docs/source/inference.mdx

@@ -117,7 +117,7 @@ lerobot-rollout \
     --strategy.num_episodes=20 \
     --policy.path=outputs/pretrain/checkpoints/last/pretrained_model \
     --robot.type=bi_openarm_follower \
-    --teleop.type=openarm_mini \
+    --teleop.type=bi_openarm_mini \
     --dataset.repo_id=${HF_USER}/rollout_hil_data \
     --dataset.single_task="Fold the T-shirt"
 ```
@@ -157,6 +157,44 @@ Foot pedal input is also supported via `--strategy.input_device=pedal`. Configur
 | `--strategy.input_device`            | Input device: `keyboard` or `pedal` (default: keyboard) |
 | `--teleop.type`                      | **Required.** Teleoperator type                         |
 
+### Episodic (`--strategy.type=episodic`)
+
+Episode-oriented recording that mirrors the behavior of `lerobot-record`. The policy drives the robot for each episode; an optional teleoperator can drive the robot during the reset phase between episodes.
+
+```bash
+lerobot-rollout \
+    --strategy.type=episodic \
+    --policy.path=${HF_USER}/my_policy \
+    --robot.type=so100_follower \
+    --robot.port=/dev/ttyACM0 \
+    --teleop.type=so100_leader \
+    --teleop.port=/dev/ttyACM1 \
+    --dataset.repo_id=${HF_USER}/my_eval_data \
+    --dataset.num_episodes=20 \
+    --dataset.episode_time_s=30 \
+    --dataset.reset_time_s=10 \
+    --dataset.single_task="Pick up the red cube"
+```
+
+Teleop is optional — if omitted the robot holds its position during the reset phase.
+
+**Keyboard controls:**
+
+| Key         | Action                           |
+| ----------- | -------------------------------- |
+| `→` (right) | End the current episode early    |
+| `←` (left)  | Discard episode and re-record it |
+| `ESC`       | Stop the recording session       |
+
+| Flag                                            | Description                                                                |
+| ----------------------------------------------- | -------------------------------------------------------------------------- |
+| `--dataset.num_episodes`                        | Number of episodes to record                                               |
+| `--dataset.episode_time_s`                      | Duration of each recording episode in seconds                              |
+| `--dataset.reset_time_s`                        | Duration of the reset phase between episodes in seconds                    |
+| `--teleop.type`                                 | Optional. Teleoperator to drive the robot during resets                    |
+| `--strategy.reset_to_initial_position`          | Whether to reset the robot to its initial position between episodes        |
+| `--strategy.smooth_leader_to_follower_handover` | Whether to turn on or off the leader -> follower smooth handover behavior. |
+
 ---
 
 ## Inference Backends

docs/source/policy_groot_README.md

@@ -1,13 +1,6 @@
 ## Research Paper
 
-GR00T N1 technical report (covers the GR00T N1.x family, including N1.7): https://arxiv.org/abs/2503.14734
-
-GR00T N1.7 model card: https://huggingface.co/nvidia/GR00T-N1.7-3B
-
-GR00T N1.5 research page (earlier version): https://research.nvidia.com/labs/gear/gr00t-n1_5/
-
-> GR00T N1.5 support was removed from LeRobot; the last release supporting it is `lerobot==0.5.1`.
-> Current releases support GR00T N1.7 only.
+Paper: https://research.nvidia.com/labs/gear/gr00t-n1_5/
 
 ## Repository
 
@@ -31,103 +24,4 @@ Code: https://github.com/NVIDIA/Isaac-GR00T
 
 Blog: https://developer.nvidia.com/isaac/gr00t
 
-Hugging Face Models:
-
-- GR00T N1.7: https://huggingface.co/nvidia/GR00T-N1.7-3B
-- GR00T N1.7 LIBERO checkpoints: https://huggingface.co/nvidia/GR00T-N1.7-LIBERO
-
-## Original-vs-LeRobot parity test
-
-`tests/policies/groot/test_groot_vs_original.py` verifies this LeRobot
-reimplementation of GR00T N1.7 (Qwen3-VL backbone + flow-matching action head)
-against NVIDIA's original `gr00t` package with two comparisons, each parametrized
-over every embodiment tag present in the checkpoint:
-
-1. **Model parity** — given byte-identical pre-processed inputs and the same
-   flow-matching seed (recorded in each artifact), both implementations must produce
-   the **same raw model output** (`get_action(...)["action_pred"]`, the normalized
-   flow-matching prediction). Output shapes must match exactly; any action-horizon
-   or action-dim mismatch fails the test.
-2. **Preprocessor parity** — given the identical raw observations (per-camera
-   frames, state vectors, language instruction), LeRobot's own preprocessor pipeline
-   (real Qwen3-VL chat template / tokenizer / image packing + checkpoint-driven
-   state normalization, no mocks) must produce the **same collated model inputs**
-   (`input_ids`, `attention_mask`, `pixel_values`, `image_grid_thw`, `state`,
-   `embodiment_id`) as the original package's processor.
-
-### Why two environments
-
-The original `gr00t` package pins `transformers==4.57.3` (Python 3.10); this
-integration requires `transformers>=5.x` (Qwen3-VL). Under 5.x, `PretrainedConfig`
-is itself a defaulted dataclass, so the original config dataclasses fail to import
-(`non-default argument follows default argument`). The two implementations therefore
-**cannot be imported in the same Python process**.
-
-So the test uses a **producer / consumer** split across two venvs:
-
-1. **Producer** — `tests/policies/groot/utils/dump_original_n1_7.py`, run in the _original_
-   gr00t venv. For each embodiment it builds dummy inputs generically from the
-   checkpoint metadata (state dims from `statistics.json`; camera/language keys from
-   the processor modality configs), runs the original model, and saves to one `.npz`
-   per tag: the raw observations (`raw::` keys), the exact collated inputs
-   (`in::` keys), the seed, and the raw `action_pred`.
-2. **Consumer** — the pytest above, run in the _LeRobot_ venv. It discovers every
-   `.npz`; the model-parity case replays the byte-identical collated inputs through
-   the LeRobot model with the recorded seed and asserts the outputs match, and the
-   preprocessor-parity case replays the raw observations through LeRobot's full
-   preprocessor pipeline and asserts the collated tensors match.
-
-> Artifacts generated by older versions of the dump script contain no `raw::`
-> fields; the preprocessor-parity case then **skips** with a regeneration hint.
-> Re-run the producer to refresh them.
-
-### Fairness controls
-
-- **Same pre-processed inputs (model parity)** — the original processor's `input_ids`,
-  `pixel_values`, `image_grid_thw`, `attention_mask`, `state`, `embodiment_id` are
-  fed verbatim to the LeRobot model (no re-tokenization / re-normalization), so the
-  model comparison isolates the model. LeRobot's own tokenization / image packing is
-  covered separately by the preprocessor-parity case, which compares its output
-  against those same collated tensors from identical raw observations.
-- **Same precision + attention kernel** — both sides run **fp32 + SDPA**. The
-  original defaults to `use_flash_attention=True` (flash_attention_2 + bf16); the
-  producer forces SDPA + fp32. (With the defaults the gap is ~3e-2 — pure
-  kernel/rounding noise, not an implementation difference.)
-- **Same flow-matching seed** — fixed right before sampling on both sides; the
-  producer records it in each artifact (`--seed`, default 42) and the consumer
-  replays the recorded value.
-
-### How to run
-
-```bash
-# Resolve a local checkpoint (GR00T-N1.7-LIBERO / libero_10)
-CKPT=$(python - <<'PY'
-import os
-from huggingface_hub import snapshot_download
-print(os.path.join(snapshot_download("nvidia/GR00T-N1.7-LIBERO",
-      allow_patterns=["libero_10/*"]), "libero_10"))
-PY
-)
-
-# 1) Produce the original-side artifacts for all embodiments (original gr00t venv, CUDA)
-CUDA_VISIBLE_DEVICES=0 /path/to/Isaac-GR00T/.venv-original/bin/python \
-    tests/policies/groot/utils/dump_original_n1_7.py \
-    --ckpt "$CKPT" --out-dir tests/policies/groot/artifacts --device cuda --seed 42
-
-# 2) Run the parity test (LeRobot venv) — one parametrized case per embodiment
-CUDA_VISIBLE_DEVICES=0 GROOT_PARITY_DEVICE=cuda \
-    uv run pytest tests/policies/groot/test_groot_vs_original.py -v -s
-```
-
-The `.npz` artifacts are local-only (gitignored, ~6–10 MB each) and are regenerated by
-the producer; they are never committed. The tests **skip** (do not fail) on CI or
-when the checkpoint / artifacts are absent.
-
-#### Env knobs (all optional)
-
-| Var                                       | Default                          | Purpose                               |
-| ----------------------------------------- | -------------------------------- | ------------------------------------- |
-| `GROOT_N1_7_PARITY_DIR`                   | `tests/policies/groot/artifacts` | directory of per-tag `.npz` artifacts |
-| `GROOT_N1_7_LIBERO_CKPT`                  | auto (HF cache)                  | override checkpoint dir               |
-| `GROOT_PARITY_DEVICE`                     | `cuda` if available              | `cpu` or `cuda`                       |
-| `GROOT_PARITY_ATOL` / `GROOT_PARITY_RTOL` | `1e-3`                           | comparison tolerance                  |
+Hugging Face Model: https://huggingface.co/nvidia/GR00T-N1.5-3B

examples/annotations/run_hf_job.py

@@ -0,0 +1,77 @@
+#!/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' && "
+    "pip install --upgrade-strategy only-if-needed "
+    "datasets pyarrow av jsonlines draccus 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}")

pyproject.toml

@@ -115,8 +115,8 @@ dataset = [
 ]
 training = [
     "lerobot[dataset]",
-    "accelerate>=1.10.0,<2.0.0",
-    "wandb>=0.24.0,<0.25.0",
+    "wandb>=0.24.0,<0.28.0",
+    "lerobot[accelerate-dep]",
 ]
 hardware = [
     "lerobot[pynput-dep]",
@@ -142,7 +142,8 @@ pygame-dep = ["pygame>=2.5.1,<2.7.0"]
 # (noble ships urdfdom 3.x). Cap below 0.9.16 until system urdfdom 4.x is broadly available.
 placo-dep = ["placo>=0.9.6,<0.9.16"]
 transformers-dep = ["transformers>=5.4.0,<5.6.0"]
-grpcio-dep = ["grpcio==1.73.1", "protobuf>=6.31.1,<6.32.0"]
+grpcio-dep = ["grpcio>=1.73.1,<2.0.0", "protobuf>=6.31.1,<8.0.0"]
+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"]
@@ -177,7 +178,12 @@ unitree_g1 = [
     "lerobot[matplotlib-dep]",
     "lerobot[pygame-dep]",
 ]
-reachy2 = ["reachy2_sdk>=1.0.15,<1.1.0"]
+# reachy2-sdk caps grpcio<=1.73.1 and protobuf<=6.32.0; quarantined here so downstream users aren't held back. reachy2-sdk is unlikely to release new versions.
+reachy2 = [
+    "reachy2_sdk>=1.0.15,<1.1.0",
+    "grpcio<=1.73.1",
+    "protobuf<=6.32.0",
+]
 # Seeed Studio reBot B601-DM follower (motorbridge / CAN) + StarArm102 / reBot Arm 102
 # leader (motorbridge-smart-servo / FashionStar UART servos).
 rebot = ["lerobot[motorbridge-dep]", "lerobot[motorbridge-smart-servo-dep]"]
@@ -199,7 +205,7 @@ wallx = [
 ]
 pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"]
 molmoact2 = ["lerobot[transformers-dep]", "lerobot[peft-dep]", "lerobot[scipy-dep]"]
-smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "accelerate>=1.7.0,<2.0.0"]
+smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "lerobot[accelerate-dep]"]
 multi_task_dit = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]"]
 groot = [
     "lerobot[transformers-dep]",
@@ -208,30 +214,47 @@ groot = [
     "dm-tree>=0.1.8,<1.0.0",
     "timm>=1.0.0,<1.1.0",
     "decord>=0.6.0,<1.0.0; (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
+    "ninja>=1.11.1,<2.0.0",
+    "flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'"
 ]
 sarm = ["lerobot[transformers-dep]", "pydantic>=2.0.0,<3.0.0", "faker>=33.0.0,<35.0.0", "lerobot[matplotlib-dep]", "lerobot[qwen-vl-utils-dep]"]
 robometer = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]", "lerobot[peft-dep]"]
 topreward = ["lerobot[transformers-dep]"]
 xvla = ["lerobot[transformers-dep]"]
 eo1 = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]"]
-hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
+hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.14,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
 vla_jepa = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[qwen-vl-utils-dep]"]
 
 # Features
 async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"]
 peft = ["lerobot[transformers-dep]", "lerobot[peft-dep]"]
 
+# Annotation pipeline (lerobot-annotate). The only backend is ``openai``,
+# which talks to any OpenAI-compatible server (``vllm serve`` /
+# ``transformers serve`` / hosted). Distributed runs use Hugging Face Jobs
+# (see examples/annotations/run_hf_job.py).
+annotations = [
+    "lerobot[dataset]",
+    "lerobot[transformers-dep]",
+    "openai>=1.40,<2.0",
+    # ``vllm`` is intentionally NOT a hard dep: it pins an older torch, and
+    # uv's single unified lock would then cap ``torch`` for every extra
+    # (e.g. forcing 2.8 while ``torchcodec`` in [dataset] needs 2.11 -> ABI
+    # break in CI). The HF Jobs image (``vllm/vllm-openai``) provides vLLM;
+    # install it locally only if you run your own ``vllm serve``.
+]
+
 # Development
-dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1", "mypy>=1.19.1", "ruff>=0.14.1", "lerobot[notebook]"]
+dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools>=1.73.1,<2.0.0", "mypy>=1.19.1", "ruff>=0.14.1", "lerobot[notebook]"]
 notebook = ["jupyter>=1.0.0,<2.0.0", "ipykernel>=6.0.0,<7.0.0"]
 test = ["pytest>=8.1.0,<9.0.0", "pytest-timeout>=2.4.0,<3.0.0", "pytest-cov>=5.0.0,<8.0.0", "mock-serial>=0.0.1,<0.1.0 ; sys_platform != 'win32'"]
 video_benchmark = ["scikit-image>=0.23.2,<0.26.0", "pandas>=2.2.2,<2.4.0"]
 
 # Simulation
 # NOTE: Explicitly listing scipy helps flatten the dependecy tree.
-aloha = ["lerobot[dataset]", "gym-aloha>=0.1.2,<0.2.0", "lerobot[scipy-dep]"]
+aloha = ["lerobot[dataset]", "gym-aloha>=0.1.4,<0.2.0", "lerobot[scipy-dep]"]
 pusht = ["lerobot[dataset]", "gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
-libero = ["lerobot[dataset]", "lerobot[transformers-dep]", "hf-libero>=0.1.3,<0.2.0; sys_platform == 'linux'", "lerobot[scipy-dep]"]
+libero = ["lerobot[dataset]", "lerobot[transformers-dep]", "hf-libero>=0.1.4,<0.2.0; sys_platform == 'linux'", "lerobot[scipy-dep]"]
 metaworld = ["lerobot[dataset]", "metaworld==3.0.0", "lerobot[scipy-dep]"]
 # NOTE: vlabench is NOT exposed as a `lerobot` extra. Its only distribution
 # is the OpenMOSS/VLABench GitHub repo (package name `VLABench`, no PyPI
@@ -278,7 +301,7 @@ all = [
     "lerobot[pi]",
     "lerobot[molmoact2]",
     "lerobot[smolvla]",
-    "lerobot[groot]",
+    # "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn
     "lerobot[xvla]",
     "lerobot[hilserl]",
     "lerobot[vla_jepa]",
@@ -315,6 +338,7 @@ lerobot-find-joint-limits="lerobot.scripts.lerobot_find_joint_limits:main"
 lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
 lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main"
 lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main"
+lerobot-annotate="lerobot.scripts.lerobot_annotate:main"
 lerobot-rollout="lerobot.scripts.lerobot_rollout:main"
 
 # ---------------- Tool Configurations ----------------
@@ -333,7 +357,7 @@ torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 
 [tool.setuptools.package-data]
-lerobot = ["envs/*.json"]
+lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/lerobot/annotations/__init__.py

@@ -0,0 +1,15 @@
+#!/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.

src/lerobot/annotations/steerable_pipeline/__init__.py

@@ -0,0 +1,36 @@
+#!/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.
+"""Steerable annotation pipeline producing ``language_persistent`` and
+``language_events`` columns for LeRobot datasets.
+
+The pipeline is decomposed into three independently runnable modules whose
+outputs are staged per-episode before a final parquet rewrite:
+
+- :mod:`.modules.plan_subtasks_memory` (the ``plan`` module) — persistent styles
+- :mod:`.modules.interjections_and_speech` (the ``interjections`` module) — event styles + speech
+- :mod:`.modules.general_vqa` (the ``vqa`` module) — event-style VQA pairs
+"""
+
+from .config import AnnotationPipelineConfig
+from .validator import StagingValidator, ValidationReport
+from .writer import LanguageColumnsWriter
+
+__all__ = [
+    "AnnotationPipelineConfig",
+    "LanguageColumnsWriter",
+    "StagingValidator",
+    "ValidationReport",
+]

src/lerobot/annotations/steerable_pipeline/config.py

@@ -0,0 +1,211 @@
+#!/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.
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class PlanConfig:
+    """``plan`` module: subtasks + plan + memory + task augmentation."""
+
+    enabled: bool = True
+
+    # ``task_aug`` rephrasings at t=0 (renderer rotates ${task} among them); 0 disables.
+    n_task_rephrasings: int = 10
+
+    # Derive the task from video instead of episode_task: off / if_short / always.
+    # Affects prompts only; ``meta/tasks.parquet`` is untouched.
+    derive_task_from_video: str = "if_short"
+    derive_task_min_words: int = 3
+
+    # --- Frame input: timestamped contact sheets (always on) ---------------
+    # The subtask describe/segment passes ALWAYS render the episode as
+    # macrodata/refiner-style contact sheets: sampled frames packed into JPEG
+    # grids with each frame's timestamp burned into its corner, so the VLM
+    # cites the exact source time of a boundary directly. This is far cheaper
+    # in vision tokens than one image per frame (≈2× faster subtask generation
+    # in practice), which is why the sampling is dense by default.
+    #
+    # ``frames_per_second`` is the sampling rate: 2.0 = one frame every 0.5s.
+    frames_per_second: float = 2.0
+    # Frame budget per VLM call (= columns × rows × sheets). When a whole
+    # episode sampled at ``frames_per_second`` exceeds this, the episode is
+    # AUTOMATICALLY split into consecutive windows of
+    # ``max_frames_per_prompt`` frames each (one describe→segment call per
+    # window, still at the full ``frames_per_second`` density), and the
+    # per-window spans are merged + stitched into one contiguous cover. So an
+    # episode of any length is always covered at the full sampling density.
+    max_frames_per_prompt: int = 60
+    contact_sheet_columns: int = 5
+    contact_sheet_frames_per_sheet: int = 20
+    contact_sheet_frame_width: int = 224
+    contact_sheet_quality: int = 84
+
+    min_subtask_seconds: float = 1.5
+    plan_max_steps: int = 8
+
+    # Narrate-only grounding pass before segmenting — best defense against subtasks
+    # invented from the task text (+1 VLM call/episode).
+    subtask_describe_first: bool = True
+
+    # Emit ``style="plan"`` rows at each boundary; False = subtasks + memory only.
+    emit_plan: bool = True
+
+    # Emit ``style="memory"`` rows at each boundary; False = subtasks (+ plan) only.
+    # Symmetric counterpart of ``emit_plan``.
+    emit_memory: bool = True
+
+    # (subtask spans are always stitched to a contiguous full-episode cover; not configurable.)
+
+    # Optional EgoMimic-style 5-axis task augmentation; replaces n_task_rephrasings.
+    task_aug_axes: TaskAugAxesConfig = field(default_factory=lambda: TaskAugAxesConfig())
+
+
+@dataclass
+class TaskAugAxesConfig:
+    """5-axis t=0 task augmentation (EgoMimic-style): synonym / omit_arm /
+    omit_orientation / omit_grasp_method / combined. Replaces n_task_rephrasings
+    when enabled; each variant becomes a ``task_aug`` row. Axes with nothing to
+    omit emit fewer entries. Defaults (3+3+2+2+2) match EgoMimic."""
+
+    enabled: bool = False
+
+    synonym_paraphrase: int = 3
+    omit_arm: int = 3
+    omit_orientation: int = 2
+    omit_grasp_method: int = 2
+    combined_omissions: int = 2
+
+
+@dataclass
+class InterjectionsConfig:
+    """``interjections`` module: interjections + paired speech."""
+
+    enabled: bool = True
+
+    # Each emits a paired (interjection, speech) row + a plan refresh at that ts.
+    max_interjections_per_episode: int = 3
+    interjection_min_t: float = 2.0
+
+    # Frame window centered on the timestamp so the VLM sees motion, not one frame.
+    interjection_window_seconds: float = 2.0
+    interjection_window_frames: int = 4
+
+
+@dataclass
+class VqaConfig:
+    """``vqa`` module: general VQA."""
+
+    enabled: bool = True
+    vqa_emission_hz: float = 1.0
+    K: int = 1
+    """Consecutive frames per emission tick. The VLM grounds on the FIRST frame,
+    so K>1 smears stale labels onto moved frames. Default 1 (no smear)."""
+    question_types: tuple[str, ...] = ("bbox", "keypoint", "count", "attribute", "spatial")
+
+    # True: ground VQA only on --vlm.camera_key (default: every camera).
+    restrict_to_default_camera: bool = False
+
+
+@dataclass
+class VlmConfig:
+    """Shared Qwen-VL client configuration."""
+
+    # Only ``openai`` (OpenAI-compatible vLLM server, auto-spawned when
+    # auto_serve=True); ``stub`` is for tests.
+    backend: str = "openai"
+    model_id: str = "Qwen/Qwen3.6-27B"
+
+    # OpenAI-compatible endpoint; ``EMPTY`` key works for local servers.
+    api_base: str = "http://localhost:8000/v1"
+    api_key: str = "EMPTY"
+
+    # Spawn a server if none answers api_base; False = fail fast on a remote.
+    auto_serve: bool = True
+    serve_port: int = 8000
+    # Override the auto-serve command; ``{port}`` substituted per replica.
+    serve_command: str | None = None
+
+    # Independent servers for round-robin routing (one per GPU). num_gpus=0 = one each.
+    parallel_servers: int = 1
+    num_gpus: int = 0
+    client_concurrency: int = 16
+    serve_ready_timeout_s: float = 600.0
+
+    max_new_tokens: int = 512
+    temperature: float = 0.2
+
+    # Auto-serve context length (None → 32768); other vLLM flags go in serve_command.
+    max_model_len: int | None = None
+
+    # Camera for keyframes; None → first ``observation.images.*`` key.
+    camera_key: str | None = None
+    # Forwarded as extra_body.chat_template_kwargs (e.g. {"enable_thinking": false}).
+    chat_template_kwargs: dict[str, Any] | None = None
+
+
+@dataclass
+class ExecutorConfig:
+    """Executor settings (intra-process episode concurrency; distribution via HF Jobs)."""
+
+    # Episodes processed concurrently per phase; main knob for saturating the servers.
+    episode_parallelism: int = 16
+
+
+@dataclass
+class AnnotationPipelineConfig:
+    """Top-level config for ``lerobot-annotate`` (rewrites data shards in place)."""
+
+    # Hub dataset: download source when ``root`` unset; push target when push_to_hub
+    # is on and ``new_repo_id`` unset.
+    repo_id: str | None = None
+
+    # Separate push target (matches the LeRobot edit tools). Unset → push in place.
+    new_repo_id: str | None = None
+
+    root: Path | None = None
+
+    # Defaults to ``<root>/.annotate_staging/``.
+    staging_dir: Path | None = None
+
+    seed: int = 1729
+
+    plan: PlanConfig = field(default_factory=PlanConfig)
+    interjections: InterjectionsConfig = field(default_factory=InterjectionsConfig)
+    vqa: VqaConfig = field(default_factory=VqaConfig)
+
+    vlm: VlmConfig = field(default_factory=VlmConfig)
+    executor: ExecutorConfig = field(default_factory=ExecutorConfig)
+
+    skip_validation: bool = False
+    only_episodes: tuple[int, ...] | None = None
+
+    # Keyframe decode backend forwarded to ``decode_video_frames``. None →
+    # library default (torchcodec when available, else PyAV). Or pin
+    # ``"torchcodec"`` / ``"pyav"`` explicitly.
+    video_backend: str | None = None
+
+    # Upload to the Hub (new_repo_id if set, else repo_id; one must be set).
+    push_to_hub: bool = False
+    push_private: bool = False
+    push_commit_message: str | None = None
+
+    def resolved_staging_dir(self, root: Path) -> Path:
+        return self.staging_dir if self.staging_dir is not None else root / ".annotate_staging"

src/lerobot/annotations/steerable_pipeline/executor.py

@@ -0,0 +1,253 @@
+#!/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.
+"""In-process executor that runs the annotation phases.
+
+The executor runs **six phases** in dependency order:
+
+    phase 1: ``plan`` module (plan + subtasks + memory)
+    phase 2: ``interjections`` module (interjections + speech)
+    phase 3: ``plan`` plan-update pass — re-runs plan emission at every
+             interjection timestamp produced by phase 2
+    phase 4: ``vqa`` module (VQA)
+    phase 5: validator
+    phase 6: writer
+
+Phase 3 is why the ``plan`` module must be re-entered after the
+``interjections`` module — to refresh ``plan`` rows at interjection
+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.
+Episode-level concurrency is controlled by
+``ExecutorConfig.episode_parallelism``.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from .config import AnnotationPipelineConfig
+from .reader import EpisodeRecord, iter_episodes
+from .staging import EpisodeStaging
+from .validator import StagingValidator
+from .writer import LanguageColumnsWriter
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PhaseResult:
+    """Summary of one pipeline phase across all episodes."""
+
+    name: str
+    episodes_processed: int
+    episodes_skipped: int
+
+
+@dataclass
+class PipelineRunSummary:
+    """Aggregated result returned by :meth:`Executor.run`."""
+
+    phases: list[PhaseResult]
+    written_paths: list[Path]
+    validation_report: Any  # ValidationReport, kept Any to avoid import cycle
+
+
+@dataclass
+class Executor:
+    """Run all six phases over a dataset root in-process.
+
+    Episode-level concurrency comes from ``ExecutorConfig.episode_parallelism``
+    (a thread pool); cluster-level concurrency comes from running this
+    executor inside a Hugging Face Job. Tests construct the executor
+    directly with stub modules.
+    """
+
+    config: AnnotationPipelineConfig
+    plan: Any  # PlanSubtasksMemoryModule
+    interjections: Any  # InterjectionsAndSpeechModule
+    vqa: Any  # GeneralVqaModule
+    writer: LanguageColumnsWriter
+    validator: StagingValidator
+
+    def run(self, root: Path) -> PipelineRunSummary:
+        records = list(iter_episodes(root, only_episodes=self.config.only_episodes))
+        n = len(records)
+        if n == 0:
+            raise ValueError(f"No episodes found under {root}/data/")
+
+        print(f"[annotate] {n} episodes total", flush=True)
+
+        staging_dir = self.config.resolved_staging_dir(root)
+        staging_dir.mkdir(parents=True, exist_ok=True)
+
+        phases: list[PhaseResult] = []
+
+        # Phase 1: ``plan`` module (plan + subtasks + memory)
+        phases.append(self._run_module_phase("plan", records, staging_dir, self.plan))
+        # Phase 2: ``interjections`` module (interjections + speech). It
+        # reads the ``plan`` module's subtask rows from the same staging
+        # tree to ground the interjection prompt in the correct local subtask.
+        phases.append(self._run_module_phase("interjections", records, staging_dir, self.interjections))
+        # Phase 3: ``plan`` plan-update pass at interjection timestamps.
+        phases.append(self._run_plan_update_phase(records, staging_dir))
+        # Phase 4: ``vqa`` module (VQA)
+        phases.append(self._run_module_phase("vqa", records, staging_dir, self.vqa))
+
+        print("[annotate] running validator...", flush=True)
+        report = self.validator.validate(records, staging_dir)
+        if not report.ok and not self.config.skip_validation:
+            raise RuntimeError(f"Staging validation failed: {report.summary()}")
+        print(f"[annotate] validator: {report.summary()}", flush=True)
+
+        print(f"[annotate] writing parquet shards into {root}/data/...", flush=True)
+        written = self.writer.write_all(records, staging_dir, root)
+        print(f"[annotate] wrote {len(written)} shard(s); pipeline complete", flush=True)
+
+        # Keep meta/info.json aligned with the parquet schema we just wrote.
+        # Idempotent and additive: existing user metadata is preserved.
+        self._ensure_annotation_metadata_in_info(root)
+
+        return PipelineRunSummary(phases=phases, written_paths=written, validation_report=report)
+
+    @staticmethod
+    def _ensure_annotation_metadata_in_info(root: Path) -> None:
+        """Write language features and canonical tools to ``meta/info.json``.
+
+        ``LanguageColumnsWriter`` adds ``language_persistent`` and
+        ``language_events`` to parquet shards. The metadata must advertise
+        those columns too, otherwise non-streaming ``LeRobotDataset`` loads
+        cast against the old schema and fail on the extra parquet columns.
+        """
+        from lerobot.datasets.io_utils import load_info, write_info  # noqa: PLC0415
+        from lerobot.datasets.language import SAY_TOOL_SCHEMA, language_feature_info  # noqa: PLC0415
+
+        info_path = root / "meta" / "info.json"
+        if not info_path.exists():
+            return
+        try:
+            info = load_info(root)
+        except Exception as exc:  # noqa: BLE001
+            print(f"[annotate] could not read {info_path}: {exc}", flush=True)
+            return
+
+        changed = False
+
+        merged_features = {**info.features, **language_feature_info()}
+        if merged_features != info.features:
+            info.features = merged_features
+            changed = True
+
+        existing = info.tools or []
+        names = {(t.get("function") or {}).get("name") for t in existing if isinstance(t, dict)}
+        if SAY_TOOL_SCHEMA["function"]["name"] not in names:
+            info.tools = [*existing, SAY_TOOL_SCHEMA]
+            changed = True
+
+        if changed:
+            write_info(info, root)
+            print(
+                "[annotate] meta/info.json: "
+                f"language_features={list(language_feature_info())}, "
+                f"tools={[t['function']['name'] for t in (info.tools or [])]}",
+                flush=True,
+            )
+
+    def _run_module_phase(
+        self,
+        name: str,
+        records: list[EpisodeRecord],
+        staging_dir: Path,
+        module: Any,
+    ) -> PhaseResult:
+        if not module.enabled:
+            print(f"[annotate] phase={name} skipped (module disabled)", flush=True)
+            return PhaseResult(name=name, episodes_processed=0, episodes_skipped=len(records))
+        n = len(records)
+        parallelism = max(1, min(self.config.executor.episode_parallelism, n))
+        print(
+            f"[annotate] phase={name} starting on {n} episode(s) (parallelism={parallelism})",
+            flush=True,
+        )
+        t0 = time.time()
+
+        def _do(idx_record: tuple[int, EpisodeRecord]) -> tuple[int, int, float]:
+            i, record = idx_record
+            ep_start = time.time()
+            staging = EpisodeStaging(staging_dir, record.episode_index)
+            module.run_episode(record, staging)
+            return i, record.episode_index, time.time() - ep_start
+
+        processed = 0
+        if parallelism == 1:
+            for i, record in enumerate(records, 1):
+                _, ep_idx, elapsed = _do((i, record))
+                processed += 1
+                print(
+                    f"[annotate]   {name} episode {i}/{n} (idx={ep_idx}) done in {elapsed:.1f}s",
+                    flush=True,
+                )
+        else:
+            with ThreadPoolExecutor(max_workers=parallelism) as pool:
+                futures = [pool.submit(_do, (i, r)) for i, r in enumerate(records, 1)]
+                for fut in as_completed(futures):
+                    i, ep_idx, elapsed = fut.result()
+                    processed += 1
+                    print(
+                        f"[annotate]   {name} episode {processed}/{n} "
+                        f"(idx={ep_idx}, submit_order={i}) done in {elapsed:.1f}s",
+                        flush=True,
+                    )
+        total = time.time() - t0
+        print(f"[annotate] phase={name} complete: {processed}/{n} in {total:.1f}s", flush=True)
+        return PhaseResult(name=name, episodes_processed=processed, episodes_skipped=0)
+
+    def _run_plan_update_phase(  # noqa: PLR0915
+        self, records: list[EpisodeRecord], staging_dir: Path
+    ) -> PhaseResult:
+        """Re-emit ``plan`` rows at each timestamp the ``interjections`` module produced.
+
+        The ``plan`` module owns the prompt; the ``interjections`` module
+        produced the timestamps. This phase therefore calls back into the
+        ``plan`` module with the interjection timestamps so its existing
+        prompt path is reused.
+        """
+        if not self.plan.enabled or not self.interjections.enabled:
+            return PhaseResult(name="plan_update", episodes_processed=0, episodes_skipped=len(records))
+        processed = 0
+        for record in records:
+            staging = EpisodeStaging(staging_dir, record.episode_index)
+            interjection_rows = [
+                row for row in staging.read("interjections") if row.get("style") == "interjection"
+            ]
+            interjection_times = [float(row["timestamp"]) for row in interjection_rows]
+            interjection_texts = [str(row.get("content") or "") for row in interjection_rows]
+            if interjection_times:
+                self.plan.run_plan_updates(record, staging, interjection_times, interjection_texts)
+                processed += 1
+        # Episodes without any interjections are skipped (no plan refresh
+        # needed); count them so the summary's processed+skipped == total.
+        return PhaseResult(
+            name="plan_update",
+            episodes_processed=processed,
+            episodes_skipped=len(records) - processed,
+        )

src/lerobot/annotations/steerable_pipeline/frames.py

@@ -0,0 +1,481 @@
+#!/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.
+"""Keyframe extraction for the annotation pipeline.
+
+Modules attach decoded camera frames to their VLM prompts so the model can
+ground subtask decomposition, interjection scenarios, and VQA in actual
+visual content. The pipeline shares one provider across modules and one
+episode at a time, with a small per-episode cache so multiple modules
+querying the same timestamp pay decode cost once.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import math
+import threading
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Protocol
+
+import PIL.Image
+import torch
+
+from lerobot.configs.video import VideoEncoderConfig
+from lerobot.datasets.video_utils import decode_video_frames, reencode_video
+
+from .reader import EpisodeRecord, snap_to_frame
+
+logger = logging.getLogger(__name__)
+
+
+class FrameProvider(Protocol):
+    """Decodes camera frames at episode-relative timestamps."""
+
+    @property
+    def camera_keys(self) -> list[str]:
+        """All ``observation.images.*`` feature keys this provider can decode."""
+
+    def frames_at(
+        self,
+        record: EpisodeRecord,
+        timestamps: list[float],
+        camera_key: str | None = None,
+    ) -> list[Any]:
+        """Return one decoded frame per timestamp from ``camera_key`` (or default).
+
+        Frames are ``torch.Tensor`` (``C, H, W`` uint8) — the shape
+        :func:`lerobot.datasets.video_utils.decode_video_frames` returns.
+        :func:`to_image_blocks` converts them to PIL only at the VLM-message
+        boundary.
+
+        Empty list if the camera is unavailable. ``camera_key=None`` falls back
+        to the provider's default camera so existing single-camera callers
+        (the ``plan`` and ``interjections`` modules) keep working unchanged.
+        """
+
+    def video_for_episode(
+        self,
+        record: EpisodeRecord,
+        max_frames: int,
+        camera_key: str | None = None,
+    ) -> list[Any]:
+        """Return up to ``max_frames`` decoded frames covering the whole episode.
+
+        Sampling is uniform across the episode duration. Frames are
+        ``torch.Tensor`` (``C, H, W`` uint8); :func:`to_video_block` wraps
+        them into one ``{"type":"video", "video":<list>}`` block for a
+        Qwen-VL-compatible model that pools temporally itself. Empty list if
+        no camera available.
+        """
+
+
+@dataclass
+class _NullProvider:
+    """No-op provider used when the dataset has no video keys or in tests."""
+
+    @property
+    def camera_keys(self) -> list[str]:
+        return []
+
+    def frames_at(
+        self,
+        record: EpisodeRecord,
+        timestamps: list[float],
+        camera_key: str | None = None,
+    ) -> list[Any]:
+        return []
+
+    def video_for_episode(
+        self,
+        record: EpisodeRecord,
+        max_frames: int,
+        camera_key: str | None = None,
+    ) -> list[Any]:
+        return []
+
+
+def null_provider() -> FrameProvider:
+    return _NullProvider()
+
+
+@dataclass
+class VideoFrameProvider:
+    """Decodes frames from the dataset's ``observation.images.*`` streams.
+
+    By default the *first* camera key is used for the ``plan`` module
+    (subtask decomposition) and the ``interjections`` module (interjection
+    scenarios) — those prompts care about *what is happening*, not which
+    angle. The ``vqa`` module instead iterates over every camera in
+    :attr:`camera_keys` so each frame's
+    grounded answer (bbox/keypoint/...) is tagged with the camera it was
+    grounded against.
+
+    ``camera_key`` overrides the default-camera choice but does not restrict
+    :attr:`camera_keys`. Pass ``camera_key`` explicitly to ``frames_at`` /
+    ``video_for_episode`` to read a non-default stream.
+
+    Caches up to ``cache_size`` decoded frames per process to keep
+    co-timestamped ``interjections`` + ``plan`` plan-update calls cheap.
+    """
+
+    root: Path
+    camera_key: str | None = None
+    tolerance_s: float = 1e-2
+    cache_size: int = 256
+    # Keyframe decode backend forwarded to
+    # :func:`lerobot.datasets.video_utils.decode_video_frames`. ``None``
+    # uses the library default (torchcodec when available, else PyAV).
+    video_backend: str | None = None
+    _meta: Any = field(default=None, init=False, repr=False)
+    _cache: dict = field(default_factory=dict, init=False, repr=False)
+    _camera_keys: list[str] = field(default_factory=list, init=False, repr=False)
+    # Pipeline runs the three module phases under a ThreadPoolExecutor (see
+    # ``ExecutorConfig.episode_parallelism``); guard the dict cache and the
+    # one-shot warn flag against concurrent updates from worker threads.
+    _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
+    # Serializes decode_video_frames calls: torchcodec hands out one
+    # ``VideoDecoder`` per file from a process-wide cache, and the decoder
+    # is not safe to drive from multiple threads at once.
+    _decode_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
+    _warned_decode_fail: bool = field(default=False, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata  # noqa: PLC0415
+
+        self._meta = LeRobotDatasetMetadata(repo_id="local", root=self.root)
+        # Only ``video_keys`` are decodable here: the clip/decode paths read
+        # ``videos/<key>/from_timestamp`` from episode metadata, which exists
+        # only for video-stored cameras. Image-stored cameras (also in
+        # ``camera_keys``) would KeyError, so restrict the list — and the
+        # default — to video keys.
+        keys = list(self._meta.video_keys)
+        # Last-resort fallback: if metadata didn't surface any video keys but
+        # the caller explicitly named a camera (``--vlm.camera_key=...``),
+        # trust them — the key is by definition known to exist on the dataset.
+        if not keys and self.camera_key:
+            keys = [self.camera_key]
+        self._camera_keys = keys
+        if self.camera_key is None:
+            self.camera_key = keys[0] if keys else None
+
+    @property
+    def camera_keys(self) -> list[str]:
+        """All ``observation.images.*`` keys available on this dataset."""
+        return list(self._camera_keys)
+
+    def frames_at(
+        self,
+        record: EpisodeRecord,
+        timestamps: list[float],
+        camera_key: str | None = None,
+    ) -> list[Any]:
+        target = camera_key if camera_key is not None else self.camera_key
+        if not timestamps or target is None:
+            return []
+        # Snap each request to the nearest real frame timestamp: callers
+        # sample uniform grids whose points land mid-frame, and
+        # ``decode_video_frames`` rejects queries farther than
+        # ``tolerance_s`` from a decodable frame. Snapping also dedupes
+        # repeat queries through the cache.
+        if record.frame_timestamps:
+            timestamps = [snap_to_frame(float(ts), record.frame_timestamps) for ts in timestamps]
+
+        out: list[Any] = []
+        misses: list[float] = []
+        miss_indices: list[int] = []
+        with self._lock:
+            for i, ts in enumerate(timestamps):
+                key = (record.episode_index, target, round(float(ts), 6))
+                cached = self._cache.get(key)
+                if cached is not None:
+                    out.append(cached)
+                else:
+                    out.append(None)
+                    misses.append(float(ts))
+                    miss_indices.append(i)
+
+        if misses:
+            decoded = self._decode(record.episode_index, misses, target)
+            # ``_decode`` returns exactly one frame per requested timestamp,
+            # or an empty list if decoding failed wholesale. A partial list
+            # would mean a frame/timestamp misalignment, so only pair them up
+            # when the counts match (``strict=True`` then guards regressions).
+            if len(decoded) == len(miss_indices):
+                with self._lock:
+                    for i, frame in zip(miss_indices, decoded, strict=True):
+                        out[i] = frame
+                        key = (record.episode_index, target, round(float(timestamps[i]), 6))
+                        if len(self._cache) >= self.cache_size:
+                            self._cache.pop(next(iter(self._cache)))
+                        self._cache[key] = frame
+        # filter out any None left over from decode failures
+        return [frame for frame in out if frame is not None]
+
+    def video_for_episode(
+        self,
+        record: EpisodeRecord,
+        max_frames: int,
+        camera_key: str | None = None,
+    ) -> list[Any]:
+        """Return up to ``max_frames`` frames uniformly sampled across the episode.
+
+        The whole episode duration is covered; the model picks subtask
+        boundaries from the temporal pooling it does internally. Frames are
+        ``torch.Tensor`` (see :meth:`frames_at`).
+        """
+        target = camera_key if camera_key is not None else self.camera_key
+        if max_frames <= 0 or target is None or not record.frame_timestamps:
+            return []
+        n_frames = min(max_frames, len(record.frame_timestamps))
+        if n_frames == len(record.frame_timestamps):
+            timestamps = list(record.frame_timestamps)
+        else:
+            t0 = record.frame_timestamps[0]
+            t_last = record.frame_timestamps[-1]
+            if t_last <= t0:
+                timestamps = [float(t0)] * n_frames
+            else:
+                step = (t_last - t0) / (n_frames - 1) if n_frames > 1 else 0.0
+                timestamps = [float(t0 + i * step) for i in range(n_frames)]
+        return self.frames_at(record, timestamps, camera_key=target)
+
+    def episode_clip_path(self, record: EpisodeRecord, cache_dir: Path) -> Path | None:
+        """Extract the episode's subclip to ``cache_dir/ep_{idx:06d}.mp4``.
+
+        Returns ``None`` if the dataset has no video tracks or extraction
+        failed. Skips re-extract when the cached clip already exists.
+        Re-encodes to H.264 via
+        :func:`lerobot.datasets.video_utils.reencode_video` so the resulting
+        mp4 is decodable by every downstream video processor — stream-copy
+        would inherit the source codec (often AV1 in modern LeRobot
+        datasets), which vllm's libav build cannot decode.
+        """
+        if self.camera_key is None:
+            return None
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        out_path = cache_dir / f"ep_{record.episode_index:06d}.mp4"
+        if out_path.exists() and out_path.stat().st_size > 0:
+            return out_path
+        ep = self._meta.episodes[record.episode_index]
+        from_timestamp = float(ep[f"videos/{self.camera_key}/from_timestamp"])
+        to_timestamp = float(ep[f"videos/{self.camera_key}/to_timestamp"])
+        src = self.root / self._meta.get_video_file_path(record.episode_index, self.camera_key)
+        encoder = VideoEncoderConfig(vcodec="h264", pix_fmt="yuv420p", g=None, crf=23, preset="ultrafast")
+        try:
+            reencode_video(
+                src,
+                out_path,
+                camera_encoder=encoder,
+                overwrite=True,
+                start_time_s=from_timestamp,
+                end_time_s=to_timestamp,
+            )
+        except Exception:
+            logger.warning(
+                "clip extraction failed for episode %s (%s)", record.episode_index, src, exc_info=True
+            )
+            return None
+        return out_path if out_path.exists() and out_path.stat().st_size > 0 else None
+
+    def _decode(self, episode_index: int, timestamps: list[float], camera_key: str) -> list[Any]:
+        """Decode ``timestamps`` from the episode's video as ``(C, H, W)`` tensors.
+
+        Delegates to :func:`lerobot.datasets.video_utils.decode_video_frames`
+        (torchcodec when available, PyAV otherwise; ``video_backend`` pins
+        one explicitly). Returns one frame per requested timestamp, or ``[]``
+        if decoding failed — callers treat ``[]`` as "no frames available".
+        """
+        ep = self._meta.episodes[episode_index]
+        from_timestamp = ep[f"videos/{camera_key}/from_timestamp"]
+        shifted = [from_timestamp + ts for ts in timestamps]
+        video_path = self.root / self._meta.get_video_file_path(episode_index, camera_key)
+
+        try:
+            # The module phases decode under a ThreadPoolExecutor (see
+            # ``ExecutorConfig.episode_parallelism``) but torchcodec's cached
+            # per-file decoder is single-threaded, so serialize decodes on a
+            # dedicated lock. Frame extraction is a small fraction of episode
+            # wall time (VLM calls dominate), so the contention is cheap.
+            with self._decode_lock:
+                # Stacked ``(N, C, H, W)`` uint8 tensor; one row per timestamp.
+                decoded = decode_video_frames(
+                    video_path, shifted, self.tolerance_s, backend=self.video_backend, return_uint8=True
+                )
+            return list(decoded)
+        except Exception as exc:
+            # Log loudly the first time so a silent vqa-module no-op (every
+            # prompt skipped because frames_at returned []) is debuggable from
+            # the job log instead of post-hoc parquet inspection. Subsequent
+            # failures stay quiet.
+            with self._lock:
+                already_warned = self._warned_decode_fail
+                if not already_warned:
+                    self._warned_decode_fail = True
+            if not already_warned:
+                logger.warning(
+                    "VideoFrameProvider._decode failed for episode=%s camera=%s video_path=%s backend=%s: %s",
+                    episode_index,
+                    camera_key,
+                    video_path,
+                    self.video_backend,
+                    exc,
+                    exc_info=exc,
+                )
+            return []
+
+
+def make_frame_provider(
+    root: Path, camera_key: str | None = None, video_backend: str | None = None
+) -> FrameProvider:
+    """Build a :class:`VideoFrameProvider` if videos are present, else null."""
+    try:
+        provider = VideoFrameProvider(root=root, camera_key=camera_key, video_backend=video_backend)
+    except Exception:
+        return null_provider()
+    if provider.camera_key is None:
+        return null_provider()
+    return provider
+
+
+def _frame_to_pil(frame: Any) -> Any:
+    """Materialise a decoded frame as a ``PIL.Image`` for the VLM message.
+
+    Frames flow through the provider as ``torch.Tensor`` (``C, H, W`` uint8,
+    straight from :func:`decode_video_frames`); PIL is only created here, at
+    the VLM-message boundary, because the chat backends expect PIL images /
+    data URLs. Non-tensor inputs (e.g. test stubs) pass through untouched.
+    """
+    if not isinstance(frame, torch.Tensor):
+        return frame
+    array = frame.detach().cpu()
+    if array.ndim == 3 and array.shape[0] in (1, 3):
+        array = array.permute(1, 2, 0)  # (C, H, W) -> (H, W, C)
+    if array.shape[-1] == 1:
+        array = array.squeeze(-1)
+    return PIL.Image.fromarray(array.to(torch.uint8).numpy())
+
+
+def to_image_blocks(frames: list[Any]) -> list[dict[str, Any]]:
+    """Convert decoded frames to Qwen-VL-compatible image content blocks."""
+    return [{"type": "image", "image": _frame_to_pil(frame)} for frame in frames]
+
+
+def to_video_block(frames: list[Any]) -> list[dict[str, Any]]:
+    """Wrap a list of decoded frames as one Qwen-VL video block.
+
+    Returns ``[]`` when the list is empty, so the caller can splat the result
+    into a content array without a separate emptiness check.
+    """
+    if not frames:
+        return []
+    return [{"type": "video", "video": [_frame_to_pil(frame) for frame in frames]}]
+
+
+def to_video_url_block(url: str | None, fps: float = 2.0) -> list[dict[str, Any]]:
+    """Wrap a video file URL as one ``video_url`` block.
+
+    Used by the ``openai`` backend (transformers serve / vllm serve /
+    ktransformers serve), where the server handles frame sampling.
+    Returns ``[]`` when ``url`` is ``None`` so the caller can splat.
+    """
+    if not url:
+        return []
+    return [{"type": "video_url", "video_url": {"url": url}, "fps": fps}]
+
+
+def _draw_timestamp_badge(image: PIL.Image.Image, timestamp: float) -> PIL.Image.Image:
+    """Burn ``timestamp`` (seconds) into the top-left corner of ``image``.
+
+    A solid black badge with white text, so a VLM reading a contact sheet can
+    cite the exact source time of each tile (e.g. ``012.50s``) directly,
+    instead of the caller having to map tile position back to time. Mirrors
+    the macrodata/refiner contact-sheet convention.
+    """
+    from PIL import ImageDraw, ImageFont
+
+    result = image.copy()
+    draw = ImageDraw.Draw(result)
+    font = ImageFont.load_default()
+    label = f"{timestamp:06.2f}s"
+    left, top, right, bottom = draw.textbbox((0, 0), label, font=font)
+    text_w, text_h = right - left, bottom - top
+    pad = max(3, round(min(image.width, image.height) * 0.018))
+    draw.rectangle((0, 0, text_w + pad * 2, text_h + pad * 2), fill=(0, 0, 0))
+    draw.text((pad - left, pad - top), label, fill=(255, 255, 255), font=font)
+    return result
+
+
+def to_contact_sheet_blocks(
+    frames: Sequence[Any],
+    timestamps: Sequence[float],
+    *,
+    columns: int = 5,
+    frames_per_sheet: int = 20,
+    frame_width: int = 224,
+    quality: int = 84,
+) -> list[dict[str, Any]]:
+    """Pack decoded frames into timestamped JPEG contact-sheet image blocks.
+
+    Each frame is resized to ``frame_width`` wide, stamped with its
+    episode-relative timestamp, and tiled row-major into grids of
+    ``frames_per_sheet`` (``columns`` wide). One ``{"type":"image", ...}``
+    block is returned per grid; many frames collapse into a few images, so a
+    long episode's temporal coverage stays dense at a fraction of the vision
+    tokens N separate frames would cost. ``frames`` and ``timestamps`` must be
+    aligned and equal length. Returns ``[]`` for empty input.
+    """
+    from PIL import Image
+
+    if not frames:
+        return []
+    columns = max(1, columns)
+    frames_per_sheet = max(1, frames_per_sheet)
+    rows_per_sheet = math.ceil(frames_per_sheet / columns)
+
+    tiles: list[PIL.Image.Image] = []
+    for ts, frame in zip(timestamps, frames, strict=False):
+        img = _frame_to_pil(frame)
+        if not isinstance(img, PIL.Image.Image):
+            continue
+        img = img.convert("RGB")
+        if img.width != frame_width:
+            height = max(1, round(img.height * frame_width / img.width))
+            img = img.resize((frame_width, height), resample=Image.Resampling.BILINEAR)
+        tiles.append(_draw_timestamp_badge(img, float(ts)))
+    if not tiles:
+        return []
+
+    blocks: list[dict[str, Any]] = []
+    for start in range(0, len(tiles), frames_per_sheet):
+        chunk = tiles[start : start + frames_per_sheet]
+        cell_w = max(tile.width for tile in chunk)
+        cell_h = max(tile.height for tile in chunk)
+        sheet = Image.new("RGB", (cell_w * columns, cell_h * rows_per_sheet), color=(0, 0, 0))
+        for i, tile in enumerate(chunk):
+            x = (i % columns) * cell_w
+            y = (i // columns) * cell_h
+            sheet.paste(tile, (x, y))
+        # JPEG round-trip at ``quality`` to match the refiner convention and
+        # shrink the wire payload; vision-token count is set by resolution, so
+        # the real saving is the grid packing, not the codec.
+        buf = io.BytesIO()
+        sheet.save(buf, format="JPEG", quality=quality)
+        buf.seek(0)
+        blocks.append({"type": "image", "image": Image.open(buf).convert("RGB")})
+    return blocks

src/lerobot/annotations/steerable_pipeline/modules/__init__.py

@@ -0,0 +1,25 @@
+#!/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.
+
+from .general_vqa import GeneralVqaModule
+from .interjections_and_speech import InterjectionsAndSpeechModule
+from .plan_subtasks_memory import PlanSubtasksMemoryModule
+
+__all__ = [
+    "GeneralVqaModule",
+    "InterjectionsAndSpeechModule",
+    "PlanSubtasksMemoryModule",
+]

src/lerobot/annotations/steerable_pipeline/modules/general_vqa.py

@@ -0,0 +1,248 @@
+#!/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.
+"""``vqa`` module: general VQA at a timed cadence.
+
+Every ``1/hz`` seconds an emission tick fires; each tick anchors ``K``
+consecutive frames, and every anchored frame gets its own VQA pair. Each
+pair is grounded on that single anchor frame — there is no per-pair frame
+window. For datasets with multiple cameras, every anchored frame produces
+one ``(vqa, user)`` + ``(vqa, assistant)`` pair *per camera*: each pair is
+generated against that camera's frame and stamped with the matching
+``camera`` field on the emitted rows. The resolver disambiguates via
+``camera=...``; recipes that consume VQA do so through one sub-recipe
+per camera (see ``recipes/pi05_hirobot.yaml``).
+
+Within a single (frame, camera) we still emit at most one ``(vqa, user)``
+and one ``(vqa, assistant)`` row, so the resolver contract stays scalar.
+
+Question types covered (per the plan's ``vqa`` table): bbox, keypoint,
+count, attribute, spatial. The assistant's ``content`` is a JSON string
+whose schema depends on the question type. Malformed JSON triggers one
+retry inside :meth:`VlmClient.generate_json`.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import random
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..config import VqaConfig
+from ..frames import FrameProvider, null_provider, to_image_blocks
+from ..prompts import load as load_prompt
+from ..reader import EpisodeRecord
+from ..staging import EpisodeStaging
+from ..validator import classify_vqa_answer
+from ..vlm_client import VlmClient
+
+
+def _emission_anchor_indices(frame_timestamps: Sequence[float], hz: float, k: int) -> list[int]:
+    """Return the relative frame indices to anchor VQA emissions to.
+
+    For each emission tick (every ``1/hz`` seconds), we anchor ``k``
+    consecutive frames starting at the tick. Ticks fall on the nearest
+    available source frame timestamp.
+    """
+    if hz <= 0 or k <= 0 or not frame_timestamps:
+        return []
+    t0 = frame_timestamps[0]
+    t_last = frame_timestamps[-1]
+    period = 1.0 / hz
+    indices: list[int] = []
+    t = t0
+    while t <= t_last + 1e-9:
+        # find the index of the nearest frame to t
+        nearest_i = min(range(len(frame_timestamps)), key=lambda i: abs(frame_timestamps[i] - t))
+        for offset in range(k):
+            j = nearest_i + offset
+            if j >= len(frame_timestamps):
+                break
+            if not indices or indices[-1] != j:
+                indices.append(j)
+        t += period
+    # dedupe while preserving order
+    seen: set[int] = set()
+    deduped: list[int] = []
+    for i in indices:
+        if i in seen:
+            continue
+        seen.add(i)
+        deduped.append(i)
+    return deduped
+
+
+@dataclass
+class GeneralVqaModule:
+    """Emit grounded VQA pairs at a timed cadence."""
+
+    vlm: VlmClient
+    config: VqaConfig
+    seed: int = 1729
+    frame_provider: FrameProvider = field(default_factory=null_provider)
+    _warned_no_camera: bool = field(default=False, init=False, repr=False)
+
+    @property
+    def enabled(self) -> bool:
+        return self.config.enabled
+
+    def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
+        if not record.frame_timestamps:
+            staging.write("vqa", [])
+            return
+        rng = random.Random(f"{self.seed}:{record.episode_index}:vqa")
+        anchor_idx = _emission_anchor_indices(
+            record.frame_timestamps, self.config.vqa_emission_hz, self.config.K
+        )
+        cameras = self._target_cameras()
+        if not cameras:
+            # No camera available — emit nothing rather than producing
+            # untagged rows that would fail validation. Surface a loud one-
+            # time warning so this is never silently a no-op.
+            if not self._warned_no_camera:
+                logging.getLogger(__name__).warning(
+                    "vqa module found no cameras on the frame provider — "
+                    "every episode will emit zero VQA rows. Check that the "
+                    "dataset declares observation.images.* features in "
+                    "meta/info.json; passing --vlm.camera_key=<key> at the "
+                    "CLI now also seeds the cameras list as a fallback."
+                )
+                self._warned_no_camera = True
+            staging.write("vqa", [])
+            return
+
+        # Build all messages first (one per (frame, camera)), then issue them
+        # as a single batched generate_json call so the client can fan them
+        # out concurrently.
+        per_call: list[tuple[float, str, str, list[dict[str, Any]]]] = []
+        for idx in anchor_idx:
+            ts = float(record.frame_timestamps[idx])
+            qtype = rng.choice(self.config.question_types)
+            for camera in cameras:
+                messages = self._build_messages(record, qtype, ts, camera)
+                # Skip cameras that decoded to zero frames at this ts: no point
+                # asking the VLM to ground a bbox without an image.
+                if not _has_image_block(messages):
+                    continue
+                per_call.append((ts, camera, qtype, messages))
+
+        if not per_call:
+            staging.write("vqa", [])
+            return
+
+        results = self.vlm.generate_json([m for _, _, _, m in per_call])
+
+        rows: list[dict[str, Any]] = []
+        for (ts, camera, _qtype, _messages), result in zip(per_call, results, strict=True):
+            qa = self._postprocess(result)
+            if qa is None:
+                continue
+            question, answer = qa
+            rows.append(
+                {
+                    "role": "user",
+                    "content": question,
+                    "style": "vqa",
+                    "timestamp": ts,
+                    "camera": camera,
+                    "tool_calls": None,
+                }
+            )
+            rows.append(
+                {
+                    "role": "assistant",
+                    "content": json.dumps(answer, sort_keys=True),
+                    "style": "vqa",
+                    "timestamp": ts,
+                    "camera": camera,
+                    "tool_calls": None,
+                }
+            )
+        staging.write("vqa", rows)
+
+    def _target_cameras(self) -> list[str]:
+        """Return the cameras the ``vqa`` module should iterate per anchored frame.
+
+        Defaults to every camera the provider exposes. Datasets with no
+        cameras (or test/null providers) yield an empty list, which makes
+        ``run_episode`` a no-op.
+
+        When ``config.restrict_to_default_camera`` is set, VQA grounds on
+        only the provider's default camera (the single ``--vlm.camera_key``
+        stream), matching the plan / interjection modules so the whole
+        pipeline focuses on one view.
+        """
+        all_cameras = list(getattr(self.frame_provider, "camera_keys", []) or [])
+        if getattr(self.config, "restrict_to_default_camera", False):
+            default = getattr(self.frame_provider, "camera_key", None)
+            if default and default in all_cameras:
+                return [default]
+            # ``restrict_to_default_camera`` is set but the configured default
+            # isn't one the provider exposes. Returning it anyway would make
+            # ``_decode`` raise a KeyError deep in frame extraction, so warn and
+            # fall through to every available camera instead.
+            if default:
+                logging.getLogger(__name__).warning(
+                    "restrict_to_default_camera is set but camera_key=%r is not in the "
+                    "provider's cameras %s; grounding VQA on all available cameras instead.",
+                    default,
+                    all_cameras,
+                )
+        return all_cameras
+
+    def _build_messages(
+        self,
+        record: EpisodeRecord,
+        question_type: str,
+        frame_timestamp: float,
+        camera_key: str,
+    ) -> list[dict[str, Any]]:
+        prompt = load_prompt("vqa").format(
+            episode_task=record.episode_task,
+            question_type=question_type,
+        )
+        images = self.frame_provider.frames_at(record, [frame_timestamp], camera_key=camera_key)
+        content = [*to_image_blocks(images), {"type": "text", "text": prompt}]
+        return [{"role": "user", "content": content}]
+
+    def _postprocess(self, result: Any) -> tuple[str, dict[str, Any]] | None:
+        if not isinstance(result, dict):
+            return None
+        question = result.get("question")
+        answer = result.get("answer")
+        if not isinstance(question, str) or not question.strip():
+            return None
+        if not isinstance(answer, dict):
+            return None
+        # The validator will enforce shape; here we just sanity-check that the
+        # answer matches *some* known shape so we can drop garbage early.
+        if classify_vqa_answer(answer) is None:
+            return None
+        return question.strip(), answer
+
+
+def _has_image_block(messages: list[dict[str, Any]]) -> bool:
+    """Return True if any user content block is a populated image block."""
+    for msg in messages:
+        content = msg.get("content")
+        if not isinstance(content, list):
+            continue
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "image":
+                return True
+    return False

src/lerobot/annotations/steerable_pipeline/modules/interjections_and_speech.py

@@ -0,0 +1,211 @@
+#!/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.
+"""``interjections`` module: interjections + paired speech (EVENT styles + speech atoms).
+
+Two sub-passes:
+
+1. At ``t=0``, emit ONLY a speech tool-call atom (acknowledgement of the
+   canonical task). No interjection row — the canonical task is already the
+   user utterance from ``meta/tasks.parquet``.
+
+2. For mid-episode interruptions, emit a co-timestamped pair:
+       {role:user, style:interjection, content:<text>}
+       speech atom (role:assistant, style:None, tool_calls=[say(...)])
+   Both rows go in ``language_events`` at the same timestamp.
+
+The ``plan`` module's :meth:`run_plan_updates` reuses this module's
+interjection timestamps to refresh the ``plan`` row at the same instant.
+"""
+
+from __future__ import annotations
+
+import random
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..config import InterjectionsConfig
+from ..frames import FrameProvider, null_provider, to_image_blocks
+from ..prompts import load as load_prompt
+from ..reader import EpisodeRecord, reconstruct_subtask_spans, snap_to_frame
+from ..staging import EpisodeStaging
+from ..vlm_client import VlmClient
+from ..writer import speech_atom
+
+
+@dataclass
+class InterjectionsAndSpeechModule:
+    """Generate task-start speech and mid-episode interjection/speech pairs."""
+
+    vlm: VlmClient
+    config: InterjectionsConfig
+    seed: int = 1729
+    frame_provider: FrameProvider = field(default_factory=null_provider)
+
+    @property
+    def enabled(self) -> bool:
+        return self.config.enabled
+
+    def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
+        rows: list[dict[str, Any]] = []
+        if record.frame_timestamps:
+            t0 = float(record.frame_timestamps[0])
+            initial = self._initial_speech(record)
+            if initial:
+                rows.append(speech_atom(t0, initial))
+        # Pull the ``plan`` module's subtask spans for this episode so the
+        # interjection prompt can ground itself in the actual current
+        # subtask at each chosen timestamp. The ``plan`` module ran first.
+        episode_end_t = float(record.frame_timestamps[-1]) if record.frame_timestamps else None
+        subtask_spans = reconstruct_subtask_spans(staging.read("plan"), episode_end_t=episode_end_t)
+        rows.extend(self._mid_episode_interjections(record, subtask_spans))
+        staging.write("interjections", rows)
+
+    @staticmethod
+    def _subtask_at(spans: Sequence[dict[str, Any]], t: float) -> str | None:
+        current: str | None = None
+        for span in spans:
+            if float(span["start"]) <= t:
+                current = span.get("text")
+            else:
+                break
+        return current
+
+    def _initial_speech(self, record: EpisodeRecord) -> str | None:
+        prompt = load_prompt("interjections_initial_speech").format(
+            episode_task=record.episode_task,
+        )
+        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
+        result = self.vlm.generate_json([messages])[0]
+        if isinstance(result, dict) and isinstance(result.get("text"), str):
+            text = result["text"].strip()
+            if text:
+                return text
+        return None
+
+    def _mid_episode_interjections(
+        self,
+        record: EpisodeRecord,
+        subtask_spans: Sequence[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        """Generate interjections aligned with the actual demo trajectory.
+
+        Teleop data is frozen — the robot already executed every step in
+        the video. A *counterfactual* interjection like "actually skip
+        the wipe" contradicts what then happens in the video, which is
+        what qwen36moe-10/11 surfaced as low-quality interjections.
+
+        Instead, anchor every interjection at a subtask boundary and
+        write it as a natural user request for the *upcoming* subtask.
+        The robot's visible next behavior IS the interjection's effect,
+        so the training signal stays consistent: interjection text →
+        plan refresh → action stream all line up.
+        """
+        if self.config.max_interjections_per_episode <= 0:
+            return []
+        if len(subtask_spans) < 2:
+            # Need at least one transition (subtask 0 → subtask 1).
+            return []
+        # Deterministic per-episode RNG so reruns are stable across SLURM jobs.
+        rng = random.Random(f"{self.seed}:{record.episode_index}:interjection")
+
+        # Boundaries: the start time of every subtask except the first
+        # (which is just t0 and is covered by the initial-task speech atom).
+        boundaries: list[tuple[float, str, str]] = []
+        for i in range(1, len(subtask_spans)):
+            ts = float(subtask_spans[i]["start"])
+            if ts < self.config.interjection_min_t:
+                continue
+            prev_text = (subtask_spans[i - 1].get("text") or "").strip()
+            next_text = (subtask_spans[i].get("text") or "").strip()
+            if not next_text:
+                continue
+            boundaries.append((ts, prev_text, next_text))
+        if not boundaries:
+            return []
+
+        n = min(self.config.max_interjections_per_episode, len(boundaries))
+        chosen = sorted(rng.sample(boundaries, n), key=lambda b: b[0])
+
+        out: list[dict[str, Any]] = []
+        for t, prev_subtask, next_subtask in chosen:
+            t_snap = snap_to_frame(t, record.frame_timestamps)
+            # Window straddles the boundary so the VLM sees the end of the
+            # previous subtask and the start of the next one — same
+            # conditioning the policy will see at training time.
+            window_ts = self._window_timestamps(t_snap, record.frame_timestamps)
+            prompt = load_prompt("interjections_interjection").format(
+                episode_task=record.episode_task,
+                prev_subtask=prev_subtask or "(starting from initial state)",
+                next_subtask=next_subtask,
+                timestamp=t_snap,
+                window_seconds=self.config.interjection_window_seconds,
+            )
+            images = self.frame_provider.frames_at(record, window_ts)
+            content = [*to_image_blocks(images), {"type": "text", "text": prompt}]
+            messages = [{"role": "user", "content": content}]
+            result = self.vlm.generate_json([messages])[0]
+            if not isinstance(result, dict):
+                continue
+            interjection_text = result.get("interjection")
+            speech_text = result.get("speech")
+            if not isinstance(interjection_text, str) or not interjection_text.strip():
+                continue
+            if not isinstance(speech_text, str) or not speech_text.strip():
+                continue
+            out.append(
+                {
+                    "role": "user",
+                    "content": interjection_text.strip(),
+                    "style": "interjection",
+                    "timestamp": t_snap,
+                    "tool_calls": None,
+                }
+            )
+            out.append(speech_atom(t_snap, speech_text.strip()))
+        return out
+
+    def _window_timestamps(self, t_anchor: float, frame_timestamps: Sequence[float]) -> list[float]:
+        """Return a small set of frame timestamps centered on ``t_anchor``.
+
+        The window straddles the subtask boundary the interjection sits
+        on: roughly half the frames cover the end of the previous
+        subtask, half cover the start of the next one. The VLM therefore
+        sees BOTH what just finished AND what's about to start, which is
+        the conditioning we need to write a natural "now please do X"
+        request that matches the visible upcoming behavior.
+        """
+        if not frame_timestamps:
+            return [t_anchor]
+        n = max(1, int(self.config.interjection_window_frames))
+        if n == 1:
+            return [t_anchor]
+        window = float(self.config.interjection_window_seconds)
+        step = window / max(1, n - 1)
+        # Center the window on the anchor so half lands before, half after.
+        start_offset = -window / 2.0
+        targets = [t_anchor + start_offset + step * i for i in range(n)]
+        first_ts = float(frame_timestamps[0])
+        last_ts = float(frame_timestamps[-1])
+        snapped: list[float] = []
+        seen: set[float] = set()
+        for tgt in targets:
+            clamped = min(last_ts, max(first_ts, tgt))
+            t = snap_to_frame(clamped, frame_timestamps)
+            if t not in seen:
+                seen.add(t)
+                snapped.append(t)
+        return snapped or [t_anchor]

src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py

@@ -0,0 +1,780 @@
+#!/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.
+"""``plan`` module: subtask decomposition + plan + memory (PERSISTENT styles)."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..config import PlanConfig
+from ..frames import (
+    FrameProvider,
+    null_provider,
+    to_contact_sheet_blocks,
+)
+from ..prompts import load as load_prompt
+from ..reader import EpisodeRecord, reconstruct_subtask_spans, snap_to_frame
+from ..staging import EpisodeStaging
+from ..vlm_client import VlmClient
+
+logger = logging.getLogger(__name__)
+
+
+# Prepended to every describe / segment prompt so the VLM knows the images are
+# timestamped contact-sheet grids, not a single video, and reads the burned-in
+# per-tile timestamp when choosing boundaries.
+def _contact_sheet_preamble(columns: int) -> str:
+    return (
+        "CONTACT SHEETS — how to read the images below:\n"
+        f"- Each image is a grid of sampled video frames, {columns} per row, "
+        "with time running left-to-right then top-to-bottom (row-major).\n"
+        "- Each frame has its timestamp burned into the top-left corner, e.g. "
+        '"012.50s". Use that printed timestamp (not the tile position) when you '
+        "choose start/end times; boundaries should land on or near a printed "
+        "timestamp.\n"
+        "- Frames continue across grids: an action may span the end of one sheet "
+        "and the start of the next, so do not place a boundary just because a new "
+        "image begins.\n\n"
+    )
+
+
+# Appended to every describe (and segment) prompt. A visual, causal definition
+# of where one event ends and the next begins — adapted from macrodata/refiner —
+# to sharpen cut points while the existing prompt keeps owning the imperative
+# phrasing.
+_CAUSAL_BOUNDARY_RULES = (
+    "EVENT BOUNDARIES — where one event ends and the next begins:\n"
+    "- Start a new event whenever the world state changes: an object becomes "
+    "held (the gripper closes on it), an object is released (the gripper opens "
+    "and it stays put), an object reaches a new location, a lid/door/drawer "
+    "changes open/closed state, a tool starts or stops affecting a surface, or "
+    "contents visibly move (e.g. poured).\n"
+    "- If a single action changes the same state gradually and continuously, "
+    "keep it as ONE event — do not split it.\n"
+    "- If the same action repeats on different objects or target locations, "
+    "treat each repetition as a separate event.\n"
+    "- Do NOT create boundaries for idle time, camera motion, hesitation, or "
+    "tiny hand adjustments."
+)
+
+
+@dataclass
+class PlanSubtasksMemoryModule:
+    """Generate subtask spans, plan, and memory rows.
+
+    All output is persistent (lives in ``language_persistent``):
+
+    - ``subtask`` rows: one per span, stamped at the span's *start* timestamp
+      (snapped to an exact frame).
+    - ``plan`` rows: emitted at ``t=0``; refreshed at every interjection
+      timestamp via :meth:`run_plan_updates` (called by the executor after
+      the ``interjections`` module completes).
+    - ``memory`` rows: emitted at each subtask boundary (= subtask start
+      timestamp from the second subtask onward).
+    """
+
+    vlm: VlmClient
+    config: PlanConfig
+    frame_provider: FrameProvider = field(default_factory=null_provider)
+
+    @property
+    def enabled(self) -> bool:
+        return self.config.enabled
+
+    def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
+        rows: list[dict[str, Any]] = []
+        # Task driving every plan-module prompt: canonical episode_task, or a
+        # video-derived one when it's empty/placeholder (see derive_task_*).
+        effective_task = self._resolve_effective_task(record)
+        # task_aug rows at t=0: phrasings the renderer rotates ${task} through.
+        # Either the structured 5-axis taxonomy (task_aug_axes.enabled) or
+        # free-form n_task_rephrasings; the effective task is always emitted
+        # first so the rotation covers the source-of-truth phrasing.
+        t0 = float(record.frame_timestamps[0]) if record.frame_timestamps else 0.0
+        variants: list[str] | None = None
+        if self.config.task_aug_axes.enabled and effective_task:
+            variants = self._generate_task_aug_by_axes(effective_task, self.config.task_aug_axes)
+        elif self.config.n_task_rephrasings > 0 and effective_task:
+            variants = self._generate_task_rephrasings(effective_task, n=self.config.n_task_rephrasings)
+        if variants is not None:
+            rows.extend(self._task_aug_rows([effective_task, *variants], t0))
+
+        subtask_spans = self._generate_subtasks(record, task=effective_task)
+
+        # subtask rows
+        for span in subtask_spans:
+            rows.append(
+                {
+                    "role": "assistant",
+                    "content": span["text"],
+                    "style": "subtask",
+                    "timestamp": snap_to_frame(span["start"], record.frame_timestamps),
+                    "tool_calls": None,
+                }
+            )
+        # Plan rows at every subtask boundary (incl. t=0). The plan is a
+        # numbered list of still-todo subtasks, so re-emitting at each
+        # boundary makes it shrink as work progresses — ${plan} at frame t is
+        # exactly what's left to do.
+        if self.config.emit_plan:
+            for span in subtask_spans:
+                boundary_t = snap_to_frame(span["start"], record.frame_timestamps)
+                plan_text = self._generate_plan(
+                    record, subtask_spans, refresh_t=boundary_t, task=effective_task
+                )
+                if plan_text is not None:
+                    rows.append(
+                        {
+                            "role": "assistant",
+                            "content": plan_text,
+                            "style": "plan",
+                            "timestamp": float(boundary_t),
+                            "tool_calls": None,
+                        }
+                    )
+        # memory rows at every subtask boundary except the very first start;
+        # skipped entirely when ``emit_memory`` is False (subtasks-only / plan-only).
+        prior_memory = ""
+        memory_boundaries = enumerate(subtask_spans[1:], start=1) if self.config.emit_memory else []
+        for i, span in memory_boundaries:
+            completed = subtask_spans[i - 1]["text"]
+            remaining = [s["text"] for s in subtask_spans[i:]]
+            mem_text = self._generate_memory(record, prior_memory, completed, remaining, task=effective_task)
+            if mem_text:
+                ts = snap_to_frame(span["start"], record.frame_timestamps)
+                rows.append(
+                    {
+                        "role": "assistant",
+                        "content": mem_text,
+                        "style": "memory",
+                        "timestamp": ts,
+                        "tool_calls": None,
+                    }
+                )
+                prior_memory = mem_text
+        staging.write("plan", rows)
+
+    # ------------------------------------------------------------------
+    # Task derivation + rephrasings
+    # ------------------------------------------------------------------
+
+    _PLACEHOLDER_TASKS: frozenset[str] = frozenset(
+        {
+            "debug",
+            "test",
+            "tbd",
+            "todo",
+            "n/a",
+            "na",
+            "untitled",
+            "unnamed",
+            "default",
+            "placeholder",
+        }
+    )
+
+    def _resolve_effective_task(self, record: EpisodeRecord) -> str:
+        """Decide which task string drives the ``plan`` module for this episode.
+
+        Returns the user-supplied ``record.episode_task`` unless
+        ``derive_task_from_video`` says otherwise (see config docstring).
+        Falls back gracefully to the canonical task if video derivation
+        fails.
+        """
+        canonical = (record.episode_task or "").strip()
+        mode = (self.config.derive_task_from_video or "off").strip().lower()
+        if mode == "always":
+            derived = self._derive_task_from_video(record)
+            return derived or canonical
+        if mode == "if_short" and self._task_seems_bad(canonical):
+            derived = self._derive_task_from_video(record)
+            if derived:
+                return derived
+        return canonical
+
+    def _task_seems_bad(self, task: str) -> bool:
+        if not task:
+            return True
+        if len(task.split()) < int(self.config.derive_task_min_words):
+            return True
+        return task.lower() in self._PLACEHOLDER_TASKS
+
+    @staticmethod
+    def _task_aug_rows(phrasings: Sequence[str], t0: float) -> list[dict[str, Any]]:
+        """Build deduplicated ``task_aug`` rows (role=user) at ``t0``."""
+        seen: set[str] = set()
+        rows: list[dict[str, Any]] = []
+        for phrasing in phrasings:
+            key = phrasing.strip()
+            if not key or key in seen:
+                continue
+            seen.add(key)
+            rows.append(
+                {"role": "user", "content": key, "style": "task_aug", "timestamp": t0, "tool_calls": None}
+            )
+        return rows
+
+    # ------------------------------------------------------------------
+    # VLM call helpers — every plan-module prompt follows the same shape:
+    # build messages → single VLM call → pull a named field.
+    # ------------------------------------------------------------------
+
+    def _vlm_field(self, messages: list[dict[str, Any]], field: str) -> Any:
+        """Run a single VLM call and return ``result[field]`` or ``None``.
+
+        Centralizes the ``vlm.generate_json([m])[0]`` + ``isinstance(dict)``
+        dance every prompt-call site needs.
+        """
+        result = self.vlm.generate_json([messages])[0]
+        if isinstance(result, dict):
+            return result.get(field)
+        return None
+
+    @staticmethod
+    def _text_message(text: str) -> list[dict[str, Any]]:
+        """One-shot text-only user message wrapped for ``generate_json``."""
+        return [{"role": "user", "content": [{"type": "text", "text": text}]}]
+
+    def _video_message(
+        self,
+        record: EpisodeRecord,
+        prompt: str,
+        window: tuple[float, float] | None = None,
+    ) -> list[dict[str, Any]]:
+        """User message combining the (optionally windowed) contact sheets with ``prompt``.
+
+        The prompt is always prefixed with a short explanation of how to read
+        the timestamped grids, so the model treats them as one ordered
+        sequence of frames rather than unrelated images.
+        """
+        prompt = _contact_sheet_preamble(self.config.contact_sheet_columns) + prompt
+        content = [*self._episode_video_block(record, window=window), {"type": "text", "text": prompt}]
+        return [{"role": "user", "content": content}]
+
+    def _derive_task_from_video(self, record: EpisodeRecord) -> str | None:
+        """Ask the VLM "what is this video about" with no task hint at all."""
+        text = self._vlm_field(self._video_message(record, load_prompt("plan_video_task")), "task")
+        return text.strip() if isinstance(text, str) and text.strip() else None
+
+    def _generate_task_rephrasings(self, base_task: str, *, n: int) -> list[str]:
+        """Generate ``n`` text-only paraphrases of ``base_task``."""
+        if n <= 0 or not base_task:
+            return []
+        prompt = load_prompt("plan_task_rephrasings").format(base_task=base_task, n=n)
+        raw = self._vlm_field(self._text_message(prompt), "rephrasings")
+        if not isinstance(raw, list):
+            return []
+        out = [item.strip().strip('"').strip("'") for item in raw if isinstance(item, str)]
+        return [s for s in out if s][:n]
+
+    # ------------------------------------------------------------------
+    # Structured 5-axis task augmentation (EgoMimic-style taxonomy)
+    # ------------------------------------------------------------------
+
+    def _generate_task_aug_by_axes(self, base_task: str, axes_cfg: Any) -> list[str]:
+        """One VLM call → variants along the 5-axis taxonomy.
+
+        Variants from all axes are flattened into a single list (the
+        downstream pipeline doesn't need to know about the per-axis
+        bucketing — every variant becomes a ``task_aug`` row). Order
+        is preserved for reproducibility: synonym_paraphrase first,
+        then omit_arm, then omit_orientation, then omit_grasp_method,
+        then combined_omissions.
+        """
+        if not base_task:
+            return []
+        prompt = load_prompt("plan_task_aug_axes").format(
+            base_task=base_task,
+            n_synonym=axes_cfg.synonym_paraphrase,
+            n_omit_arm=axes_cfg.omit_arm,
+            n_omit_orientation=axes_cfg.omit_orientation,
+            n_omit_grasp_method=axes_cfg.omit_grasp_method,
+            n_combined=axes_cfg.combined_omissions,
+        )
+        result = self.vlm.generate_json([self._text_message(prompt)])[0]
+        if not isinstance(result, dict):
+            return []
+        ordered_axes = (
+            "synonym_paraphrase",
+            "omit_arm",
+            "omit_orientation",
+            "omit_grasp_method",
+            "combined_omissions",
+        )
+        flat: list[str] = []
+        seen: set[str] = set()
+        for axis in ordered_axes:
+            entries = result.get(axis)
+            if not isinstance(entries, list):
+                continue
+            for item in entries:
+                if not isinstance(item, str):
+                    continue
+                key = item.strip().strip('"').strip("'")
+                if not key or key in seen:
+                    continue
+                seen.add(key)
+                flat.append(key)
+        return flat
+
+    def _episode_video_block(
+        self, record: EpisodeRecord, window: tuple[float, float] | None = None
+    ) -> list[dict[str, Any]]:
+        """Timestamped contact sheets for the describe / segmentation prompts.
+
+        Always renders the (optionally windowed) episode as contact sheets:
+        frames sampled at ``frames_per_second`` and packed into timestamped
+        JPEG grids. ``max_frames_per_prompt`` caps the frame count; whole
+        episodes that exceed it are windowed upstream in
+        :meth:`_generate_subtasks` so each call stays within budget while the
+        full episode keeps its sampling density.
+
+        When ``window=(w0, w1)`` is given the badges are WINDOW-RELATIVE
+        (``ts - w0``) to match the window-relative time frame the
+        segmentation prompt works in (spans are offset back to absolute time
+        afterwards).
+        """
+        if not record.frame_timestamps:
+            return []
+        if window is not None:
+            w0, w1 = float(window[0]), float(window[1])
+            dur = max(0.0, w1 - w0)
+            n = max(1, int(round(dur * self.config.frames_per_second)) + 1)
+            n = min(n, self.config.max_frames_per_prompt)
+            if n <= 1 or dur <= 0.0:
+                timestamps = [0.5 * (w0 + w1)]
+            else:
+                step = dur / (n - 1)
+                timestamps = [w0 + i * step for i in range(n)]
+            frames = self.frame_provider.frames_at(record, timestamps)
+            rel = [ts - w0 for ts in timestamps[: len(frames)]]
+            return self._contact_sheet_blocks(frames, rel)
+        episode_duration = record.frame_timestamps[-1] - record.frame_timestamps[0]
+        n = max(1, int(round(episode_duration * self.config.frames_per_second)) + 1)
+        n = min(n, self.config.max_frames_per_prompt)
+        timestamps = self._uniform_episode_timestamps(record, n)
+        frames = self.frame_provider.frames_at(record, timestamps)
+        return self._contact_sheet_blocks(frames, timestamps[: len(frames)])
+
+    @staticmethod
+    def _uniform_episode_timestamps(record: EpisodeRecord, n: int) -> list[float]:
+        """``n`` episode-relative timestamps spanning ``[t0, t_last]`` uniformly."""
+        ts = record.frame_timestamps
+        if n >= len(ts):
+            return [float(t) for t in ts]
+        t0, t_last = float(ts[0]), float(ts[-1])
+        if t_last <= t0 or n <= 1:
+            return [t0] * max(1, n)
+        step = (t_last - t0) / (n - 1)
+        return [t0 + i * step for i in range(n)]
+
+    def _contact_sheet_blocks(self, frames: list[Any], timestamps: list[float]) -> list[dict[str, Any]]:
+        """Build timestamped contact-sheet image blocks from decoded frames."""
+        return to_contact_sheet_blocks(
+            frames,
+            timestamps,
+            columns=self.config.contact_sheet_columns,
+            frames_per_sheet=self.config.contact_sheet_frames_per_sheet,
+            frame_width=self.config.contact_sheet_frame_width,
+            quality=self.config.contact_sheet_quality,
+        )
+
+    def run_plan_updates(
+        self,
+        record: EpisodeRecord,
+        staging: EpisodeStaging,
+        interjection_times: Sequence[float],
+        interjection_texts: Sequence[str] | None = None,
+    ) -> None:
+        """Append additional ``plan`` rows at every interjection timestamp.
+
+        Plans refresh ONLY on user interjections (event-driven). The
+        interjection text is forwarded into the prompt so the refreshed plan
+        reflects the user's correction.
+        """
+        if not self.config.emit_plan:
+            return
+        existing = staging.read("plan")
+        # Pass the last frame timestamp so the final span is closed (else its
+        # end == start, zero duration, and a refresh inside it is missed).
+        episode_end_t = float(record.frame_timestamps[-1]) if record.frame_timestamps else None
+        spans = reconstruct_subtask_spans(existing, episode_end_t=episode_end_t)
+        already_planned: set[float] = {float(r["timestamp"]) for r in existing if r.get("style") == "plan"}
+        new_rows = list(existing)
+
+        texts: list[str | None] = (
+            [None] * len(interjection_times)
+            if interjection_texts is None
+            else [str(t) if t else None for t in interjection_texts]
+        )
+        for raw_t, inter_text in zip(interjection_times, texts, strict=True):
+            t = snap_to_frame(raw_t, record.frame_timestamps)
+            if t in already_planned:
+                continue
+            already_planned.add(t)
+            plan_text = self._generate_plan(record, spans, refresh_t=t, interjection=inter_text)
+            if plan_text is not None:
+                new_rows.append(
+                    {
+                        "role": "assistant",
+                        "content": plan_text,
+                        "style": "plan",
+                        "timestamp": t,
+                        "tool_calls": None,
+                    }
+                )
+        staging.write("plan", new_rows)
+
+    def _generate_subtasks(self, record: EpisodeRecord, *, task: str | None = None) -> list[dict[str, Any]]:
+        """Generate subtask spans, optionally via a multi-call quality chain.
+
+        Single call (default): watch video → emit subtask JSON.
+
+        Multi-call (opt-in, higher quality, more VLM calls):
+          1. ``subtask_describe_first`` — a grounding pass that narrates
+             ONLY what is visible (no JSON commitment to subtasks yet);
+             its description is injected into the segmentation prompt so
+             the model segments its own grounded observations instead of
+             pattern-matching the task text.
+          2. segmentation — emit subtask JSON (as before).
+        """
+        if record.row_count == 0 or not record.frame_timestamps:
+            return []
+        episode_duration = record.frame_timestamps[-1] - record.frame_timestamps[0]
+        effective_task = task if task is not None else record.episode_task
+
+        # ---- Auto-windowing (keeps the full sampling density) --------
+        # Contact sheets are cheap, but a whole long episode sampled at
+        # ``frames_per_second`` can still exceed ``max_frames_per_prompt``.
+        # When it does, split into consecutive windows of exactly that many
+        # frames (one describe→segment call each, still at the full sampling
+        # density), then merge + stitch — so an episode of any length is
+        # covered at full density rather than subsampled into one sparse call.
+        fps = max(1e-6, float(self.config.frames_per_second))
+        n_whole = int(round(episode_duration * fps)) + 1
+        if n_whole > self.config.max_frames_per_prompt:
+            window_s = self.config.max_frames_per_prompt / fps
+            return self._generate_subtasks_windowed(record, effective_task, window_s)
+
+        # ---- Pass 1 (optional): grounding description ----------------
+        observation_block = ""
+        if getattr(self.config, "subtask_describe_first", False):
+            description = self._describe_episode(record, effective_task)
+            if description:
+                observation_block = (
+                    "You watched this video and described, chronologically, "
+                    "ONLY what the robot actually does:\n"
+                    f'"""{description}"""\n\n'
+                    "Segment THAT grounded description (cross-checked against "
+                    "the video) into atomic subtasks. Do not introduce any "
+                    "action that is not in your description above.\n\n"
+                )
+
+        # ---- Pass 2: segmentation ------------------------------------
+        prompt = self._with_causal_rules(
+            load_prompt("plan_subtasks").format(
+                episode_task=effective_task,
+                min_subtask_seconds=self.config.min_subtask_seconds,
+                max_steps=self.config.plan_max_steps,
+                episode_duration=f"{episode_duration:.3f}",
+                observation_block=observation_block,
+            )
+        )
+        spans = self._vlm_field(self._video_message(record, prompt), "subtasks")
+        cleaned = self._clean_spans(spans, record)
+        if not cleaned:
+            return []
+
+        # ---- Full-episode coverage stitch ----------------------------
+        # The VLM can start after t0 or leave gaps, so frames fall through
+        # with no active subtask. Always stitch into a contiguous
+        # [t0, t_last] cover.
+        cleaned = self._stitch_full_coverage(cleaned, record)
+
+        return cleaned
+
+    def _generate_subtasks_windowed(
+        self, record: EpisodeRecord, task: str, window_s: float
+    ) -> list[dict[str, Any]]:
+        """Subtask generation in fixed-length windows at constant fps.
+
+        Splits ``[t0, t_last]`` into consecutive windows of ``window_s``
+        seconds, runs the describe -> segment chain on each window's own
+        frames (sampled at ``frames_per_second``), offsets
+        each window's spans back to absolute episode time, then merges +
+        stitches into a contiguous whole-episode cover.
+        """
+        t0 = float(record.frame_timestamps[0])
+        t_last = float(record.frame_timestamps[-1])
+        all_spans: list[dict[str, Any]] = []
+        w0 = t0
+        n_windows = 0
+        while w0 < t_last - 1e-6:
+            w1 = min(w0 + window_s, t_last)
+            all_spans.extend(self._subtasks_for_window(record, task, w0, w1))
+            n_windows += 1
+            w0 = w1
+        logger.info(
+            "episode %d: windowed subtask gen over %d window(s) of %.1fs -> %d raw spans",
+            record.episode_index,
+            n_windows,
+            window_s,
+            len(all_spans),
+        )
+        # Merge across windows: clamp to the absolute episode, sort, and
+        # frame-snap to distinct starts (handles any boundary collisions).
+        cleaned = self._clean_spans(all_spans, record)
+        if not cleaned:
+            return []
+        return self._stitch_full_coverage(cleaned, record)
+
+    def _subtasks_for_window(
+        self, record: EpisodeRecord, task: str, w0: float, w1: float
+    ) -> list[dict[str, Any]]:
+        """Run describe -> segment on one ``[w0, w1]`` window.
+
+        The model works in window-RELATIVE time ``[0, L]`` (it perceives
+        the window as a clip starting at 0); spans are offset back to
+        absolute ``[w0, w1]`` before returning.
+        """
+        window = (w0, w1)
+        win_len = max(0.0, w1 - w0)
+
+        observation_block = ""
+        if getattr(self.config, "subtask_describe_first", False):
+            description = self._describe_episode(record, task, window=window)
+            if description:
+                observation_block = (
+                    "You watched this video clip and described, chronologically, "
+                    "ONLY what the robot actually does:\n"
+                    f'"""{description}"""\n\n'
+                    "Segment THAT grounded description (cross-checked against "
+                    "the clip) into atomic subtasks. Do not introduce any "
+                    "action that is not in your description above.\n\n"
+                )
+
+        prompt = self._with_causal_rules(
+            load_prompt("plan_subtasks").format(
+                episode_task=task,
+                min_subtask_seconds=self.config.min_subtask_seconds,
+                max_steps=self.config.plan_max_steps,
+                episode_duration=f"{win_len:.3f}",
+                observation_block=observation_block,
+            )
+        )
+        spans = self._vlm_field(self._video_message(record, prompt, window=window), "subtasks")
+        # Window-relative clamp; no frame-snap dedupe yet (done on the
+        # merged absolute set).
+        cleaned = self._clean_spans(spans, record, bounds=(0.0, win_len), dedupe=False)
+        if not cleaned:
+            return []
+
+        # Offset window-relative spans back to absolute episode time.
+        for s in cleaned:
+            s["start"] = w0 + float(s["start"])
+            s["end"] = w0 + float(s["end"])
+        return cleaned
+
+    def _stitch_full_coverage(
+        self, spans: list[dict[str, Any]], record: EpisodeRecord
+    ) -> list[dict[str, Any]]:
+        """Make subtask spans tile the full episode with no gaps.
+
+        * The first subtask starts at the episode's first frame ``t0``
+          (any idle / approach before the first labelled action is folded
+          into it), so every early frame has an active subtask.
+        * Each subtask's ``end`` is snapped to the next subtask's
+          ``start`` (gaps between spans are closed), and the final
+          subtask's ``end`` extends to the last frame ``t_last``.
+
+        Starts are otherwise left as the (already frame-snapped, distinct)
+        values the VLM produced — only the FIRST start is pulled
+        back to ``t0``, which can't collide with a later span because it
+        was already the earliest. Purely deterministic; runs after the
+        VLM passes.
+        """
+        if not spans or not record.frame_timestamps:
+            return spans
+        t0 = float(record.frame_timestamps[0])
+        t_last = float(record.frame_timestamps[-1])
+        spans = sorted(spans, key=lambda s: float(s["start"]))
+        spans[0]["start"] = t0
+        for i in range(len(spans) - 1):
+            spans[i]["end"] = float(spans[i + 1]["start"])
+        spans[-1]["end"] = t_last
+        for s in spans:
+            if float(s["end"]) < float(s["start"]):
+                s["end"] = float(s["start"])
+        return spans
+
+    @staticmethod
+    def _with_causal_rules(prompt: str) -> str:
+        """Append the causal event-boundary rules to a describe/segment prompt."""
+        return f"{prompt}\n\n{_CAUSAL_BOUNDARY_RULES}"
+
+    def _clean_spans(
+        self,
+        spans: Any,
+        record: EpisodeRecord,
+        bounds: tuple[float, float] | None = None,
+        dedupe: bool = True,
+    ) -> list[dict[str, Any]]:
+        """Clamp / sort / (optionally) dedupe raw VLM subtask spans into valid rows.
+
+        ``bounds`` overrides the clamp range — pass the window's
+        ``(w_lo, w_hi)`` when cleaning window-relative spans, or leave
+        ``None`` to clamp to the whole episode ``[t0, t_last]``.
+        ``dedupe`` runs the frame-snap distinct-start step; skip it for
+        window-relative spans (frame snapping is done once on the merged,
+        absolute-time set).
+        """
+        if not spans:
+            return []
+        if bounds is not None:
+            lo, hi = float(bounds[0]), float(bounds[1])
+        else:
+            lo = record.frame_timestamps[0]
+            hi = record.frame_timestamps[-1]
+        cleaned: list[dict[str, Any]] = []
+        for span in spans:
+            try:
+                start = float(span["start"])
+                end = float(span["end"])
+                text = str(span["text"]).strip()
+            except (KeyError, ValueError, TypeError):
+                continue
+            start = max(lo, min(start, hi))
+            end = max(lo, min(end, hi))
+            if end < start:
+                start, end = end, start
+            if not text:
+                continue
+            cleaned.append({"text": text, "start": start, "end": end})
+        cleaned.sort(key=lambda s: s["start"])
+        if dedupe:
+            return self._dedupe_starts_to_distinct_frames(cleaned, record)
+        return cleaned
+
+    def _describe_episode(
+        self, record: EpisodeRecord, task: str, window: tuple[float, float] | None = None
+    ) -> str:
+        """Grounding pass: free-form chronological description of the (windowed) video."""
+        prompt = self._with_causal_rules(load_prompt("plan_subtask_describe").format(episode_task=task))
+        text = self._vlm_field(self._video_message(record, prompt, window=window), "description")
+        return text.strip() if isinstance(text, str) and text.strip() else ""
+
+    @staticmethod
+    def _dedupe_starts_to_distinct_frames(
+        spans: list[dict[str, Any]], record: EpisodeRecord
+    ) -> list[dict[str, Any]]:
+        """Bump same-frame subtask starts onto distinct frames.
+
+        Two consecutive VLM spans whose ``start`` rounds to the same
+        source frame (after :func:`snap_to_frame`) would otherwise emit
+        two ``style=subtask`` rows at the identical persistent
+        timestamp. The training-time renderer's ``active_at(t,
+        style=subtask)`` resolver can't disambiguate that and raises
+        ``Ambiguous resolver for style='subtask'``.
+
+        Walk the (sorted-by-start) spans, snap each to its frame, and
+        if the snapped frame is already taken push the span onto the
+        next unused frame so both subtasks survive on distinct
+        timestamps. If the episode ends before a free frame is found,
+        the trailing span is dropped with a warning — better than
+        poisoning the render.
+        """
+        if not spans:
+            return spans
+        frames = record.frame_timestamps
+        if not frames:
+            return spans
+        used: set[float] = set()
+        out: list[dict[str, Any]] = []
+        for span in spans:
+            ts = snap_to_frame(span["start"], frames)
+            if ts in used:
+                next_ts = next((f for f in frames if f > ts and f not in used), None)
+                if next_ts is None:
+                    logger.warning(
+                        "episode %d: subtask %r snapped to occupied frame "
+                        "%.3f and no free later frame exists — dropping",
+                        record.episode_index,
+                        span.get("text"),
+                        ts,
+                    )
+                    continue
+                ts = next_ts
+            used.add(ts)
+            new_span = {**span, "start": ts}
+            if float(new_span.get("end", ts)) < ts:
+                new_span["end"] = ts
+            out.append(new_span)
+        return out
+
+    def _generate_plan(
+        self,
+        record: EpisodeRecord,  # noqa: ARG002  (kept for signature stability)
+        subtask_spans: Sequence[dict[str, Any]],
+        *,
+        refresh_t: float | None = None,
+        interjection: str | None = None,  # noqa: ARG002
+        task: str | None = None,  # noqa: ARG002
+    ) -> str | None:
+        """Deterministic plan = numbered list of *still-todo* subtasks.
+
+        No VLM call: a plain numbered list keeps the plan aligned with the
+        upcoming subtasks (the old VLM "compact hierarchical plan" prompt
+        cost a round-trip per episode/refresh and could diverge).
+
+            1. <subtask 1>
+            2. <subtask 2>
+
+        On a refresh at ``refresh_t`` (from ``run_plan_updates`` on
+        interjections, and ``run_episode`` at each boundary), only subtasks
+        starting at or after ``refresh_t`` are included — so it always
+        describes what's left.
+        """
+        if not subtask_spans:
+            return None
+        remaining = [
+            s for s in subtask_spans if refresh_t is None or float(s.get("start", 0.0)) >= float(refresh_t)
+        ]
+        if not remaining:
+            # Past the last subtask boundary on a late refresh — nothing
+            # left to plan; emit None so the caller skips the row.
+            return None
+        return "\n".join(f"{i}. {span.get('text', '').strip()}" for i, span in enumerate(remaining, start=1))
+
+    def _generate_memory(
+        self,
+        record: EpisodeRecord,
+        prior_memory: str,
+        completed: str,
+        remaining: Sequence[str],
+        *,
+        task: str | None = None,
+    ) -> str:
+        prompt = load_prompt("plan_memory").format(
+            episode_task=(task if task is not None else record.episode_task),
+            prior_memory=prior_memory or "(none)",
+            completed_subtask=completed,
+            remaining_subtasks=", ".join(remaining) if remaining else "(none)",
+        )
+        memory = self._vlm_field(self._text_message(prompt), "memory")
+        return memory.strip() if isinstance(memory, str) else ""

src/lerobot/annotations/steerable_pipeline/prompts/__init__.py

@@ -0,0 +1,33 @@
+#!/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.
+"""Prompt templates loaded as plain text.
+
+One file per use site. Templates use ``str.format(**vars)`` substitution; we
+intentionally avoid jinja2 here so the templates remain inspectable in
+plain editors and roundtrip cleanly through ``ruff format``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+_DIR = Path(__file__).parent
+
+
+def load(name: str) -> str:
+    """Read prompt template ``name.txt`` from the ``prompts/`` directory."""
+    path = _DIR / f"{name}.txt"
+    return path.read_text(encoding="utf-8")

src/lerobot/annotations/steerable_pipeline/prompts/interjections_initial_speech.txt

@@ -0,0 +1,12 @@
+The user just asked the robot: "{episode_task}".
+
+Generate a short verbal acknowledgement the robot would speak back before
+beginning the task. Style: compact, confident, friendly.
+
+Examples (Hi Robot, Shi 2025): "Sure, I won't put cheese on it.",
+"OK, starting with the sponge.", "Got it.".
+
+Prefer very short replies: "Got it.", "On it.", "OK."
+
+Output strictly valid JSON:
+  {{ "text": "<the spoken acknowledgement>" }}

src/lerobot/annotations/steerable_pipeline/prompts/interjections_interjection.txt

@@ -0,0 +1,46 @@
+You are generating training data for a Hi Robot-style hierarchical
+robot policy. The robot in this demonstration has ALREADY executed
+every step shown in the video — we cannot retroactively change the
+action stream. To keep training data consistent with the video, the
+"interjection" must align with what the robot is *about to do next* in
+the demonstration, framed as a natural mid-task user request.
+
+The episode's overall task: "{episode_task}".
+
+The images above show roughly {window_seconds:.1f} seconds straddling a
+subtask boundary in the demonstration:
+
+- Subtask the robot just finished: "{prev_subtask}"
+- Subtask the robot is about to start: "{next_subtask}"
+- Time into episode: {timestamp:.2f}s
+
+Write ONE compact interjection the user would naturally say at this
+moment to prompt / confirm / encourage the robot to do "{next_subtask}".
+Keep it like a mid-task coaching cue, not a full instruction paragraph.
+Also write the robot's compact verbal acknowledgement.
+
+Hard rules:
+
+- The interjection MUST be consistent with the next subtask. The user
+  cannot ask for something different from what the robot then does in
+  the video. If you're tempted to say "actually skip X" or "do Y
+  instead", DO NOT — those would contradict the demonstration.
+- The interjection must reference an object, location, or action that
+  is plausible given the visible scene and the next subtask text.
+- One short phrase or sentence each. Conversational, not robotic.
+- Prefer direct cues: "{next_subtask}, please."; "Now {next_subtask}."
+- Keep robot speech very short: "OK.", "On it.", "Doing that."
+
+Style examples (vary the phrasing — don't reuse these verbatim):
+  - "Now go ahead and {next_subtask}."
+  - "Great, can you {next_subtask} next?"
+  - "{next_subtask}, please."
+  - "Before you continue, please {next_subtask}."
+  - "Looking good — {next_subtask} now."
+  - "Okay, {next_subtask}."
+
+Output strictly valid JSON:
+  {{
+    "interjection": "<short cue from the user, asking for the next subtask>",
+    "speech":       "<short robot acknowledgement>"
+  }}

src/lerobot/annotations/steerable_pipeline/prompts/plan_memory.txt

@@ -0,0 +1,36 @@
+You are updating the robot's compressed semantic memory at the boundary of
+a completed subtask.
+
+Reference (verbatim from MEM, Torne 2026):
+"Remove or compress information in the language memory whenever
+appropriate. Keep ONLY the minimal set of relevant information for future
+task execution. Specific object attributes (colors, precise quantities of
+each item) get discarded when their details won't affect subsequent
+actions. Functional outcomes (where items went, how many) are preserved."
+
+Episode task: "{episode_task}"
+Previous memory: {prior_memory}
+Just-completed subtask: "{completed_subtask}"
+Remaining subtasks (for relevance judgement only): {remaining_subtasks}
+
+Write the memory as a short FIRST-PERSON, PAST-TENSE narrative of what the
+robot has accomplished so far — the running story it would tell itself.
+
+Authoring rules:
+- First person, past tense. Every sentence starts with "I": "I picked
+  up...", "I opened...", "I moved to...".
+- One or two short sentences. Extend the previous memory with the
+  just-completed subtask; do not rewrite it from scratch.
+- Keep WHAT happened (functional outcomes — where items went, how many),
+  drop HOW (grasp details, motions).
+- Compress completed steps and drop object attributes (colors, exact
+  counts) once they no longer affect the remaining subtasks.
+
+Example (MEM, Torne 2026):
+  Before: "I prepared the pot and got the potatoes, milk, and butter. I
+           moved to the drawer."
+  After:  "I prepared the pot and got the ingredients. I opened the
+           drawer with the masher."
+
+Output strictly valid JSON:
+  {{ "memory": "<one or two short first-person past-tense sentences>" }}

src/lerobot/annotations/steerable_pipeline/prompts/plan_subtask_describe.txt

@@ -0,0 +1,27 @@
+You are watching a teleoperated robot demonstration from a single
+camera. The user asked the robot to: "{episode_task}"
+
+This is an OBSERVATION pass. Watch the entire clip and describe, in
+chronological order, ONLY what the robot physically does — the concrete
+motions, approaches, contacts, grasps, releases, and relocations you can
+actually SEE in the frames.
+
+Hard rules:
+- Describe only motion visible in the video. Do NOT use the task
+  instruction to guess steps that aren't shown. The instruction is the
+  goal; the video is ground truth.
+- Do NOT segment into named subtasks yet and do NOT output JSON beyond
+  the single field below. Just narrate what happens.
+- Give an approximate timestamp (in seconds) for each distinct event,
+  e.g. "0.0-1.4s: the base drives forward toward the stove".
+- Do NOT invent objects, grasps, destinations, or steps. If the robot
+  only does one thing (e.g. it just navigates and the clip ends), say
+  exactly that and nothing more.
+- Be concrete and literal. "the gripper closes on the mug" — not "the
+  robot prepares to make coffee".
+
+Output strictly valid JSON:
+
+  {{
+    "description": "<chronological, timestamped description of ONLY what is visible>"
+  }}

src/lerobot/annotations/steerable_pipeline/prompts/plan_subtasks.txt

@@ -0,0 +1,112 @@
+You are labeling a teleoperated robot demonstration.
+
+The user originally asked: "{episode_task}"
+
+You are shown the entire demonstration as a single video. Watch the
+whole clip, then segment it into a list of consecutive atomic subtasks
+the robot performs.
+
+{observation_block}GROUNDING — read this first, it overrides everything below:
+- Label ONLY what the robot actually does in the video. Every subtask
+  you emit must correspond to motion you can SEE in specific frames.
+- Do NOT invent, anticipate, or pad. If the robot only does one thing
+  (e.g. it just navigates to a location and the clip ends), emit
+  EXACTLY ONE subtask. Many demonstrations are a single atomic skill.
+- ``max_steps`` below is a hard CEILING, not a target. Emitting fewer
+  subtasks than the ceiling is not just allowed, it is expected for
+  short / atomic demonstrations. One correct subtask is far better
+  than several invented ones.
+- If the video does not clearly show the action implied by the task,
+  describe what you actually see — do NOT fabricate the task's steps
+  from the instruction text. The instruction tells you the goal; the
+  VIDEO is the ground truth for what happened.
+
+Authoring rules — Hi Robot atom granularity, pi0.7-style short prompts:
+
+- Each subtask = one COMPOSITE atomic skill the low-level policy can
+  execute end-to-end. A "skill" bundles its own approach motion with
+  its terminal action — do NOT split the approach off as its own
+  subtask. The whole-arm policy already learns to reach as part of
+  every manipulation primitive.
+- Write each subtask as an IMPERATIVE COMMAND, starting with one of
+  these verbs (extend only when none fits):
+    pick up <obj>           — approach + grasp + lift in one subtask
+    put <obj> on/in <loc>   — transport + release in one subtask
+    place <obj> on/in <loc> — synonym of "put"; pick one and stay consistent
+    push <obj>              — contact + linear shove
+    pull <obj>              — contact + linear retract
+    turn <knob/dial/handle> — rotary actuation
+    press <button>          — single-press contact
+    open <drawer/door/lid>  — full open motion
+    close <drawer/door/lid> — full close motion
+    pour <src> into <dst>   — tilt + flow
+    insert <obj> into <slot>— alignment + push-fit
+    go to <loc>             — ONLY when no grasp / actuation follows
+                             (e.g. a pure relocation between phases).
+                             If the next subtask grasps something at
+                             that location, drop "go to ..." and just
+                             write "pick up ..." instead.
+- Forbidden ultra-fine splits — the VLM is NOT allowed to emit these
+  as standalone subtasks; fold them into the parent composite:
+    "move to X"   → fold into "pick up X" (or whatever follows)
+    "reach for X" → fold into "pick up X"
+    "grasp X"     → fold into "pick up X"
+    "lift X"      → fold into "pick up X" (or "put X on Y" if it's
+                    the transport phase of a place)
+    "release X"   → fold into "put X on Y" (or "place X in Y")
+- Keep it SHORT — a verb phrase, not a sentence. Drop articles
+  ("the", "a") and adverbs ("carefully", "slowly"). Add a "how"
+  detail (which hand, which grasp point) ONLY when it is needed to
+  disambiguate. Every subtask must begin with one of the verbs
+  above (no leading nouns, no "then", no "first").
+- NEVER use third person. Never write "the robot", "the arm", "the
+  gripper moves", "it picks up" — the robot is implied. Command it,
+  do not describe it.
+- Use the exact object nouns from the task above. If the task says
+  "cube", every subtask says "cube" — never switch to "block". If it
+  says "box", never switch to "bin"/"container". Keep vocabulary
+  consistent across the whole episode.
+- Good: "pick up blue cube", "put blue cube in box", "open drawer",
+  "turn red knob", "press start button", "go to sink".
+- Bad: "move to blue cube" (approach as its own subtask — forbidden,
+  must be folded into "pick up blue cube"); "the robot arm moves
+  towards the blue cube" (third person, too long); "carefully pick
+  up the cube" (adverb, article); "release the yellow block"
+  ("block" when the task said "cube", and "release" must be folded
+  into a "put"/"place" subtask).
+- Subtasks are non-overlapping and cover the full episode in order.
+  Choose the cut points yourself based on what you see in the video
+  (gripper open/close events, contact, regrasps, transitions).
+- Each subtask spans at least {min_subtask_seconds} seconds. If a
+  candidate span would be shorter, merge it into its neighbour
+  rather than emitting it.
+- Do not exceed {max_steps} subtasks total. Fewer, larger composites
+  are preferred over many micro-steps.
+- Every subtask's [start_time, end_time] must lie within
+  [0.0, {episode_duration}] seconds.
+
+SPECIAL CASES — verb disambiguation (each rule is narrowly visual and
+fires ONLY on the spatial situation it names; it must not change how you
+label any other situation):
+- STACK vs PUT: if an object is placed ON TOP OF another specific object
+  (not on a flat table / shelf / counter), use "stack ... on ...", not
+  "put". "stack blue book on green book", NOT "put blue book on table".
+- INSERT vs PUT: if an object goes INTO a fitted slot / hole / socket /
+  receptacle (push-fit), use "insert ... into ...", not "put".
+- RETRIEVE/PICK-UP vs PUT (direction): watch the gripper. If it CLOSES
+  on the object and the object moves WITH the hand, it is "pick up" /
+  "retrieve" (object leaves its location). If the gripper OPENS and the
+  object stays where the hand left it, it is "put" / "place" (object
+  arrives at a location). Decide by which way the object moves, not by
+  where the hand ends up.
+- POUR vs PUT: only use "pour" when the source is tilted and contents
+  flow out; moving a full container without tilting is "put"/"place".
+
+Output strictly valid JSON of shape:
+
+  {{
+    "subtasks": [
+      {{"text": "<short imperative verb phrase>", "start": <float>, "end": <float>}},
+      ...
+    ]
+  }}

src/lerobot/annotations/steerable_pipeline/prompts/plan_task_aug_axes.txt

@@ -0,0 +1,67 @@
+You are generating structured augmentations of a robot task instruction
+for training a language-conditioned policy. Unlike free-form rephrasing,
+your variants follow a NAMED 5-axis taxonomy — each axis omits or varies
+a specific element of the task while preserving its meaning.
+
+Original task: "{base_task}"
+
+Produce variants along five named axes. Each axis has a target count.
+The whole batch should expose the policy to maximum linguistic diversity
+WITHOUT changing what the robot is supposed to do.
+
+Axes and target counts:
+
+  synonym_paraphrase ({n_synonym}):
+    Different wording / verbs / sentence structure. ALL information
+    from the original task is preserved — same object, same arm
+    specification if present, same orientation if present, same grasp
+    if present.
+
+  omit_arm ({n_omit_arm}):
+    Drop the left/right/both arm specification from the task. Skip
+    entirely (emit 0 entries) if the original task does NOT mention an
+    arm. Do not invent an arm specification just to omit it.
+
+  omit_orientation ({n_omit_orientation}):
+    Drop orientation cues (upright, sideways, facing the user,
+    long-edge-first, etc.). Skip entirely if no orientation cue is
+    present in the original task.
+
+  omit_grasp_method ({n_omit_grasp_method}):
+    Drop the grip / grasp method specification (pinch, wrap, hold by
+    the rim, etc.). Skip entirely if no grasp method is mentioned.
+
+  combined_omissions ({n_combined}):
+    Combine TWO of the above omissions simultaneously (e.g. drop both
+    arm and orientation). Skip entirely if fewer than two of (arm,
+    orientation, grasp_method) appear in the original task.
+
+Hard rules:
+- Each variant MUST preserve the core action, the target object, AND
+  the goal / destination. Do not change which object is involved, where
+  it goes, or the high-level action. "Navigate to the stove" may become
+  "go to the stove" or "head over to the stove" — it must NEVER become
+  "wander around the kitchen", "explore the room", or anything that
+  drops or generalises the stove destination. If you cannot vary the
+  wording without changing the goal, emit fewer variants.
+- Only the FIVE listed elements (wording, arm, orientation, grasp
+  method, or a combination) may be varied or omitted. The verb's
+  meaning, the object, and the destination are fixed.
+- Each variant is plain prose, no markdown, no quotes, no list numbers.
+- Each variant must be DISTINCT from every other variant in the entire
+  output, both within and across axes. Near-duplicates are not allowed.
+- If an axis cannot reach its target count because the original task
+  lacks the omittable element, emit fewer entries — do NOT pad the
+  axis with paraphrases that belong to a different axis.
+- Variants should not all start with verbs — vary sentence structure
+  (some imperative, some polite request, some question).
+
+Output strictly valid JSON of shape:
+
+  {{
+    "synonym_paraphrase": ["<v1>", "<v2>", ...],
+    "omit_arm": ["<v1>", "<v2>", ...],
+    "omit_orientation": ["<v1>", ...],
+    "omit_grasp_method": ["<v1>", ...],
+    "combined_omissions": ["<v1>", ...]
+  }}

src/lerobot/annotations/steerable_pipeline/prompts/plan_task_rephrasings.txt

@@ -0,0 +1,32 @@
+You are generating training data for a Hi Robot-style policy. We need
+{n} alternative phrasings of the same robot task so the policy sees
+diverse user prompts during training instead of the same canonical
+string repeated every frame.
+
+Original task:
+"{base_task}"
+
+Generate exactly {n} alternative phrasings of the same task. Vary:
+
+- formality (casual / polite / curt)
+- verbosity (mostly short imperative; occasional polite request)
+- word choice (synonyms, different verbs)
+- sentence structure (imperative / question / suggestion)
+
+Hard rules:
+- Each phrasing MUST preserve the exact meaning of the original task.
+  Do not change which object is involved, the destination, or the
+  action. Do not add extra steps. Do not invent new objects.
+- Each phrasing must be a short phrase or sentence, plain prose, no
+  markdown, no quotes, no list numbers.
+- Phrasings must be distinct — no near-duplicates.
+- Output exactly {n} entries.
+
+Output strictly valid JSON:
+  {{
+    "rephrasings": [
+      "<phrasing 1>",
+      "<phrasing 2>",
+      ...
+    ]
+  }}

src/lerobot/annotations/steerable_pipeline/prompts/plan_video_task.txt

@@ -0,0 +1,17 @@
+The video above shows a robot manipulation episode in full. Look at
+the entire video and describe in ONE concise sentence what the robot
+is doing.
+
+Rules:
+- One sentence, in natural English, like a user instruction.
+- Capture the goal of the demonstration, not low-level motions.
+  Example: "place the yellow cube into the red bin" — not "move the
+  end-effector down 5cm and close the gripper".
+- 4 to 15 words. Plain prose, no markdown, no bullets, no quotes.
+- Do not invent objects or actions that aren't visible.
+- Do not output anything other than the JSON object below.
+
+Output strictly valid JSON:
+  {{
+    "task": "<single concise sentence describing what the robot does in this video>"
+  }}

src/lerobot/annotations/steerable_pipeline/prompts/vqa.txt

@@ -0,0 +1,32 @@
+You are generating a frame-grounded visual question/answer pair for
+chain-of-thought training. Reference: ECoT (Zawalski 2024) and Steerable
+Policies — both train policies on grounded features such as bounding box
+pixel coordinates, keypoints, counts, attributes, and spatial relations.
+
+The frame shows a robot working on: "{episode_task}".
+
+Question types and the EXACT answer JSON shape required for each:
+
+  bbox       => {{"detections": [{{"label": "<obj>", "bbox_format": "xyxy",
+                                    "bbox": [x1, y1, x2, y2]}}, ...]}}
+                bbox is in pixel coordinates (x_min, y_min, x_max, y_max).
+                ECoT example: "a white cup [124, 25, 176, 113]".
+
+  keypoint   => {{"label": "<point>", "point_format": "xy",
+                  "point": [x, y]}}
+
+  count      => {{"label": "<obj>", "count": <int>,
+                  "note": "<optional short note>"}}
+
+  attribute  => {{"label": "<obj>", "attribute": "<color|shape|state|...>",
+                  "value": "<observed value>"}}
+
+  spatial    => {{"subject": "<obj>", "relation": "<left_of|right_of|on|in|"
+                  "above|below|near>", "object": "<obj>"}}
+
+Generate a question of type "{question_type}". Output strictly valid JSON:
+
+  {{
+    "question": "<short, frame-grounded question>",
+    "answer":   <object whose shape matches the schema above>
+  }}

src/lerobot/annotations/steerable_pipeline/reader.py

@@ -0,0 +1,216 @@
+#!/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.
+"""Datatrove-shaped reader.
+
+The reader walks ``data/chunk-*/file-*.parquet`` and yields one record per
+episode containing:
+
+- ``episode_index``: int
+- ``frame_timestamps``: tuple[float, ...]
+- ``frame_indices``: tuple[int, ...]
+- ``episode_task``: str (canonical task from ``meta/tasks.parquet``)
+- ``data_path``: pathlib.Path of the source parquet shard
+- ``frames_df``: pandas.DataFrame slice for the episode (only loaded on demand)
+
+This shape lets each module operate per-episode without loading all parquet
+rows into memory at once.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import pyarrow.parquet as pq
+
+from lerobot.datasets.io_utils import load_tasks
+from lerobot.datasets.utils import DEFAULT_TASKS_PATH
+
+
+@dataclass
+class EpisodeRecord:
+    """Per-episode record yielded by the reader."""
+
+    episode_index: int
+    episode_task: str
+    frame_timestamps: tuple[float, ...]
+    frame_indices: tuple[int, ...]
+    data_path: Path
+    row_offset: int  # row offset within the parquet file where this episode starts
+    row_count: int  # number of rows for this episode
+
+    # Memoized parquet slice — populated on first ``frames_df()`` call so
+    # repeat queries from different modules don't re-read the whole shard.
+    _frames_df_cache: Any = field(default=None, init=False, repr=False, compare=False)
+
+    def frames_df(self):  # type: ignore[no-untyped-def]
+        """Lazy-load the pandas slice for this episode (memoized)."""
+        if self._frames_df_cache is None:
+            import pandas as pd  # noqa: PLC0415  - deferred for optional dataset extra
+
+            table = pq.read_table(self.data_path)
+            df: pd.DataFrame = table.to_pandas()
+            self._frames_df_cache = df.iloc[self.row_offset : self.row_offset + self.row_count].reset_index(
+                drop=True
+            )
+        return self._frames_df_cache
+
+
+def reconstruct_subtask_spans(
+    rows: Sequence[dict[str, Any]],
+    *,
+    episode_end_t: float | None = None,
+) -> list[dict[str, Any]]:
+    """Turn ``style="subtask"`` rows into ``{text, start, end}`` spans.
+
+    Each span's ``end`` is the next span's ``start``. The final span's
+    ``end`` defaults to its own ``start`` (zero-duration) — pass
+    ``episode_end_t`` to extend it to the episode's last frame instead,
+    which is what downstream consumers (memory, interjection boundary
+    selection) expect.
+
+    Used by the ``plan`` module (plan-update pass) and the
+    ``interjections`` module (interjection anchoring), which both need the
+    same span shape.
+    """
+    sorted_rows = sorted(
+        (r for r in rows if r.get("style") == "subtask"),
+        key=lambda r: float(r["timestamp"]),
+    )
+    spans: list[dict[str, Any]] = []
+    for r in sorted_rows:
+        t = float(r["timestamp"])
+        if spans:
+            spans[-1]["end"] = t
+        spans.append({"text": r.get("content") or "", "start": t, "end": t})
+    if spans and episode_end_t is not None and float(episode_end_t) > spans[-1]["start"]:
+        spans[-1]["end"] = float(episode_end_t)
+    return spans
+
+
+def snap_to_frame(t: float, frame_timestamps: Sequence[float]) -> float:
+    """Snap an arbitrary float to the nearest exact source frame timestamp.
+
+    Modules use this when emitting event-style rows so the row's
+    timestamp matches a real parquet frame: event rows must land on an
+    exact frame, otherwise the per-frame event lookup the writer does
+    would never match them.
+    """
+    if not frame_timestamps:
+        return float(t)
+    nearest = min(frame_timestamps, key=lambda f: abs(f - t))
+    return float(nearest)
+
+
+def _load_tasks_lookup(root: Path) -> dict[int, str]:
+    """Map ``task_index -> task`` from ``meta/tasks.parquet``.
+
+    Returns an empty dict when the file is absent — the task description is
+    derived later from the video if needed. Reuses the library-level
+    :func:`lerobot.datasets.io_utils.load_tasks`, which returns the tasks
+    frame indexed by task string with a ``task_index`` column.
+    """
+    if not (root / DEFAULT_TASKS_PATH).exists():
+        return {}
+    tasks = load_tasks(root)
+    return {int(idx): str(task) for task, idx in zip(tasks.index, tasks["task_index"], strict=True)}
+
+
+def iter_episodes(root: Path, *, only_episodes: tuple[int, ...] | None = None) -> Iterator[EpisodeRecord]:
+    """Yield :class:`EpisodeRecord` for every episode under ``root/data/``.
+
+    Episodes are yielded in ascending ``episode_index`` order. The reader does
+    not assume a specific chunk/file layout: it scans every ``*.parquet``
+    under ``data/`` and groups by ``episode_index``.
+    """
+    tasks = _load_tasks_lookup(root)
+    data_dir = root / "data"
+    parquet_files = sorted(data_dir.rglob("*.parquet"))
+
+    only_set = set(only_episodes) if only_episodes is not None else None
+
+    for path in parquet_files:
+        yield from _iter_one_path(path, tasks, only_set)
+
+
+def _iter_one_path(path: Path, tasks: dict[int, str], only_set: set[int] | None) -> Iterator[EpisodeRecord]:
+    table = pq.read_table(path)
+    names = table.column_names
+    if "episode_index" not in names:
+        return
+    episode_col = table.column("episode_index").to_pylist()
+    timestamp_col = (
+        table.column("timestamp").to_pylist() if "timestamp" in names else [0.0] * len(episode_col)
+    )
+    frame_col = (
+        table.column("frame_index").to_pylist() if "frame_index" in names else list(range(len(episode_col)))
+    )
+    task_col = table.column("task_index").to_pylist() if "task_index" in names else None
+
+    def _build(
+        ep: int,
+        start: int,
+        end: int,
+        task_idx: int | None,
+        ts_buf: list[float],
+        fi_buf: list[int],
+    ) -> EpisodeRecord | None:
+        if only_set is not None and ep not in only_set:
+            return None
+        task = tasks.get(task_idx, "") if task_idx is not None else ""
+        return EpisodeRecord(
+            episode_index=ep,
+            episode_task=task,
+            frame_timestamps=tuple(ts_buf),
+            frame_indices=tuple(fi_buf),
+            data_path=path,
+            row_offset=start,
+            row_count=end - start,
+        )
+
+    cur_ep: int | None = None
+    start_offset = 0
+    ts_buf: list[float] = []
+    fi_buf: list[int] = []
+    cur_task_idx: int | None = None
+
+    for i, ep in enumerate(episode_col):
+        if cur_ep is None:
+            cur_ep = ep
+            start_offset = i
+            ts_buf = [timestamp_col[i]]
+            fi_buf = [frame_col[i]]
+            cur_task_idx = task_col[i] if task_col is not None else None
+            continue
+        if ep != cur_ep:
+            rec = _build(cur_ep, start_offset, i, cur_task_idx, ts_buf, fi_buf)
+            if rec is not None:
+                yield rec
+            cur_ep = ep
+            start_offset = i
+            ts_buf = [timestamp_col[i]]
+            fi_buf = [frame_col[i]]
+            cur_task_idx = task_col[i] if task_col is not None else None
+        else:
+            ts_buf.append(timestamp_col[i])
+            fi_buf.append(frame_col[i])
+
+    if cur_ep is not None:
+        rec = _build(cur_ep, start_offset, len(episode_col), cur_task_idx, ts_buf, fi_buf)
+        if rec is not None:
+            yield rec

src/lerobot/annotations/steerable_pipeline/staging.py

@@ -0,0 +1,92 @@
+#!/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.
+"""Per-episode staging.
+
+Each module writes its raw output as a JSONL file under
+``<staging_dir>/episode_{ep:06d}/<module>.jsonl``. The writer reads back this
+staging tree and partitions rows into the two language columns.
+
+JSONL is preferred over parquet here because the staging artifact is meant to
+be human-inspectable, easy to diff between prompt iterations, and trivially
+appended to. The final dataset format is parquet; staging is just an
+intermediate.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+ModuleName = str
+
+_MODULES: tuple[ModuleName, ...] = (
+    "plan",
+    "interjections",
+    "vqa",
+)
+
+
+@dataclass
+class EpisodeStaging:
+    """Filesystem layout for a single episode's staged module outputs."""
+
+    root: Path
+    episode_index: int
+
+    @property
+    def episode_dir(self) -> Path:
+        return self.root / f"episode_{self.episode_index:06d}"
+
+    def path_for(self, module: ModuleName) -> Path:
+        if module not in _MODULES:
+            raise ValueError(f"Unknown module {module!r}; expected one of {_MODULES}")
+        return self.episode_dir / f"{module}.jsonl"
+
+    def write(self, module: ModuleName, rows: Iterable[dict[str, Any]]) -> Path:
+        path = self.path_for(module)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        # Atomic replace: a crash mid-write would otherwise leave a
+        # half-written JSONL file that ``read()`` would then fail to
+        # parse. Write to a sibling .tmp and rename so the target path
+        # only ever points at a complete file.
+        tmp_path = path.with_suffix(path.suffix + ".tmp")
+        with tmp_path.open("w", encoding="utf-8") as f:
+            for row in rows:
+                f.write(json.dumps(row, ensure_ascii=False, sort_keys=True))
+                f.write("\n")
+        tmp_path.replace(path)
+        return path
+
+    def read(self, module: ModuleName) -> list[dict[str, Any]]:
+        path = self.path_for(module)
+        if not path.exists():
+            return []
+        out: list[dict[str, Any]] = []
+        with path.open(encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    out.append(json.loads(line))
+        return out
+
+    def read_all(self) -> dict[ModuleName, list[dict[str, Any]]]:
+        return {m: self.read(m) for m in _MODULES}
+
+    def has(self, module: ModuleName) -> bool:
+        return self.path_for(module).exists()

src/lerobot/annotations/steerable_pipeline/validator.py

@@ -0,0 +1,332 @@
+#!/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.
+"""Pre-write validation against staged outputs.
+
+Runs after all three modules have written their per-episode artifacts but
+*before* the writer rewrites parquet shards. The validator never touches
+parquet; it only inspects the staging tree and the source frame timestamps
+exposed by :class:`EpisodeRecord`.
+
+Checks (per the plan's "Intermediate staging and validation" section):
+
+- exact timestamp alignment against source frame timestamps
+- no orphan speech / interjection pairs
+- plan / memory emission consistency (events have a paired persistent row)
+- VQA assistant ``content`` is valid JSON (one of bbox / keypoint / count /
+  attribute / spatial)
+- every row maps to its correct column under :func:`column_for_style`
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from lerobot.datasets.language import (
+    LANGUAGE_EVENTS,
+    LANGUAGE_PERSISTENT,
+    column_for_style,
+    is_view_dependent_style,
+    validate_camera_field,
+)
+
+from .reader import EpisodeRecord
+from .staging import EpisodeStaging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ValidationReport:
+    """Outcome of one validation pass across all episodes."""
+
+    errors: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+    episodes_checked: int = 0
+
+    @property
+    def ok(self) -> bool:
+        return not self.errors
+
+    def add_error(self, message: str) -> None:
+        self.errors.append(message)
+
+    def add_warning(self, message: str) -> None:
+        self.warnings.append(message)
+
+    def summary(self) -> str:
+        return f"checked={self.episodes_checked} errors={len(self.errors)} warnings={len(self.warnings)}"
+
+
+VQA_ANSWER_SHAPES: dict[str, set[str]] = {
+    "bbox": {"detections"},
+    "keypoint": {"label", "point_format", "point"},
+    "count": {"label", "count"},
+    "attribute": {"label", "attribute", "value"},
+    "spatial": {"subject", "relation", "object"},
+}
+
+
+def classify_vqa_answer(payload: Any) -> str | None:
+    """Best-effort classification of a VQA answer payload to a question type."""
+    if not isinstance(payload, dict):
+        return None
+    keys = set(payload.keys())
+    for kind, required in VQA_ANSWER_SHAPES.items():
+        if required.issubset(keys):
+            return kind
+    return None
+
+
+@dataclass
+class StagingValidator:
+    """Walks the staging tree and produces a :class:`ValidationReport`."""
+
+    timestamp_atol: float = 0.0  # exact-match by default
+    dataset_camera_keys: tuple[str, ...] | None = None
+    """Known ``observation.images.*`` keys on the dataset. When set, the
+    validator additionally enforces that every view-dependent row's
+    ``camera`` field references one of these keys. Pass ``None`` (default)
+    to skip that cross-check (e.g. in unit tests with no real dataset)."""
+
+    def validate(
+        self,
+        records: Sequence[EpisodeRecord],
+        staging_dir: Path,
+    ) -> ValidationReport:
+        report = ValidationReport()
+        for record in records:
+            self._validate_episode(record, staging_dir, report)
+            report.episodes_checked += 1
+        return report
+
+    def _validate_episode(
+        self,
+        record: EpisodeRecord,
+        staging_dir: Path,
+        report: ValidationReport,
+    ) -> None:
+        staging = EpisodeStaging(staging_dir, record.episode_index)
+        staged = staging.read_all()
+        all_rows: list[dict[str, Any]] = []
+        for module_name, rows in staged.items():
+            for row in rows:
+                row = {**row, "_module": module_name}
+                all_rows.append(row)
+
+        frame_ts = set(record.frame_timestamps)
+
+        events: list[dict[str, Any]] = []
+        persistent: list[dict[str, Any]] = []
+        for row in all_rows:
+            self._check_column_routing(row, report, record.episode_index)
+            self._check_camera_field(row, report, record.episode_index, self.dataset_camera_keys)
+            # ``_check_column_routing`` already recorded any unknown-style error;
+            # don't let the same ``column_for_style`` lookup raise here uncaught.
+            try:
+                column = column_for_style(row.get("style"))
+            except ValueError:
+                continue
+            if column == LANGUAGE_PERSISTENT:
+                persistent.append(row)
+            else:
+                events.append(row)
+
+        for row in events:
+            self._check_event_timestamp_alignment(row, frame_ts, report, record.episode_index)
+
+        self._check_speech_interjection_pairs(events, report, record.episode_index)
+        self._check_plan_memory_consistency(persistent, events, report, record.episode_index)
+        self._check_vqa_json(events, report, record.episode_index)
+        self._check_vqa_uniqueness_per_frame_camera(events, report, record.episode_index)
+
+    def _check_camera_field(
+        self,
+        row: dict[str, Any],
+        report: ValidationReport,
+        episode_index: int,
+        dataset_camera_keys: Sequence[str] | None,
+    ) -> None:
+        """Enforce the camera invariant + that the key matches the dataset's cameras."""
+        style = row.get("style")
+        camera = row.get("camera")
+        try:
+            validate_camera_field(style, camera)
+        except ValueError as exc:
+            report.add_error(f"ep={episode_index} module={row.get('_module')}: {exc}")
+            return
+        if is_view_dependent_style(style) and dataset_camera_keys and camera not in dataset_camera_keys:
+            report.add_error(
+                f"ep={episode_index} module={row.get('_module')}: camera {camera!r} on style "
+                f"{style!r} is not one of the dataset's video keys {sorted(dataset_camera_keys)!r}"
+            )
+
+    def _check_vqa_uniqueness_per_frame_camera(
+        self,
+        events: Iterable[dict[str, Any]],
+        report: ValidationReport,
+        episode_index: int,
+    ) -> None:
+        """Ensure at most one (vqa, user) and one (vqa, assistant) per (t, camera)."""
+        counts: dict[tuple[float, str, str], int] = {}
+        for row in events:
+            if row.get("style") != "vqa":
+                continue
+            ts = row.get("timestamp")
+            camera = row.get("camera")
+            role = row.get("role")
+            if ts is None or camera is None or role is None:
+                continue  # other validators flag these
+            key = (float(ts), str(camera), str(role))
+            counts[key] = counts.get(key, 0) + 1
+        for (ts, camera, role), n in counts.items():
+            if n > 1:
+                report.add_error(
+                    f"ep={episode_index}: {n} duplicate vqa rows at t={ts} "
+                    f"camera={camera!r} role={role!r}; expected at most one per (t, camera, role)"
+                )
+
+    def _check_column_routing(
+        self,
+        row: dict[str, Any],
+        report: ValidationReport,
+        episode_index: int,
+    ) -> None:
+        style = row.get("style")
+        module = row.get("_module")
+        try:
+            target_col = column_for_style(style)
+        except ValueError:
+            report.add_error(f"ep={episode_index} module={module}: unknown style {style!r}")
+            return
+        if module == "plan" and target_col != LANGUAGE_PERSISTENT:
+            report.add_error(
+                f"ep={episode_index} module=plan emitted style {style!r} that routes to {target_col} (must be persistent)"
+            )
+        if module in {"interjections", "vqa"} and target_col != LANGUAGE_EVENTS:
+            report.add_error(
+                f"ep={episode_index} module={module} emitted style {style!r} that routes to {target_col} (must be events)"
+            )
+
+    def _check_event_timestamp_alignment(
+        self,
+        row: dict[str, Any],
+        frame_ts: set[float],
+        report: ValidationReport,
+        episode_index: int,
+    ) -> None:
+        ts = row.get("timestamp")
+        if ts is None:
+            report.add_error(f"ep={episode_index}: event row missing timestamp: {row!r}")
+            return
+        if self.timestamp_atol == 0.0:
+            if float(ts) not in frame_ts:
+                report.add_error(
+                    f"ep={episode_index}: event row timestamp {ts!r} does not match any source frame timestamp"
+                )
+        else:
+            if not any(abs(float(ts) - f) <= self.timestamp_atol for f in frame_ts):
+                report.add_error(
+                    f"ep={episode_index}: event row timestamp {ts!r} not within {self.timestamp_atol}s of any frame"
+                )
+
+    def _check_speech_interjection_pairs(
+        self,
+        events: Iterable[dict[str, Any]],
+        report: ValidationReport,
+        episode_index: int,
+    ) -> None:
+        speech_ts: dict[float, int] = {}
+        interjection_ts: dict[float, int] = {}
+        for row in events:
+            ts = row.get("timestamp")
+            if ts is None:
+                continue
+            ts_f = float(ts)
+            if row.get("style") is None and row.get("role") == "assistant":
+                speech_ts[ts_f] = speech_ts.get(ts_f, 0) + 1
+            if row.get("style") == "interjection":
+                interjection_ts[ts_f] = interjection_ts.get(ts_f, 0) + 1
+
+        for ts in interjection_ts:
+            if ts not in speech_ts:
+                report.add_error(f"ep={episode_index}: interjection at t={ts} has no paired speech atom")
+
+    def _check_plan_memory_consistency(
+        self,
+        persistent: Sequence[dict[str, Any]],
+        events: Sequence[dict[str, Any]],
+        report: ValidationReport,
+        episode_index: int,
+    ) -> None:
+        plan_ts = sorted({float(r["timestamp"]) for r in persistent if r.get("style") == "plan"})
+        memory_ts = sorted({float(r["timestamp"]) for r in persistent if r.get("style") == "memory"})
+        subtask_ts = sorted({float(r["timestamp"]) for r in persistent if r.get("style") == "subtask"})
+        interjection_ts = sorted(
+            {
+                float(r["timestamp"])
+                for r in events
+                if r.get("style") == "interjection" and r.get("timestamp") is not None
+            }
+        )
+
+        if persistent and not plan_ts:
+            report.add_warning(f"ep={episode_index}: persistent rows present but no plan emitted")
+        # every interjection should have a same-timestamp plan refresh
+        for ts in interjection_ts:
+            if ts not in set(plan_ts):
+                report.add_error(
+                    f"ep={episode_index}: interjection at t={ts} has no co-timestamped plan update"
+                )
+        # memory should be emitted at subtask boundaries (subset relation)
+        if memory_ts and subtask_ts:
+            mem_set = set(memory_ts)
+            sub_set = set(subtask_ts)
+            stray = sorted(mem_set - sub_set)
+            if stray:
+                report.add_warning(f"ep={episode_index}: memory rows at {stray} not at any subtask boundary")
+
+    def _check_vqa_json(
+        self,
+        events: Iterable[dict[str, Any]],
+        report: ValidationReport,
+        episode_index: int,
+    ) -> None:
+        for row in events:
+            if row.get("style") != "vqa" or row.get("role") != "assistant":
+                continue
+            content = row.get("content")
+            if content is None:
+                report.add_error(
+                    f"ep={episode_index}: VQA assistant row at t={row.get('timestamp')} has null content"
+                )
+                continue
+            try:
+                payload = json.loads(content)
+            except (TypeError, ValueError) as exc:
+                report.add_error(
+                    f"ep={episode_index}: VQA assistant content not valid JSON at t={row.get('timestamp')}: {exc}"
+                )
+                continue
+            shape = classify_vqa_answer(payload)
+            if shape is None:
+                report.add_error(
+                    f"ep={episode_index}: VQA assistant payload at t={row.get('timestamp')} does not match any known shape: keys={list(payload) if isinstance(payload, dict) else type(payload).__name__}"
+                )

src/lerobot/annotations/steerable_pipeline/vlm_client.py

@@ -0,0 +1,617 @@
+#!/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.
+"""Shared Qwen-VL client.
+
+The pipeline uses a single shared VLM across modules. vLLM is preferred when
+available (high throughput, JSON-guided decoding); transformers is the
+fallback. A ``stub`` backend is used for unit tests so fixtures never call
+into a real model.
+
+The client speaks one method, :meth:`VlmClient.generate_json`, which:
+
+- accepts a list of OpenAI/HF-style multimodal messages,
+- requests JSON output from the server,
+- batches requests transparently,
+- and reprompts once on a JSON parse failure with an inline correction
+  message before raising.
+"""
+
+from __future__ import annotations
+
+import atexit
+import base64
+import io
+import json
+import os
+import shlex
+import signal
+import subprocess
+import sys
+import threading
+import time
+import urllib.request
+from collections.abc import Callable, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from .config import VlmConfig
+
+
+class VlmClient(Protocol):
+    """Protocol every backend must implement."""
+
+    def generate_json(
+        self,
+        messages_batch: Sequence[Sequence[dict[str, Any]]],
+        *,
+        max_new_tokens: int | None = None,
+        temperature: float | None = None,
+    ) -> list[Any]:
+        """Generate one JSON-decoded response per messages list."""
+
+
+@dataclass
+class StubVlmClient:
+    """Deterministic stub used in unit tests.
+
+    A test passes a callable that maps the *last user message text* (or, if
+    that is empty, the full message list) to a JSON-serializable response.
+    """
+
+    responder: Callable[[Sequence[dict[str, Any]]], Any]
+
+    def generate_json(
+        self,
+        messages_batch: Sequence[Sequence[dict[str, Any]]],
+        *,
+        max_new_tokens: int | None = None,
+        temperature: float | None = None,
+    ) -> list[Any]:
+        return [self.responder(list(messages)) for messages in messages_batch]
+
+
+def _strip_to_json(text: str) -> Any:
+    text = text.strip()
+    # Strip <think>...</think> blocks (Qwen3 Thinking style)
+    while "<think>" in text and "</think>" in text:
+        start = text.find("<think>")
+        end = text.find("</think>", start) + len("</think>")
+        text = (text[:start] + text[end:]).strip()
+    # Strip ```json ... ``` fences from chat-tuned backbones
+    if text.startswith("```"):
+        first = text.find("\n")
+        last = text.rfind("```")
+        if first != -1 and last != -1 and last > first:
+            text = text[first + 1 : last].strip()
+    try:
+        return json.loads(text)
+    except (ValueError, json.JSONDecodeError):
+        pass
+    # Fall back to extracting the first balanced {...} block.
+    obj_text = _extract_first_json_object(text)
+    if obj_text is None:
+        raise json.JSONDecodeError("No JSON object found", text, 0)
+    return json.loads(obj_text)
+
+
+def _extract_first_json_object(text: str) -> str | None:
+    """Return the first balanced ``{...}`` substring, ignoring braces in
+    string literals. Returns ``None`` if no balanced block is found."""
+    start = text.find("{")
+    if start < 0:
+        return None
+    depth = 0
+    in_string = False
+    escape = False
+    for i in range(start, len(text)):
+        ch = text[i]
+        if escape:
+            escape = False
+            continue
+        if ch == "\\":
+            escape = True
+            continue
+        # Note: ``escape`` is always False here — the ``if escape`` branch
+        # above already handled and reset it.
+        if ch == '"':
+            in_string = not in_string
+            continue
+        if in_string:
+            continue
+        if ch == "{":
+            depth += 1
+        elif ch == "}":
+            depth -= 1
+            if depth == 0:
+                return text[start : i + 1]
+    return None
+
+
+@dataclass
+class _GenericTextClient:
+    """Wraps any text-generation callable in JSON-mode + one-retry semantics."""
+
+    generate_text: Callable[[Sequence[Sequence[dict[str, Any]]], int, float], list[str]]
+    config: VlmConfig
+
+    def generate_json(
+        self,
+        messages_batch: Sequence[Sequence[dict[str, Any]]],
+        *,
+        max_new_tokens: int | None = None,
+        temperature: float | None = None,
+    ) -> list[Any]:
+        max_tok = max_new_tokens if max_new_tokens is not None else self.config.max_new_tokens
+        temp = temperature if temperature is not None else self.config.temperature
+        raw = self.generate_text(messages_batch, max_tok, temp)
+        out: list[Any] = []
+        for messages, text in zip(messages_batch, raw, strict=True):
+            try:
+                out.append(_strip_to_json(text))
+                continue
+            except (ValueError, json.JSONDecodeError):
+                pass
+            retry = list(messages) + [
+                {"role": "assistant", "content": text},
+                {
+                    "role": "user",
+                    "content": (
+                        "Your previous reply was not valid JSON. "
+                        "Reply with strictly valid JSON, no prose, no fences."
+                    ),
+                },
+            ]
+            retry_text = self.generate_text([retry], max_tok, temp)[0]
+            try:
+                out.append(_strip_to_json(retry_text))
+            except (ValueError, json.JSONDecodeError):
+                # After retry: log preview and return None instead of crashing
+                # the whole pipeline. Modules treat None as "skip".
+                preview = retry_text.strip().replace("\n", " ")[:200]
+                print(
+                    f"[vlm] WARNING: failed to parse JSON after retry; preview: {preview!r}",
+                    flush=True,
+                )
+                out.append(None)
+        return out
+
+
+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.
+
+    For ``stub``, construct :class:`StubVlmClient` directly with a responder
+    callable; it is rejected here to make accidental misuse obvious.
+    """
+    if config.backend == "openai":
+        return _make_openai_client(config)
+    if config.backend == "stub":
+        raise ValueError(
+            "Use StubVlmClient(...) directly for the stub backend; make_vlm_client builds real clients."
+        )
+    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 "
+            "vllm/vllm-openai image and talks to it over the OpenAI-compatible API."
+        )
+    raise ValueError(f"Unknown VLM backend: {config.backend!r}")
+
+
+def _make_openai_client(config: VlmConfig) -> VlmClient:
+    """Backend that talks to any OpenAI-compatible server.
+
+    Compatible with ``vllm serve``, ``transformers serve``,
+    ``ktransformers serve``, and hosted endpoints. By default the server
+    is expected to be already running. Set ``auto_serve=True`` to have
+    this client spawn one (default: ``transformers serve``), wait until
+    it's ready, and tear it down on process exit.
+
+    Image blocks ``{"type":"image", "image":<PIL.Image>}`` are
+    auto-converted to ``image_url`` data-URLs. Video blocks
+    ``{"type":"video", "video":[<PIL>...]}`` are forwarded as
+    multi-frame ``video_url`` items where supported.
+    """
+    try:
+        from openai import OpenAI  # type: ignore[import-not-found]
+    except ImportError as exc:
+        raise ImportError(
+            "openai package is required for backend='openai'. Install with `pip install openai`."
+        ) from exc
+
+    api_base = config.api_base
+    api_key = config.api_key
+    auto_serve = config.auto_serve
+    api_bases: list[str] = [api_base]
+
+    print(
+        f"[lerobot-annotate] backend=openai model={config.model_id} "
+        f"api_base={api_base} auto_serve={auto_serve}",
+        flush=True,
+    )
+    if auto_serve:
+        if config.parallel_servers > 1:
+            print(
+                f"[lerobot-annotate] spawning {config.parallel_servers} parallel servers",
+                flush=True,
+            )
+            api_bases = _spawn_parallel_inference_servers(config)
+        elif _server_is_up(api_base):
+            print(f"[lerobot-annotate] reusing server already up at {api_base}", flush=True)
+        else:
+            print("[lerobot-annotate] no server reachable; spawning one", flush=True)
+            api_base = _spawn_inference_server(config)
+            api_bases = [api_base]
+            print(f"[lerobot-annotate] server ready at {api_base}", flush=True)
+
+    clients = [OpenAI(base_url=base, api_key=api_key) for base in api_bases]
+    # round-robin counter for parallel mode
+    rr_counter = {"i": 0}
+
+    # ``mm_processor_kwargs`` is a vllm-specific extra; transformers serve
+    # rejects it with HTTP 422. Send it only when explicitly opted in via
+    # an env var (e.g. ``LEROBOT_OPENAI_SEND_MM_KWARGS=1`` for vllm).
+    send_mm_kwargs = os.environ.get("LEROBOT_OPENAI_SEND_MM_KWARGS", "").lower() in {"1", "true", "yes"}
+
+    rr_lock = threading.Lock()
+
+    def _one_call(messages: Sequence[dict[str, Any]], max_tok: int, temp: float) -> str:
+        api_messages, mm_kwargs = _to_openai_messages(messages)
+        kwargs: dict[str, Any] = {
+            "model": config.model_id,
+            "messages": api_messages,
+            "max_tokens": max_tok,
+            "temperature": temp,
+        }
+        extra_body: dict[str, Any] = {}
+        if send_mm_kwargs and mm_kwargs:
+            extra_body["mm_processor_kwargs"] = {**mm_kwargs, "do_sample_frames": True}
+        if config.chat_template_kwargs:
+            extra_body["chat_template_kwargs"] = config.chat_template_kwargs
+        if extra_body:
+            kwargs["extra_body"] = extra_body
+        with rr_lock:
+            chosen = clients[rr_counter["i"] % len(clients)]
+            rr_counter["i"] += 1
+        response = chosen.chat.completions.create(**kwargs)
+        return response.choices[0].message.content or ""
+
+    def _gen(batch: Sequence[Sequence[dict[str, Any]]], max_tok: int, temp: float) -> list[str]:
+        if len(batch) <= 1 or config.client_concurrency <= 1:
+            return [_one_call(messages, max_tok, temp) for messages in batch]
+        # Parallel fan-out — vllm batches these on the server side.
+        max_workers = min(config.client_concurrency, len(batch))
+        with ThreadPoolExecutor(max_workers=max_workers) as pool:
+            futures = [pool.submit(_one_call, messages, max_tok, temp) for messages in batch]
+            return [f.result() for f in futures]
+
+    return _GenericTextClient(_gen, config)
+
+
+def _bind_serve_port(cmd: str, port: int) -> str:
+    """Bind a serve command to ``port``: substitute a ``{port}`` placeholder
+    if present, else append ``--port`` when the command omits it (leaving an
+    explicit ``--port`` untouched). Shared by the single- and parallel-server
+    paths so a serve_command never reaches the server with a literal
+    ``{port}``."""
+    if "{port}" in cmd:
+        return cmd.replace("{port}", str(port))
+    if "--port" not in cmd:
+        return f"{cmd} --port {port}"
+    return cmd
+
+
+def _spawn_parallel_inference_servers(config: VlmConfig) -> list[str]:
+    """Spawn ``config.parallel_servers`` independent vllm replicas.
+
+    Each replica:
+    - is pinned to a single GPU via ``CUDA_VISIBLE_DEVICES``
+    - listens on ``serve_port + i``
+    - is shut down via the same atexit hook as the single-server path
+
+    Returns the list of ``api_base`` URLs the client should round-robin
+    across.
+    """
+    n = config.parallel_servers
+    api_bases: list[str] = []
+    procs: list[subprocess.Popen] = []
+    ready_events: list[threading.Event] = []
+    # Multiple readiness signals — uvicorn's own banner is suppressed at
+    # ``--uvicorn-log-level warning``, so we also accept vllm's own
+    # "Starting vLLM API server" line and the route-listing line. The
+    # HTTP probe below is the ultimate fallback.
+    ready_markers = (
+        "Uvicorn running",
+        "Application startup complete",
+        "Starting vLLM API server",
+        "Available routes are",
+    )
+    # Single lock for all server-stream threads so multibyte chars from
+    # different servers don't interleave and tear UTF-8 sequences.
+    print_lock = threading.Lock()
+
+    base_cmd = config.serve_command or (
+        f"vllm serve {shlex.quote(config.model_id)} "
+        f"--tensor-parallel-size 1 "
+        f"--max-model-len {config.max_model_len or 32768} "
+        f"--uvicorn-log-level warning"
+    )
+
+    num_gpus = config.num_gpus if config.num_gpus > 0 else n
+    for i in range(n):
+        port = config.serve_port + i
+        gpu = i % num_gpus
+        env = os.environ.copy()
+        env["CUDA_VISIBLE_DEVICES"] = str(gpu)
+        cmd = _bind_serve_port(base_cmd, port)
+        api_base = f"http://localhost:{port}/v1"
+        api_bases.append(api_base)
+        print(f"[server-{i}] launching on GPU {gpu} port {port}: {cmd}", flush=True)
+        proc = subprocess.Popen(
+            shlex.split(cmd),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            bufsize=1,
+            env=env,
+        )
+        procs.append(proc)
+        ready = threading.Event()
+        ready_events.append(ready)
+
+        def _stream(idx: int, p: subprocess.Popen, ev: threading.Event) -> None:
+            # Read whole lines and emit each line atomically under the
+            # shared print_lock so output from N servers stays readable.
+            assert p.stdout is not None
+            for line in iter(p.stdout.readline, ""):
+                with print_lock:
+                    sys.stdout.write(f"[server-{idx}] {line}")
+                    if not line.endswith(("\n", "\r")):
+                        sys.stdout.write("\n")
+                    sys.stdout.flush()
+                if any(m in line for m in ready_markers):
+                    ev.set()
+
+        threading.Thread(target=_stream, args=(i, proc, ready), daemon=True).start()
+
+        def _probe(idx: int, base: str, ev: threading.Event, p: subprocess.Popen) -> None:
+            while not ev.is_set() and p.poll() is None:
+                if _server_is_up(base):
+                    print(f"[server-{idx}] ready (http probe)", flush=True)
+                    ev.set()
+                    return
+                time.sleep(2)
+
+        threading.Thread(target=_probe, args=(i, api_base, ready, proc), daemon=True).start()
+
+    def _shutdown() -> None:
+        for i, p in enumerate(procs):
+            if p.poll() is None:
+                print(f"[server-{i}] stopping pid={p.pid}", flush=True)
+                p.send_signal(signal.SIGINT)
+        for p in procs:
+            try:
+                p.wait(timeout=15)
+            except subprocess.TimeoutExpired:
+                p.kill()
+                p.wait(timeout=5)
+
+    atexit.register(_shutdown)
+
+    deadline = time.monotonic() + config.serve_ready_timeout_s
+    while any(not ev.is_set() for ev in ready_events) and time.monotonic() < deadline:
+        for i, p in enumerate(procs):
+            if p.poll() is not None:
+                raise RuntimeError(
+                    f"[server-{i}] inference server exited unexpectedly with rc={p.returncode}"
+                )
+        time.sleep(2)
+    if any(not ev.is_set() for ev in ready_events):
+        raise RuntimeError(f"[server] not all replicas became ready within {config.serve_ready_timeout_s}s")
+    print(f"[lerobot-annotate] all {n} servers ready: {api_bases}", flush=True)
+    return api_bases
+
+
+def _server_is_up(api_base: str) -> bool:
+    """Return True if ``api_base/models`` answers 200 within 2 seconds."""
+    url = api_base.rstrip("/") + "/models"
+    # ``api_base`` is the user-configured local-server URL we just spawned
+    # or the user passed in via ``--vlm.api_base``; the bandit B310 warning
+    # is for arbitrary user-controlled URLs with file:/ schemes which
+    # cannot reach this code path.
+    try:
+        with urllib.request.urlopen(url, timeout=2) as resp:  # noqa: S310  # nosec B310
+            return resp.status == 200
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def _spawn_inference_server(config: VlmConfig) -> str:
+    """Spawn ``transformers serve`` (or ``serve_command``), wait until it
+    accepts ``/v1/models``, and register a shutdown hook.
+
+    Streams the server's stdout/stderr to the parent terminal in
+    real-time on a background thread so users can see model-load
+    progress and errors as they happen.
+
+    Returns the full ``api_base`` URL the OpenAI client should use.
+    """
+    cmd = config.serve_command
+    if not cmd:
+        cmd = (
+            f"transformers serve {shlex.quote(config.model_id)} "
+            f"--port {config.serve_port} --continuous-batching"
+        )
+    # Bind the single server to ``serve_port`` (what ``api_base`` below
+    # targets): substitute a literal ``{port}`` placeholder, else append
+    # ``--port``. Without this a serve_command carrying ``{port}`` would
+    # reach the server unsubstituted and fail to parse.
+    cmd = _bind_serve_port(cmd, config.serve_port)
+    api_base = f"http://localhost:{config.serve_port}/v1"
+    print(f"[server] launching: {cmd}", flush=True)
+    proc = subprocess.Popen(
+        shlex.split(cmd),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+    )
+
+    # Watch the server output for the uvicorn readiness banner. This is
+    # more reliable than polling /v1/models because transformers serve
+    # rescans its cache on every model-list request, which can exceed
+    # the urllib timeout and trigger an infinite probe loop.
+    ready_event = threading.Event()
+    # See _spawn_parallel_inference_servers for why we accept these.
+    ready_markers = (
+        "Uvicorn running",
+        "Application startup complete",
+        "Starting vLLM API server",
+        "Available routes are",
+    )
+
+    def _probe() -> None:
+        while not ready_event.is_set() and proc.poll() is None:
+            if _server_is_up(api_base):
+                print("[server] ready (http probe)", flush=True)
+                ready_event.set()
+                return
+            time.sleep(2)
+
+    threading.Thread(target=_probe, daemon=True).start()
+
+    def _stream_output() -> None:
+        # Read raw chunks instead of iterating lines so tqdm progress
+        # bars (which overwrite using \r) flush in real time.
+        assert proc.stdout is not None
+        buf = ""
+        prefix_started = False
+        while True:
+            ch = proc.stdout.read(1)
+            if ch == "":
+                # process exited; flush any tail
+                if buf:
+                    sys.stdout.write(buf)
+                    sys.stdout.flush()
+                return
+            if not prefix_started:
+                sys.stdout.write("[server] ")
+                prefix_started = True
+            sys.stdout.write(ch)
+            sys.stdout.flush()
+            buf += ch
+            if ch in ("\n", "\r"):
+                if any(marker in buf for marker in ready_markers):
+                    ready_event.set()
+                buf = ""
+                prefix_started = False
+
+    threading.Thread(target=_stream_output, daemon=True).start()
+
+    def _shutdown() -> None:
+        if proc.poll() is None:
+            print(f"[server] stopping pid={proc.pid}", flush=True)
+            proc.send_signal(signal.SIGINT)
+            try:
+                proc.wait(timeout=15)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+                proc.wait(timeout=5)
+
+    atexit.register(_shutdown)
+
+    deadline = time.monotonic() + config.serve_ready_timeout_s
+    while time.monotonic() < deadline:
+        if proc.poll() is not None:
+            raise RuntimeError(
+                f"[server] inference server exited unexpectedly with rc={proc.returncode}. "
+                f"See [server] log lines above for the cause."
+            )
+        if ready_event.wait(timeout=2):
+            return api_base
+    proc.terminate()
+    raise RuntimeError(f"[server] did not become ready within {config.serve_ready_timeout_s}s")
+
+
+def _to_openai_messages(
+    messages: Sequence[dict[str, Any]],
+) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+    """Convert internal messages to OpenAI chat format.
+
+    Returns ``(api_messages, mm_kwargs)``. Multimodal-processor kwargs
+    (``fps`` from ``video_url`` blocks) are extracted out so the caller
+    can pass them via ``extra_body.mm_processor_kwargs`` rather than
+    inside the content blocks (which transformers serve rejects).
+
+    File-URL video blocks are inlined as base64 data URLs.
+    """
+    out_messages: list[dict[str, Any]] = []
+    mm_kwargs: dict[str, Any] = {}
+    for message in messages:
+        content = message.get("content")
+        if not isinstance(content, list):
+            out_messages.append({"role": message["role"], "content": content})
+            continue
+        out_blocks: list[dict[str, Any]] = []
+        for block in content:
+            block_type = block.get("type") if isinstance(block, dict) else None
+            if block_type == "text":
+                out_blocks.append({"type": "text", "text": block.get("text", "")})
+            elif block_type == "image":
+                out_blocks.append(
+                    {"type": "image_url", "image_url": {"url": _pil_to_data_url(block["image"])}}
+                )
+            elif block_type == "video":
+                frames = block.get("video", [])
+                for img in frames:
+                    out_blocks.append({"type": "image_url", "image_url": {"url": _pil_to_data_url(img)}})
+            elif block_type == "video_url":
+                video_url = dict(block["video_url"])
+                url = video_url.get("url", "")
+                if url.startswith("file://"):
+                    video_url["url"] = _file_to_data_url(url[len("file://") :])
+                out_blocks.append({"type": "video_url", "video_url": video_url})
+                fps = block.get("fps")
+                if fps is not None:
+                    mm_kwargs["fps"] = fps
+            else:
+                out_blocks.append(block)
+        out_messages.append({"role": message["role"], "content": out_blocks})
+    return out_messages, mm_kwargs
+
+
+def _file_to_data_url(path: str) -> str:
+    """Read a local video file and return a base64 ``data:video/mp4`` URL."""
+    with open(path, "rb") as f:
+        b64 = base64.b64encode(f.read()).decode("ascii")
+    return f"data:video/mp4;base64,{b64}"
+
+
+def _pil_to_data_url(image: Any) -> str:
+    """Encode a PIL.Image as a base64 data URL."""
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+    return f"data:image/png;base64,{b64}"

src/lerobot/annotations/steerable_pipeline/writer.py

@@ -0,0 +1,341 @@
+#!/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.
+"""Final parquet rewrite.
+
+For every episode the writer:
+
+1. reads the staged module outputs,
+2. partitions them into a persistent slice (PERSISTENT_STYLES) and an event
+   slice (EVENT_ONLY_STYLES + style=None tool-call atoms),
+3. sorts each slice deterministically,
+4. broadcasts the persistent slice across every frame in the episode,
+5. for each frame, materializes the sublist of event rows whose timestamp
+   exactly equals that frame's timestamp,
+6. drops the legacy ``subtask_index`` column,
+7. writes the parquet shard back in place.
+
+The writer does NOT add a dataset-level ``tools`` column. Tool *calls* are
+emitted per-row via the existing ``tool_calls`` field on the v3.1 row
+struct for every speech atom. The tool *schema* (the description
+of the ``say`` function and its parameters) is a fixed code constant —
+``SAY_TOOL_SCHEMA`` below — and downstream chat-template consumers import
+it directly rather than reading a redundant per-row column.
+
+Invariants enforced here (and re-checked by the validator):
+
+- per-episode persistent slice is byte-identical across every frame;
+- ``language_events`` rows on a frame all have ``timestamp == frame_ts``
+  (timestamps come straight from the source parquet — never recomputed);
+- every row passes ``column_for_style(style)``.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from lerobot.datasets.io_utils import write_table_one_row_group_per_episode
+from lerobot.datasets.language import (
+    EVENT_ONLY_STYLES,
+    LANGUAGE_EVENTS,
+    LANGUAGE_PERSISTENT,
+    PERSISTENT_STYLES,
+    column_for_style,
+    validate_camera_field,
+)
+
+from .reader import EpisodeRecord
+from .staging import EpisodeStaging
+
+logger = logging.getLogger(__name__)
+
+
+# Tool schema constants live in lerobot.datasets.language — single
+# source of truth. Re-exported here so existing imports
+# (``from lerobot.annotations.steerable_pipeline.writer import SAY_TOOL_SCHEMA``)
+# keep working.
+from lerobot.datasets.language import DEFAULT_TOOLS, SAY_TOOL_SCHEMA  # noqa: F401, E402
+
+
+def _row_persistent_sort_key(row: dict[str, Any]) -> tuple:
+    return (float(row["timestamp"]), row.get("style") or "", row.get("role") or "")
+
+
+def _row_event_sort_key(row: dict[str, Any]) -> tuple:
+    # events are bucketed per-frame, but within a frame we still want determinism
+    return (
+        row.get("style") or "",
+        row.get("role") or "",
+        row.get("camera") or "",
+    )
+
+
+def _normalize_row(row: dict[str, Any], style: str | None, *, with_timestamp: bool) -> dict[str, Any]:
+    """Coerce a staged row into the language-column struct shape.
+
+    Key order matches ``PERSISTENT_ROW_FIELDS`` / ``EVENT_ROW_FIELDS`` — the
+    writer infers the parquet struct schema from insertion order, so
+    ``timestamp`` (persistent rows only) sits between ``style`` and ``camera``.
+    """
+    camera = row.get("camera")
+    validate_camera_field(style, camera)
+    out: dict[str, Any] = {
+        "role": str(row["role"]),
+        "content": None if row.get("content") is None else str(row["content"]),
+        "style": style,
+    }
+    if with_timestamp:
+        out["timestamp"] = float(row["timestamp"])
+    out["camera"] = None if camera is None else str(camera)
+    out["tool_calls"] = _normalize_tool_calls(row.get("tool_calls"))
+    return out
+
+
+def _normalize_persistent_row(row: dict[str, Any]) -> dict[str, Any]:
+    """Coerce a staged row into the persistent column's struct shape."""
+    style = row.get("style")
+    if style not in PERSISTENT_STYLES:
+        raise ValueError(
+            f"persistent slice contains row with non-persistent style {style!r}; "
+            "row would be misrouted under column_for_style()"
+        )
+    if "timestamp" not in row:
+        raise ValueError(f"persistent row missing timestamp: {row!r}")
+    if "role" not in row:
+        # Friendly error from the writer instead of a raw KeyError below;
+        # the validator doesn't check ``role`` yet.
+        raise ValueError(f"persistent row missing role: {row!r}")
+    return _normalize_row(row, style, with_timestamp=True)
+
+
+def _normalize_event_row(row: dict[str, Any]) -> dict[str, Any]:
+    """Coerce a staged row into the event column's struct shape (no timestamp)."""
+    style = row.get("style")
+    if style is not None and style not in EVENT_ONLY_STYLES:
+        raise ValueError(
+            f"event slice contains row with style {style!r}; expected None or one of {EVENT_ONLY_STYLES}"
+        )
+    if column_for_style(style) != LANGUAGE_EVENTS:
+        raise ValueError(f"event row with style {style!r} would not route to language_events")
+    if "role" not in row:
+        raise ValueError(f"event row missing role: {row!r}")
+    return _normalize_row(row, style, with_timestamp=False)
+
+
+def _normalize_tool_calls(value: Any) -> list[Any] | None:
+    if value is None:
+        return None
+    if not isinstance(value, list):
+        raise ValueError(f"tool_calls must be a list or None, got {type(value).__name__}")
+    return list(value)
+
+
+def _validate_atom_invariants(row: dict[str, Any]) -> None:
+    """At-least-one of content/tool_calls; style=None implies tool_calls."""
+    has_content = row.get("content") is not None
+    has_tools = row.get("tool_calls") is not None
+    if not (has_content or has_tools):
+        raise ValueError(f"row has neither content nor tool_calls: {row!r}")
+    if row.get("style") is None and not has_tools:
+        raise ValueError(f"style=None requires tool_calls: {row!r}")
+
+
+def _validate_speech_atom(row: dict[str, Any]) -> None:
+    """Speech atoms: role=assistant, style=None, content=None, say tool call."""
+    if row.get("style") is not None:
+        return  # not a speech atom
+    if row.get("role") != "assistant":
+        raise ValueError(f"speech atom must have role=assistant: {row!r}")
+    if row.get("content") is not None:
+        raise ValueError(f"speech atom must have content=null: {row!r}")
+    tool_calls = row.get("tool_calls")
+    if not tool_calls or not isinstance(tool_calls, list):
+        raise ValueError(f"speech atom must have non-empty tool_calls list: {row!r}")
+    first = tool_calls[0]
+    if not isinstance(first, dict):
+        raise ValueError(f"speech atom tool_calls[0] must be a dict: {row!r}")
+    if first.get("type") != "function":
+        raise ValueError(f"speech atom tool_calls[0].type must be 'function': {row!r}")
+    fn = first.get("function") or {}
+    if fn.get("name") != "say":
+        raise ValueError(f"speech atom tool_calls[0].function.name must be 'say': {row!r}")
+    args = fn.get("arguments") or {}
+    if not isinstance(args, dict) or "text" not in args or not isinstance(args["text"], str):
+        raise ValueError(f"speech atom must carry 'text' string in arguments: {row!r}")
+
+
+@dataclass
+class LanguageColumnsWriter:
+    """Rewrite ``data/chunk-*/file-*.parquet`` with the two language columns."""
+
+    drop_existing_subtask_index: bool = True
+
+    def write_all(
+        self,
+        records: Sequence[EpisodeRecord],
+        staging_dir: Path,
+        root: Path,
+    ) -> list[Path]:
+        episodes_by_path: dict[Path, list[EpisodeRecord]] = defaultdict(list)
+        for record in records:
+            episodes_by_path[record.data_path].append(record)
+
+        written: list[Path] = []
+        for path, eps in episodes_by_path.items():
+            self._rewrite_one(path, eps, staging_dir, root)
+            written.append(path)
+        return written
+
+    def _rewrite_one(
+        self,
+        path: Path,
+        episodes: Sequence[EpisodeRecord],
+        staging_dir: Path,
+        root: Path,
+    ) -> None:
+        table = pq.read_table(path)
+        n_rows = table.num_rows
+
+        # Ensure we cover every episode in the file. Episodes that don't have
+        # staging artifacts are passed through with empty annotation lists —
+        # this keeps the writer idempotent and safe for partial reruns.
+        staged_per_ep: dict[int, dict[str, list[dict[str, Any]]]] = {}
+        for record in episodes:
+            staging = EpisodeStaging(staging_dir, record.episode_index)
+            staged_per_ep[record.episode_index] = staging.read_all()
+
+        persistent_by_ep: dict[int, list[dict[str, Any]]] = {}
+        events_by_ep_ts: dict[int, dict[float, list[dict[str, Any]]]] = {}
+
+        for ep_index, ep_staged in staged_per_ep.items():
+            persistent_rows: list[dict[str, Any]] = []
+            event_rows: list[dict[str, Any]] = []  # carry timestamp until bucketed
+            for _module_name, rows in ep_staged.items():
+                for row in rows:
+                    style = row.get("style")
+                    if column_for_style(style) == LANGUAGE_PERSISTENT:
+                        persistent_rows.append(row)
+                    else:
+                        event_rows.append(row)
+
+            persistent_rows.sort(key=_row_persistent_sort_key)
+            normalized_persistent = []
+            for r in persistent_rows:
+                _validate_atom_invariants(r)
+                _validate_speech_atom(r)
+                normalized_persistent.append(_normalize_persistent_row(r))
+            persistent_by_ep[ep_index] = normalized_persistent
+
+            buckets: dict[float, list[dict[str, Any]]] = defaultdict(list)
+            for r in event_rows:
+                _validate_atom_invariants(r)
+                _validate_speech_atom(r)
+                ts = float(r["timestamp"])
+                buckets[ts].append(_normalize_event_row(r))
+            for ts in list(buckets.keys()):
+                buckets[ts].sort(key=_row_event_sort_key)
+            events_by_ep_ts[ep_index] = buckets
+
+        episode_col = (
+            table.column("episode_index").to_pylist() if "episode_index" in table.column_names else None
+        )
+        ts_col = table.column("timestamp").to_pylist() if "timestamp" in table.column_names else None
+        if episode_col is None or ts_col is None:
+            raise ValueError(f"{path} is missing 'episode_index' or 'timestamp' — required by the writer.")
+
+        per_row_persistent: list[list[dict[str, Any]]] = []
+        per_row_events: list[list[dict[str, Any]]] = []
+        for i in range(n_rows):
+            ep = episode_col[i]
+            ts = float(ts_col[i])
+            per_row_persistent.append(persistent_by_ep.get(ep, []))
+            buckets = events_by_ep_ts.get(ep, {})
+            per_row_events.append(buckets.get(ts, []))
+
+        new_table = self._materialize_table(
+            table, per_row_persistent, per_row_events, drop_old=self.drop_existing_subtask_index
+        )
+        # Re-emit one row group per episode (a bulk pq.write_table would collapse
+        # them into one). Write to a sibling tmp path and atomically rename so a
+        # crash mid-write can't leave a half-written shard.
+        tmp_path = path.with_suffix(path.suffix + ".tmp")
+        write_table_one_row_group_per_episode(new_table, tmp_path)
+        tmp_path.replace(path)
+
+    def _materialize_table(
+        self,
+        table: pa.Table,
+        persistent: list[list[dict[str, Any]]],
+        events: list[list[dict[str, Any]]],
+        *,
+        drop_old: bool,
+    ) -> pa.Table:
+        cols = []
+        names = []
+        for name in table.column_names:
+            if drop_old and name == "subtask_index":
+                continue
+            if name in (LANGUAGE_PERSISTENT, LANGUAGE_EVENTS):
+                continue  # we'll re-add canonical versions
+            # Strip any legacy ``tools`` column previously emitted by older
+            # writers — the schema no longer uses it (constant lives in
+            # SAY_TOOL_SCHEMA / DEFAULT_TOOLS).
+            if name == "tools":
+                continue
+            cols.append(table.column(name))
+            names.append(name)
+
+        # We let pyarrow infer struct/list schema rather than passing the
+        # canonical type from `lerobot.datasets.language` directly: that type
+        # uses `pa.json_()` for the `tool_calls` element type, which
+        # `pa.array(..., type=...)` cannot materialize from Python lists on
+        # current pyarrow versions. The inferred schema round-trips through
+        # parquet and `LeRobotDataset` correctly — `tests/datasets/test_language.py`
+        # exercises the same flow.
+        persistent_arr = pa.array(persistent)
+        events_arr = pa.array(events)
+
+        cols.extend([persistent_arr, events_arr])
+        names.extend([LANGUAGE_PERSISTENT, LANGUAGE_EVENTS])
+
+        return pa.Table.from_arrays(cols, names=names)
+
+
+def speech_atom(timestamp: float, text: str) -> dict[str, Any]:
+    """Build a canonical speech tool-call atom for the events column."""
+    return {
+        "role": "assistant",
+        "content": None,
+        "style": None,
+        "timestamp": float(timestamp),
+        "camera": None,
+        "tool_calls": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "say",
+                    "arguments": {"text": text},
+                },
+            }
+        ],
+    }

src/lerobot/common/control_utils.py

@@ -18,6 +18,7 @@ from __future__ import annotations
 # Utilities
 ########################################################################################
 import logging
+import time
 import traceback
 from contextlib import nullcontext
 from copy import copy
@@ -243,3 +244,72 @@ def sanity_check_dataset_robot_compatibility(
         raise ValueError(
             "Dataset metadata compatibility check failed with mismatches:\n" + "\n".join(mismatches)
         )
+
+
+########################################################################################
+# Teleoperator smooth handover helpers
+# NOTE(Maxime): These functions use minimal type hints to maintain compatibility with utils
+# being a root module.
+########################################################################################
+
+
+def teleop_supports_feedback(teleop) -> bool:
+    """Return True when the teleop can receive position feedback (is actuated).
+
+    Actuated teleops (e.g. SO-101, OpenArmMini) have non-empty ``feedback_features``
+    and expose ``enable_torque`` / ``disable_torque`` motor-control methods.
+
+    TODO(Maxime): See if it is possible to unify this interface across teleops instead of duck-typing.
+    """
+    return (
+        bool(teleop.feedback_features)
+        and hasattr(teleop, "disable_torque")
+        and hasattr(teleop, "enable_torque")
+    )
+
+
+def teleop_smooth_move_to(teleop, target_pos: dict, duration_s: float = 2.0, fps: int = 30) -> None:
+    """Smoothly move an actuated teleop to ``target_pos`` via linear interpolation.
+
+    Requires the teleoperator to support feedback (i.e. have non-empty
+    ``feedback_features`` and implement ``disable_torque`` / ``enable_torque``).
+
+    ``target_pos`` is expected to be in the teleop's action/feedback key space.
+    For homogeneous setups (e.g. SO-101 leader + SO-101 follower) this matches
+    the robot action key space directly.
+
+    TODO(Maxime): This blocks up to ``duration_s`` seconds; during this time the
+    follower robot does not receive new actions, which could be an issue on LeKiwi.
+    """
+    teleop.enable_torque()
+    current = teleop.get_action()
+    steps = max(int(duration_s * fps), 1)
+
+    for step in range(steps + 1):
+        t = step / steps
+        interp = {
+            k: current[k] * (1 - t) + target_pos[k] * t if k in target_pos else current[k] for k in current
+        }
+        teleop.send_feedback(interp)
+        time.sleep(1 / fps)
+
+
+def follower_smooth_move_to(
+    robot, current: dict, target: dict, duration_s: float = 1.0, fps: int = 30
+) -> None:
+    """Smoothly move the follower robot from ``current`` to ``target`` action.
+
+    Used when the teleop is non-actuated: instead of driving the leader arm to
+    the follower, the follower is brought to the teleop's current pose so the
+    robot meets the operator's hand rather than jumping to it on the first frame.
+
+    Both ``current`` and ``target`` must be in the robot action key space
+    (i.e. the output of ``robot_action_processor``).
+    """
+    steps = max(int(duration_s * fps), 1)
+
+    for step in range(steps + 1):
+        t = step / steps
+        interp = {k: current[k] * (1 - t) + target[k] * t if k in target else current[k] for k in current}
+        robot.send_action(interp)
+        time.sleep(1 / fps)

src/lerobot/common/train_utils.py

@@ -49,8 +49,19 @@ def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Pa
     return output_dir / CHECKPOINTS_DIR / step_identifier
 
 
-def save_training_step(step: int, save_dir: Path) -> None:
-    write_json({"step": step}, save_dir / TRAINING_STEP)
+def save_training_step(
+    step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None
+) -> None:
+    state: dict = {"step": step}
+    # num_processes and batch_size are recorded so a resumed run can detect a changed world size or
+    # batch size: the sampler's resume offset is computed from the (num_processes, batch_size) that
+    # produced `step`, since both scale how many sampler positions a step consumes (see
+    # compute_sampler_state).
+    if num_processes is not None:
+        state["num_processes"] = num_processes
+    if batch_size is not None:
+        state["batch_size"] = batch_size
+    write_json(state, save_dir / TRAINING_STEP)
 
 
 def load_training_step(save_dir: Path) -> int:
@@ -58,6 +69,16 @@ def load_training_step(save_dir: Path) -> int:
     return training_step["step"]
 
 
+def load_training_num_processes(checkpoint_dir: Path) -> int | None:
+    """World size recorded at checkpoint time, or None for checkpoints written before it was stored."""
+    return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("num_processes")
+
+
+def load_training_batch_size(checkpoint_dir: Path) -> int | None:
+    """Per-process batch size recorded at checkpoint time, or None for older checkpoints."""
+    return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("batch_size")
+
+
 def update_last_checkpoint(checkpoint_dir: Path) -> Path:
     last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK
     if last_checkpoint_dir.is_symlink():
@@ -75,6 +96,8 @@ def save_checkpoint(
     scheduler: LRScheduler | None = None,
     preprocessor: PolicyProcessorPipeline | None = None,
     postprocessor: PolicyProcessorPipeline | None = None,
+    num_processes: int | None = None,
+    batch_size: int | None = None,
 ) -> None:
     """This function creates the following directory structure:
 
@@ -100,6 +123,10 @@ def save_checkpoint(
         scheduler (LRScheduler | None, optional): The scheduler to save the state from. Defaults to None.
         preprocessor: The preprocessor/pipeline to save. Defaults to None.
         postprocessor: The postprocessor/pipeline to save. Defaults to None.
+        num_processes (int | None, optional): Distributed world size to record for sample-exact
+            resume. Defaults to None (not recorded).
+        batch_size (int | None, optional): Per-process batch size to record for sample-exact
+            resume. Defaults to None (not recorded).
     """
     pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
     policy.save_pretrained(pretrained_dir)
@@ -112,7 +139,9 @@ def save_checkpoint(
         preprocessor.save_pretrained(pretrained_dir)
     if postprocessor is not None:
         postprocessor.save_pretrained(pretrained_dir)
-    save_training_state(checkpoint_dir, step, optimizer, scheduler)
+    save_training_state(
+        checkpoint_dir, step, optimizer, scheduler, num_processes=num_processes, batch_size=batch_size
+    )
 
 
 def save_training_state(
@@ -120,6 +149,8 @@ def save_training_state(
     train_step: int,
     optimizer: Optimizer | None = None,
     scheduler: LRScheduler | None = None,
+    num_processes: int | None = None,
+    batch_size: int | None = None,
 ) -> None:
     """
     Saves the training step, optimizer state, scheduler state, and rng state.
@@ -131,10 +162,12 @@ def save_training_state(
             Defaults to None.
         scheduler (LRScheduler | None, optional): The scheduler from which to save the state_dict.
             Defaults to None.
+        num_processes (int | None, optional): Distributed world size to record. Defaults to None.
+        batch_size (int | None, optional): Per-process batch size to record. Defaults to None.
     """
     save_dir = checkpoint_dir / TRAINING_STATE_DIR
     save_dir.mkdir(parents=True, exist_ok=True)
-    save_training_step(train_step, save_dir)
+    save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size)
     save_rng_state(save_dir)
     if optimizer is not None:
         save_optimizer_state(optimizer, save_dir)

src/lerobot/configs/policies.py

@@ -79,6 +79,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
     # Either the repo ID of a model hosted on the Hub or a path to a directory containing weights
     # saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch.
     pretrained_path: Path | None = None
+    # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained model version.
+    pretrained_revision: str | None = None
 
     def __post_init__(self) -> None:
         if not self.device or not is_torch_device_available(self.device):

src/lerobot/configs/rewards.py

@@ -56,6 +56,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     device: str | None = None
 
     pretrained_path: str | None = None
+    # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained reward model version.
+    pretrained_revision: str | None = None
 
     push_to_hub: bool = False
     repo_id: str | None = None

src/lerobot/datasets/__init__.py

@@ -50,7 +50,7 @@ from .lerobot_dataset import LeRobotDataset
 from .multi_dataset import MultiLeRobotDataset
 from .pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
 from .pyav_utils import check_video_encoder_parameters_pyav, detect_available_encoders_pyav
-from .sampler import EpisodeAwareSampler
+from .sampler import EpisodeAwareSampler, compute_sampler_state
 from .streaming_dataset import StreamingLeRobotDataset
 from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card
 from .video_utils import VideoEncodingManager
@@ -82,6 +82,7 @@ __all__ = [
     "aggregate_stats",
     "convert_image_to_video_dataset",
     "create_initial_features",
+    "compute_sampler_state",
     "create_lerobot_dataset_card",
     "column_for_style",
     "delete_episodes",

src/lerobot/datasets/aggregate.py

@@ -32,6 +32,7 @@ from .feature_utils import features_equal_for_merge, get_hf_features_from_featur
 from .io_utils import (
     get_file_size_in_mb,
     get_parquet_file_size_in_mb,
+    to_parquet_one_row_group_per_episode,
     to_parquet_with_hf_images,
     write_info,
     write_stats,
@@ -286,6 +287,8 @@ def aggregate_datasets(
     data_files_size_in_mb: int | None = None,
     video_files_size_in_mb: int | None = None,
     chunk_size: int | None = None,
+    concatenate_videos: bool = True,
+    concatenate_data: bool = True,
 ):
     """Aggregates multiple LeRobot datasets into a single unified dataset.
 
@@ -303,6 +306,8 @@ def aggregate_datasets(
         data_files_size_in_mb: Maximum size for data files in MB (defaults to DEFAULT_DATA_FILE_SIZE_IN_MB)
         video_files_size_in_mb: Maximum size for video files in MB (defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB)
         chunk_size: Maximum number of files per chunk (defaults to DEFAULT_CHUNK_SIZE)
+        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")
 
@@ -351,8 +356,12 @@ def aggregate_datasets(
     dst_meta.episodes = {}
 
     for src_meta in tqdm.tqdm(all_metadata, desc="Copy data and videos"):
-        videos_idx = aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size)
-        data_idx = aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size)
+        videos_idx = aggregate_videos(
+            src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos
+        )
+        data_idx = aggregate_data(
+            src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data
+        )
 
         meta_idx = aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx)
 
@@ -367,7 +376,9 @@ def aggregate_datasets(
     logging.info("Aggregation complete.")
 
 
-def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size):
+def aggregate_videos(
+    src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
+):
     """Aggregates video chunks from a source dataset into the destination dataset.
 
     Handles video file concatenation and rotation based on file size limits.
@@ -379,6 +390,7 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
         videos_idx: Dictionary tracking video chunk and file indices.
         video_files_size_in_mb: Maximum size for video files in MB (defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB)
         chunk_size: Maximum number of files per chunk (defaults to DEFAULT_CHUNK_SIZE)
+        concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
     Returns:
         dict: Updated videos_idx with current chunk and file indices.
     """
@@ -439,7 +451,7 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
             src_size = get_file_size_in_mb(src_path)
             dst_size = get_file_size_in_mb(dst_path)
 
-            if dst_size + src_size >= video_files_size_in_mb:
+            if not concatenate_videos or dst_size + src_size >= video_files_size_in_mb:
                 # Rotate to a new file - offset is 0
                 chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, chunk_size)
                 dst_key = (chunk_idx, file_idx)
@@ -477,7 +489,7 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
     return videos_idx
 
 
-def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size):
+def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
     """Aggregates data chunks from a source dataset into the destination dataset.
 
     Reads source data files, updates indices to match the aggregated dataset,
@@ -493,6 +505,7 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
         data_idx: Dictionary tracking data chunk and file indices.
         data_files_size_in_mb: Maximum size for data files in MB.
         chunk_size: Maximum number of files per chunk.
+        concatenate_data: When False, keep one parquet per source file instead of packing into shards.
 
     Returns:
         dict: Updated data_idx with current chunk and file indices.
@@ -538,6 +551,8 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
             contains_images=contains_images,
             aggr_root=dst_meta.root,
             hf_features=hf_features,
+            concatenate=concatenate_data,
+            one_row_group_per_episode=True,
         )
 
         # Record the mapping from source to actual destination
@@ -614,6 +629,8 @@ def append_or_create_parquet_file(
     contains_images: bool = False,
     aggr_root: Path = None,
     hf_features: datasets.Features | None = None,
+    concatenate: bool = True,
+    one_row_group_per_episode: bool = False,
 ) -> tuple[dict[str, int], tuple[int, int]]:
     """Appends data to an existing parquet file or creates a new one based on size constraints.
 
@@ -630,6 +647,9 @@ def append_or_create_parquet_file(
         contains_images: Whether the data contains images requiring special handling.
         aggr_root: Root path for the aggregated dataset.
         hf_features: Optional HuggingFace Features schema for proper image typing.
+        concatenate: When False, always rotate to a new file instead of appending to the current one.
+        one_row_group_per_episode: True for DATA parquet (emit one row group per episode); False for
+            the episodes-metadata parquet (already one row per episode).
 
     Returns:
         tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
@@ -642,6 +662,8 @@ def append_or_create_parquet_file(
         dst_path.parent.mkdir(parents=True, exist_ok=True)
         if contains_images:
             to_parquet_with_hf_images(df, dst_path, features=hf_features)
+        elif one_row_group_per_episode:
+            to_parquet_one_row_group_per_episode(df, dst_path)
         else:
             df.to_parquet(dst_path)
         return idx, (dst_chunk, dst_file)
@@ -649,7 +671,7 @@ def append_or_create_parquet_file(
     src_size = get_parquet_file_size_in_mb(src_path)
     dst_size = get_parquet_file_size_in_mb(dst_path)
 
-    if dst_size + src_size >= max_mb:
+    if not concatenate or dst_size + src_size >= max_mb:
         idx["chunk"], idx["file"] = update_chunk_file_indices(idx["chunk"], idx["file"], chunk_size)
         dst_chunk, dst_file = idx["chunk"], idx["file"]
         new_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
@@ -668,6 +690,8 @@ def append_or_create_parquet_file(
 
     if contains_images:
         to_parquet_with_hf_images(final_df, target_path, features=hf_features)
+    elif one_row_group_per_episode:
+        to_parquet_one_row_group_per_episode(final_df, target_path)
     else:
         final_df.to_parquet(target_path)
 

src/lerobot/datasets/compute_stats.py

@@ -59,6 +59,8 @@ class RunningQuantileStats:
             batch: An array where all dimensions except the last are batch dimensions.
         """
         batch = batch.reshape(-1, batch.shape[-1])
+        # Promote integer and low-precision inputs before computing squared statistics.
+        batch = batch.astype(np.result_type(batch.dtype, np.float32), copy=False)
         num_elements, vector_length = batch.shape
 
         if self._count == 0:

src/lerobot/datasets/dataset_tools.py

@@ -261,6 +261,8 @@ def merge_datasets(
     datasets: list[LeRobotDataset],
     output_repo_id: str,
     output_dir: str | Path | None = None,
+    concatenate_videos: bool = True,
+    concatenate_data: bool = True,
 ) -> LeRobotDataset:
     """Merge multiple LeRobotDatasets into a single dataset.
 
@@ -270,6 +272,8 @@ def merge_datasets(
         datasets: List of LeRobotDatasets to merge.
         output_repo_id: Merged dataset identifier.
         output_dir: Root directory where the merged dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/output_repo_id.
+        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.
     """
     if not datasets:
         raise ValueError("No datasets to merge")
@@ -284,6 +288,8 @@ def merge_datasets(
         aggr_repo_id=output_repo_id,
         roots=roots,
         aggr_root=output_dir,
+        concatenate_videos=concatenate_videos,
+        concatenate_data=concatenate_data,
     )
 
     merged_dataset = LeRobotDataset(

src/lerobot/datasets/io_utils.py

@@ -20,6 +20,7 @@ import datasets
 import numpy as np
 import pandas
 import pandas as pd
+import pyarrow as pa
 import pyarrow.dataset as pa_ds
 import pyarrow.parquet as pq
 import torch
@@ -270,21 +271,49 @@ def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[to
     return items_dict
 
 
+def write_table_one_row_group_per_episode(table: pa.Table, path: Path) -> None:
+    """Write ``table`` with one parquet row group per episode (in episode order).
+
+    Keeps shards random-access friendly (``read_row_group(i)`` fetches episode i),
+    mirroring the recording writer. ``table`` must carry a contiguous
+    ``episode_index`` column.
+    """
+    episode_index = table.column("episode_index").to_numpy(zero_copy_only=False)
+    starts = np.concatenate(([0], np.nonzero(np.diff(episode_index))[0] + 1))
+    writer = pq.ParquetWriter(str(path), table.schema, compression="snappy", use_dictionary=True)
+    try:
+        for start, stop in zip(starts, np.append(starts[1:], len(episode_index)), strict=True):
+            writer.write_table(table.slice(start, stop - start))  # one episode -> one row group
+    finally:
+        writer.close()
+
+
 def to_parquet_with_hf_images(
     df: pandas.DataFrame, path: Path, features: datasets.Features | None = None
 ) -> None:
-    """This function correctly writes to parquet a panda DataFrame that contains images encoded by HF dataset.
-    This way, it can be loaded by HF dataset and correctly formatted images are returned.
+    """Write a DataFrame with HF-encoded images to parquet, one row group per episode.
 
-    Args:
-        df: DataFrame to write to parquet.
-        path: Path to write the parquet file.
-        features: Optional HuggingFace Features schema. If provided, ensures image columns
-                  are properly typed as Image() in the parquet schema.
+    Images are embedded into the arrow table first (``ParquetWriter.write_table``
+    does not embed external image files like ``Dataset.to_parquet`` does).
+    ``features`` types image columns as ``Image()`` in the parquet schema.
     """
-    # TODO(qlhoest): replace this weird synthax by `df.to_parquet(path)` only
     ds = datasets.Dataset.from_dict(df.to_dict(orient="list"), features=features)
-    ds.to_parquet(path)
+    ds = embed_images(ds)
+    table = ds.with_format("arrow")[:]
+    if "episode_index" in table.column_names:
+        write_table_one_row_group_per_episode(table, path)
+    else:
+        # No episode boundaries to align row groups to — keep a single write.
+        pq.write_table(table, str(path))
+
+
+def to_parquet_one_row_group_per_episode(df: pandas.DataFrame, path: Path) -> None:
+    """Write a (non-image) DataFrame to parquet with one row group per episode."""
+    table = pa.Table.from_pandas(df, preserve_index=False)
+    if "episode_index" in table.column_names:
+        write_table_one_row_group_per_episode(table, path)
+    else:
+        pq.write_table(table, str(path))
 
 
 def item_to_torch(item: dict) -> dict:

src/lerobot/datasets/pipeline_features.py

@@ -70,19 +70,21 @@ def aggregate_pipeline_dataset_features(
     initial_features: dict[PipelineFeatureType, dict[str, Any]],
     *,
     use_videos: bool = True,
+    exclude_images: bool = False,
     patterns: Sequence[str] | None = None,
 ) -> dict[str, dict]:
     """
     Aggregates and filters pipeline features to create a dataset-ready features dictionary.
 
     This function transforms initial features using the pipeline, categorizes them as action or observations
-    (image or state), filters them based on `use_videos` and `patterns`, and finally
+    (image or state), filters them based on `exclude_images` and `patterns`, and finally
     formats them for use with a Hugging Face LeRobot Dataset.
 
     Args:
         pipeline: The DataProcessorPipeline to apply.
         initial_features: A dictionary of raw feature specs for actions and observations.
-        use_videos: If False, image features are excluded.
+        use_videos: Controls the storage dtype for image features. If True, images are stored as "video"; if False, they are stored as "image".
+        exclude_images: If True, image features are dropped entirely from the output.
         patterns: A sequence of regex patterns to filter action and state features.
                   Image features are not affected by this filter.
 
@@ -120,7 +122,7 @@ def aggregate_pipeline_dataset_features(
             )
 
             # 2. Apply filtering rules.
-            if is_image and not use_videos:
+            if is_image and exclude_images:
                 continue
             if not is_image and not should_keep(key, compiled_patterns):
                 continue

src/lerobot/datasets/sampler.py

@@ -14,14 +14,36 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+import math
 from collections.abc import Iterator
 
+import numpy as np
 import torch
 
 logger = logging.getLogger(__name__)
 
 
 class EpisodeAwareSampler:
+    """Sampler over episode frames that stores only per-episode boundaries.
+
+    Logical positions map to frame indices on the fly (O(num_episodes) construction memory)
+    instead of materializing a Python list of every frame index.
+
+    Each epoch is shuffled with a `torch.randperm` seeded from `(seed, epoch)`, so the data order
+    is a pure function of `(seed, epoch)`: it reproduces on every rank without synchronizing the
+    global RNG (no `generator` to sync across distributed ranks), and `state_dict` /
+    `load_state_dict` resume a run sample-exactly by regenerating the epoch's permutation and
+    continuing from the saved offset. Each call to `__iter__` advances the epoch. During a
+    resumed epoch, `__len__` still reports the full length.
+
+    Epoch advancement: `__iter__` eagerly advances the epoch, and `set_epoch` / `load_state_dict`
+    set it explicitly. Within a single run callers should rely on exactly one of these mechanisms,
+    not both: advancing the epoch by hand *and* letting `__iter__` auto-advance over the same
+    iterations would skip or repeat epochs. The training loop drives it purely through `__iter__`
+    (via `cycle`); `set_epoch` / `load_state_dict` are used only to (re)position before iteration
+    starts (e.g. on resume or in tests).
+    """
+
     def __init__(
         self,
         dataset_from_indices: list[int],
@@ -30,57 +52,125 @@ class EpisodeAwareSampler:
         drop_n_first_frames: int = 0,
         drop_n_last_frames: int = 0,
         shuffle: bool = False,
+        seed: int = 0,
     ):
-        """Sampler that optionally incorporates episode boundary information.
-
+        """
         Args:
-            dataset_from_indices: List of indices containing the start of each episode in the dataset.
-            dataset_to_indices: List of indices containing the end of each episode in the dataset.
-            episode_indices_to_use: List of episode indices to use. If None, all episodes are used.
-                                    Assumes that episodes are indexed from 0 to N-1.
-            drop_n_first_frames: Number of frames to drop from the start of each episode.
-            drop_n_last_frames: Number of frames to drop from the end of each episode.
+            dataset_from_indices: Start index of each episode in the dataset.
+            dataset_to_indices: End index of each episode in the dataset.
+            episode_indices_to_use: Episode indices to use; None means all.
+            drop_n_first_frames: Frames to drop from the start of each episode.
+            drop_n_last_frames: Frames to drop from the end of each episode.
             shuffle: Whether to shuffle the indices.
+            seed: Seed the permutation is derived from (together with the epoch).
         """
         if drop_n_first_frames < 0:
             raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}")
         if drop_n_last_frames < 0:
             raise ValueError(f"drop_n_last_frames must be >= 0, got {drop_n_last_frames}")
 
-        indices = []
-        for episode_idx, (start_index, end_index) in enumerate(
-            zip(dataset_from_indices, dataset_to_indices, strict=True)
-        ):
-            if episode_indices_to_use is None or episode_idx in episode_indices_to_use:
-                ep_length = end_index - start_index
-                if drop_n_first_frames + drop_n_last_frames >= ep_length:
-                    logger.warning(
-                        "Episode %d has %d frames but drop_n_first_frames=%d and "
-                        "drop_n_last_frames=%d removes all frames. Skipping.",
-                        episode_idx,
-                        ep_length,
-                        drop_n_first_frames,
-                        drop_n_last_frames,
-                    )
-                    continue
-                indices.extend(range(start_index + drop_n_first_frames, end_index - drop_n_last_frames))
+        from_indices = np.asarray(dataset_from_indices, dtype=np.int64)
+        to_indices = np.asarray(dataset_to_indices, dtype=np.int64)
+        if from_indices.shape != to_indices.shape:
+            raise ValueError(
+                f"dataset_from_indices and dataset_to_indices must have the same length, "
+                f"got {len(from_indices)} and {len(to_indices)}"
+            )
 
-        if not indices:
+        used = np.ones(len(from_indices), dtype=bool)
+        if episode_indices_to_use is not None:
+            used = np.zeros(len(from_indices), dtype=bool)
+            used[np.asarray(episode_indices_to_use, dtype=np.int64)] = True
+
+        starts = from_indices + drop_n_first_frames
+        lengths = to_indices - drop_n_last_frames - starts
+        for episode_idx in np.flatnonzero(used & (lengths <= 0)):
+            logger.warning(
+                "Episode %d has %d frames but drop_n_first_frames=%d and "
+                "drop_n_last_frames=%d removes all frames. Skipping.",
+                episode_idx,
+                to_indices[episode_idx] - from_indices[episode_idx],
+                drop_n_first_frames,
+                drop_n_last_frames,
+            )
+        used &= lengths > 0
+        if not used.any():
             raise ValueError(
                 "No valid frames remain after applying drop_n_first_frames and drop_n_last_frames. "
                 "All episodes were either filtered out or had too few frames."
             )
 
-        self.indices = indices
+        self._starts = starts[used]
+        self._cum_lengths = np.cumsum(lengths[used])
+        self._num_frames = int(self._cum_lengths[-1])
         self.shuffle = shuffle
+        self.seed = seed
+        self._epoch = 0
+        self._start_index = 0
+
+    @property
+    def indices(self) -> list[int]:
+        """Materialized frame indices in unshuffled order; O(num_frames), introspection only."""
+        return [self._frame_index(k) for k in range(self._num_frames)]
+
+    def set_epoch(self, epoch: int) -> None:
+        self._epoch = epoch
+
+    def state_dict(self) -> dict:
+        return {"epoch": self._epoch, "start_index": self._start_index}
+
+    def load_state_dict(self, state: dict) -> None:
+        self._epoch = state["epoch"]
+        self._start_index = state["start_index"]
+
+    def _epoch_generator(self, epoch: int) -> torch.Generator:
+        # Derive a per-epoch seed from (seed, epoch) so the permutation is a pure function of both
+        # and reproduces identically on every rank without touching the global RNG.
+        epoch_seed = int(np.random.SeedSequence([self.seed, epoch]).generate_state(1, dtype=np.uint64)[0])
+        return torch.Generator().manual_seed(epoch_seed)
+
+    def _frame_index(self, position: int) -> int:
+        episode = int(np.searchsorted(self._cum_lengths, position, side="right"))
+        position_in_episode = position - (int(self._cum_lengths[episode - 1]) if episode > 0 else 0)
+        return int(self._starts[episode]) + position_in_episode
 
     def __iter__(self) -> Iterator[int]:
+        # Advance epoch state eagerly, not on first consumption of the generator.
+        epoch, start = self._epoch, self._start_index
+        self._epoch += 1
+        self._start_index = 0
+        return self._iter_epoch(epoch, start)
+
+    def _iter_epoch(self, epoch: int, start: int) -> Iterator[int]:
         if self.shuffle:
-            for i in torch.randperm(len(self.indices)):
-                yield self.indices[i]
+            order = torch.randperm(self._num_frames, generator=self._epoch_generator(epoch))
+            for k in range(start, self._num_frames):
+                yield self._frame_index(int(order[k]))
         else:
-            for i in self.indices:
-                yield i
+            for k in range(start, self._num_frames):
+                yield self._frame_index(k)
 
     def __len__(self) -> int:
-        return len(self.indices)
+        return self._num_frames
+
+
+def compute_sampler_state(step: int, num_frames: int, batch_size: int, num_processes: int) -> dict:
+    """Map an optimization step to an `EpisodeAwareSampler` state for sample-exact resume.
+
+    Under accelerate's batch sharding, one step consumes `batch_size * num_processes` sampler
+    positions and each rank sees `ceil(ceil(num_frames / batch_size) / num_processes)` batches
+    per epoch (`even_batches` padding included). The start index provably stays below
+    `num_frames`; the `min` is defensive.
+
+    Assumptions (resume is only sample-exact when they hold):
+        - `num_processes` and `batch_size` match the run that wrote the checkpoint. Both scale how
+          many positions a step consumes, so the epoch/offset are wrong if either changed. The
+          caller passes the checkpoint's `num_processes` and `batch_size` and warns on a mismatch.
+        - accelerate uses `even_batches=True` (its default). The `ceil(... / num_processes)` term
+          mirrors that padding; with `even_batches=False` the per-epoch batch count differs and
+          the boundary is off.
+    """
+    batches_per_epoch = math.ceil(math.ceil(num_frames / batch_size) / num_processes)
+    epoch, batches_into_epoch = divmod(step, batches_per_epoch)
+    start_index = min(batches_into_epoch * batch_size * num_processes, num_frames)
+    return {"epoch": epoch, "start_index": start_index}

src/lerobot/datasets/video_utils.py

@@ -481,8 +481,10 @@ def reencode_video(
     encoder_threads: int | None = None,
     log_level: int | None = av.logging.WARNING,
     overwrite: bool = False,
+    start_time_s: float | None = None,
+    end_time_s: float | None = None,
 ) -> None:
-    """Re-encode a video file using the given encoder configuration.
+    """Re-encode a video file, optionally trimming it to ``[start_time_s, end_time_s)``.
 
     Args:
         input_video_path: Existing video file to read.
@@ -491,10 +493,17 @@ def reencode_video(
         encoder_threads: Optional thread count forwarded to :meth:`VideoEncoderConfig.get_codec_options`.
         log_level: libav log level while encoding, or ``None`` to leave logging unchanged. Defaults to WARNING.
         overwrite: When ``False`` and ``output_video_path`` already exists, skip and log a warning.
+        start_time_s: When set, trim the output to start at this timestamp (seconds).
+        end_time_s: When set, trim the output to end at this timestamp (seconds, exclusive).
     """
 
     camera_encoder = camera_encoder or camera_encoder_defaults()
 
+    if (start_time_s is not None and start_time_s < 0) or (end_time_s is not None and end_time_s < 0):
+        raise ValueError(f"Trim times must be non-negative, got start={start_time_s}, end={end_time_s}.")
+    if start_time_s is not None and end_time_s is not None and end_time_s <= start_time_s:
+        raise ValueError(f"end_time_s ({end_time_s}) must be greater than start_time_s ({start_time_s}).")
+
     output_video_path = Path(output_video_path)
 
     if output_video_path.exists() and not overwrite:
@@ -526,6 +535,10 @@ def reencode_video(
             width = int(in_stream.width)
             height = int(in_stream.height)
 
+            # Seek to the keyframe at or before start_time_s to avoid reading from the start.
+            if start_time_s is not None:
+                src.seek(int(start_time_s * av.time_base), backward=True)
+
             with av.open(
                 tmp_output_video_path,
                 mode="w",
@@ -539,7 +552,14 @@ def reencode_video(
                 out_stream.height = height
 
                 for frame in src.decode(in_stream):
+                    frame_time_s = frame.time
+                    if start_time_s is not None and frame_time_s < start_time_s:
+                        continue
+                    if end_time_s is not None and frame_time_s >= end_time_s:
+                        break
                     frame = frame.reformat(width=width, height=height, format=pix_fmt)
+                    if start_time_s is not None:
+                        frame.pts = None  # reset timestamps so the trimmed output starts at t=0
                     packet = out_stream.encode(frame)
                     if packet:
                         dst.mux(packet)

src/lerobot/policies/factory.py

@@ -252,6 +252,7 @@ class ProcessorConfigKwargs(TypedDict, total=False):
 def make_pre_post_processors(
     policy_cfg: PreTrainedConfig,
     pretrained_path: str | None = None,
+    pretrained_revision: str | None = None,
     **kwargs: Unpack[ProcessorConfigKwargs],
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
@@ -280,22 +281,26 @@ def make_pre_post_processors(
             policy configuration type.
     """
     if pretrained_path:
+        # TODO(Steven): Temporary patch, implement correctly the processors for Gr00t
         if isinstance(policy_cfg, GrootConfig):
-            from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained
+            # GROOT handles normalization in groot_pack_inputs_v3 step
+            # Need to override both stats AND normalize_min_max since saved config might be empty
+            preprocessor_overrides = {}
+            postprocessor_overrides = {}
+            preprocessor_overrides["groot_pack_inputs_v3"] = {
+                "stats": kwargs.get("dataset_stats"),
+                "normalize_min_max": True,
+            }
 
-            return make_groot_pre_post_processors_from_pretrained(
-                config=policy_cfg,
-                pretrained_path=pretrained_path,
-                dataset_stats=kwargs.get("dataset_stats"),
-                preprocessor_overrides=kwargs.get("preprocessor_overrides"),
-                postprocessor_overrides=kwargs.get("postprocessor_overrides"),
-                preprocessor_config_filename=kwargs.get(
-                    "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
-                ),
-                postprocessor_config_filename=kwargs.get(
-                    "postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json"
-                ),
-            )
+            # Also ensure postprocessing slices to env action dim and unnormalizes with dataset stats
+            env_action_dim = policy_cfg.output_features[ACTION].shape[0]
+            postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = {
+                "stats": kwargs.get("dataset_stats"),
+                "normalize_min_max": True,
+                "env_action_dim": env_action_dim,
+            }
+            kwargs["preprocessor_overrides"] = preprocessor_overrides
+            kwargs["postprocessor_overrides"] = postprocessor_overrides
 
         preprocessor = PolicyProcessorPipeline.from_pretrained(
             pretrained_model_name_or_path=pretrained_path,
@@ -305,6 +310,7 @@ def make_pre_post_processors(
             overrides=kwargs.get("preprocessor_overrides", {}),
             to_transition=batch_to_transition,
             to_output=transition_to_batch,
+            revision=pretrained_revision,
         )
         postprocessor = PolicyProcessorPipeline.from_pretrained(
             pretrained_model_name_or_path=pretrained_path,
@@ -314,6 +320,7 @@ def make_pre_post_processors(
             overrides=kwargs.get("postprocessor_overrides", {}),
             to_transition=policy_action_to_transition,
             to_output=transition_to_policy_action,
+            revision=pretrained_revision,
         )
         _reconnect_relative_absolute_steps(preprocessor, postprocessor)
         return preprocessor, postprocessor
@@ -553,6 +560,7 @@ def make_policy(
         # Load a pretrained policy and override the config if needed (for example, if there are inference-time
         # hyperparameters that we want to vary).
         kwargs["pretrained_name_or_path"] = cfg.pretrained_path
+        kwargs["revision"] = cfg.pretrained_revision
         policy = policy_cls.from_pretrained(**kwargs)
     elif cfg.pretrained_path and cfg.use_peft:
         # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo

src/lerobot/policies/groot/__init__.py

@@ -18,12 +18,4 @@ from .configuration_groot import GrootConfig
 from .modeling_groot import GrootPolicy
 from .processor_groot import make_groot_pre_post_processors
 
-__all__ = ["GR00TN17", "GR00TN17Config", "GrootConfig", "GrootPolicy", "make_groot_pre_post_processors"]
-
-
-def __getattr__(name: str):
-    if name in {"GR00TN17", "GR00TN17Config"}:
-        from .groot_n1_7 import GR00TN17, GR00TN17Config
-
-        return {"GR00TN17": GR00TN17, "GR00TN17Config": GR00TN17Config}[name]
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+__all__ = ["GrootConfig", "GrootPolicy", "make_groot_pre_post_processors"]

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

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+
+
+def swish(x):
+    return x * torch.sigmoid(x)
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    """
+    Produces a sinusoidal encoding of shape (B, T, w)
+    given timesteps of shape (B, T).
+    """
+
+    def __init__(self, embedding_dim):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+
+    def forward(self, timesteps):
+        # timesteps: shape (B, T)
+        # We'll compute sin/cos frequencies across dim T
+        timesteps = timesteps.float()  # ensure float
+
+        b, t = timesteps.shape
+        device = timesteps.device
+
+        half_dim = self.embedding_dim // 2
+        # typical log space frequencies for sinusoidal encoding
+        exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * (
+            torch.log(torch.tensor(10000.0)) / half_dim
+        )
+        # Expand timesteps to (B, T, 1) then multiply
+        freqs = timesteps.unsqueeze(-1) * exponent.exp()  # (B, T, half_dim)
+
+        sin = torch.sin(freqs)
+        cos = torch.cos(freqs)
+        enc = torch.cat([sin, cos], dim=-1)  # (B, T, w)
+
+        return enc

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

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 
-import logging
 from typing import TYPE_CHECKING
 
 import torch
@@ -43,9 +42,6 @@ else:
     Timesteps = None
 
 
-logger = logging.getLogger(__name__)
-
-
 class TimestepEncoder(nn.Module):
     def __init__(self, embedding_dim, compute_dtype=torch.float32):
         require_package("diffusers", extra="groot")
@@ -185,7 +181,8 @@ class BasicTransformerBlock(nn.Module):
         attn_output = self.attn1(
             norm_hidden_states,
             encoder_hidden_states=encoder_hidden_states,
-            attention_mask=encoder_attention_mask if encoder_hidden_states is not None else attention_mask,
+            attention_mask=attention_mask,
+            # encoder_attention_mask=encoder_attention_mask,
         )
         if self.final_dropout:
             attn_output = self.final_dropout(attn_output)
@@ -269,8 +266,8 @@ class DiT(ModelMixin, ConfigMixin):
         self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
         self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
         self.proj_out_2 = nn.Linear(self.inner_dim, self.config.output_dim)
-        logger.debug(
-            "Total number of DiT parameters: %d",
+        print(
+            "Total number of DiT parameters: ",
             sum(p.numel() for p in self.parameters() if p.requires_grad),
         )
 
@@ -321,71 +318,6 @@ class DiT(ModelMixin, ConfigMixin):
             return self.proj_out_2(hidden_states)
 
 
-class AlternateVLDiT(DiT):
-    """N1.7 DiT variant that alternates cross-attention over image and text tokens."""
-
-    def __init__(self, *args, attend_text_every_n_blocks: int = 2, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.attend_text_every_n_blocks = attend_text_every_n_blocks
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.LongTensor | None = None,
-        encoder_attention_mask: torch.Tensor | None = None,
-        return_all_hidden_states: bool = False,
-        image_mask: torch.Tensor | None = None,
-        backbone_attention_mask: torch.Tensor | None = None,
-    ):
-        if image_mask is None:
-            raise ValueError("image_mask is required for AlternateVLDiT.")
-        if backbone_attention_mask is None:
-            raise ValueError("backbone_attention_mask is required for AlternateVLDiT.")
-
-        temb = self.timestep_encoder(timestep)
-        hidden_states = hidden_states.contiguous()
-        encoder_hidden_states = encoder_hidden_states.contiguous()
-
-        image_attention_mask = image_mask & backbone_attention_mask
-        non_image_attention_mask = (~image_mask) & backbone_attention_mask
-
-        all_hidden_states = [hidden_states]
-        if not self.config.interleave_self_attention:
-            raise ValueError("AlternateVLDiT requires interleave_self_attention=True.")
-
-        for idx, block in enumerate(self.transformer_blocks):
-            if idx % 2 == 1:
-                hidden_states = block(
-                    hidden_states,
-                    attention_mask=None,
-                    encoder_hidden_states=None,
-                    encoder_attention_mask=None,
-                    temb=temb,
-                )
-            else:
-                curr_encoder_attention_mask = (
-                    non_image_attention_mask
-                    if idx % (2 * self.attend_text_every_n_blocks) == 0
-                    else image_attention_mask
-                )
-                hidden_states = block(
-                    hidden_states,
-                    attention_mask=None,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=curr_encoder_attention_mask,
-                    temb=temb,
-                )
-            all_hidden_states.append(hidden_states)
-
-        conditioning = temb
-        shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
-        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
-        if return_all_hidden_states:
-            return self.proj_out_2(hidden_states), all_hidden_states
-        return self.proj_out_2(hidden_states)
-
-
 class SelfAttentionTransformer(ModelMixin, ConfigMixin):
     _supports_gradient_checkpointing = True
 
@@ -430,8 +362,8 @@ class SelfAttentionTransformer(ModelMixin, ConfigMixin):
                 for _ in range(self.config.num_layers)
             ]
         )
-        logger.debug(
-            "Total number of SelfAttentionTransformer parameters: %d",
+        print(
+            "Total number of SelfAttentionTransformer parameters: ",
             sum(p.numel() for p in self.parameters() if p.requires_grad),
         )
 

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

@@ -0,0 +1,408 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import field
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import nn
+from torch.distributions import Beta
+
+from lerobot.utils.import_utils import _transformers_available
+
+# Conditional import for type checking and lazy loading
+if TYPE_CHECKING or _transformers_available:
+    from transformers import PretrainedConfig
+    from transformers.feature_extraction_utils import BatchFeature
+else:
+    PretrainedConfig = object
+    BatchFeature = None
+
+from .action_encoder import (
+    SinusoidalPositionalEncoding,
+    swish,
+)
+from .cross_attention_dit import DiT, SelfAttentionTransformer
+
+
+class CategorySpecificLinear(nn.Module):
+    def __init__(self, num_categories, input_dim, hidden_dim):
+        super().__init__()
+        self.num_categories = num_categories
+        # For each category, we have separate weights and biases.
+        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
+        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))
+
+    def forward(self, x, cat_ids):
+        selected_w = self.W[cat_ids]
+        selected_b = self.b[cat_ids]
+        return torch.bmm(x, selected_w) + selected_b.unsqueeze(1)
+
+
+class CategorySpecificMLP(nn.Module):
+    def __init__(self, num_categories, input_dim, hidden_dim, output_dim):
+        super().__init__()
+        self.num_categories = num_categories
+        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
+        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)
+
+    def forward(self, x, cat_ids):
+        hidden = F.relu(self.layer1(x, cat_ids))
+        return self.layer2(hidden, cat_ids)
+
+
+class MultiEmbodimentActionEncoder(nn.Module):
+    def __init__(self, action_dim, hidden_size, num_embodiments):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_embodiments = num_embodiments
+
+        # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w}
+        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)  # (d -> w)
+        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)  # (2w -> w)
+        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)  # (w -> w)
+        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)
+
+    def forward(self, actions, timesteps, cat_ids):
+        """
+        actions:   shape (B, T, action_dim)
+        timesteps: shape (B,)  -- a single scalar per batch item
+        cat_ids:   shape (B,)
+        returns:   shape (B, T, hidden_size)
+        """
+        b, t, _ = actions.shape
+
+        # 1) Expand each batch's single scalar time 'tau' across all T steps
+        #    so that shape => (B, T)
+        #    e.g. if timesteps is (B,), replicate across T
+        if timesteps.dim() == 1 and timesteps.shape[0] == b:
+            # shape (B,) => (B,T)
+            timesteps = timesteps.unsqueeze(1).expand(-1, t)
+        else:
+            raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.")
+
+        # 2) Standard action MLP step for shape => (B, T, w)
+        a_emb = self.W1(actions, cat_ids)
+
+        # 3) Get the sinusoidal encoding (B, T, w)
+        tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype)
+
+        # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish
+        x = torch.cat([a_emb, tau_emb], dim=-1)
+        x = swish(self.W2(x, cat_ids))
+
+        # 5) Finally W3 => (B, T, w)
+        x = self.W3(x, cat_ids)
+        return x
+
+
+class FlowmatchingActionHeadConfig(PretrainedConfig):
+    """NOTE: N1.5 uses XEmbFlowmatchingPolicyHeadConfig as action head"""
+
+    add_pos_embed: bool = field(default=True, metadata={"help": "Whether to add positional embedding"})
+    model_dtype: str = field(default="float32", metadata={"help": "Model data type."})
+    diffusion_model_cfg: dict = field(default=None, metadata={"help": "Diffusion model configuration."})
+    input_embedding_dim: int = field(default=1536, metadata={"help": "Input embedding channel dimension."})
+    backbone_embedding_dim: int = field(
+        default=1536, metadata={"help": "Backbone embedding channel dimension."}
+    )
+
+    hidden_size: int = field(default=1024, metadata={"help": "Input embedding dimension."})
+    max_seq_len: int = field(default=1024, metadata={"help": "Maximum Sequence Length"})
+    action_dim: int = field(default=None, metadata={"help": "Action dimension."})
+    action_horizon: int = field(default=None, metadata={"help": "Action horizon."})
+    noise_beta_alpha: float = field(default=1.5, metadata={"help": ""})
+    noise_beta_beta: float = field(default=1.0, metadata={"help": ""})
+    noise_s: float = field(default=0.999, metadata={"help": "Flow matching noise Beta distribution s."})
+    num_timestep_buckets: int = field(
+        default=1000, metadata={"help": "Number of timestep discretization buckets."}
+    )
+    num_inference_timesteps: int = field(
+        default=None,
+        metadata={"help": "Number of inference steps for noise diffusion."},
+    )
+    max_num_embodiments: int = field(default=32, metadata={"help": "Number of embodiments."})
+    tune_projector: bool = field(default=True, metadata={"help": "Whether to tune the projector."})
+    tune_diffusion_model: bool = field(
+        default=True, metadata={"help": "Whether to tune the diffusion model."}
+    )
+    load_pretrained_det_decode_layer_path: str = field(
+        default=None, metadata={"help": "Path to pretrained detection model."}
+    )
+    detection_coeff: float = field(default=1.0, metadata={"help": "Detection coefficient."})
+
+    freeze_decode_layer: bool = field(default=False)
+    expand_batch: int = field(default=None)
+    use_vlln: bool = field(default=True)
+
+    vl_self_attention_cfg: dict = field(default=None)
+    num_target_vision_tokens: int = field(default=32, metadata={"help": "Number of target vision tokens."})
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+
+class FlowmatchingActionHead(nn.Module):
+    config_class = FlowmatchingActionHeadConfig
+    supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        config: FlowmatchingActionHeadConfig,
+    ):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.input_embedding_dim = config.input_embedding_dim
+
+        self.model = DiT(**config.diffusion_model_cfg)
+        self.action_dim = config.action_dim
+        self.action_horizon = config.action_horizon
+        self.num_inference_timesteps = config.num_inference_timesteps
+
+        self.state_encoder = CategorySpecificMLP(
+            num_categories=config.max_num_embodiments,
+            input_dim=config.max_state_dim,
+            hidden_dim=self.hidden_size,
+            output_dim=self.input_embedding_dim,
+        )
+        self.action_encoder = MultiEmbodimentActionEncoder(
+            action_dim=config.action_dim,
+            hidden_size=self.input_embedding_dim,
+            num_embodiments=config.max_num_embodiments,
+        )
+        self.action_decoder = CategorySpecificMLP(
+            num_categories=config.max_num_embodiments,
+            input_dim=self.hidden_size,
+            hidden_dim=self.hidden_size,
+            output_dim=self.action_dim,
+        )
+        self.future_tokens = nn.Embedding(config.num_target_vision_tokens, self.input_embedding_dim)
+        nn.init.normal_(self.future_tokens.weight, mean=0.0, std=0.02)
+
+        self.vlln = nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity()
+        self.vl_self_attention = (
+            SelfAttentionTransformer(**config.vl_self_attention_cfg) if config.use_vlln else nn.Identity()
+        )
+
+        if config.add_pos_embed:
+            self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim)
+            nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)
+
+        self._noise_beta_alpha = config.noise_beta_alpha
+        self._noise_beta_beta = config.noise_beta_beta
+        self._beta_dist = None
+        self.num_timestep_buckets = config.num_timestep_buckets
+        self.config = config
+        self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model)
+
+    def set_trainable_parameters(self, tune_projector: bool, tune_diffusion_model: bool):
+        self.tune_projector = tune_projector
+        self.tune_diffusion_model = tune_diffusion_model
+        for p in self.parameters():
+            p.requires_grad = True
+        if not tune_projector:
+            self.state_encoder.requires_grad_(False)
+            self.action_encoder.requires_grad_(False)
+            self.action_decoder.requires_grad_(False)
+            if self.config.add_pos_embed:
+                self.position_embedding.requires_grad_(False)
+        if not tune_diffusion_model:
+            self.model.requires_grad_(False)
+        print(f"Tune action head projector: {self.tune_projector}")
+        print(f"Tune action head diffusion model: {self.tune_diffusion_model}")
+        # Check if any parameters are still trainable. If not, print a warning.
+        if not tune_projector and not tune_diffusion_model:
+            for name, p in self.named_parameters():
+                if p.requires_grad:
+                    print(f"Action head trainable parameter: {name}")
+        if not any(p.requires_grad for p in self.parameters()):
+            print("Warning: No action head trainable parameters found.")
+
+    def set_frozen_modules_to_eval_mode(self):
+        """
+        Huggingface will call model.train() at each training_step. To ensure
+        the expected behaviors for modules like dropout, batchnorm, etc., we
+        need to call model.eval() for the frozen modules.
+        """
+        if self.training:
+            if not self.tune_projector:
+                self.state_encoder.eval()
+                self.action_encoder.eval()
+                self.action_decoder.eval()
+                if self.config.add_pos_embed:
+                    self.position_embedding.eval()
+            if not self.tune_diffusion_model:
+                self.model.eval()
+
+    def sample_time(self, batch_size, device, dtype):
+        if self._beta_dist is None:
+            self._beta_dist = Beta(self._noise_beta_alpha, self._noise_beta_beta, validate_args=False)
+        sample = self._beta_dist.sample([batch_size]).to(device, dtype=dtype)
+        return (self.config.noise_s - sample) / self.config.noise_s
+
+    def prepare_input(self, batch: dict) -> BatchFeature:
+        return BatchFeature(data=batch)
+
+    def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature:
+        backbone_features = backbone_output["backbone_features"]
+        backbone_features = self.vlln(backbone_features)
+        backbone_features = self.vl_self_attention(backbone_features)
+        backbone_output["backbone_features"] = backbone_features
+        return backbone_output
+
+    def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature:
+        # Set frozen modules to eval
+        self.set_frozen_modules_to_eval_mode()
+
+        backbone_output = self.process_backbone_output(backbone_output)
+
+        if self.config.expand_batch is not None:
+            for k, v in backbone_output.items():
+                ndim = len(v.shape)
+                factors = [self.config.expand_batch]
+                while len(factors) < ndim:
+                    factors.append(1)
+                factors = tuple(factors)
+                expanded = v.repeat(*factors)
+                backbone_output[k] = expanded
+
+            for k, v in action_input.items():
+                ndim = len(v.shape)
+                factors = [self.config.expand_batch]
+                while len(factors) < ndim:
+                    factors.append(1)
+                factors = tuple(factors)
+                expanded = v.repeat(*factors)
+                action_input[k] = expanded
+
+        # Get vision and language embeddings.
+        vl_embs = backbone_output.backbone_features
+        device = vl_embs.device
+
+        # Get embodiment ID.
+        embodiment_id = action_input.embodiment_id
+
+        # Embed state.
+        state_features = self.state_encoder(action_input.state, embodiment_id)
+
+        # Embed noised action trajectory.
+        actions = action_input.action
+        noise = torch.randn(actions.shape, device=actions.device, dtype=actions.dtype)
+        t = self.sample_time(actions.shape[0], device=actions.device, dtype=actions.dtype)
+        t = t[:, None, None]  # shape (B,1,1) for broadcast
+
+        noisy_trajectory = (1 - t) * noise + t * actions
+        velocity = actions - noise
+
+        # Convert (continuous) t -> discrete if needed
+        t_discretized = (t[:, 0, 0] * self.num_timestep_buckets).long()
+        action_features = self.action_encoder(noisy_trajectory, t_discretized, embodiment_id)
+
+        # Maybe add position embedding.
+        if self.config.add_pos_embed:
+            pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
+            pos_embs = self.position_embedding(pos_ids).unsqueeze(0)
+            action_features = action_features + pos_embs
+
+        # Join vision, language, state and action embedding along sequence dimension.
+        future_tokens = self.future_tokens.weight.unsqueeze(0).expand(vl_embs.shape[0], -1, -1)
+        sa_embs = torch.cat((state_features, future_tokens, action_features), dim=1)
+
+        vl_attn_mask = backbone_output.backbone_attention_mask
+
+        model_output = self.model(
+            hidden_states=sa_embs,
+            encoder_hidden_states=vl_embs,
+            encoder_attention_mask=vl_attn_mask,
+            timestep=t_discretized,
+            return_all_hidden_states=False,  # NOTE (YL): not using flare now
+        )
+        pred = self.action_decoder(model_output, embodiment_id)
+        pred_actions = pred[:, -actions.shape[1] :]
+
+        # Slice out only the action portion of pred and target.
+        action_mask = action_input.action_mask
+        loss = F.mse_loss(pred_actions, velocity, reduction="none") * action_mask
+        loss = loss.sum() / action_mask.sum()
+        output_dict = {
+            "loss": loss,
+        }
+        return BatchFeature(data=output_dict)
+
+    @torch.no_grad()
+    def get_action(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature:
+        backbone_output = self.process_backbone_output(backbone_output)
+
+        # Get vision and language embeddings.
+        vl_embs = backbone_output.backbone_features
+        embodiment_id = action_input.embodiment_id
+
+        # Embed state.
+        state_features = self.state_encoder(action_input.state, embodiment_id)
+
+        # Set initial actions as the sampled noise.
+        batch_size = vl_embs.shape[0]
+        device = vl_embs.device
+        actions = torch.randn(
+            size=(batch_size, self.config.action_horizon, self.config.action_dim),
+            dtype=vl_embs.dtype,
+            device=device,
+        )
+
+        num_steps = self.num_inference_timesteps
+        dt = 1.0 / num_steps
+
+        # Run denoising steps.
+        for t in range(num_steps):
+            t_cont = t / float(num_steps)  # e.g. goes 0, 1/N, 2/N, ...
+            t_discretized = int(t_cont * self.num_timestep_buckets)
+
+            # Embed noised action trajectory.
+            timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device)
+            action_features = self.action_encoder(actions, timesteps_tensor, embodiment_id)
+            # Maybe add position embedding.
+            if self.config.add_pos_embed:
+                pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
+                pos_embs = self.position_embedding(pos_ids).unsqueeze(0)
+                action_features = action_features + pos_embs
+
+            # Join vision, language, state and action embedding along sequence dimension.
+            future_tokens = self.future_tokens.weight.unsqueeze(0).expand(vl_embs.shape[0], -1, -1)
+            sa_embs = torch.cat((state_features, future_tokens, action_features), dim=1)
+
+            # Run model forward.
+            model_output = self.model(
+                hidden_states=sa_embs,
+                encoder_hidden_states=vl_embs,
+                timestep=timesteps_tensor,
+            )
+            pred = self.action_decoder(model_output, embodiment_id)
+
+            pred_velocity = pred[:, -self.action_horizon :]
+
+            # Update actions using euler integration.
+            actions = actions + dt * pred_velocity
+        return BatchFeature(data={"action_pred": actions})
+
+    @property
+    def device(self):
+        return next(iter(self.parameters())).device
+
+    @property
+    def dtype(self):
+        return next(iter(self.parameters())).dtype

src/lerobot/policies/groot/configuration_groot.py

@@ -14,228 +14,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 from dataclasses import dataclass, field
-from pathlib import Path
 
 from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
 from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
 from lerobot.utils.constants import ACTION, OBS_STATE
 
-from .utils import read_json
-
-logger = logging.getLogger(__name__)
-
-GROOT_N1_7 = "n1.7"
-# Legacy GR00T N1.5 identifier. N1.5 is NOT a supported model_version (it is
-# intentionally absent from _GROOT_MODEL_VERSION_ALIASES so normalize_groot_model_version
-# still rejects it). It is retained only so that infer_groot_model_version can recognise
-# an N1.5 base path/checkpoint and the N1.7 config/loader can reject the mismatch.
-GROOT_N1_5 = "n1.5"
-# Canonical guidance appended to every error raised when an N1.5 checkpoint, config,
-# or processor pipeline is detected. Keep this message in sync with docs/source/groot.mdx.
-GROOT_N1_5_REMOVAL_GUIDANCE = (
-    "GR00T N1.5 support was removed from LeRobot. "
-    "To keep using an N1.5 checkpoint, pin the last release that supports it: "
-    "`pip install 'lerobot==0.5.1'`. To use the current release, migrate to GR00T N1.7 "
-    "(model_version='n1.7', base model nvidia/GR00T-N1.7-3B)."
-)
-GROOT_N1_7_BASE_MODEL = "nvidia/GR00T-N1.7-3B"
-GROOT_N1_7_BACKBONE_MODEL = "nvidia/Cosmos-Reason2-2B"
-# Default GR00T N1.7 training resolution. Fallback if processor_config lacks sizing. Prevents mismatched
-# full-res patchification by forcing a resize. Mirrored by GR00T_N1_7_DEFAULTS in groot_n1_7.py.
-N1_7_DEFAULT_IMAGE_TARGET_SIZE = (256, 256)
-N1_7_DEFAULT_IMAGE_CROP_SIZE = (230, 230)
-GROOT_ACTION_DECODE_TRANSFORM_LIBERO = "libero"
-# Sentinel meaning "the user did not pick an action decode transform": __post_init__ resolves it
-# to the embodiment default ('libero' for 'libero_sim', otherwise None). It is distinct from an
-# explicit 'none' (resolved to None) so an opt-out survives a draccus save/load round-trip.
-GROOT_ACTION_DECODE_TRANSFORM_AUTO = "auto"
-
-_GROOT_MODEL_VERSION_ALIASES = {
-    "n1.7": GROOT_N1_7,
-    "n1_7": GROOT_N1_7,
-    "n1d7": GROOT_N1_7,
-    "n17": GROOT_N1_7,
-    "1.7": GROOT_N1_7,
-}
-
-# Legacy N1.5 spellings, kept ONLY so they can be detected and rejected with
-# GROOT_N1_5_REMOVAL_GUIDANCE (see GROOT_N1_5 above). Never map these to a supported version.
-_GROOT_N1_5_VERSION_ALIASES = {"n1.5", "n1_5", "n1d5", "n15", "1.5"}
-
-_GROOT_ACTION_DECODE_TRANSFORM_ALIASES = {
-    GROOT_ACTION_DECODE_TRANSFORM_AUTO: GROOT_ACTION_DECODE_TRANSFORM_AUTO,
-    "none": None,
-    "": None,
-    GROOT_ACTION_DECODE_TRANSFORM_LIBERO: GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
-}
-
-
-def normalize_groot_model_version(model_version: str) -> str:
-    normalized = _GROOT_MODEL_VERSION_ALIASES.get(model_version.lower())
-    if normalized is None:
-        supported = GROOT_N1_7
-        message = f"Unsupported GR00T model_version '{model_version}'. Supported versions: {supported}."
-        if model_version.lower() in _GROOT_N1_5_VERSION_ALIASES:
-            message = f"{message} {GROOT_N1_5_REMOVAL_GUIDANCE}"
-        raise ValueError(message)
-    return normalized
-
-
-def normalize_groot_action_decode_transform(transform: str | None) -> str | None:
-    if transform is None:
-        return None
-    normalized = _GROOT_ACTION_DECODE_TRANSFORM_ALIASES.get(transform.lower())
-    if normalized is None and transform.lower() not in _GROOT_ACTION_DECODE_TRANSFORM_ALIASES:
-        supported = ", ".join(
-            sorted(key for key, value in _GROOT_ACTION_DECODE_TRANSFORM_ALIASES.items() if value is not None)
-        )
-        raise ValueError(
-            f"Unsupported GR00T N1.7 action decode transform '{transform}'. "
-            f"Supported transforms: none, {supported}."
-        )
-    return normalized
-
-
-def infer_groot_model_version(model_path: str | None) -> str | None:
-    if not model_path:
-        return None
-    model_path_lower = model_path.lower()
-    if "gr00t-n1.7" in model_path_lower or "gr00t_n1.7" in model_path_lower:
-        return GROOT_N1_7
-    # Detect legacy N1.5 paths so the N1.7 config/loader can reject the mismatch.
-    # N1.5 is unsupported, but it must still be recognised here to fail loudly
-    # rather than silently treating an N1.5 checkpoint as N1.7.
-    if "gr00t-n1.5" in model_path_lower or "gr00t_n1.5" in model_path_lower:
-        return GROOT_N1_5
-    config_version = _infer_groot_model_version_from_local_config(model_path)
-    if config_version is not None:
-        return config_version
-    return None
-
-
-def is_raw_groot_n1_7_checkpoint(model_path: str | Path | None) -> bool:
-    if model_path is None:
-        return False
-
-    path = Path(model_path).expanduser()
-    if path.is_dir():
-        config_path = path / "config.json"
-    elif path.name == "config.json":
-        config_path = path
-    else:
-        return False
-
-    config = read_json(config_path)
-    return "type" not in config and _infer_groot_model_version_from_config(config) == GROOT_N1_7
-
-
-def infer_groot_n1_7_embodiment_tag(model_path: str | Path | None) -> str | None:
-    if model_path is None:
-        return None
-
-    processor_config_path = Path(model_path).expanduser() / "processor_config.json"
-    processor_config = read_json(processor_config_path)
-
-    modality_configs = processor_config.get("processor_kwargs", {}).get("modality_configs", {})
-    if not isinstance(modality_configs, dict):
-        return None
-    if "libero_sim" in modality_configs:
-        return "libero_sim"
-    if len(modality_configs) == 1:
-        return next(iter(modality_configs))
-    return None
-
-
-def infer_groot_n1_7_action_horizon(
-    model_path: str | Path | None, embodiment_tag: str | None = None
-) -> int | None:
-    if model_path is None:
-        return None
-
-    processor_config_path = Path(model_path).expanduser() / "processor_config.json"
-    processor_config = read_json(processor_config_path)
-
-    processor_kwargs = processor_config.get("processor_kwargs", {})
-    if not isinstance(processor_kwargs, dict):
-        return None
-    modality_configs = processor_kwargs.get("modality_configs", {})
-    if not isinstance(modality_configs, dict):
-        return None
-
-    if embodiment_tag is None:
-        embodiment_tag = infer_groot_n1_7_embodiment_tag(model_path)
-    if embodiment_tag is None:
-        return None
-
-    embodiment_config = modality_configs.get(embodiment_tag, {})
-    if not isinstance(embodiment_config, dict):
-        return None
-    action_config = embodiment_config.get("action", {})
-    if not isinstance(action_config, dict):
-        return None
-    delta_indices = action_config.get("delta_indices", [])
-    if not isinstance(delta_indices, list):
-        return None
-    return len(delta_indices) or None
-
-
-def infer_groot_n1_7_action_execution_horizon(
-    model_path: str | Path | None, embodiment_tag: str | None = None
-) -> int | None:
-    action_horizon = infer_groot_n1_7_action_horizon(model_path, embodiment_tag)
-    if action_horizon is None:
-        return None
-
-    if embodiment_tag is None:
-        embodiment_tag = infer_groot_n1_7_embodiment_tag(model_path)
-    if embodiment_tag == "libero_sim":
-        # NVIDIA's N1.7 LIBERO rollout wrapper replans after 8 of the 16 decoded
-        # actions. Keeping that execution cadence avoids stale open-loop chunks.
-        return min(action_horizon, 8)
-    return action_horizon
-
-
-def _infer_groot_model_version_from_local_config(model_path: str) -> str | None:
-    path = Path(model_path).expanduser()
-    if path.is_dir():
-        config_path = path / "config.json"
-    elif path.name == "config.json":
-        config_path = path
-    else:
-        return None
-
-    return _infer_groot_model_version_from_config(read_json(config_path))
-
-
-def _infer_groot_model_version_from_config(config: dict) -> str | None:
-    model_version = config.get("model_version")
-    if isinstance(model_version, str):
-        if model_version.lower() in _GROOT_N1_5_VERSION_ALIASES:
-            return GROOT_N1_5
-        try:
-            return normalize_groot_model_version(model_version)
-        except ValueError:
-            return None
-
-    candidates = [config.get("model_type"), *(config.get("architectures") or [])]
-    for candidate in candidates:
-        if not isinstance(candidate, str):
-            continue
-        normalized = candidate.lower().replace("-", "_")
-        if normalized in {"gr00tn1d7", "gr00t_n1d7", "gr00t_n1_7"}:
-            return GROOT_N1_7
-        if normalized in {"gr00t_n1_5", "gr00tn1_5", "gr00t_n15", "gr00t_n1d5", "gr00tn1d5"}:
-            return GROOT_N1_5
-    if config.get("model_name") == GROOT_N1_7_BACKBONE_MODEL:
-        return GROOT_N1_7
-    # The Eagle VLM backbone is specific to pre-N1.7 GR00T checkpoints (N1.7 uses Cosmos/Qwen3-VL).
-    backbone_cfg = config.get("backbone_cfg")
-    if isinstance(backbone_cfg, dict) and "eagle_path" in backbone_cfg:
-        return GROOT_N1_5
-    return None
-
 
 @PreTrainedConfig.register_subclass("groot")
 @dataclass
@@ -244,44 +28,35 @@ class GrootConfig(PreTrainedConfig):
 
     # Basic policy settings
     n_obs_steps: int = 1
-    chunk_size: int = 40
-    n_action_steps: int = 40
+    chunk_size: int = 50
+    n_action_steps: int = 50
 
     # Dimension settings (must match pretrained GR00T model expectations)
     # Maximum state dimension. Shorter states will be zero-padded.
-    max_state_dim: int = 132
+    max_state_dim: int = 64
 
     # Maximum action dimension. Shorter actions will be zero-padded.
-    max_action_dim: int = 132
+    max_action_dim: int = 32
 
-    # GR00T normalizes state/action internally in its processor steps (min/max with
-    # q01/q99 percentiles, per embodiment), and the Qwen3-VL backbone's image processor
-    # handles image normalization. The policy therefore does NOT use LeRobot's
-    # NormalizerProcessorStep/UnnormalizerProcessorStep, so this mapping is intentionally
-    # IDENTITY for every feature and is not consulted by make_groot_pre_post_processors.
+    # Normalization (start with identity, adjust as needed)
     normalization_mapping: dict[str, NormalizationMode] = field(
         default_factory=lambda: {
             "VISUAL": NormalizationMode.IDENTITY,
-            "STATE": NormalizationMode.IDENTITY,
-            "ACTION": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.MEAN_STD,
+            "ACTION": NormalizationMode.MEAN_STD,
         }
     )
 
-    # Groot-specific model parameters
+    # Image preprocessing (adjust to match Groot's expected input)
+    image_size: tuple[int, int] = (224, 224)
 
-    # Path or HuggingFace model ID for the base GR00T N1.7 model whose backbone weights and
-    # checkpoint sidecars (statistics.json, processor_config.json, ...) are loaded. This is the
-    # model *source*, and is intentionally distinct from the inherited `pretrained_path`:
-    # `pretrained_path` (`--policy.path`) points at a saved LeRobot checkpoint directory whose
-    # `config.json` carries a `type` field, whereas a raw NVIDIA GR00T checkpoint has no such
-    # field and so can only be loaded through `base_model_path` (`--policy.base_model_path`).
-    # Defaults to GROOT_N1_7_BASE_MODEL when unset (resolved in __post_init__).
-    base_model_path: str | None = None
+    # Groot-specific model parameters (from groot_finetune_script.py)
 
-    # Optional named action transform applied after raw N1.7 checkpoint decoding and before env.step().
-    # 'auto' (default) resolves to the embodiment default ('libero' for 'libero_sim', otherwise no
-    # transform). Pass 'none' to explicitly disable the transform, including for 'libero_sim'.
-    action_decode_transform: str | None = GROOT_ACTION_DECODE_TRANSFORM_AUTO
+    # Path or HuggingFace model ID for the base Groot model
+    base_model_path: str = "nvidia/GR00T-N1.5-3B"
+
+    # HF repo ID (or local path) that hosts vocab.json and merges.txt for Eagle tokenizer.
+    tokenizer_assets_repo: str = "lerobot/eagle2hg-processor-groot-n1p5"
 
     # Embodiment tag to use for training (e.g. 'new_embodiment', 'gr1')
     embodiment_tag: str = "new_embodiment"
@@ -300,31 +75,20 @@ class GrootConfig(PreTrainedConfig):
     # Whether to fine-tune the diffusion model
     tune_diffusion_model: bool = True
 
-    # Whether to fine-tune the VL LayerNorm + VL self-attention projector in the action head.
-    tune_vlln: bool = True
+    # LoRA parameters (from groot_finetune_script.py)
+    # Rank for the LORA model. If 0, no LORA will be used.
+    lora_rank: int = 0
 
-    # Number of top LLM backbone layers to fine-tune (0 = none). Lets you adapt just the final
-    # language layers without unfreezing the whole backbone; independent of `tune_llm`, which tunes
-    # the entire LLM.
-    tune_top_llm_layers: int = 0
+    # Alpha value for the LORA model
+    lora_alpha: int = 16
 
-    # Inference-time knob: Number of flow-matching denoising steps used to decode an action chunk.
-    # Trades inference latency for action quality.
-    # None keeps the checkpoint value (GR00T N1.7 default: 4).
-    num_inference_timesteps: int | None = None
+    # Dropout rate for the LORA model
+    lora_dropout: float = 0.1
 
-    # Inference-time knob: Real-Time Chunking (RTC) overlap-blend ramp rate, used when the RTC engine
-    # supplies a previous-chunk prefix. Higher values blend the overlapping prefix more aggressively.
-    # None keeps the checkpoint value (GR00T N1.7 default: 6.0).
-    rtc_ramp_rate: float | None = None
+    # Whether to use the full model for LORA
+    lora_full_model: bool = False
 
-    # Inference-time knob: Whether to request the flash-attention-2 kernel for the Qwen3-VL backbone.
-    # flash-attn is an optional, user-managed optimization; when it is absent (the default),
-    # the backbone transparently falls back to SDPA, which is numerically equivalent.
-    # Set to True only after installing a flash-attn build matching your torch/CUDA env.
-    use_flash_attention: bool = False
-
-    # Training parameters
+    # Training parameters (matching groot_finetune_script.py)
     optimizer_lr: float = 1e-4
     optimizer_betas: tuple[float, float] = (0.95, 0.999)
     optimizer_eps: float = 1e-8
@@ -332,22 +96,17 @@ class GrootConfig(PreTrainedConfig):
     warmup_ratio: float = 0.05
     use_bf16: bool = True
 
-    # TODO(Steven): Remove these deprecated fields in a future release.
-    # Deprecated Isaac-GR00T runner / GR00T N1.5 fields, plus the (never-wired) LoRA fields — all
-    # unused by the LeRobot N1.7 implementation except the `tokenizer_assets_repo` N1.5 tripwire and
-    # the `image_size` legacy remap in __post_init__. They are kept ONLY so a config.json saved by an
-    # earlier lerobot release (notably a GR00T N1.5 checkpoint) still parses under draccus — which
-    # rejects unknown fields — and is then rejected with a clear N1.5 removal message rather than an
-    # opaque draccus decoding error.
-    image_size: tuple[int, int] = (256, 256)  # image sizing is handled by the backbone's image processor.
-    tokenizer_assets_repo: str | None = None
-    lora_rank: int = 0
-    lora_alpha: int = 16
-    lora_dropout: float = 0.1
-    lora_full_model: bool = False
+    # Dataset parameters
+    # Video backend to use for training ('decord' or 'torchvision_av')
     video_backend: str = "decord"
+
+    # Whether to balance dataset weights in mixture datasets
     balance_dataset_weights: bool = True
+
+    # Whether to sample trajectories weighted by their length
     balance_trajectory_weights: bool = True
+
+    # Optional dataset paths for delegating training to Isaac-GR00T runner
     dataset_paths: list[str] | None = None
     output_dir: str = "./tmp/gr00t"
     save_steps: int = 1000
@@ -358,65 +117,6 @@ class GrootConfig(PreTrainedConfig):
     resume: bool = False
 
     def __post_init__(self):
-        if self.tokenizer_assets_repo is not None:
-            raise ValueError(
-                "Config sets 'tokenizer_assets_repo', which only existed for GR00T N1.5; this looks "
-                f"like a legacy GR00T N1.5 checkpoint or config. {GROOT_N1_5_REMOVAL_GUIDANCE}"
-            )
-
-        self.action_decode_transform = normalize_groot_action_decode_transform(self.action_decode_transform)
-        if self.base_model_path is None:
-            self.base_model_path = GROOT_N1_7_BASE_MODEL
-
-        # The N1.7 LIBERO checkpoints emit a [0, 1] gripper action, but the LIBERO
-        # simulator expects the OpenVLA/[-1, 1] sign convention. NVIDIA's rollout
-        # wrapper applies this conversion; mirror it here so eval on the
-        # 'libero_sim' embodiment grasps correctly instead of scoring 0% success.
-        # This matches the embodiment-specific handling already done for the
-        # action execution horizon (see infer_groot_n1_7_action_execution_horizon).
-        # Only the 'auto' sentinel resolves to the embodiment default; an explicit
-        # 'none' (normalized to None above) keeps the transform disabled.
-        if self.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_AUTO:
-            self.action_decode_transform = (
-                GROOT_ACTION_DECODE_TRANSFORM_LIBERO if self.embodiment_tag == "libero_sim" else None
-            )
-
-        # GR00T N1.5-era default values (e.g. --policy.chunk_size=50 from old commands or
-        # stale configs) are migrated to the values the N1.7 checkpoints expect, with a
-        # warning. The dataclass defaults are already the N1.7 values, so a plain
-        # GrootConfig() never triggers this.
-        legacy_default_remaps = (
-            ("max_state_dim", 64, 132),
-            ("max_action_dim", 32, 132),
-            ("chunk_size", 50, 40),
-            ("n_action_steps", 50, 40),
-            ("image_size", (224, 224), (256, 256)),
-        )
-        for field_name, legacy_value, n1_7_value in legacy_default_remaps:
-            current_value = getattr(self, field_name)
-            if isinstance(legacy_value, tuple):
-                current_value = tuple(current_value)
-            if current_value == legacy_value:
-                logger.warning(
-                    "GrootConfig.%s=%s matches a legacy GR00T N1.5-era default; remapping it to %s, "
-                    "the value expected by GR00T N1.7 checkpoints. Set a different value explicitly "
-                    "if this is not what you want.",
-                    field_name,
-                    legacy_value,
-                    n1_7_value,
-                )
-                setattr(self, field_name, n1_7_value)
-
-        inferred_version = infer_groot_model_version(self.base_model_path)
-        if inferred_version is not None and inferred_version != GROOT_N1_7:
-            message = (
-                f"GR00T model_version '{GROOT_N1_7}' does not match base_model_path "
-                f"'{self.base_model_path}', which looks like '{inferred_version}'."
-            )
-            if inferred_version == GROOT_N1_5:
-                message = f"{message} {GROOT_N1_5_REMOVAL_GUIDANCE}"
-            raise ValueError(message)
-
         super().__post_init__()
 
         if self.n_action_steps > self.chunk_size:
@@ -424,6 +124,9 @@ class GrootConfig(PreTrainedConfig):
                 f"n_action_steps ({self.n_action_steps}) cannot exceed chunk_size ({self.chunk_size})"
             )
 
+        # groot_repo_path is now optional since we ported the components
+        # No validation needed
+
     def validate_features(self) -> None:
         """Validate and set up input/output features for Groot."""
         image_features = [key for key, feat in self.input_features.items() if feat.type == FeatureType.VISUAL]
@@ -489,10 +192,7 @@ class GrootConfig(PreTrainedConfig):
     @property
     def action_delta_indices(self) -> list[int]:
         """Return indices for delta actions."""
-        model_action_horizon = (
-            infer_groot_n1_7_action_horizon(self.base_model_path, self.embodiment_tag) or 40
-        )
-        return list(range(min(self.chunk_size, model_action_horizon)))
+        return list(range(min(self.chunk_size, 16)))
 
     @property
     def reward_delta_indices(self) -> None:

src/lerobot/policies/groot/eagle2_hg_model/configuration_eagle2_5_vl.py

@@ -0,0 +1,135 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import copy
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class Eagle25VLConfig(PretrainedConfig):
+    model_type = "eagle_2_5_vl"
+    is_composition = True
+    sub_configs = {"vision_config": SiglipVisionConfig, "text_config": Qwen2Config}
+
+    def __init__(
+        self,
+        vision_config=None,
+        text_config=None,
+        use_backbone_lora=0,
+        use_llm_lora=0,
+        pad2square=False,
+        select_layer=-4,
+        force_image_size=None,
+        downsample_ratio=0.5,
+        template=None,
+        dynamic_image_size=False,
+        use_thumbnail=False,
+        loss_version="v1",
+        min_dynamic_tiles=1,
+        max_dynamic_tiles=6,
+        mlp_checkpoint=False,
+        initializer_range=0.02,
+        _attn_implementation="flash_attention_2",
+        _attn_implementation_autoset=False,
+        llm_config=None,
+        image_token_index=None,
+        use_pixel_shuffle=True,
+        mlp_connector_layers=2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if vision_config is None:
+            vision_config = {"model_type": "siglip_vision_model"}
+            logger.info("vision_config is None. Initializing the InternVisionConfig with default values.")
+
+        if text_config is None:
+            text_config = {"architectures": ["Qwen2ForCausalLM"]}
+            logger.info(
+                "text_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`)."
+            )
+
+        if vision_config["model_type"] == "siglip_vision_model":
+            self.vision_config = SiglipVisionConfig(**vision_config)
+        else:
+            raise ValueError("Unsupported model_type: {}".format(vision_config["model_type"]))
+
+        if text_config["architectures"][0] == "LlamaForCausalLM":
+            self.text_config = LlamaConfig(**text_config)
+        elif text_config["architectures"][0] == "Qwen2ForCausalLM":
+            self.text_config = Qwen2Config(**text_config)
+        elif text_config["architectures"][0] == "Qwen3ForCausalLM":
+            self.text_config = Qwen3Config(**text_config)
+        else:
+            raise ValueError("Unsupported architecture: {}".format(text_config["architectures"][0]))
+        self.use_backbone_lora = use_backbone_lora
+        self.use_llm_lora = use_llm_lora
+        self.mlp_checkpoint = mlp_checkpoint
+        self.pad2square = pad2square
+        self.select_layer = select_layer
+        self.force_image_size = force_image_size
+        self.downsample_ratio = downsample_ratio
+        self.template = template
+        self.dynamic_image_size = dynamic_image_size
+        self.use_thumbnail = use_thumbnail
+        self.loss_version = loss_version
+        self.initializer_range = initializer_range
+        self.min_dynamic_tiles = min_dynamic_tiles
+        self.max_dynamic_tiles = max_dynamic_tiles
+        self.tie_word_embeddings = self.text_config.tie_word_embeddings
+        self._attn_implementation = _attn_implementation
+        self._attn_implementation_autoset = _attn_implementation_autoset
+        self.image_token_index = image_token_index
+        self.use_pixel_shuffle = use_pixel_shuffle
+        self.mlp_connector_layers = mlp_connector_layers
+        logger.info(f"min_dynamic_tiles: {self.min_dynamic_tiles}")
+        logger.info(f"max_dynamic_tiles: {self.max_dynamic_tiles}")
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output["vision_config"] = self.vision_config.to_dict()
+        output["text_config"] = self.text_config.to_dict()
+        output["model_type"] = self.__class__.model_type
+        output["use_backbone_lora"] = self.use_backbone_lora
+        output["use_llm_lora"] = self.use_llm_lora
+        output["pad2square"] = self.pad2square
+        output["select_layer"] = self.select_layer
+        output["force_image_size"] = self.force_image_size
+        output["downsample_ratio"] = self.downsample_ratio
+        output["template"] = self.template
+        output["dynamic_image_size"] = self.dynamic_image_size
+        output["use_thumbnail"] = self.use_thumbnail
+        output["min_dynamic_tiles"] = self.min_dynamic_tiles
+        output["max_dynamic_tiles"] = self.max_dynamic_tiles
+        output["tie_word_embeddings"] = self.tie_word_embeddings
+        output["_attn_implementation"] = self._attn_implementation
+        output["_attn_implementation_autoset"] = self._attn_implementation_autoset
+        output["use_pixel_shuffle"] = self.use_pixel_shuffle
+        output["mlp_connector_layers"] = self.mlp_connector_layers
+        return output

src/lerobot/policies/groot/eagle2_hg_model/image_processing_eagle2_5_vl_fast.py

@@ -0,0 +1,503 @@
+# --------------------------------------------------------
+# NVIDIA
+# Copyright (c) 2025 NVIDIA
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import annotations
+
+# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py
+from transformers.image_processing_utils import (
+    BatchFeature,
+    get_patch_output_size,
+)
+from transformers.image_processing_utils_fast import (
+    BaseImageProcessorFast,
+    ImagesKwargs,
+    group_images_by_shape,
+    reorder_images,
+)
+from transformers.image_utils import (
+    IMAGENET_STANDARD_MEAN,  # 0.5, 0.5, 0.5
+    IMAGENET_STANDARD_STD,  # 0.5, 0.5, 0.5
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    SizeDict,
+    get_image_size,
+    make_flat_list_of_images,
+    validate_kwargs,
+)
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    TensorType,
+    add_start_docstrings,
+    is_torch_available,
+    is_torchvision_v2_available,
+)
+from transformers.video_utils import VideoInput
+
+if is_torch_available():
+    import torch
+if is_torchvision_v2_available():
+    from torchvision.transforms.v2 import functional as F  # noqa: N812
+    from transformers.image_utils import pil_torch_interpolation_mapping
+else:
+    from torchvision.transforms import functional as F  # noqa: N812
+
+
+def crop(img: torch.Tensor, left: int, top: int, right: int, bottom: int) -> torch.Tensor:
+    """Crop the given numpy array.
+
+    Args:
+        img (torch.Tensor): Image to be cropped. Format should be (C, H, W).
+        left (int): The left coordinate of the crop box.
+        top (int): The top coordinate of the crop box.
+        right (int): The right coordinate of the crop box.
+        bottom (int): The bottom coordinate of the crop box.
+
+    Returns:
+        torch.Tensor: Cropped image.
+    """
+    if not isinstance(img, torch.Tensor):
+        raise TypeError(f"img should be torch.Tensor. Got {type(img)}")
+
+    if img.ndim not in [2, 3]:
+        raise ValueError(f"Image should have 2 or 3 dimensions. Got {img.ndim}")
+
+    img_height = img.shape[1]
+    img_width = img.shape[2]
+    if top < 0 or left < 0 or bottom > img_height or right > img_width:
+        raise ValueError("Crop coordinates out of bounds")
+
+    if top >= bottom or left >= right:
+        raise ValueError("Invalid crop coordinates")
+
+    return img[:, top:bottom, left:right]
+
+
+class Eagle25VLFastImageProcessorKwargs(ImagesKwargs):
+    max_dynamic_tiles: int | None
+    min_dynamic_tiles: int | None
+    use_thumbnail: bool | None
+    pad_during_tiling: bool | None
+    do_pad: bool | None
+
+
+@add_start_docstrings(
+    "Constructs a fast ConvNeXT image processor. Based on [`SiglipImageProcessor`] with incorporation of processing each video frame.",
+    # BASE_IMAGE_PROCESSOR_FAST_DOCSTRING, TODO: this was depreciated from transformers remove!
+    """
+        image_grid_pinpoints (`List[List[int]]`, *optional*):
+            A list of possible resolutions to use for processing high resolution images. The best resolution is selected
+            based on the original size of the image. Can be overridden by `image_grid_pinpoints` in the `preprocess`
+            method. Not used for processing videos.
+        do_pad (`bool`, *optional*):
+            Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest
+            number of patches in the batch. Padding will be applied to the bottom and right with zeros.
+    """,
+)
+class Eagle25VLImageProcessorFast(BaseImageProcessorFast):
+    resample = PILImageResampling.BICUBIC
+    image_mean = IMAGENET_STANDARD_MEAN
+    image_std = IMAGENET_STANDARD_STD
+    size = {"height": 448, "width": 448}
+    default_to_square = False
+    crop_size = None
+    do_resize = True
+    do_center_crop = None
+    do_rescale = True
+    do_normalize = True
+    do_convert_rgb = True
+    do_pad = True
+    max_dynamic_tiles = 12
+    min_dynamic_tiles = 1
+    use_thumbnail = True
+    pad_during_tiling = False
+    valid_kwargs = Eagle25VLFastImageProcessorKwargs
+    model_input_names = ["pixel_values_videos"]
+
+    def __init__(self, **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs]):
+        super().__init__(**kwargs)
+
+    @add_start_docstrings(
+        # BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS, TODO: this was depreciated from transformers remove!
+        """
+            max_dynamic_tiles (`int`, *optional*):
+                The maximum number of dynamic tiles to use for processing high resolution images.
+            min_dynamic_tiles (`int`, *optional*):
+                The minimum number of dynamic tiles to use for processing high resolution images.
+            use_thumbnail (`bool`, *optional*):
+                Whether to use a thumbnail for processing high resolution images.
+            pad_during_tiling (`bool`, *optional*):
+                Whether to pad the image during tiling.
+            do_pad (`bool`, *optional*):
+                    Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest
+                    number of patches in the batch. Padding will be applied to the bottom and right with zeros.
+        """,
+    )
+
+    # NOTE(YL): we will overload the preprocess method to add the image_flags
+    # def preprocess(
+    #     self, images: ImageInput, **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs]
+    # ) -> BatchFeature:
+    #     return super().preprocess(images, **kwargs)
+
+    def _prepare_images_structure(
+        self,
+        images: ImageInput,
+        expected_ndims: int = 3,
+    ) -> ImageInput:
+        """
+        Prepare the images structure for processing.
+
+        Args:
+            images (`ImageInput`):
+                The input images to process.
+            expected_ndims (`int`, *optional*, defaults to 3):
+                Expected number of dimensions for the images (added for transformers >=4.53.0 compatibility).
+
+        Returns:
+            `ImageInput`: The images with a valid nesting.
+        """
+        return make_flat_list_of_images(images)
+
+    def _resize_for_patching(
+        self,
+        image: torch.Tensor,
+        target_resolution: tuple,
+        interpolation: F.InterpolationMode,
+        input_data_format: ChannelDimension,
+    ) -> torch.Tensor:
+        """
+        Resizes an image to a target resolution while maintaining aspect ratio.
+
+        Args:
+            image ("torch.Tensor"):
+                The input image.
+            target_resolution (tuple):
+                The target resolution (height, width) of the image.
+            interpolation (`InterpolationMode`):
+                Resampling filter to use if resizing the image.
+            input_data_format (`ChannelDimension` or `str`):
+                The channel dimension format of the input image.
+
+        Returns:
+            "torch.Tensor": The resized and padded image.
+        """
+        new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format)
+
+        # Resize the image
+        resized_image = F.resize(image, (new_height, new_width), interpolation=interpolation)
+
+        return resized_image
+
+    def find_closest_aspect_ratio(self, aspect_ratio, target_ratios, width, height, image_size):
+        """
+        previous version mainly focus on ratio.
+        We also consider area ratio here.
+        """
+        best_factor = float("-inf")
+        best_ratio = (1, 1)
+        area = width * height
+        for ratio in target_ratios:
+            target_aspect_ratio = ratio[0] / ratio[1]
+            # ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+            # area_ratio = (ratio[0] * ratio[1] * image_size * image_size) / area
+            """
+            new area > 60% of original image area is enough.
+            """
+            factor_based_on_area_n_ratio = min(
+                (ratio[0] * ratio[1] * image_size * image_size) / area, 0.6
+            ) * min(target_aspect_ratio / aspect_ratio, aspect_ratio / target_aspect_ratio)
+
+            if factor_based_on_area_n_ratio > best_factor:
+                best_factor = factor_based_on_area_n_ratio
+                best_ratio = ratio
+
+        return best_ratio
+
+    def _pad_for_patching(
+        self, image: torch.Tensor, target_resolution: tuple, input_data_format: ChannelDimension
+    ) -> torch.Tensor:
+        """
+        Pad an image to a target resolution while maintaining aspect ratio.
+        """
+        target_height, target_width = target_resolution
+        new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format)
+
+        paste_x = (target_width - new_width) // 2
+        paste_y = (target_height - new_height) // 2
+
+        padded_image = F.pad(image, padding=[paste_x, paste_y, paste_x, paste_y])
+
+        return padded_image
+
+    def _get_image_patches(
+        self,
+        image: torch.Tensor,
+        min_num: int,
+        max_num: int,
+        size: tuple,
+        tile_size: int,
+        use_thumbnail: bool,
+        interpolation: F.InterpolationMode,
+        pad_during_tiling: bool,
+    ) -> list[torch.Tensor]:
+        image_size = get_image_size(image, channel_dim=ChannelDimension.FIRST)
+        orig_height, orig_width = image_size
+        aspect_ratio = orig_width / orig_height
+
+        # calculate the existing image aspect ratio
+        target_ratios = {
+            (i, j)
+            for n in range(min_num, max_num + 1)
+            for i in range(1, n + 1)
+            for j in range(1, n + 1)
+            if i * j <= max_num and i * j >= min_num
+        }
+        target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+        # find the closest aspect ratio to the target
+        target_aspect_ratio = self.find_closest_aspect_ratio(
+            aspect_ratio, target_ratios, orig_width, orig_height, tile_size
+        )
+
+        # calculate the target width and height
+        target_width = tile_size * target_aspect_ratio[0]
+        target_height = tile_size * target_aspect_ratio[1]
+        blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+        if pad_during_tiling:
+            resized_image = self._resize_for_patching(
+                image,
+                (target_height, target_width),
+                interpolation=interpolation,
+                input_data_format=ChannelDimension.FIRST,
+            )
+            padded_image = self._pad_for_patching(
+                resized_image,
+                (target_height, target_width),
+                input_data_format=ChannelDimension.FIRST,
+            )
+            image_used_to_split = padded_image
+        else:
+            image_used_to_split = F.resize(image, (target_height, target_width), interpolation=interpolation)
+
+        processed_tiles = []
+        for i in range(blocks):
+            box = (
+                (i % (target_width // tile_size)) * tile_size,
+                (i // (target_width // tile_size)) * tile_size,
+                ((i % (target_width // tile_size)) + 1) * tile_size,
+                ((i // (target_width // tile_size)) + 1) * tile_size,
+            )
+            # split the image
+            split_img = crop(image_used_to_split, box[0], box[1], box[2], box[3])
+            processed_tiles.append(split_img)
+        assert len(processed_tiles) == blocks
+
+        if use_thumbnail and len(processed_tiles) != 1:
+            thumbnail_img = F.resize(image, (tile_size, tile_size), interpolation=interpolation)
+            processed_tiles.append(thumbnail_img)
+
+        return processed_tiles
+
+    def _pad_for_batching(
+        self,
+        pixel_values: list[torch.Tensor],
+    ) -> list[torch.Tensor]:
+        """
+        Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.
+
+        Args:
+            pixel_values (`List[torch.Tensor]`):
+                An array of pixel values of each images of shape (`batch_size`, `num_patches`, `image_in_3D`)
+
+        Returns:
+            List[`torch.Tensor`]: The padded images.
+        """
+        max_patch = max(len(x) for x in pixel_values)
+        pixel_values = [
+            torch.nn.functional.pad(image, pad=[0, 0, 0, 0, 0, 0, 0, max_patch - image.shape[0]])
+            for image in pixel_values
+        ]
+
+        return pixel_values
+
+    def _preprocess(
+        self,
+        images: list[torch.Tensor],
+        do_resize: bool,
+        size: SizeDict,
+        max_dynamic_tiles: int,
+        min_dynamic_tiles: int,
+        use_thumbnail: bool,
+        pad_during_tiling: bool,
+        interpolation: F.InterpolationMode | None,
+        do_center_crop: bool,
+        crop_size: SizeDict,
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        do_pad: bool,
+        return_tensors: str | TensorType | None,
+        pad_size: SizeDict | None = None,  # Added for transformers >=4.53.0 compatibility
+        disable_grouping: bool | None = None,  # Added for transformers >=4.53.0 compatibility
+    ) -> BatchFeature:
+        processed_images = []
+        image_sizes = []
+        # Determine the size tuple
+        if size and size.height and size.width:
+            size_tuple = (size.height, size.width)
+        else:
+            size_tuple = (size.shortest_edge, size.shortest_edge)
+
+        # Determine the patch size
+        if crop_size and crop_size.height:
+            tile_size = crop_size.height
+        elif size and size.height:
+            tile_size = size.height
+        else:
+            tile_size = size.shortest_edge
+
+        for image in images:
+            image_patches = self._get_image_patches(
+                image,
+                min_num=min_dynamic_tiles,
+                max_num=max_dynamic_tiles,
+                size=size_tuple,
+                tile_size=tile_size,
+                use_thumbnail=use_thumbnail,
+                interpolation=interpolation,
+                pad_during_tiling=pad_during_tiling,
+            )
+
+            # Group images by size for batched processing
+            processed_image_patches_grouped = {}
+            # Added for transformers >=4.53.0 compatibility
+            grouped_image_patches, grouped_image_patches_index = group_images_by_shape(
+                image_patches,
+                disable_grouping=disable_grouping,
+            )
+
+            for shape, stacked_image_patches in grouped_image_patches.items():
+                if do_resize:
+                    stacked_image_patches = self.resize(
+                        image=stacked_image_patches,
+                        size=size,
+                        interpolation=interpolation,
+                    )
+                if do_center_crop:
+                    stacked_image_patches = self.center_crop(stacked_image_patches, crop_size)
+                # Fused rescale and normalize
+                stacked_image_patches = self.rescale_and_normalize(
+                    stacked_image_patches,
+                    do_rescale,
+                    rescale_factor,
+                    do_normalize,
+                    image_mean,
+                    image_std,
+                )
+                processed_image_patches_grouped[shape] = stacked_image_patches
+            processed_image_patches = reorder_images(
+                processed_image_patches_grouped, grouped_image_patches_index
+            )
+            processed_image_patches = (
+                torch.stack(processed_image_patches, dim=0) if return_tensors else processed_image_patches
+            )
+            processed_images.append(processed_image_patches)
+            image_sizes.append(get_image_size(image, ChannelDimension.FIRST))
+
+        if do_pad:
+            processed_images = self._pad_for_batching(processed_images)
+
+        # processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images
+        processed_images = torch.cat(processed_images, dim=0) if return_tensors else processed_images
+        return BatchFeature(
+            data={"pixel_values": processed_images, "image_sizes": image_sizes},
+            tensor_type=return_tensors,
+        )
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        videos: VideoInput = None,
+        **kwargs: Unpack[Eagle25VLFastImageProcessorKwargs],
+    ) -> BatchFeature:
+        validate_kwargs(
+            captured_kwargs=kwargs.keys(),
+            valid_processor_keys=self.valid_kwargs.__annotations__.keys(),
+        )
+        # Set default kwargs from self. This ensures that if a kwarg is not provided
+        # by the user, it gets its default value from the instance, or is set to None.
+        for kwarg_name in self.valid_kwargs.__annotations__:
+            kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None))
+
+        # Extract parameters that are only used for preparing the input images
+        do_convert_rgb = kwargs.pop("do_convert_rgb")
+        input_data_format = kwargs.pop("input_data_format")
+        device = kwargs.pop("device")
+        # Prepare input images
+        # transformers >= 4.53.0: uses _prepare_image_like_inputs instead of _prepare_input_images
+        if images is not None:
+            images = self._prepare_image_like_inputs(
+                images=images,
+                do_convert_rgb=do_convert_rgb,
+                input_data_format=input_data_format,
+                device=device,
+            )
+
+        if videos is not None:
+            videos = self._prepare_image_like_inputs(
+                images=videos,
+                do_convert_rgb=do_convert_rgb,
+                input_data_format=input_data_format,
+                device=device,
+            )
+
+        # Update kwargs that need further processing before being validated
+        kwargs = self._further_process_kwargs(**kwargs)
+
+        # Validate kwargs
+        self._validate_preprocess_kwargs(**kwargs)
+
+        # torch resize uses interpolation instead of resample
+        # Added for transformers >=4.53.0 compatibility
+        resample = kwargs.pop("resample", self.resample)
+        kwargs["interpolation"] = (
+            pil_torch_interpolation_mapping[resample]
+            if isinstance(resample, PILImageResampling | int)
+            else resample
+        )
+
+        # Filter kwargs to only include those accepted by _preprocess
+        valid_preprocess_kwargs = {
+            "do_resize",
+            "size",
+            "max_dynamic_tiles",
+            "min_dynamic_tiles",
+            "use_thumbnail",
+            "pad_during_tiling",
+            "interpolation",
+            "do_center_crop",
+            "crop_size",
+            "do_rescale",
+            "rescale_factor",
+            "do_normalize",
+            "image_mean",
+            "image_std",
+            "do_pad",
+            "return_tensors",
+            "pad_size",
+            "disable_grouping",
+        }
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_preprocess_kwargs}
+        if images is not None:
+            return self._preprocess(images, **filtered_kwargs)
+        elif videos is not None:
+            return self._preprocess(videos, **filtered_kwargs)
+
+
+__all__ = ["Eagle25VLImageProcessorFast"]

src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py

@@ -0,0 +1,396 @@
+# --------------------------------------------------------
+# NVIDIA
+# Copyright (c) 2025 NVIDIA
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import inspect
+
+import torch
+import torch.utils.checkpoint as cp
+from peft import LoraConfig, get_peft_model
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import GenerationConfig
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.llama.modeling_llama import LlamaForCausalLM
+from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM
+from transformers.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM
+from transformers.models.siglip.modeling_siglip import SiglipVisionModel
+from transformers.utils import add_start_docstrings, logging
+
+from .configuration_eagle2_5_vl import Eagle25VLConfig
+
+logger = logging.get_logger(__name__)
+
+
+# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/modeling_llava_onevision.py#L241C1-L280C1
+EAGLE2_5_VL_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`Eagle25VLConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare Eagle2_5_VL Model outputting raw hidden-states without any specific head on top.",
+    EAGLE2_5_VL_START_DOCSTRING,
+)
+class Eagle25VLPreTrainedModel(PreTrainedModel):
+    config_class = Eagle25VLConfig
+    base_model_prefix = "model"
+    main_input_name = "input_ids"
+    supports_gradient_checkpointing = True
+    _no_split_modules = [
+        "Qwen2DecoderLayer",
+        "LlamaDecoderLayer",
+        "Siglip2EncoderLayer",
+        "SiglipEncoderLayer",
+    ]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_cache_class = True
+    _supports_static_cache = True
+    _supports_quantized_cache = True
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear | nn.Conv2d):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class Eagle25VLForConditionalGeneration(Eagle25VLPreTrainedModel, GenerationMixin):
+    config_class = Eagle25VLConfig
+
+    def __init__(self, config: Eagle25VLConfig, vision_model=None, language_model=None):
+        super().__init__(config)
+
+        image_size = config.force_image_size or config.vision_config.image_size
+        patch_size = config.vision_config.patch_size
+        self.patch_size = patch_size
+        if config.use_pixel_shuffle:
+            self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio**2))
+        else:
+            self.num_image_token = int((image_size // patch_size) ** 2)
+
+        self.select_layer = config.select_layer
+        self.downsample_ratio = config.downsample_ratio
+        self.loss_version = config.loss_version
+        self.mlp_checkpoint = config.mlp_checkpoint
+        self.use_pixel_shuffle = config.use_pixel_shuffle
+        self.mlp_connector_layers = config.mlp_connector_layers
+        logger.info(f"num_image_token: {self.num_image_token}")
+        logger.info(f"mlp_checkpoint: {self.mlp_checkpoint}")
+        if vision_model is not None:
+            self.vision_model = vision_model
+        else:
+            if config.vision_config.model_type == "siglip_vision_model":
+                config.vision_config._attn_implementation = "flash_attention_2"
+                self.vision_model = SiglipVisionModel(config.vision_config)
+            else:
+                raise NotImplementedError(f"{config.vision_config.model_type} is not implemented.")
+
+        if language_model is not None:
+            self.language_model = language_model
+        else:
+            if config.text_config.architectures[0] == "LlamaForCausalLM":
+                self.language_model = LlamaForCausalLM(config.text_config)
+            elif config.text_config.architectures[0] == "Phi3ForCausalLM":
+                raise NotImplementedError("Phi3 is not implemented.")
+                # self.language_model = Phi3ForCausalLM(config.text_config)
+            elif config.text_config.architectures[0] == "Qwen2ForCausalLM":
+                assert config.text_config._attn_implementation == "flash_attention_2", (
+                    f"Qwen2 must use flash_attention_2 but got {config.text_config._attn_implementation}"
+                )
+                self.language_model = Qwen2ForCausalLM(config.text_config)
+            elif config.text_config.architectures[0] == "Qwen3ForCausalLM":
+                self.language_model = Qwen3ForCausalLM(config.text_config)
+            else:
+                raise NotImplementedError(f"{config.text_config.architectures[0]} is not implemented.")
+
+        vit_hidden_size = config.vision_config.hidden_size
+        llm_hidden_size = config.text_config.hidden_size
+
+        if config.mlp_connector_layers == 2:
+            self.mlp1 = nn.Sequential(
+                nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
+                nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
+                nn.GELU(),
+                nn.Linear(llm_hidden_size, llm_hidden_size),
+            )
+        elif config.mlp_connector_layers == 1 and config.use_pixel_shuffle:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
+            )
+        elif config.mlp_connector_layers == 1 and not config.use_pixel_shuffle:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(vit_hidden_size, llm_hidden_size),
+            )
+        else:
+            raise NotImplementedError(f"{config.mlp_connector_layers} is not implemented.")
+
+        self.image_token_index = config.image_token_index
+        self.neftune_alpha = None
+
+        if config.use_backbone_lora:
+            self.wrap_backbone_lora(r=config.use_backbone_lora, lora_alpha=2 * config.use_backbone_lora)
+
+        self.use_llm_lora = config.use_llm_lora
+        if config.use_llm_lora:
+            self.wrap_llm_lora(r=config.use_llm_lora, lora_alpha=2 * config.use_llm_lora)
+
+        self.check_forward_kwargs()
+
+    def check_forward_kwargs(self):
+        # We intentionally avoid using **kwargs in forward because Hugging Face Transformers
+        # has special handling for functions with **kwargs parameters that would affect
+        # how our model is processed during training and inference.
+        forward_params = inspect.signature(self.forward).parameters
+        assert not any(k.kind == inspect.Parameter.VAR_KEYWORD for k in forward_params.values())
+
+    def wrap_backbone_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
+        lora_config = LoraConfig(
+            r=r,
+            target_modules=[
+                "self_attn.q_proj",
+                "self_attn.k_proj",
+                "self_attn.v_proj",
+                "self_attn.out_proj",
+                "mlp.fc1",
+                "mlp.fc2",
+            ],
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+        )
+        self.vision_model = get_peft_model(self.vision_model, lora_config)
+        self.vision_model.print_trainable_parameters()
+
+    def wrap_llm_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
+        lora_config = LoraConfig(
+            r=r,
+            target_modules=[
+                "self_attn.q_proj",
+                "self_attn.k_proj",
+                "self_attn.v_proj",
+                "self_attn.o_proj",
+                "mlp.gate_proj",
+                "mlp.down_proj",
+                "mlp.up_proj",
+            ],
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+            task_type="CAUSAL_LM",
+        )
+        self.language_model = get_peft_model(self.language_model, lora_config)
+        self.language_model.enable_input_require_grads()
+        self.language_model.print_trainable_parameters()
+        self.use_llm_lora = True
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        image_flags: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        num_tiles_list: list[torch.Tensor] | None = None,
+    ) -> tuple | CausalLMOutputWithPast:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        vit_embeds = self.extract_feature(pixel_values)
+
+        if image_flags is not None:
+            image_flags = image_flags.view(-1)
+            vit_embeds = vit_embeds[image_flags == 1]
+
+        b, n, c = input_embeds.shape
+        input_embeds = input_embeds.reshape(b * n, c)
+
+        input_ids = input_ids.reshape(b * n)
+        selected = input_ids == self.image_token_index
+        try:
+            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, c)
+        except Exception as e:
+            vit_embeds = vit_embeds.reshape(-1, c)
+            print(
+                f"warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, "
+                f"vit_embeds.shape={vit_embeds.shape}"
+            )
+            n_token = selected.sum()
+            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]
+
+        input_embeds = input_embeds.reshape(b, n, c)
+
+        outputs = self.language_model(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def pixel_shuffle(self, x, scale_factor=0.5):
+        n, w, h, c = x.size()
+        # N, W, H, C --> N, W, H * scale, C // scale
+        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
+        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
+        x = x.permute(0, 2, 1, 3).contiguous()
+        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
+        x = x.view(n, int(h * scale_factor), int(w * scale_factor), int(c / (scale_factor * scale_factor)))
+
+        x = x.permute(0, 2, 1, 3).contiguous()
+        return x
+
+    def extract_feature(self, pixel_values):
+        if self.select_layer == -1:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values, output_hidden_states=False, return_dict=True
+            )
+            if hasattr(vit_embeds, "last_hidden_state"):
+                vit_embeds = vit_embeds.last_hidden_state
+
+        else:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values, output_hidden_states=True, return_dict=True
+            ).hidden_states[self.select_layer]
+
+        if self.use_pixel_shuffle:
+            h = w = int(vit_embeds.shape[1] ** 0.5)
+            vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+            vit_embeds = self.pixel_shuffle(
+                vit_embeds, scale_factor=self.downsample_ratio
+            )  # torch.Size([B, 1024, 1024]) -> torch.Size([B, 16, 16, 4096])
+            vit_embeds = vit_embeds.reshape(
+                vit_embeds.shape[0], -1, vit_embeds.shape[-1]
+            )  # torch.Size([B, 16, 16, 4096]) -> torch.Size([B, 256, 4096])
+
+        if self.mlp_checkpoint and vit_embeds.requires_grad:
+            vit_embeds = cp.checkpoint(self.mlp1, vit_embeds)
+        else:
+            vit_embeds = self.mlp1(vit_embeds)
+
+        return vit_embeds
+
+    @torch.no_grad()
+    def generate(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        input_ids: torch.FloatTensor | None = None,
+        attention_mask: torch.LongTensor | None = None,
+        visual_features: torch.FloatTensor | None = None,
+        generation_config: GenerationConfig | None = None,
+        output_hidden_states: bool | None = None,
+        image_sizes: list[tuple[int, int]] | None = None,
+        **generate_kwargs,
+    ) -> torch.LongTensor:
+        if pixel_values is not None:
+            if visual_features is not None:
+                vit_embeds = visual_features
+            else:
+                vit_embeds = self.extract_feature(pixel_values)
+
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+            b, n, c = input_embeds.shape
+            input_embeds = input_embeds.reshape(b * n, c)
+
+            input_ids = input_ids.reshape(b * n)
+            selected = input_ids == self.config.image_token_index
+            assert selected.sum() != 0
+            input_embeds[selected] = vit_embeds.reshape(-1, c).to(input_embeds.device)
+
+            input_embeds = input_embeds.reshape(b, n, c)
+        else:
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        if "use_cache" not in generate_kwargs:
+            generate_kwargs["use_cache"] = True
+
+        outputs = self.language_model.generate(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            output_hidden_states=output_hidden_states,
+            **generate_kwargs,
+        )
+
+        return outputs
+
+    # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_input_embeddings
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_input_embeddings
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_output_embeddings
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_output_embeddings
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.set_decoder
+    def set_decoder(self, decoder):
+        self.language_model.set_decoder(decoder)
+
+    # Copied from transformers.models.llava_next.modeling_llava_next.LlavaNextForConditionalGeneration.get_decoder
+    def get_decoder(self):
+        return self.language_model.get_decoder()

src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py

@@ -0,0 +1,541 @@
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Processor class for Eagle25VL.
+copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/processing_llava_onevision.py
+"""
+
+import base64
+import os
+import re
+from io import BytesIO
+
+import requests
+import torch
+from PIL import Image
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+from transformers.utils import logging
+from transformers.video_utils import VideoInput
+
+logger = logging.get_logger(__name__)
+
+
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 256
+
+
+def to_rgb(pil_image: Image.Image) -> Image.Image:
+    if pil_image.mode == "RGBA":
+        white_background = Image.new("RGB", pil_image.size, (255, 255, 255))
+        white_background.paste(pil_image, mask=pil_image.split()[3])  # Use alpha channel as mask
+        return white_background
+    else:
+        return pil_image.convert("RGB")
+
+
+def fetch_image(ele: dict[str, str | Image.Image]) -> Image.Image:
+    image = ele["image"] if "image" in ele else ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        response = requests.get(image, stream=True, timeout=10)
+        image_obj = Image.open(BytesIO(response.content))
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(
+            f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}"
+        )
+    image = to_rgb(image_obj)
+    if "scale_factor" in ele:
+        scale_factor = ele["scale_factor"]
+        image = image.resize((image.width * scale_factor, image.height * scale_factor), Image.BILINEAR)
+    return image
+
+
+class Eagle25VLProcessorKwargs(ProcessingKwargs, total=False):
+    # see processing_utils.ProcessingKwargs documentation for usage.
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+        },
+        "images_kwargs": {},
+        "videos_kwargs": {"max_dynamic_tiles": 1},
+    }
+
+
+class Eagle25VLProcessor(ProcessorMixin):
+    r"""
+    Constructs a Eagle25VL processor which wraps a Eagle25VL video processor, Eagle25VL image processor and a Eagle25VL tokenizer into a single processor.
+
+    [`Eagle25VLProcessor`] offers all the functionalities of [`Eagle25VLVideoProcessor`], [`Eagle25VLImageProcessor`] and [`Eagle25VLTokenizer`]. See the
+    [`~Eagle25VLVideoProcessor.__call__`], [`~Eagle25VLProcessor.__call__`] and [`~Eagle25VLProcessor.decode`] for more information.
+
+    Args:
+        image_processor ([`LlavaOnevisionImageProcessor`], *optional*):
+            The image processor is a required input.
+        tokenizer ([`LlamaTokenizerFast`], *optional*):
+            The tokenizer is a required input.
+        num_image_tokens (`int`, *optional*):
+            Number of image tokens for one imagethat will be returned by vision tower.
+        vision_feature_select_strategy (`str`, *optional*):
+            The feature selection strategy used to select the vision feature from the vision backbone.
+            Should be same as in model's config
+        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+            in a chat into a tokenizable string.
+        image_token (`str`, *optional*, defaults to `"<image>"`):
+            Special token used to denote image location.
+        video_token (`str`, *optional*, defaults to `"<video>"`):
+            Special token used to denote video location.
+    """
+
+    attributes = ["image_processor", "tokenizer"]
+    valid_kwargs = [
+        "chat_template",
+        "num_image_tokens",
+        "vision_feature_select_strategy",
+        "image_token",
+        "video_token",
+        "images_kwargs",
+        "videos_kwargs",
+        "text_kwargs",
+    ]
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        vision_feature_select_strategy=None,
+        chat_template=None,
+        image_token="<IMG_CONTEXT>",  # nosec: B107
+        video_token="<IMG_CONTEXT>",  # nosec: B107
+        tokens_per_tile=256,
+        image_placeholder="image",
+        video_placeholder="video",
+        image_start_token="<img>",
+        image_end_token="</img>",
+        **kwargs,
+    ):
+        self.vision_feature_select_strategy = vision_feature_select_strategy
+        self.image_token = tokenizer.image_token if hasattr(tokenizer, "image_token") else image_token
+        self.video_token = tokenizer.video_token if hasattr(tokenizer, "video_token") else video_token
+        self.image_token_id = (
+            tokenizer.image_token_id
+            if getattr(tokenizer, "image_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.image_token)
+        )
+        self.video_token_id = (
+            tokenizer.video_token_id
+            if getattr(tokenizer, "video_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.video_token)
+        )
+        self.image_placeholder = image_placeholder
+        self.video_placeholder = video_placeholder
+        self.tokens_per_tile = tokens_per_tile
+        self.image_start_token = image_start_token
+        self.image_end_token = image_end_token
+        if "auto_map" in kwargs:
+            self.auto_map = kwargs["auto_map"]
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    def replace_media_placeholder(
+        self, text, image_list, video_list, timestamps_list, fps_list, **output_kwargs
+    ):
+        num_of_images_in_this_sample = 0
+        num_of_videos_in_this_sample = 0
+        # Regular expression pattern to match formats like <image-1> or <video-2>
+        pattern = re.compile(rf"<({self.image_placeholder}|{self.video_placeholder})-(\d+)>")
+        unified_frame_list = []
+
+        # image_min_dynamic_tiles = output_kwargs["images_kwargs"].get(
+        #     "min_dynamic_tiles", self.image_processor.min_dynamic_tiles
+        # )
+        # image_max_dynamic_tiles = output_kwargs["images_kwargs"].get(
+        #     "max_dynamic_tiles", self.image_processor.max_dynamic_tiles
+        # )
+        # image_use_thumbnail = output_kwargs["images_kwargs"].get(
+        #     "use_thumbnail", self.image_processor.use_thumbnail
+        # )
+        video_min_dynamic_tiles = output_kwargs["videos_kwargs"].get(
+            "min_dynamic_tiles", self.image_processor.min_dynamic_tiles
+        )
+        video_max_dynamic_tiles = output_kwargs["videos_kwargs"].get(
+            "max_dynamic_tiles", self.image_processor.max_dynamic_tiles
+        )
+        video_use_thumbnail = output_kwargs["videos_kwargs"].get(
+            "use_thumbnail", self.image_processor.use_thumbnail
+        )
+
+        tile_size = self.image_processor.size.get("height", 448)
+
+        # Function to replace tags in a single text
+        def replace_in_text(text):
+            # repl callback function for each match replacement operation
+            def repl(match):
+                nonlocal unified_frame_list
+                nonlocal num_of_images_in_this_sample
+                nonlocal num_of_videos_in_this_sample
+                media_type = match.group(1)  # 'image' or 'video'
+                idx_in_list = int(match.group(2)) - 1  # Convert to list index (0-based)
+                # Select the corresponding path based on media type
+                idx_mapper = {
+                    0: "first",
+                    1: "second",
+                    2: "third",
+                    3: "fourth",
+                    4: "fifth",
+                    5: "sixth",
+                    6: "seventh",
+                    7: "eighth",
+                    8: "ninth",
+                    9: "tenth",
+                }
+                if media_type == "image":
+                    image_inputs = self.image_processor(
+                        images=[image_list[idx_in_list]],
+                        videos=None,
+                        **output_kwargs["images_kwargs"],
+                    )
+                    if isinstance(image_inputs["pixel_values"], list):
+                        _pv = image_inputs["pixel_values"]
+                        if _pv and isinstance(_pv[0], list):
+                            _pv = [t for sub in _pv for t in sub]
+                        image_inputs["pixel_values"] = torch.stack(
+                            [t if isinstance(t, torch.Tensor) else torch.as_tensor(t) for t in _pv]
+                        )
+                    num_all_tiles = image_inputs["pixel_values"].shape[0]
+                    special_placeholder = f"<image {idx_in_list + 1}>{self.image_start_token}{self.image_token * num_all_tiles * self.tokens_per_tile}{self.image_end_token}"
+                    unified_frame_list.append(image_inputs)
+                    num_of_images_in_this_sample += 1
+
+                elif media_type == "video":
+                    video_inputs = self.image_processor(
+                        images=None,
+                        videos=[video_list[idx_in_list]],
+                        **output_kwargs["videos_kwargs"],
+                    )
+                    if isinstance(video_inputs["pixel_values"], list):
+                        _pv = video_inputs["pixel_values"]
+                        if _pv and isinstance(_pv[0], list):
+                            _pv = [t for sub in _pv for t in sub]
+                        video_inputs["pixel_values"] = torch.stack(
+                            [t if isinstance(t, torch.Tensor) else torch.as_tensor(t) for t in _pv]
+                        )
+                    num_all_tiles = video_inputs["pixel_values"].shape[0]
+                    image_sizes = video_inputs["image_sizes"]
+                    if timestamps_list is not None and -1 not in timestamps_list:
+                        frame_timestamps = timestamps_list[idx_in_list]
+                    else:
+                        frame_timestamps = None
+                    sampled_fps = fps_list[idx_in_list] if fps_list is not None else None
+
+                    num_of_tiles_each_frame = [
+                        self.get_number_tiles_based_on_image_size(
+                            image_size,
+                            video_min_dynamic_tiles,
+                            video_max_dynamic_tiles,
+                            video_use_thumbnail,
+                            tile_size,
+                        )
+                        for image_size in image_sizes
+                    ]
+                    assert sum(num_of_tiles_each_frame) == num_all_tiles, (
+                        f"The number of tiles in each frame is not equal to the total number of tiles: {sum(num_of_tiles_each_frame)} != {num_all_tiles}"
+                    )
+
+                    if frame_timestamps is not None:
+                        assert len(frame_timestamps) == len(num_of_tiles_each_frame), (
+                            f"The number of timestamps is not equal to the number of frames: {len(frame_timestamps)} != {len(num_of_tiles_each_frame)}"
+                        )
+                        special_placeholder = [
+                            f"Frame {i + 1} sample at {frame_timestamps[i]:.2f}s: {self.image_start_token}{self.image_token * num_of_tiles * self.tokens_per_tile}{self.image_end_token}"
+                            for i, num_of_tiles in enumerate(num_of_tiles_each_frame)
+                        ]
+                    else:
+                        special_placeholder = [
+                            f"Frame {i + 1}: {self.image_start_token}{self.image_token * num_of_tiles * self.tokens_per_tile}{self.image_end_token}"
+                            for i, num_of_tiles in enumerate(num_of_tiles_each_frame)
+                        ]
+
+                    if sampled_fps is not None:
+                        special_placeholder = (
+                            f"The {idx_mapper[idx_in_list]} video sampled with {sampled_fps:.2f} fps: "
+                            + "".join(special_placeholder)
+                        )
+                    else:
+                        special_placeholder = f"The {idx_mapper[idx_in_list]} video: " + "".join(
+                            special_placeholder
+                        )
+                    unified_frame_list.append(video_inputs)
+                    num_of_videos_in_this_sample += 1
+                else:
+                    raise ValueError(f"Unknown media type: {media_type}")
+                return special_placeholder
+
+            return pattern.sub(repl, text)
+
+        text = replace_in_text(text)
+        if len(unified_frame_list) > 0:
+
+            def _to_tensor(v):
+                if isinstance(v, torch.Tensor):
+                    return v
+                if isinstance(v, list):
+                    if v and isinstance(v[0], list):
+                        v = [t for sub in v for t in sub]
+                    return torch.stack([t if isinstance(t, torch.Tensor) else torch.as_tensor(t) for t in v])
+                return torch.as_tensor(v)
+
+            pixel_values = torch.cat([_to_tensor(frame["pixel_values"]) for frame in unified_frame_list])
+            image_sizes = torch.cat([_to_tensor(frame["image_sizes"]) for frame in unified_frame_list])
+        else:
+            pixel_values = None
+            image_sizes = None
+        return (
+            text,
+            pixel_values,
+            image_sizes,
+            num_of_images_in_this_sample,
+            num_of_videos_in_this_sample,
+        )
+
+    def __call__(
+        self,
+        images: ImageInput = None,
+        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
+        audio=None,
+        videos: VideoInput = None,
+        **kwargs: Unpack[Eagle25VLProcessorKwargs],
+    ) -> BatchFeature:
+        """
+        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+        and `kwargs` arguments to LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to
+        LlavaNextImageProcessor's [`~LlavaNextImageProcessor.__call__`] if `images` is not `None`. Please refer to the docstring
+        of the above two methods for more information.
+
+        Args:
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            text (`str`, `List[str]`, `List[List[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
+
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+            - **pixel_values_videos** -- Pixel values of a video input to be fed to a model. Returned when `videos` is not `None`.
+            - **image_sizes** -- Size of each image that will be used to unpad an image. Returned when `images` is not `None`.
+        """
+
+        output_kwargs = self._merge_kwargs(
+            Eagle25VLProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+
+        if isinstance(text, str):
+            text_list = [text]
+        elif not isinstance(text, list) and not isinstance(text[0], str):
+            raise ValueError("Invalid input text. Please provide a string, or a list of strings")
+        elif isinstance(text, list) and isinstance(text[0], str):
+            text_list = text
+
+        if images is None:
+            images = []
+        if videos is None:
+            videos = []
+
+        pixel_values_list = []
+        image_sizes_list = []
+        new_sample_list = []
+        image_start_idx = 0
+        video_start_idx = 0
+        timestamps_batch = output_kwargs["videos_kwargs"].pop("timestamps", None)
+        fps_batch = output_kwargs["videos_kwargs"].pop("fps", None)
+        for sample in text_list:
+            timestamps_list = timestamps_batch[video_start_idx:] if timestamps_batch is not None else None
+            fps_list = fps_batch[video_start_idx:] if fps_batch is not None else None
+            (
+                sample,
+                pixel_values,
+                image_sizes,
+                num_of_images_in_this_sample,
+                num_of_videos_in_this_sample,
+            ) = self.replace_media_placeholder(
+                sample,
+                images[image_start_idx:],
+                videos[video_start_idx:],
+                timestamps_list,
+                fps_list,
+                **output_kwargs,
+            )
+            new_sample_list.append(sample)
+            if pixel_values is not None:
+                pixel_values_list.append(pixel_values)
+                image_sizes_list.append(image_sizes)
+            image_start_idx += num_of_images_in_this_sample
+            video_start_idx += num_of_videos_in_this_sample
+
+        if len(pixel_values_list) > 0:
+            image_inputs = {
+                "pixel_values": torch.cat(pixel_values_list),
+                "image_sizes": torch.cat(image_sizes_list),
+            }
+        else:
+            image_inputs = {}
+        video_inputs = {}
+        text_inputs = self.tokenizer(new_sample_list, **output_kwargs["text_kwargs"])
+        return BatchFeature(data={**text_inputs, **image_inputs, **video_inputs})
+
+    def get_number_tiles_based_on_image_size(
+        self, image_size: tuple, min_num: int, max_num: int, use_thumbnail: bool, tile_size: int
+    ) -> int:
+        """
+        Get the number of tiles based on the image size.
+        """
+        orig_height, orig_width = image_size
+        aspect_ratio = orig_width / orig_height
+        # calculate the existing image aspect ratio
+        target_ratios = {
+            (i, j)
+            for n in range(min_num, max_num + 1)
+            for i in range(1, n + 1)
+            for j in range(1, n + 1)
+            if i * j <= max_num and i * j >= min_num
+        }
+        target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+        # find the closest aspect ratio to the target
+        target_aspect_ratio = self.image_processor.find_closest_aspect_ratio(
+            aspect_ratio, target_ratios, orig_width, orig_height, tile_size
+        )
+        tiles_num = target_aspect_ratio[0] * target_aspect_ratio[1]
+        if use_thumbnail and tiles_num > 1:
+            tiles_num += 1
+        return tiles_num
+
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Llama
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Llama
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.model_input_names
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+
+    # override to save video-config in a separate config file
+    def save_pretrained(self, save_directory, **kwargs):
+        if os.path.isfile(save_directory):
+            raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file")
+        os.makedirs(save_directory, exist_ok=True)
+
+        outputs = super().save_pretrained(save_directory, **kwargs)
+        return outputs
+
+    # override to load video-config from a separate config file
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        processor = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        # if return_unused_kwargs a tuple is returned where the second element is 'unused_kwargs'
+        if isinstance(processor, tuple):
+            processor = processor[0]
+        return processor
+
+    # Copy from https://github.com/QwenLM/Qwen2.5-VL/blob/main/qwen-vl-utils/src/qwen_vl_utils/vision_process.py
+    def process_vision_info(
+        self,
+        conversations: list[dict] | list[list[dict]],
+        return_video_kwargs: bool = False,
+    ) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None, dict | None]:
+        vision_infos = self.extract_vision_info(conversations)
+        ## Read images or videos
+        image_inputs = []
+        video_inputs = []
+        video_sample_fps_list = []
+        video_timestamps_list = []
+        for vision_info in vision_infos:
+            if "image" in vision_info or "image_url" in vision_info:
+                image_inputs.append(fetch_image(vision_info))
+            else:
+                raise ValueError("image, image_url or video should in content.")
+        if len(image_inputs) == 0:
+            image_inputs = None
+        if len(video_inputs) == 0:
+            video_inputs = None
+        if return_video_kwargs:
+            return (
+                image_inputs,
+                video_inputs,
+                {"fps": video_sample_fps_list, "timestamps": video_timestamps_list},
+            )
+        return image_inputs, video_inputs
+
+    def extract_vision_info(self, conversations: list[dict] | list[list[dict]]) -> list[dict]:
+        vision_infos = []
+        if isinstance(conversations[0], dict):
+            conversations = [conversations]
+        for conversation in conversations:
+            for message in conversation:
+                if isinstance(message["content"], list):
+                    for ele in message["content"]:
+                        if (
+                            "image" in ele
+                            or "image_url" in ele
+                            or "video" in ele
+                            or ele["type"] in ("image", "image_url", "video")
+                        ):
+                            vision_infos.append(ele)
+        return vision_infos
+
+
+__all__ = ["Eagle25VLProcessor"]

src/lerobot/policies/groot/groot_n1.py

@@ -0,0 +1,380 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+from huggingface_hub.errors import HFValidationError, RepositoryNotFoundError
+
+from lerobot.utils.import_utils import _transformers_available
+
+# Conditional import for type checking and lazy loading
+if TYPE_CHECKING or _transformers_available:
+    from huggingface_hub.dataclasses import strict
+    from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
+    from transformers.feature_extraction_utils import BatchFeature
+else:
+
+    def strict(cls):
+        return cls
+
+    AutoConfig = None
+    AutoModel = None
+    PretrainedConfig = object
+    PreTrainedModel = object
+    BatchFeature = None
+
+try:
+    import tree
+except ImportError:
+    tree = None
+
+from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME
+
+from .action_head.flow_matching_action_head import (
+    FlowmatchingActionHead,
+    FlowmatchingActionHeadConfig,
+)
+from .utils import ensure_eagle_cache_ready
+
+DEFAULT_VENDOR_EAGLE_PATH = str((Path(__file__).resolve().parent / "eagle2_hg_model").resolve())
+DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5"
+
+
+class EagleBackbone(nn.Module):
+    def __init__(
+        self,
+        tune_llm: bool = False,
+        tune_visual: bool = False,
+        select_layer: int = -1,
+        reproject_vision: bool = False,
+        use_flash_attention: bool = False,
+        load_bf16: bool = False,
+        eagle_path: str = DEFAULT_VENDOR_EAGLE_PATH,
+        tokenizer_assets_repo: str = DEFAULT_TOKENIZER_ASSETS_REPO,
+        project_to_dim: int = 1536,
+    ):
+        """
+        Args:
+            tune_llm: whether to tune the LLM model (default: True)
+            tune_visual: whether to tune the visual model (default: False)
+        """
+        super().__init__()
+        assert not reproject_vision, "Reproject vision is not implemented here, set to False"
+
+        # Prefer loading Eagle model config from the cache directory where vendor files were copied.
+        vendor_dir = DEFAULT_VENDOR_EAGLE_PATH
+        cache_dir = HF_LEROBOT_HOME / tokenizer_assets_repo
+        try:
+            ensure_eagle_cache_ready(vendor_dir, cache_dir, tokenizer_assets_repo)
+        except Exception as exc:  # nosec: B110
+            print(f"[GROOT] Warning: failed to prepare Eagle cache for backbone: {exc}")
+
+        config = AutoConfig.from_pretrained(str(cache_dir), trust_remote_code=True)
+        self.eagle_model = AutoModel.from_config(config, trust_remote_code=True)
+
+        if project_to_dim is not None:
+            self.eagle_linear = torch.nn.Linear(2048, project_to_dim)
+        else:
+            self.eagle_linear = torch.nn.Identity()
+
+        # needed since we don't use these layers. Also saves compute
+        while len(self.eagle_model.language_model.model.layers) > select_layer:
+            self.eagle_model.language_model.model.layers.pop(-1)
+
+        self.select_layer = select_layer
+        self.set_trainable_parameters(tune_llm, tune_visual)
+
+    def set_trainable_parameters(self, tune_llm: bool, tune_visual: bool):
+        self.tune_llm = tune_llm
+        self.tune_visual = tune_visual
+        for p in self.parameters():
+            p.requires_grad = True
+        if not tune_llm:
+            self.eagle_model.language_model.requires_grad_(False)
+        if not tune_visual:
+            self.eagle_model.vision_model.requires_grad_(False)
+            self.eagle_model.mlp1.requires_grad_(False)
+        print(f"Tune backbone llm: {self.tune_llm}")
+        print(f"Tune backbone visual: {self.tune_visual}")
+        # Check if any parameters are still trainable. If not, print a warning.
+        if not tune_llm and not tune_visual:
+            for name, p in self.named_parameters():
+                if p.requires_grad:
+                    print(f"Backbone trainable parameter: {name}")
+        if not any(p.requires_grad for p in self.parameters()):
+            print("Warning: No backbone trainable parameters found.")
+
+    def set_frozen_modules_to_eval_mode(self):
+        """
+        Huggingface will call model.train() at each training_step. To ensure
+        the expected behaviors for modules like dropout, batchnorm, etc., we
+        need to call model.eval() for the frozen modules.
+        """
+        if self.training:
+            if self.eagle_model.language_model and not self.tune_llm:
+                self.eagle_model.language_model.eval()
+            if self.eagle_model.vision_model and not self.tune_visual:
+                self.eagle_model.vision_model.eval()
+
+    def prepare_input(self, batch: dict) -> BatchFeature:
+        return BatchFeature(data=batch)
+
+    def forward_eagle(self, vl_input: BatchFeature) -> BatchFeature:
+        eagle_prefix = "eagle_"
+        eagle_input = {
+            k.removeprefix(eagle_prefix): v for k, v in vl_input.items() if k.startswith(eagle_prefix)
+        }
+        del eagle_input["image_sizes"]
+
+        eagle_output = self.eagle_model(**eagle_input, output_hidden_states=True, return_dict=True)
+        eagle_features = eagle_output.hidden_states[self.select_layer]
+
+        eagle_features = self.eagle_linear(eagle_features)
+        return eagle_features, eagle_input["attention_mask"]
+
+    def forward(self, vl_input: BatchFeature) -> BatchFeature:
+        self.set_frozen_modules_to_eval_mode()
+
+        eagle_embeds, eagle_mask = self.forward_eagle(vl_input)
+
+        # YL (TODO HACK): to resolve DDP issue when tune_visual=True
+        # Ensure all trainable parameters in vision_model are used in the forward pass for DDP compatibility
+        if self.training and self.tune_visual:
+            dummy_term = torch.tensor(
+                0.0, device=eagle_embeds.device, dtype=eagle_embeds.dtype, requires_grad=True
+            )
+            for param in self.eagle_model.vision_model.parameters():
+                if param.requires_grad:
+                    dummy_term = dummy_term + 0.0 * param.sum()
+            eagle_embeds = eagle_embeds + dummy_term
+
+        return BatchFeature(
+            data={"backbone_features": eagle_embeds, "backbone_attention_mask": eagle_mask}
+        )  # [B, T2, hidden_size]
+
+
+BACKBONE_FEATURE_KEY = "backbone_features"
+ACTION_KEY = "action_pred"
+LOSS_KEY = "loss"
+ERROR_MSG = "Error: unexpected input/output"
+N_COLOR_CHANNELS = 3
+
+
+# config
+@strict
+class GR00TN15Config(PretrainedConfig):
+    model_type = "gr00t_n1_5"
+
+    backbone_cfg: dict[str, Any] | None = None
+    action_head_cfg: dict[str, Any] | None = None
+    action_horizon: int = 0
+    action_dim: int = 0
+    compute_dtype: str = "float32"
+
+    def __post_init__(self, **kwargs):
+        self.backbone_cfg = {} if self.backbone_cfg is None else self.backbone_cfg
+        self.action_head_cfg = {} if self.action_head_cfg is None else self.action_head_cfg
+        super().__post_init__(**kwargs)
+
+
+# real model
+class GR00TN15(PreTrainedModel):
+    supports_gradient_checkpointing = True
+    config_class = GR00TN15Config
+    """
+    we expect the backbone output to have a key 'backbone_features' with shape (batch_size, n, hidden_size)
+    here n is variable and can be e.g. time, 1 or user specified
+    we expect the action head output to have a key 'action_pred' with shape (batch_size, time, action_dim) during inference time
+    we expect these to have type BatchFeature, and they can of course have many other user specified keys too
+    """
+
+    def __init__(
+        self,
+        config: GR00TN15Config,
+        local_model_path: str,
+    ):
+        assert isinstance(config.backbone_cfg, dict)
+        assert isinstance(config.action_head_cfg, dict)
+
+        super().__init__(config)
+        self.local_model_path = local_model_path
+
+        self.backbone = EagleBackbone(**config.backbone_cfg)
+        action_head_cfg = FlowmatchingActionHeadConfig(**config.action_head_cfg)
+        self.action_head = FlowmatchingActionHead(action_head_cfg)
+
+        self.action_horizon = config.action_horizon
+        self.action_dim = config.action_dim
+        self.compute_dtype = config.compute_dtype
+        self.post_init()
+
+    def validate_inputs(self, inputs):
+        # NOTE -- this should be handled internally by the model
+        # however, doing that will likely be breaking changes -- so we'll need to do it after the deadline
+
+        detected_error = False
+        error_msg = ERROR_MSG
+        if ACTION in inputs:
+            action = inputs[ACTION]
+            # In inference, action may be omitted or None; validate only when it's a tensor.
+            if action is None:
+                pass  # allow None during inference
+            elif isinstance(action, torch.Tensor):
+                shape_ok = (
+                    len(action.shape) == 3
+                    and action.shape[1] == self.action_horizon
+                    and action.shape[2] == self.action_dim
+                )
+                if not shape_ok:
+                    error_msg += f"\n{action.shape=}"
+                    detected_error = True
+            else:
+                # Unexpected non-tensor type provided for action
+                error_msg += f"\nInvalid type for action: {type(action)}"
+                detected_error = True
+
+        if "video" in inputs:
+            video = inputs["video"]
+            type_ok = isinstance(video, np.ndarray)
+            dtype_ok = video.dtype == np.uint8
+            shape_ok = len(video.shape) == 6 and video.shape[3] == N_COLOR_CHANNELS
+            if not type_ok:
+                error_msg += f"\n{type(video)=}"
+                detected_error = True
+            if not dtype_ok:
+                error_msg += f"\n{video.dtype=}"
+                detected_error = True
+            if not shape_ok:
+                error_msg += f"\n{video.shape=}"
+                detected_error = True
+
+        if detected_error:
+            raise ValueError(error_msg)
+
+    def validate_data(self, action_head_outputs, backbone_outputs, is_training):
+        fail_backbone = (
+            not isinstance(backbone_outputs, BatchFeature) or BACKBONE_FEATURE_KEY not in backbone_outputs
+        )
+
+        if fail_backbone:
+            error_msg = ERROR_MSG
+            error_msg += f"\n{isinstance(backbone_outputs, BatchFeature)=}"
+            error_msg += f"\n{BACKBONE_FEATURE_KEY in backbone_outputs=}"
+            error_msg += f"\n{backbone_outputs[BACKBONE_FEATURE_KEY].shape=}"
+            raise ValueError(error_msg)
+
+        fail_action_head = (not isinstance(action_head_outputs, BatchFeature)) or not (
+            (
+                LOSS_KEY in action_head_outputs and is_training
+            )  # there might not be an action prediction during training
+            or (
+                ACTION_KEY in action_head_outputs
+                and action_head_outputs[ACTION_KEY].shape[1] == self.action_horizon
+                and action_head_outputs[ACTION_KEY].shape[2] == self.action_dim
+            )
+        )
+
+        if fail_action_head:
+            error_msg = ERROR_MSG
+            error_msg += f"\n{isinstance(action_head_outputs, BatchFeature)=}"
+            error_msg += f"\n{LOSS_KEY in action_head_outputs=}"
+            error_msg += f"\n{action_head_outputs[ACTION_KEY].shape=}"
+            error_msg += f"\n{self.action_horizon=}"
+            error_msg += f"\n{self.action_dim=}"
+            raise ValueError(error_msg)
+
+    def forward(
+        self,
+        inputs: dict,
+    ) -> BatchFeature:
+        backbone_inputs, action_inputs = self.prepare_input(inputs)
+        backbone_outputs = self.backbone(backbone_inputs)
+        action_head_outputs = self.action_head(backbone_outputs, action_inputs)
+        self.validate_data(action_head_outputs, backbone_outputs, is_training=True)
+        return action_head_outputs
+
+    def get_action(
+        self,
+        inputs: dict,
+    ) -> BatchFeature:
+        backbone_inputs, action_inputs = self.prepare_input(inputs)
+        # Because the behavior of backbones remains the same for training and inference, we can use `forward` for backbones.
+        backbone_outputs = self.backbone(backbone_inputs)
+        action_head_outputs = self.action_head.get_action(backbone_outputs, action_inputs)
+        self.validate_data(action_head_outputs, backbone_outputs, is_training=False)
+        return action_head_outputs
+
+    def prepare_input(self, inputs) -> tuple[BatchFeature, BatchFeature]:
+        self.validate_inputs(inputs)
+        backbone_inputs = self.backbone.prepare_input(inputs)
+        action_inputs = self.action_head.prepare_input(inputs)
+
+        def to_device_with_maybe_dtype(x):
+            # Cast floating tensors to a memory-efficient compute dtype when requested.
+            # Rationale: Upcasting backbone activations to fp32 significantly increases VRAM.
+            # When compute_dtype is bfloat16, prefer bf16 for activations to match AMP behavior.
+            if not isinstance(x, torch.Tensor):
+                return x
+            if torch.is_floating_point(x):
+                if getattr(self, "compute_dtype", None) == "bfloat16":
+                    return x.to(self.device, dtype=torch.bfloat16)
+                # Fallback: preserve previous behavior if not using bf16 compute
+                return x.to(self.device, dtype=self.action_head.dtype)
+            # Non-floating tensors: move device only
+            return x.to(self.device)
+
+        backbone_inputs = tree.map_structure(to_device_with_maybe_dtype, backbone_inputs)
+        action_inputs = tree.map_structure(to_device_with_maybe_dtype, action_inputs)
+        return backbone_inputs, action_inputs
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        tune_visual = kwargs.pop("tune_visual", True)
+        tune_llm = kwargs.pop("tune_llm", False)
+        tune_projector = kwargs.pop("tune_projector", True)
+        tune_diffusion_model = kwargs.pop("tune_diffusion_model", True)
+
+        print(f"Loading pretrained dual brain from {pretrained_model_name_or_path}")
+        print(f"Tune backbone vision tower: {tune_visual}")
+        print(f"Tune backbone LLM: {tune_llm}")
+        print(f"Tune action head projector: {tune_projector}")
+        print(f"Tune action head DiT: {tune_diffusion_model}")
+
+        # get the current model path being downloaded
+        try:
+            # NOTE(YL) This downloads the model to the local cache and returns the local path to the model
+            # saved in ~/.cache/huggingface/hub/
+            local_model_path = snapshot_download(pretrained_model_name_or_path, repo_type="model")
+            # HFValidationError, RepositoryNotFoundError
+        except (HFValidationError, RepositoryNotFoundError):
+            print(
+                f"Model not found or avail in the huggingface hub. Loading from local path: {pretrained_model_name_or_path}"
+            )
+            local_model_path = pretrained_model_name_or_path
+
+        pretrained_model = super().from_pretrained(
+            local_model_path, local_model_path=local_model_path, **kwargs
+        )
+
+        pretrained_model.backbone.set_trainable_parameters(tune_visual=tune_visual, tune_llm=tune_llm)
+        pretrained_model.action_head.set_trainable_parameters(
+            tune_projector=tune_projector, tune_diffusion_model=tune_diffusion_model
+        )
+        return pretrained_model

src/lerobot/policies/groot/groot_n1_7.py

@@ -1,933 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import importlib
-import logging
-from contextlib import suppress
-from copy import deepcopy
-from typing import TYPE_CHECKING, Any
-
-import torch
-import torch.nn.functional as F  # noqa: N812
-from huggingface_hub import snapshot_download
-from huggingface_hub.errors import HFValidationError, RepositoryNotFoundError
-from torch import nn
-from torch.distributions import Beta
-
-from lerobot.utils.import_utils import _transformers_available, require_package
-
-from .action_head.cross_attention_dit import AlternateVLDiT, DiT, SelfAttentionTransformer
-from .configuration_groot import N1_7_DEFAULT_IMAGE_CROP_SIZE, N1_7_DEFAULT_IMAGE_TARGET_SIZE
-
-if TYPE_CHECKING or _transformers_available:
-    from transformers import (
-        AutoConfig,
-        AutoModel,
-        PretrainedConfig,
-        PreTrainedModel,
-        Qwen3VLConfig,
-        Qwen3VLForConditionalGeneration,
-    )
-    from transformers.feature_extraction_utils import BatchFeature
-else:
-    AutoConfig = None
-    AutoModel = None
-    PretrainedConfig = object
-    PreTrainedModel = object
-    BatchFeature = None
-    Qwen3VLConfig = None
-    Qwen3VLForConditionalGeneration = None
-
-try:
-    import tree
-except ImportError:
-    tree = None
-
-logger = logging.getLogger(__name__)
-
-
-GR00T_N1_7_DEFAULTS: dict[str, Any] = {
-    "model_dtype": "bfloat16",
-    "dtype": "bfloat16",
-    "model_name": "nvidia/Cosmos-Reason2-2B",
-    "backbone_model_type": "qwen",
-    "model_revision": None,
-    "tune_top_llm_layers": 0,
-    "backbone_embedding_dim": 2048,
-    "tune_llm": False,
-    "tune_visual": False,
-    "select_layer": 16,
-    "reproject_vision": False,
-    "use_flash_attention": False,
-    "load_bf16": False,
-    "backbone_trainable_params_fp32": True,
-    "image_crop_size": N1_7_DEFAULT_IMAGE_CROP_SIZE,
-    "image_target_size": N1_7_DEFAULT_IMAGE_TARGET_SIZE,
-    "shortest_image_edge": None,
-    "crop_fraction": None,
-    "random_rotation_angle": None,
-    "color_jitter_params": None,
-    "use_albumentations_transforms": True,
-    "extra_augmentation_config": None,
-    "formalize_language": True,
-    "apply_sincos_state_encoding": False,
-    "use_percentiles": True,
-    "use_relative_action": False,
-    "max_state_dim": 132,
-    "max_action_dim": 132,
-    "action_horizon": 40,
-    "hidden_size": 1024,
-    "input_embedding_dim": 1536,
-    "state_history_length": 1,
-    "add_pos_embed": True,
-    "attn_dropout": 0.2,
-    "use_vlln": True,
-    "max_seq_len": 1024,
-    "use_alternate_vl_dit": True,
-    "attend_text_every_n_blocks": 2,
-    "diffusion_model_cfg": {
-        "positional_embeddings": None,
-        "num_layers": 32,
-        "num_attention_heads": 32,
-        "attention_head_dim": 48,
-        "norm_type": "ada_norm",
-        "dropout": 0.2,
-        "final_dropout": True,
-        "output_dim": 1024,
-        "interleave_self_attention": True,
-    },
-    "vl_self_attention_cfg": {
-        "positional_embeddings": None,
-        "num_layers": 4,
-        "num_attention_heads": 32,
-        "attention_head_dim": 64,
-        "dropout": 0.2,
-        "final_dropout": True,
-    },
-    "num_inference_timesteps": 4,
-    "noise_beta_alpha": 1.5,
-    "noise_beta_beta": 1.0,
-    "noise_s": 0.999,
-    "num_timestep_buckets": 1000,
-    "tune_projector": True,
-    "tune_diffusion_model": True,
-    "tune_vlln": True,
-    "state_dropout_prob": 0.2,
-    "exclude_state": False,
-    "use_mean_std": False,
-    "max_num_embodiments": 32,
-    "rtc_ramp_rate": 6.0,
-}
-
-
-class GR00TN17Config(PretrainedConfig):
-    """Configuration for NVIDIA GR00T N1.7.
-
-    N1.7 uses the Cosmos-Reason2-2B / Qwen3-VL backbone and a multi-embodiment
-    flow-matching action head. This mirrors the public N1.7 checkpoint config
-    while keeping it local to LeRobot and independent from the external
-    Isaac-GR00T ``gr00t`` Python package.
-    """
-
-    model_type = "Gr00tN1d7"
-
-    _defaults = GR00T_N1_7_DEFAULTS
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        for key, value in GR00T_N1_7_DEFAULTS.items():
-            setattr(self, key, deepcopy(kwargs.pop(key, value)))
-        for key, value in kwargs.items():
-            setattr(self, key, value)
-
-
-class CategorySpecificLinear(nn.Module):
-    """Linear layer with category-specific weights for multi-embodiment support."""
-
-    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int):
-        super().__init__()
-        self.num_categories = num_categories
-        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
-        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))
-
-    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
-        selected_w = self.W[cat_ids]
-        selected_b = self.b[cat_ids]
-        return torch.bmm(x, selected_w) + selected_b.unsqueeze(1)
-
-
-class CategorySpecificMLP(nn.Module):
-    """Two-layer MLP with category-specific weights."""
-
-    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int, output_dim: int):
-        super().__init__()
-        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
-        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)
-
-    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
-        hidden = F.relu(self.layer1(x, cat_ids))
-        return self.layer2(hidden, cat_ids)
-
-
-class SinusoidalPositionalEncoding(nn.Module):
-    """Sinusoidal encoding of shape ``(B, T, D)`` for timestep tensors ``(B, T)``.
-
-    The frequency scalar is intentionally created on CPU and then broadcast with
-    the device-local arange result. That mirrors Isaac-GR00T's N1.7 timestep
-    embedding and avoids tiny dtype/device construction differences in parity
-    tests.
-    """
-
-    def __init__(self, embedding_dim: int):
-        super().__init__()
-        self.embedding_dim = embedding_dim
-
-    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
-        timesteps = timesteps.float()
-        half_dim = self.embedding_dim // 2
-        exponent = -torch.arange(half_dim, dtype=torch.float, device=timesteps.device) * (
-            torch.log(torch.tensor(10000.0)) / half_dim
-        )
-        freqs = timesteps.unsqueeze(-1) * exponent.exp()
-        return torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1)
-
-
-def swish(x: torch.Tensor) -> torch.Tensor:
-    return x * torch.sigmoid(x)
-
-
-class MultiEmbodimentActionEncoder(nn.Module):
-    """Action encoder with category-specific projections and sinusoidal time encoding."""
-
-    def __init__(self, action_dim: int, hidden_size: int, num_embodiments: int):
-        super().__init__()
-        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)
-        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)
-        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)
-        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)
-
-    def forward(self, actions: torch.Tensor, timesteps: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
-        batch_size, horizon, _ = actions.shape
-        if timesteps.dim() != 1 or timesteps.shape[0] != batch_size:
-            raise ValueError("Expected `timesteps` to have shape (B,).")
-        timesteps = timesteps.unsqueeze(1).expand(-1, horizon)
-        action_emb = self.W1(actions, cat_ids)
-        time_emb = self.pos_encoding(timesteps).to(dtype=action_emb.dtype)
-        x = swish(self.W2(torch.cat([action_emb, time_emb], dim=-1), cat_ids))
-        return self.W3(x, cat_ids)
-
-
-class Qwen3Backbone(nn.Module):
-    """Cosmos-Reason2/Qwen3-VL backbone used by GR00T N1.7.
-
-    The public checkpoint stores the action head in the GR00T checkpoint but
-    uses a Hugging Face Qwen3-VL-compatible backbone interface. This wrapper
-    keeps the nested HF module layout compatible across transformer versions
-    and exposes the hidden states consumed by the action head.
-    """
-
-    def __init__(
-        self,
-        model_name: str = "nvidia/Cosmos-Reason2-2B",
-        tune_llm: bool = False,
-        tune_visual: bool = False,
-        select_layer: int = -1,
-        reproject_vision: bool = False,
-        use_flash_attention: bool = False,
-        load_bf16: bool = False,
-        tune_top_llm_layers: int = 0,
-        trainable_params_fp32: bool = False,
-        transformers_loading_kwargs: dict[str, Any] | None = None,
-        load_pretrained_weights: bool = True,
-    ):
-        require_package("transformers", extra="groot")
-        super().__init__()
-        transformers_loading_kwargs = transformers_loading_kwargs or {"trust_remote_code": True}
-
-        extra_kwargs: dict[str, Any] = {}
-        if use_flash_attention:
-            try:
-                import flash_attn  # noqa: F401
-
-                extra_kwargs["attn_implementation"] = "flash_attention_2"
-            except ImportError:
-                logger.warning("flash_attn is not installed. Falling back to SDPA attention.")
-                extra_kwargs["attn_implementation"] = "sdpa"
-        if load_bf16:
-            extra_kwargs["torch_dtype"] = torch.bfloat16
-
-        if load_pretrained_weights:
-            self.model = Qwen3VLForConditionalGeneration.from_pretrained(
-                model_name,
-                **extra_kwargs,
-                **transformers_loading_kwargs,
-            ).eval()
-        else:
-            self.model = self._from_backbone_config(
-                model_name=model_name,
-                model_kwargs=extra_kwargs,
-                config_kwargs=transformers_loading_kwargs,
-            ).eval()
-
-        while len(self.language_model.layers) > select_layer:
-            self.language_model.layers.pop(-1)
-
-        self.select_layer = select_layer
-        self.set_trainable_parameters(tune_llm, tune_visual, tune_top_llm_layers)
-        if load_bf16 and trainable_params_fp32:
-            for parameter in self.parameters():
-                if parameter.requires_grad:
-                    parameter.data = parameter.data.to(torch.float32)
-
-    def set_trainable_parameters(
-        self, tune_llm: bool, tune_visual: bool, tune_top_llm_layers: int = 0
-    ) -> None:
-        self.tune_llm = tune_llm
-        self.tune_visual = tune_visual
-        for parameter in self.parameters():
-            parameter.requires_grad = True
-        if not tune_llm:
-            self.language_model.requires_grad_(False)
-        if not tune_visual:
-            self.visual.requires_grad_(False)
-        if tune_top_llm_layers > 0:
-            for layer in self.language_model.layers[-tune_top_llm_layers:]:
-                for parameter in layer.parameters():
-                    parameter.requires_grad = True
-
-    def set_frozen_modules_to_eval_mode(self) -> None:
-        if self.training:
-            if self.language_model and not self.tune_llm:
-                self.language_model.eval()
-            if self.visual and not self.tune_visual:
-                self.visual.eval()
-
-    @property
-    def language_model(self) -> nn.Module:
-        return getattr(self.model, "model", self.model).language_model
-
-    @property
-    def visual(self) -> nn.Module:
-        return getattr(self.model, "model", self.model).visual
-
-    def _from_backbone_config(
-        self,
-        *,
-        model_name: str,
-        model_kwargs: dict[str, Any],
-        config_kwargs: dict[str, Any],
-    ) -> nn.Module:
-        if _is_cosmos_reason2_backbone(model_name):
-            backbone_config = _cosmos_reason2_qwen3_vl_config()
-        else:
-            backbone_config = AutoConfig.from_pretrained(model_name, **config_kwargs)
-        return Qwen3VLForConditionalGeneration._from_config(backbone_config, **model_kwargs)
-
-    def prepare_input(self, batch: dict[str, Any]) -> BatchFeature:
-        return BatchFeature(data=batch)
-
-    def _ensure_mm_token_type_ids(self, model_input: dict[str, torch.Tensor]) -> None:
-        if "mm_token_type_ids" in model_input:
-            return
-        if "image_grid_thw" not in model_input and "video_grid_thw" not in model_input:
-            return
-
-        input_ids = model_input.get("input_ids")
-        if input_ids is None:
-            return
-
-        mm_token_type_ids = torch.zeros(input_ids.shape, dtype=torch.int32, device=input_ids.device)
-        image_token_id = getattr(self.model.config, "image_token_id", None)
-        video_token_id = getattr(self.model.config, "video_token_id", None)
-        if image_token_id is not None:
-            mm_token_type_ids[input_ids == image_token_id] = 1
-        if video_token_id is not None:
-            mm_token_type_ids[input_ids == video_token_id] = 2
-
-        model_input["mm_token_type_ids"] = mm_token_type_ids
-
-    def _ensure_legacy_qwen3_position_ids(self, model_input: dict[str, torch.Tensor]) -> None:
-        """Restore the Qwen3-VL text position ids used by older Transformers releases.
-
-        Transformers 5.x computes 3-row multimodal RoPE ids for Qwen3-VL and then
-        drops text position ids before calling text-layer flash attention. GR00T
-        N1.7 was aligned against the older Transformers path, where a fourth text
-        position row is forwarded alongside the temporal/height/width rows. Adding
-        the row here preserves the newer multimodal position computation while
-        keeping flash attention on the legacy code path.
-        """
-
-        if "position_ids" in model_input:
-            return
-
-        qwen3_model = getattr(self.model, "model", self.model)
-        compute_3d_position_ids = getattr(qwen3_model, "compute_3d_position_ids", None)
-        if compute_3d_position_ids is None:
-            return
-
-        position_ids = compute_3d_position_ids(
-            input_ids=model_input.get("input_ids"),
-            image_grid_thw=model_input.get("image_grid_thw"),
-            video_grid_thw=model_input.get("video_grid_thw"),
-            inputs_embeds=None,
-            attention_mask=model_input.get("attention_mask"),
-            past_key_values=None,
-            mm_token_type_ids=model_input.get("mm_token_type_ids"),
-        )
-        if position_ids.ndim == 3 and position_ids.shape[0] == 3:
-            position_ids = torch.cat([position_ids[:1], position_ids], dim=0)
-
-        model_input["position_ids"] = position_ids
-
-    def _last_decoder_layer_output(self, model_input: dict[str, torch.Tensor]) -> torch.Tensor:
-        """Return the pre-final-norm decoder output consumed by the N1.7 action head.
-
-        Older Transformers releases exposed this tensor as ``hidden_states[-1]``.
-        Newer releases expose the post-final-norm tensor there instead. Capturing
-        the last decoder layer output directly keeps the N1.7 action head input
-        stable across Transformers versions.
-        """
-
-        captured: dict[str, torch.Tensor] = {}
-
-        def capture_output(_module: nn.Module, _inputs: tuple[Any, ...], output: Any) -> None:
-            if isinstance(output, torch.Tensor):
-                captured["features"] = output
-            elif isinstance(output, (tuple, list)) and output:
-                captured["features"] = output[0]
-            elif hasattr(output, "last_hidden_state"):
-                captured["features"] = output.last_hidden_state
-
-        hook = self.language_model.layers[-1].register_forward_hook(capture_output)
-        try:
-            outputs = self.model(**model_input, output_hidden_states=True)
-        finally:
-            hook.remove()
-
-        return captured.get("features", outputs.hidden_states[-1])
-
-    def forward(self, vl_input: BatchFeature) -> BatchFeature:
-        self.set_frozen_modules_to_eval_mode()
-        keys_to_use = ["input_ids", "attention_mask", "pixel_values", "image_grid_thw"]
-        optional_keys = ["mm_token_type_ids", "pixel_values_videos", "video_grid_thw"]
-        model_input = {key: vl_input[key] for key in keys_to_use}
-        model_input.update({key: vl_input[key] for key in optional_keys if key in vl_input})
-        self._ensure_mm_token_type_ids(model_input)
-        self._ensure_legacy_qwen3_position_ids(model_input)
-        features = self._last_decoder_layer_output(model_input)
-        image_mask = model_input["input_ids"] == self.model.config.image_token_id
-        attention_mask = model_input["attention_mask"] == 1
-        return BatchFeature(
-            data={
-                "backbone_features": features,
-                "backbone_attention_mask": attention_mask,
-                "image_mask": image_mask,
-            }
-        )
-
-
-class GR00TN17ActionHead(nn.Module):
-    supports_gradient_checkpointing = True
-
-    def __init__(self, config: GR00TN17Config):
-        require_package("diffusers", extra="groot")
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.input_embedding_dim = config.input_embedding_dim
-
-        if config.use_alternate_vl_dit:
-            self.model = AlternateVLDiT(
-                **config.diffusion_model_cfg,
-                cross_attention_dim=config.backbone_embedding_dim,
-                attend_text_every_n_blocks=config.attend_text_every_n_blocks,
-            )
-        else:
-            self.model = DiT(
-                **config.diffusion_model_cfg,
-                cross_attention_dim=config.backbone_embedding_dim,
-            )
-
-        self.action_dim = config.max_action_dim
-        self.action_horizon = config.action_horizon
-        self.num_inference_timesteps = config.num_inference_timesteps
-        self.state_encoder = CategorySpecificMLP(
-            num_categories=config.max_num_embodiments,
-            input_dim=config.max_state_dim * config.state_history_length,
-            hidden_dim=self.hidden_size,
-            output_dim=self.input_embedding_dim,
-        )
-        self.action_encoder = MultiEmbodimentActionEncoder(
-            action_dim=self.action_dim,
-            hidden_size=self.input_embedding_dim,
-            num_embodiments=config.max_num_embodiments,
-        )
-        self.action_decoder = CategorySpecificMLP(
-            num_categories=config.max_num_embodiments,
-            input_dim=self.hidden_size,
-            hidden_dim=self.hidden_size,
-            output_dim=self.action_dim,
-        )
-        self.vlln = nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity()
-        vl_self_attention_cfg = getattr(config, "vl_self_attention_cfg", None)
-        if vl_self_attention_cfg and vl_self_attention_cfg.get("num_layers", 0) > 0:
-            self.vl_self_attention = SelfAttentionTransformer(**vl_self_attention_cfg)
-        else:
-            self.vl_self_attention = nn.Identity()
-        if config.add_pos_embed:
-            self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim)
-            nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)
-        self.state_dropout_prob = config.state_dropout_prob
-        self._noise_beta_alpha = config.noise_beta_alpha
-        self._noise_beta_beta = config.noise_beta_beta
-        self._beta_dist = None
-        self.num_timestep_buckets = config.num_timestep_buckets
-        self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model, config.tune_vlln)
-
-    def set_trainable_parameters(
-        self, tune_projector: bool, tune_diffusion_model: bool, tune_vlln: bool
-    ) -> None:
-        self.tune_projector = tune_projector
-        self.tune_diffusion_model = tune_diffusion_model
-        self.tune_vlln = tune_vlln
-        for parameter in self.parameters():
-            parameter.requires_grad = True
-        if not tune_projector:
-            self.state_encoder.requires_grad_(False)
-            self.action_encoder.requires_grad_(False)
-            self.action_decoder.requires_grad_(False)
-            if self.config.add_pos_embed:
-                self.position_embedding.requires_grad_(False)
-        if not tune_diffusion_model:
-            self.model.requires_grad_(False)
-        if not tune_vlln:
-            self.vlln.requires_grad_(False)
-            self.vl_self_attention.requires_grad_(False)
-
-    def set_frozen_modules_to_eval_mode(self) -> None:
-        if self.training:
-            if not self.tune_projector:
-                self.state_encoder.eval()
-                self.action_encoder.eval()
-                self.action_decoder.eval()
-                if self.config.add_pos_embed:
-                    self.position_embedding.eval()
-            if not self.tune_diffusion_model:
-                self.model.eval()
-            if not self.tune_vlln:
-                self.vlln.eval()
-                self.vl_self_attention.eval()
-
-    def sample_time(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-        if self._beta_dist is None:
-            beta_alpha = torch.tensor(self._noise_beta_alpha, device="cpu", dtype=torch.float32)
-            beta_beta = torch.tensor(self._noise_beta_beta, device="cpu", dtype=torch.float32)
-            self._beta_dist = Beta(beta_alpha, beta_beta, validate_args=False)
-        sample = self._beta_dist.sample([batch_size]).to(device, dtype=dtype)
-        return (1 - sample) * self.config.noise_s
-
-    def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature:
-        backbone_features = self.vlln(backbone_output["backbone_features"])
-        backbone_output["backbone_features"] = self.vl_self_attention(backbone_features)
-        return backbone_output
-
-    def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature:
-        self.set_frozen_modules_to_eval_mode()
-        backbone_output = self.process_backbone_output(backbone_output)
-        vl_embeds = backbone_output.backbone_features
-        device = vl_embeds.device
-        embodiment_id = action_input.embodiment_id
-
-        if action_input.state.shape[1] != self.config.state_history_length:
-            raise ValueError("state history length does not match GR00T N1.7 config.")
-        state = action_input.state.view(action_input.state.shape[0], 1, -1)
-        state_features = self.state_encoder(state, embodiment_id)
-
-        if self.training and self.state_dropout_prob > 0:
-            do_dropout = (
-                torch.rand(state_features.shape[0], device=state_features.device) < self.state_dropout_prob
-            )
-            state_features = state_features * (1 - do_dropout[:, None, None].to(dtype=state_features.dtype))
-
-        actions = action_input.action
-        noise = torch.randn(actions.shape, device=actions.device, dtype=actions.dtype)
-        t = self.sample_time(actions.shape[0], device=actions.device, dtype=actions.dtype)
-        t = t[:, None, None]
-        noisy_trajectory = (1 - t) * noise + t * actions
-        velocity = actions - noise
-        t_discretized = (t[:, 0, 0] * self.num_timestep_buckets).long()
-        action_features = self.action_encoder(noisy_trajectory, t_discretized, embodiment_id)
-
-        if self.config.add_pos_embed:
-            pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
-            action_features = action_features + self.position_embedding(pos_ids).unsqueeze(0)
-
-        sa_embs = torch.cat((state_features, action_features), dim=1)
-        if self.config.use_alternate_vl_dit:
-            model_output, _ = self.model(
-                hidden_states=sa_embs,
-                encoder_hidden_states=vl_embeds,
-                encoder_attention_mask=backbone_output.backbone_attention_mask,
-                timestep=t_discretized,
-                return_all_hidden_states=True,
-                image_mask=backbone_output.image_mask,
-                backbone_attention_mask=backbone_output.backbone_attention_mask,
-            )
-        else:
-            model_output, _ = self.model(
-                hidden_states=sa_embs,
-                encoder_hidden_states=vl_embeds,
-                encoder_attention_mask=backbone_output.backbone_attention_mask,
-                timestep=t_discretized,
-                return_all_hidden_states=True,
-            )
-
-        pred = self.action_decoder(model_output, embodiment_id)
-        pred_actions = pred[:, -actions.shape[1] :]
-        action_mask = action_input.action_mask.to(dtype=pred_actions.dtype)
-        action_loss = F.mse_loss(pred_actions, velocity, reduction="none") * action_mask
-        loss = action_loss.sum() / (action_mask.sum() + 1e-6)
-        return BatchFeature(
-            data={
-                "loss": loss,
-                "action_loss": action_loss,
-                "action_mask": action_mask,
-                "backbone_features": vl_embeds,
-                "state_features": state_features,
-            }
-        )
-
-    def _encode_features(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature:
-        backbone_output = self.process_backbone_output(backbone_output)
-        state = action_input.state
-        if state.shape[1] != self.config.state_history_length:
-            raise ValueError("state history length does not match GR00T N1.7 config.")
-        state = state.view(state.shape[0], 1, -1)
-        state_features = self.state_encoder(state, action_input.embodiment_id)
-        return BatchFeature(
-            data={"backbone_features": backbone_output.backbone_features, "state_features": state_features}
-        )
-
-    @torch.no_grad()
-    def get_action_with_features(
-        self,
-        backbone_features: torch.Tensor,
-        state_features: torch.Tensor,
-        embodiment_id: torch.Tensor,
-        backbone_output: BatchFeature,
-        action_input: BatchFeature,
-        options: dict[str, Any] | None = None,
-    ) -> BatchFeature:
-        vl_embeds = backbone_features
-        batch_size = vl_embeds.shape[0]
-        device = vl_embeds.device
-        actions = torch.randn(
-            size=(batch_size, self.config.action_horizon, self.action_dim),
-            dtype=vl_embeds.dtype,
-            device=device,
-        )
-        dt = 1.0 / self.num_inference_timesteps
-        vel_strength = torch.ones_like(actions)
-
-        if "action" in action_input:
-            if options is None:
-                raise ValueError("RTC options are required when action is provided to get_action.")
-            action_horizon_before_padding = options["action_horizon"]
-            actions[:, : options["rtc_overlap_steps"], :] = action_input["action"][
-                :,
-                action_horizon_before_padding - options["rtc_overlap_steps"] : action_horizon_before_padding,
-                :,
-            ]
-            vel_strength[:, : options["rtc_frozen_steps"], :] = 0.0
-            intermediate_steps = options["rtc_overlap_steps"] - options["rtc_frozen_steps"]
-            t = torch.linspace(0.0, 1.0, intermediate_steps + 2, device=device)
-            ramp = 1 - torch.exp(-options["rtc_ramp_rate"] * t)
-            ramp = ramp / ramp[-1].clamp_min(1e-8)
-            vel_strength[:, options["rtc_frozen_steps"] : options["rtc_overlap_steps"], :] = ramp[1:-1][
-                None, :, None
-            ].to(device)
-
-        for t_step in range(self.num_inference_timesteps):
-            t_cont = t_step / float(self.num_inference_timesteps)
-            t_discretized = int(t_cont * self.num_timestep_buckets)
-            timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device)
-            action_features = self.action_encoder(actions, timesteps_tensor, embodiment_id)
-            if self.config.add_pos_embed:
-                pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
-                action_features = action_features + self.position_embedding(pos_ids).unsqueeze(0)
-            sa_embs = torch.cat((state_features, action_features), dim=1)
-
-            if self.config.use_alternate_vl_dit:
-                model_output = self.model(
-                    hidden_states=sa_embs,
-                    encoder_hidden_states=vl_embeds,
-                    timestep=timesteps_tensor,
-                    image_mask=backbone_output.image_mask,
-                    backbone_attention_mask=backbone_output.backbone_attention_mask,
-                )
-            else:
-                model_output = self.model(
-                    hidden_states=sa_embs,
-                    encoder_hidden_states=vl_embeds,
-                    timestep=timesteps_tensor,
-                )
-            pred = self.action_decoder(model_output, embodiment_id)
-            actions = actions + dt * pred[:, -self.action_horizon :] * vel_strength
-
-        return BatchFeature(
-            data={
-                "action_pred": actions,
-                "backbone_features": vl_embeds,
-                "state_features": state_features,
-            }
-        )
-
-    @torch.no_grad()
-    def get_action(
-        self,
-        backbone_output: BatchFeature,
-        action_input: BatchFeature,
-        options: dict[str, Any] | None = None,
-    ) -> BatchFeature:
-        features = self._encode_features(backbone_output, action_input)
-        return self.get_action_with_features(
-            backbone_features=features.backbone_features,
-            state_features=features.state_features,
-            embodiment_id=action_input.embodiment_id,
-            backbone_output=backbone_output,
-            action_input=action_input,
-            options=options,
-        )
-
-    @property
-    def device(self) -> torch.device:
-        return next(iter(self.parameters())).device
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return next(iter(self.parameters())).dtype
-
-    def prepare_input(self, batch: dict[str, Any]) -> BatchFeature:
-        return BatchFeature(data=batch)
-
-
-def _is_cosmos_reason2_backbone(model_name: str) -> bool:
-    return str(model_name).rstrip("/") == "nvidia/Cosmos-Reason2-2B"
-
-
-def _cosmos_reason2_qwen3_vl_config() -> PretrainedConfig:
-    return Qwen3VLConfig(
-        image_token_id=151655,
-        video_token_id=151656,
-        vision_start_token_id=151652,
-        vision_end_token_id=151653,
-        tie_word_embeddings=True,
-        text_config={
-            "attention_bias": False,
-            "attention_dropout": 0.0,
-            "bos_token_id": 151643,
-            "dtype": "bfloat16",
-            "eos_token_id": 151645,
-            "head_dim": 128,
-            "hidden_act": "silu",
-            "hidden_size": 2048,
-            "initializer_range": 0.02,
-            "intermediate_size": 6144,
-            "max_position_embeddings": 262144,
-            "model_type": "qwen3_vl_text",
-            "num_attention_heads": 16,
-            "num_hidden_layers": 28,
-            "num_key_value_heads": 8,
-            "rms_norm_eps": 1e-6,
-            "rope_scaling": {
-                "mrope_interleaved": True,
-                "mrope_section": [24, 20, 20],
-                "rope_type": "default",
-            },
-            "rope_theta": 5000000,
-            "tie_word_embeddings": True,
-            "use_cache": True,
-            "vocab_size": 151936,
-        },
-        vision_config={
-            "deepstack_visual_indexes": [5, 11, 17],
-            "depth": 24,
-            "hidden_act": "gelu_pytorch_tanh",
-            "hidden_size": 1024,
-            "in_channels": 3,
-            "initializer_range": 0.02,
-            "intermediate_size": 4096,
-            "model_type": "qwen3_vl",
-            "num_heads": 16,
-            "num_position_embeddings": 2304,
-            "out_hidden_size": 2048,
-            "patch_size": 16,
-            "spatial_merge_size": 2,
-            "temporal_patch_size": 2,
-        },
-    )
-
-
-def get_backbone_cls(config: GR00TN17Config):
-    if "nvidia/Cosmos-Reason2" in config.model_name or "Qwen/Qwen3-VL" in config.model_name:
-        return Qwen3Backbone
-    if config.backbone_model_type == "qwen":
-        logger.warning(
-            "Unrecognized GR00T N1.7 backbone model name '%s'; assuming a Qwen3-VL-compatible "
-            "backbone because backbone_model_type='qwen'.",
-            config.model_name,
-        )
-        return Qwen3Backbone
-    raise ValueError(f"Unsupported GR00T N1.7 backbone model: {config.model_name}")
-
-
-class GR00TN17(PreTrainedModel):
-    """GR00T N1.7 model with a Cosmos-Reason2/Qwen3-VL backbone."""
-
-    config_class = GR00TN17Config
-    supports_gradient_checkpointing = True
-
-    def __init__(
-        self,
-        config: GR00TN17Config,
-        transformers_loading_kwargs: dict[str, Any] | None = None,
-        load_backbone_weights: bool = True,
-    ):
-        _register_with_transformers()
-        super().__init__(config)
-        transformers_loading_kwargs = transformers_loading_kwargs or {"trust_remote_code": True}
-        self.config = config
-        backbone_cls = get_backbone_cls(config)
-        self.backbone = backbone_cls(
-            model_name=config.model_name,
-            tune_llm=config.tune_llm,
-            tune_visual=config.tune_visual,
-            select_layer=config.select_layer,
-            reproject_vision=config.reproject_vision,
-            use_flash_attention=config.use_flash_attention,
-            load_bf16=config.load_bf16,
-            tune_top_llm_layers=config.tune_top_llm_layers,
-            trainable_params_fp32=config.backbone_trainable_params_fp32,
-            transformers_loading_kwargs=transformers_loading_kwargs,
-            load_pretrained_weights=load_backbone_weights,
-        )
-        self.action_head = GR00TN17ActionHead(config)
-        self.post_init()
-
-    def prepare_input(self, inputs: dict[str, Any]) -> tuple[BatchFeature, BatchFeature]:
-        global tree
-        if tree is None:
-            require_package("dm-tree", extra="groot", import_name="tree")
-            tree = importlib.import_module("tree")
-        backbone_inputs = self.backbone.prepare_input(inputs)
-        action_inputs = self.action_head.prepare_input(inputs)
-
-        def to_device_with_dtype(x):
-            if not isinstance(x, torch.Tensor):
-                return x
-            if torch.is_floating_point(x):
-                return x.to(self.device, dtype=self.dtype)
-            return x.to(self.device)
-
-        return (
-            tree.map_structure(to_device_with_dtype, backbone_inputs),
-            tree.map_structure(to_device_with_dtype, action_inputs),
-        )
-
-    def forward(self, inputs: dict[str, Any]) -> BatchFeature:
-        backbone_inputs, action_inputs = self.prepare_input(inputs)
-        backbone_outputs = self.backbone(backbone_inputs)
-        return self.action_head(backbone_outputs, action_inputs)
-
-    def get_action(self, inputs: dict[str, Any], options: dict[str, Any] | None = None) -> BatchFeature:
-        backbone_inputs, action_inputs = self.prepare_input(inputs)
-        backbone_outputs = self.backbone(backbone_inputs)
-        return self.action_head.get_action(backbone_outputs, action_inputs, options)
-
-    @property
-    def device(self) -> torch.device:
-        return next(iter(self.parameters())).device
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return next(iter(self.parameters())).dtype
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
-        tune_visual = kwargs.pop("tune_visual", True)
-        tune_llm = kwargs.pop("tune_llm", False)
-        tune_projector = kwargs.pop("tune_projector", True)
-        tune_diffusion_model = kwargs.pop("tune_diffusion_model", True)
-        tune_vlln = kwargs.pop("tune_vlln", True)
-        transformers_loading_kwargs = kwargs.pop("transformers_loading_kwargs", None) or {
-            "trust_remote_code": True
-        }
-        load_backbone_weights = kwargs.pop("load_backbone_weights", False)
-        for key in ("cache_dir", "local_files_only", "token"):
-            if key in kwargs:
-                transformers_loading_kwargs.setdefault(key, kwargs[key])
-
-        try:
-            local_model_path = snapshot_download(
-                pretrained_model_name_or_path,
-                repo_type="model",
-                revision=kwargs.get("revision"),
-                cache_dir=kwargs.get("cache_dir"),
-                local_files_only=kwargs.get("local_files_only", False),
-                token=kwargs.get("token"),
-            )
-        except (HFValidationError, RepositoryNotFoundError):
-            local_model_path = pretrained_model_name_or_path
-
-        pretrained_model = super().from_pretrained(
-            local_model_path,
-            transformers_loading_kwargs=transformers_loading_kwargs,
-            load_backbone_weights=load_backbone_weights,
-            **kwargs,
-        )
-        pretrained_model.backbone.set_trainable_parameters(
-            tune_visual=tune_visual,
-            tune_llm=tune_llm,
-            tune_top_llm_layers=pretrained_model.config.tune_top_llm_layers,
-        )
-        pretrained_model.action_head.set_trainable_parameters(
-            tune_projector=tune_projector,
-            tune_diffusion_model=tune_diffusion_model,
-            tune_vlln=tune_vlln,
-        )
-        return pretrained_model
-
-
-def _register_with_transformers() -> None:
-    """Register GR00T N1.7 with transformers' Auto* factories.
-
-    Idempotent: ``register(..., exist_ok=True)`` makes repeat calls no-ops (with a fallback that
-    suppresses the already-registered error on transformers builds whose ``register()`` predates
-    ``exist_ok``), so no run-once guard is needed.
-    """
-    if AutoConfig is None or AutoModel is None:
-        return
-    try:
-        AutoConfig.register(GR00TN17Config.model_type, GR00TN17Config, exist_ok=True)
-    except TypeError:
-        with suppress(ValueError):
-            AutoConfig.register(GR00TN17Config.model_type, GR00TN17Config)
-    try:
-        AutoModel.register(GR00TN17Config, GR00TN17, exist_ok=True)
-    except TypeError:
-        with suppress(ValueError):
-            AutoModel.register(GR00TN17Config, GR00TN17)

src/lerobot/policies/groot/modeling_groot.py

@@ -17,22 +17,28 @@
 """
 Groot Policy Wrapper for LeRobot Integration
 
-Minimal integration that delegates to Isaac-GR00T N1.7 components where
-possible without porting their code. Dataset loading and training
-orchestration are handled by LeRobot's standard training stack.
+Minimal integration that delegates to Isaac-GR00T components where possible
+without porting their code. The intent is to:
+
+- Download and load the pretrained GR00T model via GR00TN15.from_pretrained
+- Optionally align action horizon similar to gr00t_finetune.py
+- Expose predict_action via GR00T model.get_action
+- Provide a training forward that can call the GR00T model forward if batch
+  structure matches.
+
+Notes:
+- Dataset loading and full training orchestration is handled by Isaac-GR00T
+  TrainRunner in their codebase. If you want to invoke that flow end-to-end
+  from LeRobot, see `GrootPolicy.finetune_with_groot_runner` below.
 """
 
 import builtins
-import logging
 import os
 from collections import deque
 from pathlib import Path
 from typing import TypeVar
 
 import torch
-from huggingface_hub import hf_hub_download
-from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
-from huggingface_hub.errors import HfHubHTTPError
 from torch import Tensor
 
 from lerobot.configs import FeatureType, PolicyFeature
@@ -40,19 +46,8 @@ from lerobot.utils.constants import ACTION, OBS_IMAGES
 from lerobot.utils.import_utils import require_package
 
 from ..pretrained import PreTrainedPolicy
-from ..utils import get_device_from_parameters
-from .configuration_groot import (
-    GROOT_N1_5,
-    GROOT_N1_5_REMOVAL_GUIDANCE,
-    GROOT_N1_7,
-    GrootConfig,
-    infer_groot_model_version,
-    infer_groot_n1_7_action_execution_horizon,
-    infer_groot_n1_7_action_horizon,
-)
-from .groot_n1_7 import GR00TN17
-
-logger = logging.getLogger(__name__)
+from .configuration_groot import GrootConfig
+from .groot_n1 import GR00TN15
 
 T = TypeVar("T", bound="GrootPolicy")
 
@@ -72,38 +67,37 @@ class GrootPolicy(PreTrainedPolicy):
 
         # Initialize GR00T model using ported components
         self._groot_model = self._create_groot_model()
-        self._action_queue_steps = self._resolve_action_queue_steps()
 
         self.reset()
 
     def _create_groot_model(self):
-        """Create and initialize the GR00T N1.7 model using the ported components."""
-        model_kwargs = {
-            "pretrained_model_name_or_path": self.config.base_model_path,
-            "tune_llm": self.config.tune_llm,
-            "tune_visual": self.config.tune_visual,
-            "tune_projector": self.config.tune_projector,
-            "tune_diffusion_model": self.config.tune_diffusion_model,
-            # Forwarded as a GR00TN17Config override; read back by set_trainable_parameters.
-            "tune_top_llm_layers": self.config.tune_top_llm_layers,
-            "use_flash_attention": self.config.use_flash_attention,
-        }
-        # Surface the inference-time knobs onto the model config only when the user set them; None
-        # leaves the value baked into the checkpoint untouched.
-        if self.config.num_inference_timesteps is not None:
-            model_kwargs["num_inference_timesteps"] = self.config.num_inference_timesteps
-        if self.config.rtc_ramp_rate is not None:
-            model_kwargs["rtc_ramp_rate"] = self.config.rtc_ramp_rate
+        """Create and initialize the GR00T model using Isaac-GR00T API.
 
-        return GR00TN17.from_pretrained(
-            **model_kwargs,
-            tune_vlln=self.config.tune_vlln,
-            transformers_loading_kwargs={"trust_remote_code": True},
+        This is only called when creating a NEW policy (not when loading from checkpoint).
+
+        Steps (delegating to Isaac-GR00T):
+        1) Download and load pretrained model via GR00TN15.from_pretrained
+        2) Align action horizon with data_config if provided
+        """
+        # Handle Flash Attention compatibility issues
+        self._handle_flash_attention_compatibility()
+
+        model = GR00TN15.from_pretrained(
+            pretrained_model_name_or_path=self.config.base_model_path,
+            tune_llm=self.config.tune_llm,
+            tune_visual=self.config.tune_visual,
+            tune_projector=self.config.tune_projector,
+            tune_diffusion_model=self.config.tune_diffusion_model,
         )
 
+        model.compute_dtype = "bfloat16" if self.config.use_bf16 else model.compute_dtype
+        model.config.compute_dtype = model.compute_dtype
+
+        return model
+
     def reset(self):
         """Reset policy state when environment resets."""
-        self._action_queue = deque([], maxlen=self._action_queue_steps)
+        self._action_queue = deque([], maxlen=self.config.n_action_steps)
 
     @classmethod
     def from_pretrained(
@@ -124,7 +118,7 @@ class GrootPolicy(PreTrainedPolicy):
         """Load Groot policy from pretrained model.
 
         Handles two cases:
-        1. Base GR00T N1.7 models - loads the raw model
+        1. Base GR00T models (e.g., 'nvidia/GR00T-N1.5-3B') - loads the raw model
         2. Fine-tuned LeRobot checkpoints - loads config and weights from safetensors
 
         Args:
@@ -143,11 +137,13 @@ class GrootPolicy(PreTrainedPolicy):
         Returns:
             Initialized GrootPolicy instance with loaded model
         """
-        requested_version = infer_groot_model_version(str(pretrained_name_or_path)) or GROOT_N1_7
-        logger.info(
-            "The Groot policy wraps NVIDIA's GR00T %s model. Loading pretrained model from: %s",
-            requested_version,
-            pretrained_name_or_path,
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
+        from huggingface_hub.errors import HfHubHTTPError
+
+        print(
+            "The Groot policy is a wrapper around Nvidia's GR00T N1.5 model.\n"
+            f"Loading pretrained model from: {pretrained_name_or_path}"
         )
 
         model_id = str(pretrained_name_or_path)
@@ -178,7 +174,7 @@ class GrootPolicy(PreTrainedPolicy):
 
         if is_finetuned_checkpoint:
             # This is a fine-tuned LeRobot checkpoint - use parent class loading
-            logger.info("Detected fine-tuned LeRobot checkpoint, loading with state dict...")
+            print("Detected fine-tuned LeRobot checkpoint, loading with state dict...")
             return super().from_pretrained(
                 pretrained_name_or_path=pretrained_name_or_path,
                 config=config,
@@ -194,13 +190,11 @@ class GrootPolicy(PreTrainedPolicy):
             )
 
         # This is a base GR00T model - load it fresh
-        logger.info("Detected base GR00T model, loading from HuggingFace...")
+        print("Detected base GR00T model, loading from HuggingFace...")
 
         if config is None:
             # Create default config with the pretrained path
-            config = GrootConfig(
-                base_model_path=str(pretrained_name_or_path),
-            )
+            config = GrootConfig(base_model_path=str(pretrained_name_or_path))
 
             # Add minimal visual feature required for validation
             # validate_features() will automatically add state and action features
@@ -221,15 +215,6 @@ class GrootPolicy(PreTrainedPolicy):
             if hasattr(config, key):
                 setattr(config, key, value)
 
-        inferred_version = infer_groot_model_version(config.base_model_path)
-        if inferred_version is not None and inferred_version != GROOT_N1_7:
-            message = (
-                f"GR00T model_version '{GROOT_N1_7}' does not match base_model_path "
-                f"'{config.base_model_path}', which looks like '{inferred_version}'."
-            )
-            if inferred_version == GROOT_N1_5:
-                message = f"{message} {GROOT_N1_5_REMOVAL_GUIDANCE}"
-            raise ValueError(message)
         # Create a fresh policy instance - this will automatically load the GR00T model
         # in __init__ via _create_groot_model()
         policy = cls(config)
@@ -240,160 +225,21 @@ class GrootPolicy(PreTrainedPolicy):
     def get_optim_params(self) -> dict:
         return self.parameters()
 
-    def _resolve_action_queue_steps(self) -> int:
-        n_action_steps = int(self.config.n_action_steps)
-        checkpoint_action_horizon = infer_groot_n1_7_action_horizon(
-            self.config.base_model_path,
-            self.config.embodiment_tag,
-        )
-        execution_horizon = infer_groot_n1_7_action_execution_horizon(
-            self.config.base_model_path,
-            self.config.embodiment_tag,
-        )
-        horizons = [n_action_steps]
-        if checkpoint_action_horizon is not None:
-            horizons.append(checkpoint_action_horizon)
-        if execution_horizon is not None:
-            horizons.append(execution_horizon)
-        return min(horizons)
-
-    def _resolve_prediction_horizon(self, actions: Tensor) -> int:
-        """Return the policy-facing action horizon for a native GR00T prediction."""
-
-        horizons = [actions.shape[1]]
-        checkpoint_action_horizon = infer_groot_n1_7_action_horizon(
-            self.config.base_model_path,
-            self.config.embodiment_tag,
-        )
-        if checkpoint_action_horizon is not None:
-            horizons.append(checkpoint_action_horizon)
-
-        for horizon in (self.config.chunk_size, self.config.n_action_steps):
-            horizon = int(horizon)
-            if horizon > 0:
-                horizons.append(horizon)
-
-        return max(1, min(horizons))
-
-    def _filter_groot_inputs(self, batch: dict[str, Tensor], *, include_action: bool) -> dict[str, Tensor]:
-        allowed_base = {"state", "state_mask", "embodiment_id"}
-        if include_action:
-            allowed_base.update({"action", "action_mask"})
-
-        allowed_base.update(
-            {
-                "input_ids",
-                "attention_mask",
-                "pixel_values",
-                "image_grid_thw",
-                "mm_token_type_ids",
-                "pixel_values_videos",
-                "video_grid_thw",
-            }
-        )
-        allowed_base.add("action_mask")
-
-        return {
-            k: v for k, v in batch.items() if k in allowed_base and not (k.startswith("next.") or k == "info")
-        }
-
-    def _prepare_n1_7_rtc_inputs(
-        self,
-        inputs: dict[str, Tensor],
-        *,
-        inference_delay: object,
-        prev_chunk_left_over: object,
-    ) -> tuple[dict[str, Tensor], dict[str, object] | None]:
-        if prev_chunk_left_over is None:
-            return inputs, None
-        if not isinstance(prev_chunk_left_over, torch.Tensor):
-            raise TypeError("prev_chunk_left_over must be a torch.Tensor for GR00T N1.7 RTC.")
-        if prev_chunk_left_over.numel() == 0:
-            return inputs, None
-
-        prev_actions = prev_chunk_left_over
-        if prev_actions.ndim == 2:
-            prev_actions = prev_actions.unsqueeze(0)
-        elif prev_actions.ndim != 3:
-            raise ValueError("prev_chunk_left_over must have shape (T, A) or (B, T, A) for GR00T N1.7 RTC.")
-
-        state = inputs.get("state")
-        if state is None:
-            raise ValueError("GR00T N1.7 RTC requires `state` in the preprocessed batch.")
-        batch_size = state.shape[0]
-        if prev_actions.shape[0] == 1 and batch_size > 1:
-            prev_actions = prev_actions.expand(batch_size, -1, -1).clone()
-        elif prev_actions.shape[0] != batch_size:
-            raise ValueError("prev_chunk_left_over batch size must match the current GR00T N1.7 batch size.")
-
-        # The generic LeRobot RTC engine pads short leftovers with exact zero
-        # rows for fixed-shape policy calls. Native GR00T N1.7 RTC treats every
-        # provided prefix row as a real action constraint, so strip that padding
-        # before constructing the native overlap options.
-        valid_prefix_rows = prev_actions.detach().abs().sum(dim=(0, 2)) > 0
-        if valid_prefix_rows.any():
-            valid_prefix_steps = int(valid_prefix_rows.nonzero()[-1].item()) + 1
-            prev_actions = prev_actions[:, :valid_prefix_steps, :]
-        else:
-            return inputs, None
-
-        model_action_horizon = int(
-            getattr(self._groot_model.config, "action_horizon", self.config.chunk_size)
-        )
-        max_action_dim = int(getattr(self._groot_model.config, "max_action_dim", self.config.max_action_dim))
-        if prev_actions.shape[1] > model_action_horizon:
-            prev_actions = prev_actions[:, -model_action_horizon:, :]
-
-        action_horizon = int(prev_actions.shape[1])
-        if action_horizon <= 0:
-            return inputs, None
-
-        if prev_actions.shape[2] > max_action_dim:
-            prev_actions = prev_actions[:, :, :max_action_dim]
-        elif prev_actions.shape[2] < max_action_dim:
-            pad = torch.zeros(
-                prev_actions.shape[0],
-                prev_actions.shape[1],
-                max_action_dim - prev_actions.shape[2],
-                dtype=prev_actions.dtype,
-                device=prev_actions.device,
-            )
-            prev_actions = torch.cat([prev_actions, pad], dim=2)
-
-        prev_actions = prev_actions.to(device=state.device, dtype=state.dtype)
-
-        rtc_config = getattr(self.config, "rtc_config", None)
-        execution_horizon = int(getattr(rtc_config, "execution_horizon", action_horizon))
-        overlap_steps = max(0, min(action_horizon, execution_horizon))
-        if overlap_steps == 0:
-            return inputs, None
-
-        try:
-            frozen_steps = int(inference_delay or 0)
-        except (TypeError, ValueError):
-            frozen_steps = 0
-        frozen_steps = max(0, min(frozen_steps, overlap_steps))
-
-        options = {
-            "action_horizon": action_horizon,
-            "rtc_overlap_steps": overlap_steps,
-            "rtc_frozen_steps": frozen_steps,
-            "rtc_ramp_rate": float(getattr(self._groot_model.config, "rtc_ramp_rate", 6.0)),
-        }
-
-        inputs = dict(inputs)
-        inputs["action"] = prev_actions
-        return inputs, options
-
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Training forward pass.
 
         Delegates to Isaac-GR00T model.forward when inputs are compatible.
         """
-        groot_inputs = self._filter_groot_inputs(batch, include_action=True)
+        # Build a clean input dict for GR00T: keep only tensors GR00T consumes
+        allowed_base = {"state", "state_mask", "action", "action_mask", "embodiment_id"}
+        groot_inputs = {
+            k: v
+            for k, v in batch.items()
+            if (k in allowed_base or k.startswith("eagle_")) and not (k.startswith("next.") or k == "info")
+        }
 
         # Get device from model parameters
-        device = get_device_from_parameters(self)
+        device = next(self.parameters()).device
 
         # Run GR00T forward under bf16 autocast when enabled to reduce activation memory
         # Rationale: Matches original GR00T finetuning (bf16 compute, fp32 params) and avoids fp32 upcasts.
@@ -402,52 +248,38 @@ class GrootPolicy(PreTrainedPolicy):
 
         # Isaac-GR00T returns a BatchFeature; loss key is typically 'loss'
         loss = outputs.get("loss")
-        if loss is None:
-            raise RuntimeError(
-                "GR00T model.forward did not return a 'loss'. Training batches must include "
-                "'action' and 'action_mask'; check the preprocessor output."
-            )
 
         loss_dict = {"loss": loss.item()}
 
         return loss, loss_dict
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: object) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         """Predict a chunk of actions for inference by delegating to Isaac-GR00T.
 
         Returns a tensor of shape (B, n_action_steps, action_dim).
-
-        For N1.7, LeRobot's RTC leftovers are converted into the native GR00T
-        action-overlap options before calling the underlying model.
         """
         self.eval()
 
-        # Preprocessing is handled by the processor pipeline, so we just filter the batch.
-        # During inference, we do not pass action because it is predicted.
-        # N1.7 still carries a 2-D action horizon mask from its checkpoint processor.
-        groot_inputs = self._filter_groot_inputs(batch, include_action=False)
-        groot_inputs, groot_options = self._prepare_n1_7_rtc_inputs(
-            groot_inputs,
-            inference_delay=kwargs.get("inference_delay"),
-            prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
-        )
+        # Build a clean input dict for GR00T: keep only tensors GR00T consumes
+        # Preprocessing is handled by the processor pipeline, so we just filter the batch
+        # NOTE: During inference, we should NOT pass action/action_mask (that's what we're predicting)
+        allowed_base = {"state", "state_mask", "embodiment_id"}
+        groot_inputs = {
+            k: v
+            for k, v in batch.items()
+            if (k in allowed_base or k.startswith("eagle_")) and not (k.startswith("next.") or k == "info")
+        }
 
         # Get device from model parameters
-        device = get_device_from_parameters(self)
+        device = next(self.parameters()).device
 
         # Use bf16 autocast for inference to keep memory low and match backbone dtype
         with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=self.config.use_bf16):
-            if groot_options is not None:
-                outputs = self._groot_model.get_action(groot_inputs, options=groot_options)
-            else:
-                outputs = self._groot_model.get_action(groot_inputs)
+            outputs = self._groot_model.get_action(groot_inputs)
 
         actions = outputs.get("action_pred")
 
-        prediction_horizon = self._resolve_prediction_horizon(actions)
-        actions = actions[:, :prediction_horizon]
-
         original_action_dim = self.config.output_features[ACTION].shape[0]
         actions = actions[:, :, :original_action_dim]
 
@@ -460,5 +292,40 @@ class GrootPolicy(PreTrainedPolicy):
 
         if len(self._action_queue) == 0:
             actions = self.predict_action_chunk(batch)
-            self._action_queue.extend(actions[:, : self._action_queue_steps].transpose(0, 1))
+            self._action_queue.extend(actions.transpose(0, 1))
         return self._action_queue.popleft()
+
+    # -------------------------
+    # Internal helpers
+    # -------------------------
+    def _handle_flash_attention_compatibility(self) -> None:
+        """Handle Flash Attention compatibility issues by setting environment variables.
+
+        This addresses the common 'undefined symbol' error that occurs when Flash Attention
+        is compiled against a different PyTorch version than what's currently installed.
+        """
+
+        # Set environment variables to handle Flash Attention compatibility
+        # These help with symbol resolution issues
+        os.environ.setdefault("FLASH_ATTENTION_FORCE_BUILD", "0")
+        os.environ.setdefault("FLASH_ATTENTION_SKIP_CUDA_BUILD", "0")
+
+        # Try to import flash_attn and handle failures gracefully
+        try:
+            import flash_attn
+
+            print(f"[GROOT] Flash Attention version: {flash_attn.__version__}")
+        except ImportError as e:
+            print(f"[GROOT] Flash Attention not available: {e}")
+            print("[GROOT] Will use fallback attention mechanism")
+        except Exception as e:
+            if "undefined symbol" in str(e):
+                print(f"[GROOT] Flash Attention compatibility issue detected: {e}")
+                print("[GROOT] This is likely due to PyTorch/Flash Attention version mismatch")
+                print("[GROOT] Consider reinstalling Flash Attention with compatible version:")
+                print("  pip uninstall flash-attn")
+                print("  pip install --no-build-isolation flash-attn==2.6.3")
+                print("[GROOT] Continuing with fallback attention mechanism")
+            else:
+                print(f"[GROOT] Flash Attention error: {e}")
+                print("[GROOT] Continuing with fallback attention mechanism")

src/lerobot/policies/groot/utils.py

@@ -1,256 +1,47 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Shared, side-effect-free utilities for the GR00T N1.7 policy.
-
-These helpers are consumed by both the config layer (checkpoint sidecar
-inspection) and the processor layer (stat flattening, action decoding, language
-and image packing). They are pure functions with no GR00T-specific state so they
-can be unit-tested in isolation and reused without importing the heavier
-config/processor modules.
-"""
-
-from __future__ import annotations
-
-import json
 from pathlib import Path
-from typing import Any
+from shutil import copytree
 
-import numpy as np
-import torch
+from huggingface_hub import hf_hub_download
 
 
-def read_json(path: Path) -> dict[str, Any]:
-    """Read a JSON object from ``path``, returning ``{}`` on any read/parse error."""
-    try:
-        with path.open() as f:
-            data = json.load(f)
-    except (OSError, json.JSONDecodeError):
-        return {}
-    return data if isinstance(data, dict) else {}
+def ensure_eagle_cache_ready(vendor_dir: Path, cache_dir: Path, assets_repo: str) -> None:
+    """Populate the Eagle processor directory in cache and ensure tokenizer assets exist.
 
-
-def as_int_pair(value: Any) -> list[int] | None:
-    if not isinstance(value, (list, tuple)) or len(value) != 2:
-        return None
-    try:
-        return [int(value[0]), int(value[1])]
-    except (TypeError, ValueError):
-        return None
-
-
-def as_optional_int(value: Any) -> int | None:
-    if value is None:
-        return None
-    try:
-        return int(value)
-    except (TypeError, ValueError):
-        return None
-
-
-def as_optional_float(value: Any) -> float | None:
-    if value is None:
-        return None
-    try:
-        return float(value)
-    except (TypeError, ValueError):
-        return None
-
-
-def as_float_list(values: Any) -> list[float]:
-    if values is None:
-        return []
-    if isinstance(values, torch.Tensor):
-        return values.detach().cpu().reshape(-1).float().tolist()
-    if isinstance(values, np.ndarray):
-        return values.reshape(-1).astype(np.float32).tolist()
-    if isinstance(values, (list, tuple)):
-        flattened: list[float] = []
-        for value in values:
-            flattened.extend(as_float_list(value))
-        return flattened
-    return [float(values)]
-
-
-def config_value(value: Any) -> str:
-    if hasattr(value, "value"):
-        value = value.value
-    text = str(value).lower()
-    return {
-        "relative": "relative",
-        "absolute": "absolute",
-        "delta": "delta",
-        "eef": "eef",
-        "non_eef": "non_eef",
-        "default": "default",
-        "xyz_rot6d": "xyz+rot6d",
-        "xyz+rot6d": "xyz+rot6d",
-        "xyz_rotvec": "xyz+rotvec",
-        "xyz+rotvec": "xyz+rotvec",
-    }.get(text, text)
-
-
-def has_modality_stats(stats: dict[str, dict[str, Any]] | None) -> bool:
-    if not stats:
-        return False
-    return any(bool(modality_stats) for modality_stats in stats.values())
-
-
-def stat_dim_from_entry(entry: dict[str, Any]) -> int:
-    for stat_name in ("mean", "q01", "min", "max", "std"):
-        value = entry.get(stat_name)
-        if isinstance(value, list) and len(value) > 0:
-            return len(value)
-    return 0
-
-
-def flatten_n1_7_modality_stats(
-    *,
-    embodiment_stats: dict[str, Any],
-    embodiment_config: dict[str, Any],
-    modality: str,
-    use_percentiles: bool,
-    use_relative_action: bool,
-) -> dict[str, list[float]]:
-    """Flatten one N1.7 modality's grouped statistics in checkpoint order.
-
-    When checkpoints request percentile normalization, q01/q99 replace min/max
-    for regular groups. Relative action groups read from ``relative_action``
-    stats and keep min/max, matching Isaac-GR00T's processor override.
+    - Copies the vendored Eagle files into cache_dir (overwriting when needed).
+    - Downloads vocab.json and merges.txt into the same cache_dir if missing.
     """
+    cache_dir = Path(cache_dir)
+    vendor_dir = Path(vendor_dir)
 
-    source_stats = embodiment_stats.get(modality, {})
-    modality_config = embodiment_config.get(modality, {})
-    if not isinstance(source_stats, dict) or not isinstance(modality_config, dict):
-        return {}
-    modality_keys = modality_config.get("modality_keys", [])
-    if not isinstance(modality_keys, list):
-        return {}
+    try:
+        # Populate/refresh cache with vendor files to ensure a complete processor directory
+        print(f"[GROOT] Copying vendor Eagle files to cache: {vendor_dir} -> {cache_dir}")
+        copytree(vendor_dir, cache_dir, dirs_exist_ok=True)
+    except Exception as exc:  # nosec: B110
+        print(f"[GROOT] Warning: Failed to copy vendor Eagle files to cache: {exc}")
 
-    flattened: dict[str, list[float]] = {}
-    action_configs = modality_config.get("action_configs", []) if modality == "action" else []
-    if not isinstance(action_configs, list):
-        action_configs = []
-    relative_stats = embodiment_stats.get("relative_action", {})
-    if not isinstance(relative_stats, dict):
-        relative_stats = {}
+    required_assets = [
+        "vocab.json",
+        "merges.txt",
+        "added_tokens.json",
+        "chat_template.json",
+        "special_tokens_map.json",
+        "config.json",
+        "generation_config.json",
+        "preprocessor_config.json",
+        "processor_config.json",
+        "tokenizer_config.json",
+    ]
 
-    for stat_name in ("min", "max", "mean", "std"):
-        values: list[float] = []
-        source_stat_name = stat_name
-        if use_percentiles and stat_name == "min":
-            source_stat_name = "q01"
-        elif use_percentiles and stat_name == "max":
-            source_stat_name = "q99"
+    print(f"[GROOT] Assets repo: {assets_repo} \n Cache dir: {cache_dir}")
 
-        for idx, modality_key in enumerate(modality_keys):
-            if not isinstance(modality_key, str):
-                continue
-            key_source_stats = source_stats
-            key_stat_name = source_stat_name
-            if modality == "action" and use_relative_action and idx < len(action_configs):
-                action_config = action_configs[idx]
-                if isinstance(action_config, dict) and config_value(action_config.get("rep")) == "relative":
-                    key_source_stats = relative_stats
-                    key_stat_name = stat_name
-            key_stats = key_source_stats.get(modality_key, {})
-            if not isinstance(key_stats, dict):
-                raise KeyError(f"Missing statistics for {modality}.{modality_key}")
-            raw_values = key_stats.get(key_stat_name)
-            if raw_values is None:
-                raise KeyError(f"Missing '{key_stat_name}' statistics for {modality}.{modality_key}")
-            values.extend(as_float_list(raw_values))
-        if values:
-            flattened[stat_name] = values
-
-    return flattened
-
-
-def rot6d_to_matrix(rot6d: np.ndarray) -> np.ndarray:
-    rows = rot6d.reshape(2, 3).astype(np.float64)
-    row1 = rows[0] / np.linalg.norm(rows[0])
-    row2 = rows[1] - np.dot(row1, rows[1]) * row1
-    row2 = row2 / np.linalg.norm(row2)
-    row3 = np.cross(row1, row2)
-    return np.vstack([row1, row2, row3])
-
-
-def xyz_rot6d_to_homogeneous(xyz_rot6d: np.ndarray) -> np.ndarray:
-    transform = np.eye(4, dtype=np.float64)
-    transform[:3, :3] = rot6d_to_matrix(xyz_rot6d[3:])
-    transform[:3, 3] = xyz_rot6d[:3]
-    return transform
-
-
-def homogeneous_to_xyz_rot6d(transform: np.ndarray) -> np.ndarray:
-    return np.concatenate([transform[:3, 3], transform[:2, :3].reshape(-1)], axis=0)
-
-
-def relative_eef_to_absolute(action: np.ndarray, reference_state: np.ndarray) -> np.ndarray:
-    """Convert relative EEF deltas in xyz+rot6d format to absolute EEF poses."""
-
-    out = np.empty_like(action, dtype=np.float64)
-    for batch_idx in range(action.shape[0]):
-        reference = xyz_rot6d_to_homogeneous(reference_state[batch_idx])
-        for timestep in range(action.shape[1]):
-            relative = xyz_rot6d_to_homogeneous(action[batch_idx, timestep])
-            out[batch_idx, timestep] = homogeneous_to_xyz_rot6d(reference @ relative)
-    return out.astype(np.float32)
-
-
-def infer_n1_7_batch_size_and_device(
-    obs: dict[str, Any], action: torch.Tensor | None
-) -> tuple[int, torch.device]:
-    for value in list(obs.values()) + [action]:
-        if isinstance(value, torch.Tensor):
-            return value.shape[0], value.device
-    video = obs.get("video")
-    if isinstance(video, np.ndarray):
-        return video.shape[0], torch.device("cpu")
-    return 1, torch.device("cpu")
-
-
-def prepare_n1_7_language_batch(
-    language: Any,
-    batch_size: int,
-    *,
-    formalize_language: bool,
-) -> list[str]:
-    default_language = "Perform the task."
-    if language is None or (isinstance(language, str) and language == ""):
-        languages = [default_language] * batch_size
-    elif isinstance(language, str):
-        languages = [language] * batch_size
-    elif isinstance(language, (list, tuple)):
-        languages = list(language)
-        if len(languages) == 0:
-            languages = [default_language] * batch_size
-        elif len(languages) == 1 and batch_size > 1:
-            languages = languages * batch_size
-        elif len(languages) != batch_size:
-            raise ValueError(
-                f"language batch has {len(languages)} entries, but GR00T N1.7 input batch has {batch_size}."
+    for fname in required_assets:
+        dst = cache_dir / fname
+        if not dst.exists():
+            print(f"[GROOT] Fetching {fname}")
+            hf_hub_download(
+                repo_id=assets_repo,
+                filename=fname,
+                repo_type="model",
+                local_dir=str(cache_dir),
             )
-    else:
-        languages = [str(language)] * batch_size
-
-    formatted = []
-    for item in languages:
-        text = str(item) if item else default_language
-        if formalize_language:
-            text = text.lower()
-            text = "".join(ch for ch in text if ch.isalnum() or ch.isspace() or ch == "_")
-        formatted.append(text)
-    return formatted

src/lerobot/policies/pretrained.py

@@ -29,6 +29,7 @@ from huggingface_hub.errors import HfHubHTTPError
 from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
 from torch import Tensor, nn
 
+from lerobot.__version__ import __version__
 from lerobot.configs import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.utils.hub import HubMixin
@@ -38,6 +39,67 @@ from .utils import log_model_loading_keys
 T = TypeVar("T", bound="PreTrainedPolicy")
 
 
+def _build_card_context(
+    cfg: TrainPipelineConfig | None,
+    dataset_repo_id: str | None,
+    input_features: dict | None,
+    output_features: dict | None,
+) -> dict:
+    """Collect optional data for the model-card template.
+
+    Returns plain values only (no Markdown) — the template in
+    ``lerobot/templates/lerobot_modelcard_template.md`` decides how and whether to show
+    each one. Everything is best-effort: anything unavailable is left empty/None and the
+    template simply skips that section, so this never breaks a Hub push.
+    """
+    context = {
+        "training": None,
+        "input_features": input_features or {},
+        "output_features": output_features or {},
+        "dataset": None,
+        "robot_type": None,
+        "cameras": [],
+    }
+
+    if cfg is not None:
+        optimizer = getattr(cfg, "optimizer", None)
+        context["training"] = {
+            "steps": cfg.steps,
+            "batch_size": cfg.batch_size,
+            "seed": cfg.seed,
+            "optimizer": getattr(optimizer, "type", None) if optimizer else None,
+            "lr": getattr(optimizer, "lr", None) if optimizer else None,
+            "lerobot_version": __version__,
+        }
+
+    if dataset_repo_id:
+        dataset_cfg = getattr(cfg, "dataset", None)
+        try:
+            from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+
+            meta = LeRobotDatasetMetadata(
+                dataset_repo_id,
+                root=getattr(dataset_cfg, "root", None),
+                revision=getattr(dataset_cfg, "revision", None),
+            )
+            context["dataset"] = {
+                "repo_id": dataset_repo_id,
+                "episodes": meta.total_episodes,
+                "frames": meta.total_frames,
+                "fps": meta.fps,
+                "tasks": [str(task) for task in meta.tasks.index],
+            }
+            context["robot_type"] = meta.robot_type
+            context["cameras"] = [key.split(".")[-1] for key in meta.camera_keys]
+        except Exception as e:  # noqa: BLE001 — dataset details are optional, never fail the push
+            logging.warning(
+                f"Could not load dataset metadata for '{dataset_repo_id}'; those sections will be "
+                f"omitted from the model card. ({e})"
+            )
+
+    return context
+
+
 class ActionSelectKwargs(TypedDict, total=False):
     noise: Tensor | None
 
@@ -228,7 +290,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
                 self.save_pretrained(saved_path)  # Calls _save_pretrained and stores model tensors
 
             card = self.generate_model_card(
-                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
+                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg
             )
             card.save(str(saved_path / "README.md"))
 
@@ -246,9 +308,20 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             logging.info(f"Model pushed to {commit_info.repo_url.url}")
 
     def generate_model_card(
-        self, dataset_repo_id: str, model_type: str, license: str | None, tags: list[str] | None
+        self,
+        dataset_repo_id: str,
+        model_type: str,
+        license: str | None,
+        tags: list[str] | None,
+        cfg: TrainPipelineConfig | None = None,
     ) -> ModelCard:
-        base_model = "lerobot/smolvla_base" if model_type == "smolvla" else None  # Set a base model
+        base_model_mapping = {
+            "smolvla": "lerobot/smolvla_base",
+            "pi0": "lerobot/pi0_base",
+            "pi05": "lerobot/pi05_base",
+            "pi0_fast": "lerobot/pi0fast-base",
+            "xvla": "lerobot/xvla-base",
+        }
 
         card_data = ModelCardData(
             license=license or "apache-2.0",
@@ -257,13 +330,20 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             tags=list(set(tags or []).union({"robotics", "lerobot", model_type})),
             model_name=model_type,
             datasets=dataset_repo_id,
-            base_model=base_model,
+            base_model=base_model_mapping.get(model_type),
         )
 
+        context = _build_card_context(
+            cfg, dataset_repo_id, self.config.input_features, self.config.output_features
+        )
+        # Used by the template to pre-fill commands and the "Fine-tuned from" line.
+        context["policy_repo_id"] = getattr(self.config, "repo_id", None)
+        context["base_model"] = base_model_mapping.get(model_type)
+
         template_card = (
             files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text(encoding="utf-8")
         )
-        card = ModelCard.from_template(card_data, template_str=template_card)
+        card = ModelCard.from_template(card_data, template_str=template_card, **context)
         card.validate()
         return card
 

src/lerobot/processor/pipeline.py

@@ -32,7 +32,6 @@ from __future__ import annotations
 
 import importlib
 import json
-import os
 import re
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Sequence
@@ -281,6 +280,11 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
 
     before_step_hooks: list[Callable[[int, EnvTransition], None]] = field(default_factory=list, repr=False)
     after_step_hooks: list[Callable[[int, EnvTransition], None]] = field(default_factory=list, repr=False)
+    _serialized_state_filenames: tuple[str | None, ...] | None = field(
+        default=None,
+        init=False,
+        repr=False,
+    )
 
     def __call__(self, data: TInput) -> TOutput:
         """Processes input data through the full pipeline.
@@ -338,30 +342,108 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
             transition = processor_step(transition)
             yield transition
 
-    def _save_pretrained(self, save_directory: Path, **kwargs):
-        """Internal method to comply with `HubMixin`'s saving mechanism.
+    def _get_sanitized_name(self) -> str:
+        """Return a filename-safe version of the pipeline name.
 
-        This method does the actual saving work and is called by HubMixin.save_pretrained.
+        Returns:
+            The lower-cased pipeline name with non-alphanumeric characters replaced by underscores.
         """
-        config_filename = kwargs.pop("config_filename", None)
+        return re.sub(r"[^a-zA-Z0-9_]", "_", self.name.lower())
 
-        # Sanitize the pipeline name to create a valid filename prefix.
-        sanitized_name = re.sub(r"[^a-zA-Z0-9_]", "_", self.name.lower())
+    @staticmethod
+    def _get_state_filename(
+        *,
+        step_index: int,
+        registry_name: str | None,
+        sanitized_name: str,
+    ) -> str:
+        """Return the safetensors filename for one stateful processor step.
 
-        if config_filename is None:
-            config_filename = f"{sanitized_name}.json"
+        Args:
+            step_index: The index of the processor step in this pipeline.
+            registry_name: The registered processor step name, if available.
+            sanitized_name: The filename-safe pipeline name.
 
-        config: dict[str, Any] = {
+        Returns:
+            The state filename used by the existing disk serialization format.
+        """
+        if registry_name:
+            return f"{sanitized_name}_step_{step_index}_{registry_name}.safetensors"
+
+        return f"{sanitized_name}_step_{step_index}.safetensors"
+
+    @staticmethod
+    def _get_state_key(state_filename: str) -> str:
+        """Return the in-memory state key for a serialized state filename.
+
+        Args:
+            state_filename: The `.safetensors` filename from the serialized config.
+
+        Returns:
+            The state key used by the in-memory pipeline state dictionary.
+        """
+        return state_filename.removesuffix(".safetensors")
+
+    @staticmethod
+    def _get_state_filenames_from_config(loaded_config: dict[str, Any]) -> tuple[str | None, ...]:
+        """Return serialized state filenames in step order.
+
+        Args:
+            loaded_config: A validated processor pipeline config.
+
+        Returns:
+            A tuple containing each step's serialized state filename, or None for stateless steps.
+        """
+        return tuple(step_entry.get("state_file") for step_entry in loaded_config["steps"])
+
+    def _get_state_filenames_for_loading(self) -> tuple[str | None, ...]:
+        """Return expected state filenames in step order for `load_state_dict()`.
+
+        Returns:
+            The preserved serialized state filenames when available, otherwise filenames derived from
+            current non-empty step state.
+        """
+        if self._serialized_state_filenames is not None and len(self._serialized_state_filenames) == len(
+            self.steps
+        ):
+            return self._serialized_state_filenames
+
+        sanitized_name = self._get_sanitized_name()
+        state_filenames: list[str | None] = []
+
+        for step_index, processor_step in enumerate(self.steps):
+            step_state_dict = processor_step.state_dict()
+            if not step_state_dict:
+                state_filenames.append(None)
+                continue
+
+            registry_name = getattr(processor_step.__class__, "_registry_name", None)
+            state_filenames.append(
+                self._get_state_filename(
+                    step_index=step_index,
+                    registry_name=registry_name,
+                    sanitized_name=sanitized_name,
+                )
+            )
+
+        return tuple(state_filenames)
+
+    def get_config(self) -> dict[str, Any]:
+        """Return the JSON-serializable pipeline configuration.
+
+        Returns:
+            A dictionary with the same content that `save_pretrained()` writes as JSON.
+        """
+        sanitized_name = self._get_sanitized_name()
+        pipeline_config: dict[str, Any] = {
             "name": self.name,
             "steps": [],
         }
 
-        # Iterate through each step to build its configuration entry.
         for step_index, processor_step in enumerate(self.steps):
             registry_name = getattr(processor_step.__class__, "_registry_name", None)
-
             step_entry: dict[str, Any] = {}
-            # Prefer registry name for portability, otherwise fall back to full class path.
+
             if registry_name:
                 step_entry["registry_name"] = registry_name
             else:
@@ -369,31 +451,110 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
                     f"{processor_step.__class__.__module__}.{processor_step.__class__.__name__}"
                 )
 
-            # Save step configuration if `get_config` is implemented.
-            if hasattr(processor_step, "get_config"):
-                step_entry["config"] = processor_step.get_config()
+            step_entry["config"] = processor_step.get_config()
 
-            # Save step state if `state_dict` is implemented and returns a non-empty dict.
-            if hasattr(processor_step, "state_dict"):
-                state = processor_step.state_dict()
-                if state:
-                    # Clone tensors to avoid modifying the original state.
-                    cloned_state = {key: tensor.clone() for key, tensor in state.items()}
+            step_state_dict = processor_step.state_dict()
+            if step_state_dict:
+                step_entry["state_file"] = self._get_state_filename(
+                    step_index=step_index,
+                    registry_name=registry_name,
+                    sanitized_name=sanitized_name,
+                )
 
-                    # Create a unique filename for the state file.
-                    if registry_name:
-                        state_filename = f"{sanitized_name}_step_{step_index}_{registry_name}.safetensors"
-                    else:
-                        state_filename = f"{sanitized_name}_step_{step_index}.safetensors"
+            pipeline_config["steps"].append(step_entry)
 
-                    save_file(cloned_state, os.path.join(str(save_directory), state_filename))
-                    step_entry["state_file"] = state_filename
+        return pipeline_config
 
-            config["steps"].append(step_entry)
+    def state_dict(self) -> dict[str, dict[str, torch.Tensor]]:
+        """Return pipeline state tensors grouped by state key.
 
-        # Write the main configuration JSON file.
-        with open(os.path.join(str(save_directory), config_filename), "w") as file_pointer:
-            json.dump(config, file_pointer, indent=2)
+        Returns:
+            A dictionary mapping suffixless state keys to cloned step state dictionaries.
+        """
+        sanitized_name = self._get_sanitized_name()
+        pipeline_state_dict: dict[str, dict[str, torch.Tensor]] = {}
+
+        for step_index, processor_step in enumerate(self.steps):
+            step_state_dict = processor_step.state_dict()
+            if not step_state_dict:
+                continue
+
+            registry_name = getattr(processor_step.__class__, "_registry_name", None)
+            state_filename = self._get_state_filename(
+                step_index=step_index,
+                registry_name=registry_name,
+                sanitized_name=sanitized_name,
+            )
+            state_key = self._get_state_key(state_filename)
+            pipeline_state_dict[state_key] = {
+                tensor_name: tensor.clone() for tensor_name, tensor in step_state_dict.items()
+            }
+
+        return pipeline_state_dict
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, dict[str, torch.Tensor]],
+    ) -> None:
+        """Load pipeline state tensors into the existing steps.
+
+        Args:
+            state_dict: A dictionary mapping suffixless state keys to step state dictionaries.
+
+        Raises:
+            KeyError: If loading finds missing expected state or unexpected extra state.
+        """
+        expected_state_filenames = self._get_state_filenames_for_loading()
+        used_state_keys: set[str] = set()
+
+        for step_index, (processor_step, state_filename) in enumerate(
+            zip(self.steps, expected_state_filenames, strict=True)
+        ):
+            if state_filename is None:
+                continue
+
+            state_key = self._get_state_key(state_filename)
+            if state_key not in state_dict:
+                raise KeyError(
+                    f"Missing state key '{state_key}' for processor step {step_index}. "
+                    f"Available state keys: {sorted(state_dict.keys())}"
+                )
+
+            processor_step.load_state_dict(state_dict[state_key])
+            used_state_keys.add(state_key)
+
+        unexpected_state_keys = set(state_dict) - used_state_keys
+        if unexpected_state_keys:
+            expected_state_key_set = {
+                self._get_state_key(state_filename)
+                for state_filename in expected_state_filenames
+                if state_filename is not None
+            }
+            raise KeyError(
+                f"Unexpected processor state keys: {sorted(unexpected_state_keys)}. "
+                f"Expected state keys: {sorted(expected_state_key_set)}"
+            )
+
+    def _save_pretrained(self, save_directory: Path, **kwargs) -> None:
+        """Internal method to comply with `HubMixin`'s saving mechanism.
+
+        This method does the actual saving work and is called by HubMixin.save_pretrained.
+        """
+        config_filename = kwargs.pop("config_filename", None)
+        sanitized_name = self._get_sanitized_name()
+
+        if config_filename is None:
+            config_filename = f"{sanitized_name}.json"
+
+        pipeline_config = self.get_config()
+        pipeline_state_dict = self.state_dict()
+
+        for state_key, step_state_dict in pipeline_state_dict.items():
+            state_filename = f"{state_key}.safetensors"
+            save_file(step_state_dict, save_directory / state_filename)
+
+        with open(save_directory / config_filename, "w") as file_pointer:
+            json.dump(pipeline_config, file_pointer, indent=2)
 
     def save_pretrained(
         self,
@@ -577,12 +738,54 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         cls._validate_overrides_used(validated_overrides, loaded_config)
 
         # 5. Construct and return the final pipeline instance
-        return cls(
+        pipeline = cls(
             steps=steps,
             name=loaded_config.get("name", "DataProcessorPipeline"),
             to_transition=to_transition or cast(Callable[[TInput], EnvTransition], batch_to_transition),
             to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch),
         )
+        pipeline._serialized_state_filenames = cls._get_state_filenames_from_config(loaded_config)
+        return pipeline
+
+    @classmethod
+    def from_config(
+        cls,
+        config: dict[str, Any],
+        *,
+        state_dict: dict[str, dict[str, torch.Tensor]] | None = None,
+        overrides: dict[str, Any] | None = None,
+        to_transition: Callable[[TInput], EnvTransition] | None = None,
+        to_output: Callable[[EnvTransition], TOutput] | None = None,
+    ) -> DataProcessorPipeline[TInput, TOutput]:
+        """Build a pipeline from an in-memory config and optional state tensors.
+
+        Args:
+            config: A config dictionary with the same structure as the saved processor JSON.
+            state_dict: Optional in-memory pipeline state grouped by suffixless state key.
+            overrides: Optional constructor overrides keyed by registry name or class name.
+            to_transition: Optional converter from input data to `EnvTransition`.
+            to_output: Optional converter from `EnvTransition` to output data.
+
+        Returns:
+            A processor pipeline built from the config and optional state.
+        """
+        cls._validate_loaded_config("<in-memory config>", config, "<in-memory config>")
+
+        steps, remaining_override_keys = cls._build_steps_from_config(config, overrides or {})
+        cls._validate_overrides_used(remaining_override_keys, config)
+
+        pipeline = cls(
+            steps=steps,
+            name=config.get("name", "DataProcessorPipeline"),
+            to_transition=to_transition or cast(Callable[[TInput], EnvTransition], batch_to_transition),
+            to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch),
+        )
+        pipeline._serialized_state_filenames = cls._get_state_filenames_from_config(config)
+
+        if state_dict is not None:
+            pipeline.load_state_dict(state_dict)
+
+        return pipeline
 
     @classmethod
     def _load_config(
@@ -666,9 +869,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
                 ) from e
 
     @classmethod
-    def _validate_loaded_config(
-        cls, model_id: str, loaded_config: dict[str, Any], config_filename: str
-    ) -> None:
+    def _validate_loaded_config(cls, model_id: str, loaded_config: Any, config_filename: str) -> None:
         """Validate that a config was loaded and is a valid processor config.
 
         This method validates processor config format with intelligent migration detection:
@@ -688,7 +889,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
 
         Args:
             model_id: The model identifier (used for migration detection)
-            loaded_config: The loaded config dictionary (guaranteed non-None)
+            loaded_config: The loaded config value to validate (may be non-dict)
             config_filename: The config filename that was loaded (for error messages)
 
         Raises:
@@ -702,9 +903,14 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
                     model_id,
                     f"Config file '{config_filename}' is not a valid processor configuration",
                 )
+            loaded_config_description = (
+                list(loaded_config.keys())
+                if isinstance(loaded_config, dict)
+                else type(loaded_config).__name__
+            )
             raise ValueError(
                 f"Config file '{config_filename}' is not a valid processor configuration. "
-                f"Expected a config with 'steps' field, but got: {list(loaded_config.keys())}"
+                f"Expected a config with 'steps' field, but got: {loaded_config_description}"
             )
 
     @classmethod
@@ -766,26 +972,41 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
             ImportError: If a step class cannot be imported or found in registry
             ValueError: If a step cannot be instantiated with its configuration
         """
-        steps: list[ProcessorStep] = []
-        override_keys = set(overrides.keys())
+        steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
 
-        for step_entry in loaded_config["steps"]:
-            # 1. Get step class and key
-            step_class, step_key = cls._resolve_step_class(step_entry)
-
-            # 2. Instantiate step with overrides
-            step_instance = cls._instantiate_step(step_entry, step_class, step_key, overrides)
-
-            # 3. Load step state if available
+        for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
             cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
 
-            # 4. Track used overrides
-            if step_key in override_keys:
-                override_keys.discard(step_key)
+        return steps, remaining_override_keys
 
-            steps.append(step_instance)
+    @classmethod
+    def _build_steps_from_config(
+        cls,
+        loaded_config: dict[str, Any],
+        overrides: dict[str, Any],
+    ) -> tuple[list[ProcessorStep], set[str]]:
+        """Build processor steps from config without loading tensor state.
 
-        return steps, override_keys
+        Args:
+            loaded_config: The loaded processor configuration.
+            overrides: User-provided constructor overrides keyed by step key.
+
+        Returns:
+            A tuple containing instantiated steps and override keys that did not match a step.
+        """
+        processor_steps: list[ProcessorStep] = []
+        remaining_override_keys = set(overrides.keys())
+
+        for step_entry in loaded_config["steps"]:
+            step_class, step_key = cls._resolve_step_class(step_entry)
+            processor_step = cls._instantiate_step(step_entry, step_class, step_key, overrides)
+
+            if step_key in remaining_override_keys:
+                remaining_override_keys.discard(step_key)
+
+            processor_steps.append(processor_step)
+
+        return processor_steps, remaining_override_keys
 
     @classmethod
     def _resolve_step_class(cls, step_entry: dict[str, Any]) -> tuple[type[ProcessorStep], str]:
@@ -1096,7 +1317,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         return True
 
     @classmethod
-    def _is_processor_config(cls, config: dict) -> bool:
+    def _is_processor_config(cls, config: Any) -> bool:
         """Check if config follows DataProcessorPipeline format.
 
         This method validates the processor configuration structure:
@@ -1147,6 +1368,9 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         Returns:
             True if config follows valid DataProcessorPipeline format, False otherwise
         """
+        if not isinstance(config, dict):
+            return False
+
         # Must have a "steps" field with a list of step configurations
         if not isinstance(config.get("steps"), list):
             return False

src/lerobot/rewards/factory.py

@@ -124,6 +124,7 @@ def make_reward_model(cfg: RewardModelConfig, **kwargs) -> PreTrainedRewardModel
 
     if cfg.pretrained_path:
         kwargs["pretrained_name_or_path"] = cfg.pretrained_path
+        kwargs["revision"] = cfg.pretrained_revision
         reward_model = reward_cls.from_pretrained(**kwargs)
     else:
         reward_model = reward_cls(**kwargs)

src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py

@@ -18,7 +18,8 @@ import logging
 from functools import cached_property
 
 from lerobot.types import RobotAction, RobotObservation
-from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
+from lerobot.utils.bimanual import BimanualMixin
+from lerobot.utils.decorators import check_if_not_connected
 
 from ..openarm_follower import OpenArmFollower, OpenArmFollowerConfig
 from ..robot import Robot
@@ -27,7 +28,7 @@ from .config_bi_openarm_follower import BiOpenArmFollowerConfig
 logger = logging.getLogger(__name__)
 
 
-class BiOpenArmFollower(Robot):
+class BiOpenArmFollower(BimanualMixin, Robot):
     """
     Bimanual OpenArm Follower Arms
     """
@@ -39,15 +40,17 @@ class BiOpenArmFollower(Robot):
         super().__init__(config)
         self.config = config
 
-        # Top-level cameras are distributed evenly: each arm's OpenArmFollower
-        # will only open the cameras assigned to it. Per-arm cameras are used
-        # as fallback when top-level cameras are empty.
-        if config.cameras:
-            left_cameras = config.cameras
-            right_cameras = {}
-        else:
-            left_cameras = config.left_arm_config.cameras
-            right_cameras = config.right_arm_config.cameras
+        # Top-level cameras are opened by `left_arm` for convenience, but their
+        # keys stay unprefixed in observations (tracked via `_top_level_cam_keys`).
+        self._top_level_cam_keys = set(config.cameras)
+        _collisions = self._top_level_cam_keys & set(
+            config.left_arm_config.cameras
+        ) | self._top_level_cam_keys & set(config.right_arm_config.cameras)
+        if _collisions:
+            raise ValueError(
+                f"Top-level camera names collide with per-arm camera names: {sorted(_collisions)}"
+            )
+        left_arm_cameras = {**config.left_arm_config.cameras, **config.cameras}
 
         left_arm_config = OpenArmFollowerConfig(
             id=f"{config.id}_left" if config.id else None,
@@ -56,7 +59,7 @@ class BiOpenArmFollower(Robot):
             disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect,
             use_velocity_and_torque=config.left_arm_config.use_velocity_and_torque,
             max_relative_target=config.left_arm_config.max_relative_target,
-            cameras=left_cameras,
+            cameras=left_arm_cameras,
             side=config.left_arm_config.side,
             can_interface=config.left_arm_config.can_interface,
             use_can_fd=config.left_arm_config.use_can_fd,
@@ -75,7 +78,7 @@ class BiOpenArmFollower(Robot):
             disable_torque_on_disconnect=config.right_arm_config.disable_torque_on_disconnect,
             use_velocity_and_torque=config.right_arm_config.use_velocity_and_torque,
             max_relative_target=config.right_arm_config.max_relative_target,
-            cameras=right_cameras,
+            cameras=config.right_arm_config.cameras,
             side=config.right_arm_config.side,
             can_interface=config.right_arm_config.can_interface,
             use_can_fd=config.right_arm_config.use_can_fd,
@@ -95,22 +98,19 @@ class BiOpenArmFollower(Robot):
 
     @property
     def _motors_ft(self) -> dict[str, type]:
-        left_arm_motors_ft = self.left_arm._motors_ft
-        right_arm_motors_ft = self.right_arm._motors_ft
-
-        # Right first, then left — matches the teleoperator (OpenArmMini) ordering
-        # and the dataset feature names recorded during data collection.
         return {
-            **{f"right_{k}": v for k, v in right_arm_motors_ft.items()},
-            **{f"left_{k}": v for k, v in left_arm_motors_ft.items()},
+            **{f"left_{k}": v for k, v in self.left_arm._motors_ft.items()},
+            **{f"right_{k}": v for k, v in self.right_arm._motors_ft.items()},
         }
 
     @property
     def _cameras_ft(self) -> dict[str, tuple]:
-        # Cameras already have unique user-chosen names (e.g. "left_wrist", "base",
-        # "right_wrist"), so we merge them directly — unlike motors which need the
-        # left_/right_ prefix to disambiguate identical per-arm joint names.
-        return {**self.left_arm._cameras_ft, **self.right_arm._cameras_ft}
+        out: dict[str, tuple] = {}
+        for k, v in self.left_arm._cameras_ft.items():
+            out[k if k in self._top_level_cam_keys else f"left_{k}"] = v
+        for k, v in self.right_arm._cameras_ft.items():
+            out[f"right_{k}"] = v
+        return out
 
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
@@ -120,27 +120,6 @@ class BiOpenArmFollower(Robot):
     def action_features(self) -> dict[str, type]:
         return self._motors_ft
 
-    @property
-    def is_connected(self) -> bool:
-        return self.left_arm.is_connected and self.right_arm.is_connected
-
-    @check_if_already_connected
-    def connect(self, calibrate: bool = True) -> None:
-        self.left_arm.connect(calibrate)
-        self.right_arm.connect(calibrate)
-
-    @property
-    def is_calibrated(self) -> bool:
-        return self.left_arm.is_calibrated and self.right_arm.is_calibrated
-
-    def calibrate(self) -> None:
-        self.left_arm.calibrate()
-        self.right_arm.calibrate()
-
-    def configure(self) -> None:
-        self.left_arm.configure()
-        self.right_arm.configure()
-
     def setup_motors(self) -> None:
         raise NotImplementedError(
             "Motor ID configuration is typically done via manufacturer tools for CAN motors."
@@ -148,21 +127,15 @@ class BiOpenArmFollower(Robot):
 
     @check_if_not_connected
     def get_observation(self) -> RobotObservation:
-        obs_dict = {}
+        obs_dict: RobotObservation = {}
 
-        # Camera keys that should NOT get the arm prefix (they already have unique names)
-        left_cam_keys = set(self.left_arm.cameras.keys())
-        right_cam_keys = set(self.right_arm.cameras.keys())
+        # Add "left_" prefix to per-arm keys; keep top-level camera keys unprefixed.
+        for key, value in self.left_arm.get_observation().items():
+            obs_dict[key if key in self._top_level_cam_keys else f"left_{key}"] = value
 
-        # Right first, then left — matches the teleoperator (OpenArmMini) ordering
-        # and the dataset feature names recorded during data collection.
-        right_obs = self.right_arm.get_observation()
-        for key, value in right_obs.items():
-            obs_dict[key if key in right_cam_keys else f"right_{key}"] = value
-
-        left_obs = self.left_arm.get_observation()
-        for key, value in left_obs.items():
-            obs_dict[key if key in left_cam_keys else f"left_{key}"] = value
+        # Add "right_" prefix
+        for key, value in self.right_arm.get_observation().items():
+            obs_dict[f"right_{key}"] = value
 
         return obs_dict
 
@@ -189,9 +162,4 @@ class BiOpenArmFollower(Robot):
         prefixed_sent_action_left = {f"left_{key}": value for key, value in sent_action_left.items()}
         prefixed_sent_action_right = {f"right_{key}": value for key, value in sent_action_right.items()}
 
-        return {**prefixed_sent_action_right, **prefixed_sent_action_left}
-
-    @check_if_not_connected
-    def disconnect(self):
-        self.left_arm.disconnect()
-        self.right_arm.disconnect()
+        return {**prefixed_sent_action_left, **prefixed_sent_action_right}

src/lerobot/robots/bi_openarm_follower/config_bi_openarm_follower.py

@@ -32,5 +32,7 @@ class BiOpenArmFollowerConfig(RobotConfig):
     left_arm_config: OpenArmFollowerConfigBase
     right_arm_config: OpenArmFollowerConfigBase
 
-    # Top-level cameras shared across both arms.
+    # Top-level cameras not attached to a specific side. Keys are kept as-is in
+    # observations (no `left_`/`right_` prefix). Per-arm cameras (declared on
+    # `{left,right}_arm_config.cameras`) are prefixed.
     cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py

@@ -18,7 +18,8 @@ import logging
 from functools import cached_property
 
 from lerobot.types import RobotAction, RobotObservation
-from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
+from lerobot.utils.bimanual import BimanualMixin
+from lerobot.utils.decorators import check_if_not_connected
 
 from ..rebot_b601_follower import RebotB601Follower, RebotB601FollowerRobotConfig
 from ..robot import Robot
@@ -27,7 +28,7 @@ from .config_bi_rebot_b601_follower import BiRebotB601FollowerConfig
 logger = logging.getLogger(__name__)
 
 
-class BiRebotB601Follower(Robot):
+class BiRebotB601Follower(BimanualMixin, Robot):
     """Bimanual Seeed Studio reBot B601-DM follower.
 
     Composes two single-arm :class:`RebotB601Follower` instances. Observation and
@@ -41,6 +42,18 @@ class BiRebotB601Follower(Robot):
         super().__init__(config)
         self.config = config
 
+        # Top-level cameras are opened by `left_arm` for convenience, but their
+        # keys stay unprefixed in observations (tracked via `_top_level_cam_keys`).
+        self._top_level_cam_keys = set(config.cameras)
+        _collisions = self._top_level_cam_keys & set(
+            config.left_arm_config.cameras
+        ) | self._top_level_cam_keys & set(config.right_arm_config.cameras)
+        if _collisions:
+            raise ValueError(
+                f"Top-level camera names collide with per-arm camera names: {sorted(_collisions)}"
+            )
+        left_arm_cameras = {**config.left_arm_config.cameras, **config.cameras}
+
         left_arm_config = RebotB601FollowerRobotConfig(
             id=f"{config.id}_left" if config.id else None,
             calibration_dir=config.calibration_dir,
@@ -49,7 +62,7 @@ class BiRebotB601Follower(Robot):
             dm_serial_baud=config.left_arm_config.dm_serial_baud,
             disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect,
             max_relative_target=config.left_arm_config.max_relative_target,
-            cameras=config.left_arm_config.cameras,
+            cameras=left_arm_cameras,
             motor_can_ids=config.left_arm_config.motor_can_ids,
             pos_vel_velocity=config.left_arm_config.pos_vel_velocity,
             gripper_torque_ratio=config.left_arm_config.gripper_torque_ratio,
@@ -86,10 +99,12 @@ class BiRebotB601Follower(Robot):
 
     @property
     def _cameras_ft(self) -> dict[str, tuple]:
-        return {
-            **{f"left_{k}": v for k, v in self.left_arm._cameras_ft.items()},
-            **{f"right_{k}": v for k, v in self.right_arm._cameras_ft.items()},
-        }
+        out: dict[str, tuple] = {}
+        for k, v in self.left_arm._cameras_ft.items():
+            out[k if k in self._top_level_cam_keys else f"left_{k}"] = v
+        for k, v in self.right_arm._cameras_ft.items():
+            out[f"right_{k}"] = v
+        return out
 
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
@@ -99,32 +114,13 @@ class BiRebotB601Follower(Robot):
     def action_features(self) -> dict[str, type]:
         return self._motors_ft
 
-    @property
-    def is_connected(self) -> bool:
-        return self.left_arm.is_connected and self.right_arm.is_connected
-
-    @check_if_already_connected
-    def connect(self, calibrate: bool = True) -> None:
-        self.left_arm.connect(calibrate)
-        self.right_arm.connect(calibrate)
-
-    @property
-    def is_calibrated(self) -> bool:
-        return self.left_arm.is_calibrated and self.right_arm.is_calibrated
-
-    def calibrate(self) -> None:
-        self.left_arm.calibrate()
-        self.right_arm.calibrate()
-
-    def configure(self) -> None:
-        self.left_arm.configure()
-        self.right_arm.configure()
-
     @check_if_not_connected
     def get_observation(self) -> RobotObservation:
-        obs_dict = {}
-        obs_dict.update({f"left_{k}": v for k, v in self.left_arm.get_observation().items()})
-        obs_dict.update({f"right_{k}": v for k, v in self.right_arm.get_observation().items()})
+        obs_dict: RobotObservation = {}
+        for k, v in self.left_arm.get_observation().items():
+            obs_dict[k if k in self._top_level_cam_keys else f"left_{k}"] = v
+        for k, v in self.right_arm.get_observation().items():
+            obs_dict[f"right_{k}"] = v
         return obs_dict
 
     @check_if_not_connected
@@ -143,8 +139,3 @@ class BiRebotB601Follower(Robot):
             **{f"left_{k}": v for k, v in sent_action_left.items()},
             **{f"right_{k}": v for k, v in sent_action_right.items()},
         }
-
-    @check_if_not_connected
-    def disconnect(self) -> None:
-        self.left_arm.disconnect()
-        self.right_arm.disconnect()

src/lerobot/robots/bi_rebot_b601_follower/config_bi_rebot_b601_follower.py

@@ -14,7 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+
+from lerobot.cameras import CameraConfig
 
 from ..config import RobotConfig
 from ..rebot_b601_follower import RebotB601FollowerConfig
@@ -27,3 +29,8 @@ class BiRebotB601FollowerConfig(RobotConfig):
 
     left_arm_config: RebotB601FollowerConfig
     right_arm_config: RebotB601FollowerConfig
+
+    # Top-level cameras not attached to a specific side. Keys are kept as-is in
+    # observations (no `left_`/`right_` prefix). Per-arm cameras (declared on
+    # `{left,right}_arm_config.cameras`) are prefixed.
+    cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/bi_so_follower/bi_so_follower.py

@@ -18,7 +18,8 @@ import logging
 from functools import cached_property
 
 from lerobot.types import RobotAction, RobotObservation
-from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
+from lerobot.utils.bimanual import BimanualMixin
+from lerobot.utils.decorators import check_if_not_connected
 
 from ..robot import Robot
 from ..so_follower import SOFollower, SOFollowerRobotConfig
@@ -27,7 +28,7 @@ from .config_bi_so_follower import BiSOFollowerConfig
 logger = logging.getLogger(__name__)
 
 
-class BiSOFollower(Robot):
+class BiSOFollower(BimanualMixin, Robot):
     """
     [Bimanual SO Follower Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio
     """
@@ -39,6 +40,18 @@ class BiSOFollower(Robot):
         super().__init__(config)
         self.config = config
 
+        # Top-level cameras are opened by `left_arm` for convenience, but their
+        # keys stay unprefixed in observations (tracked via `_top_level_cam_keys`).
+        self._top_level_cam_keys = set(config.cameras)
+        _collisions = self._top_level_cam_keys & set(
+            config.left_arm_config.cameras
+        ) | self._top_level_cam_keys & set(config.right_arm_config.cameras)
+        if _collisions:
+            raise ValueError(
+                f"Top-level camera names collide with per-arm camera names: {sorted(_collisions)}"
+            )
+        left_arm_cameras = {**config.left_arm_config.cameras, **config.cameras}
+
         left_arm_config = SOFollowerRobotConfig(
             id=f"{config.id}_left" if config.id else None,
             calibration_dir=config.calibration_dir,
@@ -46,7 +59,7 @@ class BiSOFollower(Robot):
             disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect,
             max_relative_target=config.left_arm_config.max_relative_target,
             use_degrees=config.left_arm_config.use_degrees,
-            cameras=config.left_arm_config.cameras,
+            cameras=left_arm_cameras,
         )
 
         right_arm_config = SOFollowerRobotConfig(
@@ -77,13 +90,12 @@ class BiSOFollower(Robot):
 
     @property
     def _cameras_ft(self) -> dict[str, tuple]:
-        left_arm_cameras_ft = self.left_arm._cameras_ft
-        right_arm_cameras_ft = self.right_arm._cameras_ft
-
-        return {
-            **{f"left_{k}": v for k, v in left_arm_cameras_ft.items()},
-            **{f"right_{k}": v for k, v in right_arm_cameras_ft.items()},
-        }
+        out: dict[str, tuple] = {}
+        for k, v in self.left_arm._cameras_ft.items():
+            out[k if k in self._top_level_cam_keys else f"left_{k}"] = v
+        for k, v in self.right_arm._cameras_ft.items():
+            out[f"right_{k}"] = v
+        return out
 
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
@@ -93,42 +105,21 @@ class BiSOFollower(Robot):
     def action_features(self) -> dict[str, type]:
         return self._motors_ft
 
-    @property
-    def is_connected(self) -> bool:
-        return self.left_arm.is_connected and self.right_arm.is_connected
-
-    @check_if_already_connected
-    def connect(self, calibrate: bool = True) -> None:
-        self.left_arm.connect(calibrate)
-        self.right_arm.connect(calibrate)
-
-    @property
-    def is_calibrated(self) -> bool:
-        return self.left_arm.is_calibrated and self.right_arm.is_calibrated
-
-    def calibrate(self) -> None:
-        self.left_arm.calibrate()
-        self.right_arm.calibrate()
-
-    def configure(self) -> None:
-        self.left_arm.configure()
-        self.right_arm.configure()
-
     def setup_motors(self) -> None:
         self.left_arm.setup_motors()
         self.right_arm.setup_motors()
 
     @check_if_not_connected
     def get_observation(self) -> RobotObservation:
-        obs_dict = {}
+        obs_dict: RobotObservation = {}
 
-        # Add "left_" prefix
-        left_obs = self.left_arm.get_observation()
-        obs_dict.update({f"left_{key}": value for key, value in left_obs.items()})
+        # Add "left_" prefix to per-arm keys; keep top-level camera keys unprefixed.
+        for key, value in self.left_arm.get_observation().items():
+            obs_dict[key if key in self._top_level_cam_keys else f"left_{key}"] = value
 
         # Add "right_" prefix
-        right_obs = self.right_arm.get_observation()
-        obs_dict.update({f"right_{key}": value for key, value in right_obs.items()})
+        for key, value in self.right_arm.get_observation().items():
+            obs_dict[f"right_{key}"] = value
 
         return obs_dict
 
@@ -151,8 +142,3 @@ class BiSOFollower(Robot):
         prefixed_sent_action_right = {f"right_{key}": value for key, value in sent_action_right.items()}
 
         return {**prefixed_sent_action_left, **prefixed_sent_action_right}
-
-    @check_if_not_connected
-    def disconnect(self):
-        self.left_arm.disconnect()
-        self.right_arm.disconnect()

src/lerobot/robots/bi_so_follower/config_bi_so_follower.py

@@ -14,7 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+
+from lerobot.cameras import CameraConfig
 
 from ..config import RobotConfig
 from ..so_follower import SOFollowerConfig
@@ -27,3 +29,8 @@ class BiSOFollowerConfig(RobotConfig):
 
     left_arm_config: SOFollowerConfig
     right_arm_config: SOFollowerConfig
+
+    # Top-level cameras not attached to a specific side. Keys are kept as-is in
+    # observations (no `left_`/`right_` prefix). Per-arm cameras (declared on
+    # `{left,right}_arm_config.cameras`) are prefixed.
+    cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/rollout/__init__.py

@@ -23,6 +23,7 @@ from .configs import (
     DAggerKeyboardConfig,
     DAggerPedalConfig,
     DAggerStrategyConfig,
+    EpisodicStrategyConfig,
     HighlightStrategyConfig,
     RolloutConfig,
     RolloutStrategyConfig,
@@ -49,6 +50,7 @@ from .inference import (
 from .strategies import (
     BaseStrategy,
     DAggerStrategy,
+    EpisodicStrategy,
     HighlightStrategy,
     RolloutStrategy,
     SentryStrategy,
@@ -66,6 +68,8 @@ __all__ = [
     "HardwareContext",
     "HighlightStrategy",
     "HighlightStrategyConfig",
+    "EpisodicStrategy",
+    "EpisodicStrategyConfig",
     "InferenceEngine",
     "InferenceEngineConfig",
     "PolicyContext",

src/lerobot/rollout/configs.py

@@ -121,6 +121,35 @@ class DAggerPedalConfig:
     upload: str = "KEY_C"
 
 
+@RolloutStrategyConfig.register_subclass("episodic")
+@dataclass
+class EpisodicStrategyConfig(RolloutStrategyConfig):
+    """Episode-oriented recording that mirrors the behavior of ``lerobot-record``.
+
+    Records ``dataset.num_episodes`` episodes of maximum ``dataset.episode_time_s`` each.
+    After each episode, runs ``dataset.reset_time_s`` seconds of reset time.
+
+    Keyboard controls:
+        Right arrow  — end current episode or reset phase early
+        Left arrow   — discard current episode and re-record
+        Escape       — stop recording session
+
+    In between episodes:
+    - if there is no teleop leader, the robot is held at its initial joint positions captured at startup.
+    - else, the robot is moved smoothly to the position of the teleop leader.
+    """
+
+    # This only applies if there are no teleop leaders specified.
+    # When True (default), moves the robot back to the joint positions captured at startup.
+    # Otherwise, leave the robot in its current position.
+    reset_to_initial_position: bool = True
+
+    # Whether to turn on or off the leader -> follower smooth handover behavior.
+    # When False, fallback to follower -> leader handover.
+    # Note that leader -> follower handover is only supported when the leader has `send_feedback` capability.
+    smooth_leader_to_follower_handover: bool = True
+
+
 @RolloutStrategyConfig.register_subclass("dagger")
 @dataclass
 class DAggerStrategyConfig(RolloutStrategyConfig):
@@ -229,7 +258,13 @@ class RolloutConfig:
 
         # TODO(Steven): DAgger shouldn't require a dataset (user may want to just rollout+intervene without recording), but for now we require it to simplify the implementation.
         needs_dataset = isinstance(
-            self.strategy, (SentryStrategyConfig, HighlightStrategyConfig, DAggerStrategyConfig)
+            self.strategy,
+            (
+                SentryStrategyConfig,
+                HighlightStrategyConfig,
+                DAggerStrategyConfig,
+                EpisodicStrategyConfig,
+            ),
         )
         if needs_dataset and (self.dataset is None or not self.dataset.repo_id):
             raise ValueError(f"{self.strategy.type} strategy requires --dataset.repo_id to be set")

src/lerobot/rollout/strategies/__init__.py

@@ -17,6 +17,7 @@
 from .base import BaseStrategy
 from .core import RolloutStrategy, estimate_max_episode_seconds, safe_push_to_hub, send_next_action
 from .dagger import DAggerEvents, DAggerPhase, DAggerStrategy
+from .episodic import EpisodicStrategy
 from .factory import create_strategy
 from .highlight import HighlightStrategy
 from .sentry import SentryStrategy
@@ -27,6 +28,7 @@ __all__ = [
     "DAggerPhase",
     "DAggerStrategy",
     "HighlightStrategy",
+    "EpisodicStrategy",
     "RolloutStrategy",
     "SentryStrategy",
     "create_strategy",

src/lerobot/rollout/strategies/dagger.py

@@ -56,10 +56,14 @@ from typing import Any
 
 import numpy as np
 
-from lerobot.common.control_utils import is_headless
+from lerobot.common.control_utils import (
+    follower_smooth_move_to,
+    is_headless,
+    teleop_smooth_move_to,
+    teleop_supports_feedback,
+)
 from lerobot.datasets import VideoEncodingManager
 from lerobot.datasets.utils import DEFAULT_VIDEO_FILE_SIZE_IN_MB
-from lerobot.teleoperators import Teleoperator
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame
 from lerobot.utils.import_utils import _pynput_available
@@ -69,7 +73,6 @@ from lerobot.utils.utils import log_say
 
 from ..configs import DAggerKeyboardConfig, DAggerPedalConfig, DAggerStrategyConfig
 from ..context import RolloutContext
-from ..robot_wrapper import ThreadSafeRobot
 from .core import RolloutStrategy, estimate_max_episode_seconds, safe_push_to_hub, send_next_action
 
 PYNPUT_AVAILABLE = _pynput_available
@@ -171,64 +174,6 @@ class DAggerEvents:
         self.upload_requested.clear()
 
 
-# ---------------------------------------------------------------------------
-# Teleoperator helpers
-# ---------------------------------------------------------------------------
-
-
-def _teleop_supports_feedback(teleop: Teleoperator) -> bool:
-    """Return True when the teleop can receive position feedback (is actuated).
-    TODO(Maxime): See if it is possible to unify this interface across teleops instead of duck-typing.
-    """
-    return (
-        bool(teleop.feedback_features)
-        and hasattr(teleop, "disable_torque")
-        and hasattr(teleop, "enable_torque")
-    )
-
-
-def _teleop_smooth_move_to(
-    teleop: Teleoperator, target_pos: dict, duration_s: float = 2.0, fps: int = 30
-) -> None:
-    """Smoothly move an actuated teleop to ``target_pos`` via linear interpolation.
-
-    Requires the teleoperator to support feedback
-    (i.e. have non-empty ``feedback_features`` and implement ``disable_torque`` / ``enable_torque``).
-
-    TODO(Maxime): This blocks up to ``duration_s`` seconds, during this time
-    the follower robot doesn't receive new actions, this could be an issue on LeKiwi.
-    """
-    teleop.enable_torque()
-    current = teleop.get_action()
-    steps = max(int(duration_s * fps), 1)
-
-    for step in range(steps + 1):
-        t = step / steps
-        interp = {
-            k: current[k] * (1 - t) + target_pos[k] * t if k in target_pos else current[k] for k in current
-        }
-        teleop.send_feedback(interp)
-        time.sleep(1 / fps)
-
-
-def _follower_smooth_move_to(
-    robot: ThreadSafeRobot, current: dict, target: dict, duration_s: float = 1.0, fps: int = 30
-) -> None:
-    """Smoothly move the follower robot from ``current`` to ``target`` action.
-
-    Used when the teleop is non-actuated: instead of driving the leader arm
-    to the follower, we bring the follower to the teleop's current pose.
-    Both ``current`` and ``target`` must be in robot-action key space.
-    """
-    steps = max(int(duration_s * fps), 1)
-
-    for step in range(steps + 1):
-        t = step / steps
-        interp = {k: current[k] * (1 - t) + target[k] * t if k in target else current[k] for k in current}
-        robot.send_action(interp)
-        time.sleep(1 / fps)
-
-
 # ---------------------------------------------------------------------------
 # Input device handlers
 # ---------------------------------------------------------------------------
@@ -756,31 +701,31 @@ class DAggerStrategy(RolloutStrategy):
             logger.info("Pausing engine - robot holds position")
             engine.pause()
 
-            if _teleop_supports_feedback(teleop) and prev_action is not None:
+            if teleop_supports_feedback(teleop) and prev_action is not None:
                 # TODO(Maxime): prev_action is in robot action key space (output of robot_action_processor).
                 # send_feedback expects teleop feedback key space. For homogeneous setups (e.g. SO-101
                 # leader + SO-101 follower) the keys are identical so this works. If the processor pipeline
                 # does non-trivial key renaming (e.g. a rename_map on action keys), the interpolation in
-                # _teleop_smooth_move_to silently no-ops and the arm doesn't move.
+                # teleop_smooth_move_to silently no-ops and the arm doesn't move.
                 logger.info("Smooth handover: moving leader arm to follower position")
-                _teleop_smooth_move_to(teleop, prev_action)
+                teleop_smooth_move_to(teleop, prev_action)
 
         elif old_phase == DAggerPhase.PAUSED and new_phase == DAggerPhase.CORRECTING:
             logger.info("Entering correction mode - human teleop control")
-            if not _teleop_supports_feedback(teleop) and prev_action is not None:
+            if not teleop_supports_feedback(teleop) and prev_action is not None:
                 logger.info("Smooth handover: sliding follower to teleop position")
                 obs = robot.get_observation()
                 teleop_action = teleop.get_action()
                 processed = ctx.processors.teleop_action_processor((teleop_action, obs))
                 target = ctx.processors.robot_action_processor((processed, obs))
-                _follower_smooth_move_to(robot, prev_action, target)
+                follower_smooth_move_to(robot, prev_action, target)
 
             # unlock the teleop for human control
-            if _teleop_supports_feedback(teleop):
+            if teleop_supports_feedback(teleop):
                 teleop.disable_torque()
 
         elif old_phase == DAggerPhase.CORRECTING and new_phase == DAggerPhase.PAUSED:
-            if _teleop_supports_feedback(teleop):
+            if teleop_supports_feedback(teleop):
                 teleop.enable_torque()
 
         elif new_phase == DAggerPhase.AUTONOMOUS:
@@ -790,7 +735,7 @@ class DAggerStrategy(RolloutStrategy):
             engine.resume()
 
             # release teleop before resuming the policy
-            if _teleop_supports_feedback(teleop):
+            if teleop_supports_feedback(teleop):
                 teleop.disable_torque()
 
     # ------------------------------------------------------------------

src/lerobot/rollout/strategies/episodic.py

@@ -0,0 +1,335 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Episodic rollout strategy: mirrors the behavior of ``lerobot-record``.
+
+- Policy drives the robot during each recording episode.
+- An optional teleoperator can drive the robot during reset phases so the
+  operator can bring the environment back to its starting configuration.
+  If no teleop is connected the robot stays in its current position.
+- Keyboard controls:
+
+      Right arrow  — end the current episode or reset phase early
+      Left arrow   — discard the current episode and re-record it
+      Escape       — stop the recording session
+
+Dataset naming follows the rollout convention: repo names must start with ``rollout_``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import time
+
+from lerobot.common.control_utils import (
+    follower_smooth_move_to,
+    init_keyboard_listener,
+    is_headless,
+    teleop_smooth_move_to,
+    teleop_supports_feedback,
+)
+from lerobot.datasets import VideoEncodingManager
+from lerobot.utils.constants import ACTION, OBS_STR
+from lerobot.utils.feature_utils import build_dataset_frame
+from lerobot.utils.robot_utils import precise_sleep
+from lerobot.utils.utils import log_say
+from lerobot.utils.visualization_utils import log_rerun_data
+
+from ..configs import EpisodicStrategyConfig
+from ..context import RolloutContext
+from .core import RolloutStrategy, safe_push_to_hub, send_next_action
+
+logger = logging.getLogger(__name__)
+
+
+class EpisodicStrategy(RolloutStrategy):
+    """Policy-driven multi-episode recording, mirrors the behavior of ``lerobot-record``.
+
+    Each recording episode runs the policy for maximum ``dataset.episode_time_s``
+    seconds, recording every frame.  A reset phase of ``dataset.reset_time_s``
+    follows every episode (except the last) so the operator can manually
+    reset the environment.  During the reset phase, an optional teleoperator
+    drives the robot; if none is present the robot returns to its initial joint positions captured at startup.
+
+    The policy state (hidden state, RTC queue, interpolator) is reset at
+    the start of each recording episode.
+
+    Keyboard events:
+        right arrow  → end current episode or reset phase early
+        left arrow   → discard & re-record current episode
+        ESC          → stop the session
+    """
+
+    config: EpisodicStrategyConfig
+
+    def __init__(self, config: EpisodicStrategyConfig) -> None:
+        super().__init__(config)
+        self._listener = None
+        self._events: dict | None = None
+
+    def setup(self, ctx: RolloutContext) -> None:
+        """Start the inference engine and attach the keyboard listener."""
+        self._init_engine(ctx)
+        self._listener, self._events = init_keyboard_listener()
+        logger.info("Episodic strategy ready")
+
+    def run(self, ctx: RolloutContext) -> None:
+        """Main multi-episode recording loop."""
+        cfg = ctx.runtime.cfg
+        dataset_cfg = cfg.dataset
+        robot = ctx.hardware.robot_wrapper
+        teleop = ctx.hardware.teleop
+        dataset = ctx.data.dataset
+        events = self._events
+        features = ctx.data.dataset_features
+
+        fps = cfg.fps
+        episode_time_s = dataset_cfg.episode_time_s
+        reset_time_s = dataset_cfg.reset_time_s
+        num_episodes = dataset_cfg.num_episodes
+        single_task = dataset_cfg.single_task or cfg.task
+        play_sounds = cfg.play_sounds
+
+        display_compressed = (
+            True
+            if (cfg.display_data and cfg.display_ip is not None and cfg.display_port is not None)
+            else cfg.display_compressed_images
+        )
+
+        with VideoEncodingManager(dataset):
+            try:
+                recorded_episodes = 0
+                while recorded_episodes < num_episodes and not events["stop_recording"]:
+                    if ctx.runtime.shutdown_event.is_set():
+                        break
+
+                    # Reset policy state at episode start (discard leftover hidden state / queue)
+                    self._engine.reset()
+                    self._interpolator.reset()
+                    self._engine.resume()
+
+                    log_say(f"Recording episode {dataset.num_episodes}", play_sounds)
+                    self._policy_loop(
+                        ctx=ctx,
+                        robot=robot,
+                        events=events,
+                        features=features,
+                        fps=fps,
+                        control_time_s=episode_time_s,
+                        dataset=dataset,
+                        single_task=single_task,
+                    )
+
+                    # Reset phase, skip after the last episode (but run when re-recording)
+                    if not events["stop_recording"] and (
+                        recorded_episodes < num_episodes - 1 or events["rerecord_episode"]
+                    ):
+                        log_say("Reset the environment", play_sounds)
+
+                        if teleop:
+                            # Smooth handover so the transition to teleop control is jerk-free.
+                            # For actuated teleops: drive the leader arm to the follower's current
+                            # position so the operator takes over without fighting the arm.
+                            # For non-actuated teleops: slide the follower to the teleop's current
+                            # pose instead, since the leader cannot be driven.
+                            obs = robot.get_observation()
+                            current_pos = {k: v for k, v in obs.items() if k.endswith(".pos")}
+                            if (
+                                teleop_supports_feedback(teleop)
+                                and self.config.smooth_leader_to_follower_handover
+                            ):
+                                logger.info("Smooth handover: moving leader arm to follower position")
+                                teleop_smooth_move_to(teleop, current_pos, duration_s=2)
+                                teleop.disable_torque()
+                            else:
+                                logger.info("Smooth handover: sliding follower to teleop position")
+                                teleop_action = teleop.get_action()
+                                processed = ctx.processors.teleop_action_processor((teleop_action, obs))
+                                target = ctx.processors.robot_action_processor((processed, obs))
+                                follower_smooth_move_to(robot, current_pos, target, duration_s=1)
+
+                        elif self.config.reset_to_initial_position:
+                            # No teleop: return the robot to its startup position.
+                            self._return_to_initial_position(hw=ctx.hardware, duration_s=1)
+
+                        self._reset_loop(
+                            ctx=ctx,
+                            robot=robot,
+                            teleop=teleop,
+                            events=events,
+                            fps=fps,
+                            control_time_s=reset_time_s,
+                            display_data=cfg.display_data,
+                            display_compressed=display_compressed,
+                        )
+
+                    if events["rerecord_episode"]:
+                        log_say("Re-record episode", play_sounds)
+                        events["rerecord_episode"] = False
+                        events["exit_early"] = False
+                        dataset.clear_episode_buffer()
+
+                        # returns to its initial joint positions captured at startup
+                        if not teleop and self.config.reset_to_initial_position:
+                            self._return_to_initial_position(hw=ctx.hardware, duration_s=1)
+
+                        continue
+
+                    dataset.save_episode()
+                    recorded_episodes += 1
+            finally:
+                # Save any frames buffered in the current episode so an unexpected
+                # exception or KeyboardInterrupt does not silently drop recorded data.
+                # suppress: save_episode raises if the buffer is empty (nothing to lose).
+                logger.info("Episodic control loop ended — saving any in-progress episode")
+                with contextlib.suppress(Exception):
+                    dataset.save_episode()
+
+    def _policy_loop(
+        self,
+        ctx: RolloutContext,
+        robot,
+        events: dict,
+        features: dict,
+        fps: float,
+        control_time_s: float,
+        dataset,
+        single_task: str,
+    ) -> None:
+        """Policy-driven recording loop for a single episode."""
+        interpolator = self._interpolator
+        control_interval = interpolator.get_control_interval(fps)
+
+        timestamp = 0.0
+        start_t = time.perf_counter()
+
+        while timestamp < control_time_s:
+            loop_start = time.perf_counter()
+
+            if events["exit_early"]:
+                events["exit_early"] = False
+                break
+
+            if ctx.runtime.shutdown_event.is_set():
+                break
+
+            obs = robot.get_observation()
+            obs_processed = self._process_observation_and_notify(ctx.processors, obs)
+
+            if self._handle_warmup(ctx.runtime.cfg.use_torch_compile, loop_start, control_interval):
+                continue
+
+            action_dict = send_next_action(obs_processed, obs, ctx, interpolator)
+
+            if action_dict is not None:
+                obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
+                action_frame = build_dataset_frame(features, action_dict, prefix=ACTION)
+                dataset.add_frame({**obs_frame, **action_frame, "task": single_task})
+                self._log_telemetry(obs_processed, action_dict, ctx.runtime)
+
+            dt = time.perf_counter() - loop_start
+            sleep_t = control_interval - dt
+            if sleep_t < 0:
+                logger.warning(
+                    f"Record loop is running slower ({1 / dt:.1f} Hz) than the target FPS ({fps} Hz). "
+                    "Dataset frames might be dropped and robot control might be unstable. "
+                    "Common causes are: 1) Camera FPS not keeping up 2) Policy inference taking too long "
+                    "3) CPU starvation"
+                )
+            precise_sleep(max(sleep_t, 0.0))
+            timestamp = time.perf_counter() - start_t
+
+    def _reset_loop(
+        self,
+        ctx: RolloutContext,
+        robot,
+        teleop,
+        events: dict,
+        fps: float,
+        control_time_s: float,
+        display_data: bool,
+        display_compressed: bool,
+    ) -> None:
+        """Reset-phase loop: teleop drives the robot if available, no recording."""
+        processors = ctx.processors
+        control_interval = 1.0 / fps
+
+        timestamp = 0.0
+        start_t = time.perf_counter()
+
+        while timestamp < control_time_s:
+            loop_start = time.perf_counter()
+
+            if events["exit_early"]:
+                events["exit_early"] = False
+                break
+
+            if ctx.runtime.shutdown_event.is_set():
+                break
+
+            obs = robot.get_observation()
+
+            if teleop is not None:
+                act = teleop.get_action()
+                act_teleop = processors.teleop_action_processor((act, obs))
+                robot_action = processors.robot_action_processor((act_teleop, obs))
+                robot.send_action(robot_action)
+
+                if display_data:
+                    obs_processed = processors.robot_observation_processor(obs)
+                    log_rerun_data(
+                        observation=obs_processed,
+                        action=act_teleop,
+                        compress_images=display_compressed,
+                    )
+
+            dt = time.perf_counter() - loop_start
+            sleep_t = control_interval - dt
+            precise_sleep(max(sleep_t, 0.0))
+            timestamp = time.perf_counter() - start_t
+
+    def teardown(self, ctx: RolloutContext) -> None:
+        """Finalise dataset, stop listener, push to hub, and disconnect hardware."""
+        cfg = ctx.runtime.cfg
+        play_sounds = cfg.play_sounds
+
+        log_say("Stop recording", play_sounds, blocking=True)
+
+        if not is_headless() and self._listener is not None:
+            self._listener.stop()
+
+        if ctx.data.dataset is not None:
+            logger.info("Finalizing dataset...")
+            ctx.data.dataset.finalize()
+
+        if (
+            cfg.dataset is not None
+            and cfg.dataset.push_to_hub
+            and ctx.data.dataset is not None
+            and safe_push_to_hub(
+                ctx.data.dataset,
+                tags=cfg.dataset.tags,
+                private=cfg.dataset.private,
+            )
+        ):
+            logger.info("Dataset uploaded to hub")
+            log_say("Dataset uploaded to hub", play_sounds)
+
+        self._teardown_hardware(
+            ctx.hardware,
+            return_to_initial_position=cfg.return_to_initial_position,
+        )
+        log_say("Exiting", play_sounds)
+        logger.info("Episodic strategy teardown complete")

src/lerobot/rollout/strategies/factory.py

@@ -21,6 +21,7 @@ from typing import TYPE_CHECKING
 from .base import BaseStrategy
 from .core import RolloutStrategy
 from .dagger import DAggerStrategy
+from .episodic import EpisodicStrategy
 from .highlight import HighlightStrategy
 from .sentry import SentryStrategy
 
@@ -42,4 +43,8 @@ def create_strategy(config: RolloutStrategyConfig) -> RolloutStrategy:
         return HighlightStrategy(config)
     if config.type == "dagger":
         return DAggerStrategy(config)
-    raise ValueError(f"Unknown strategy type '{config.type}'. Available: base, sentry, highlight, dagger")
+    if config.type == "episodic":
+        return EpisodicStrategy(config)
+    raise ValueError(
+        f"Unknown strategy type '{config.type}'. Available: base, sentry, highlight, dagger, episodic"
+    )

src/lerobot/scripts/lerobot_annotate.py

@@ -0,0 +1,206 @@
+#!/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.
+"""``lerobot-annotate`` — populate ``language_persistent`` and
+``language_events`` columns on a LeRobot dataset.
+
+Annotations live directly in ``data/chunk-*/file-*.parquet``.
+
+Example:
+
+  uv run lerobot-annotate \\
+      --root=/path/to/dataset \\
+      --vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct
+
+For distributed runs, see ``examples/annotations/run_hf_job.py``.
+"""
+
+import logging
+from pathlib import Path
+
+from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
+from lerobot.annotations.steerable_pipeline.executor import Executor
+from lerobot.annotations.steerable_pipeline.frames import make_frame_provider
+from lerobot.annotations.steerable_pipeline.modules import (
+    GeneralVqaModule,
+    InterjectionsAndSpeechModule,
+    PlanSubtasksMemoryModule,
+)
+from lerobot.annotations.steerable_pipeline.validator import StagingValidator
+from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
+from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
+from lerobot.configs import parser
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
+    if cfg.root is not None:
+        return Path(cfg.root)
+    if cfg.repo_id is not None:
+        from huggingface_hub import snapshot_download
+
+        return Path(snapshot_download(repo_id=cfg.repo_id, repo_type="dataset"))
+    raise ValueError("Either --root or --repo_id must be provided.")
+
+
+@parser.wrap()
+def annotate(cfg: AnnotationPipelineConfig) -> None:
+    """Run the steerable annotation pipeline against a dataset."""
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+    root = _resolve_root(cfg)
+    logger.info("annotate: root=%s", root)
+
+    vlm = make_vlm_client(cfg.vlm)
+    frame_provider = make_frame_provider(root, camera_key=cfg.vlm.camera_key, video_backend=cfg.video_backend)
+    # Surface the resolved cameras up front so a silent vqa-module no-op
+    # is obvious in job output rather than discovered post-hoc by counting
+    # parquet rows.
+    cam_keys = list(getattr(frame_provider, "camera_keys", []) or [])
+    logger.info(
+        "annotate: frame_provider default camera=%r, all cameras=%s",
+        getattr(frame_provider, "camera_key", None),
+        cam_keys,
+    )
+    if cfg.vqa.enabled and not cam_keys:
+        logger.warning(
+            "annotate: the vqa module is enabled but no cameras were "
+            "resolved — it will produce zero VQA rows. Check "
+            "meta/info.json for observation.images.* features, or pass "
+            "--vlm.camera_key=<key> to seed the cameras list."
+        )
+    plan = PlanSubtasksMemoryModule(vlm=vlm, config=cfg.plan, frame_provider=frame_provider)
+    interjections = InterjectionsAndSpeechModule(
+        vlm=vlm, config=cfg.interjections, seed=cfg.seed, frame_provider=frame_provider
+    )
+    vqa = GeneralVqaModule(vlm=vlm, config=cfg.vqa, seed=cfg.seed, frame_provider=frame_provider)
+    writer = LanguageColumnsWriter()
+    validator = StagingValidator(
+        dataset_camera_keys=tuple(getattr(frame_provider, "camera_keys", []) or []) or None,
+    )
+
+    executor = Executor(
+        config=cfg,
+        plan=plan,
+        interjections=interjections,
+        vqa=vqa,
+        writer=writer,
+        validator=validator,
+    )
+    summary = executor.run(root)
+    logger.info("annotate: wrote %d shard(s)", len(summary.written_paths))
+    for phase in summary.phases:
+        logger.info(
+            "annotate: phase=%s processed=%d skipped=%d",
+            phase.name,
+            phase.episodes_processed,
+            phase.episodes_skipped,
+        )
+    if summary.validation_report.warnings:
+        for w in summary.validation_report.warnings:
+            logger.warning(w)
+
+    if cfg.push_to_hub:
+        if cfg.repo_id is None and cfg.new_repo_id is None:
+            raise ValueError(
+                "--push_to_hub requires --repo_id or --new_repo_id (the dataset repo to push to)."
+            )
+        _push_to_hub(root, cfg)
+
+
+def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
+    """Upload the annotated dataset directory to the Hub.
+
+    Pushes to ``cfg.new_repo_id`` when set, otherwise back to ``cfg.repo_id``.
+    """
+    from huggingface_hub import HfApi  # noqa: PLC0415
+
+    repo_id = cfg.new_repo_id or cfg.repo_id
+    commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
+    api = HfApi()
+    print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type="dataset",
+        private=cfg.push_private,
+        exist_ok=True,
+    )
+    print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
+    commit_info = api.upload_folder(
+        folder_path=str(root),
+        repo_id=repo_id,
+        repo_type="dataset",
+        commit_message=commit_message,
+        ignore_patterns=[".annotate_staging/**", "**/.DS_Store"],
+    )
+    print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
+
+    # Tag the upload with the codebase version. ``LeRobotDatasetMetadata``
+    # resolves the dataset revision via ``get_safe_version`` which scans
+    # for tags like ``v3.0``; without a tag it raises
+    # ``RevisionNotFoundError``. Read the version straight from the
+    # dataset's own ``meta/info.json`` so we tag whatever the writer
+    # actually wrote (no accidental drift if the codebase floor moves).
+    from lerobot.datasets.dataset_metadata import CODEBASE_VERSION  # noqa: PLC0415
+
+    info_path = root / "meta" / "info.json"
+    version_tag = CODEBASE_VERSION
+    if info_path.exists():
+        try:
+            from lerobot.utils.io_utils import load_json  # noqa: PLC0415
+
+            info = load_json(info_path)
+            ds_version = info.get("codebase_version")
+            if isinstance(ds_version, str) and ds_version.startswith("v"):
+                version_tag = ds_version
+        except Exception as exc:  # noqa: BLE001
+            print(
+                f"[lerobot-annotate] could not read codebase_version from info.json ({exc}); falling back to {version_tag}",
+                flush=True,
+            )
+    revision = getattr(commit_info, "oid", None)
+    tag_kwargs = {
+        "repo_id": repo_id,
+        "tag": version_tag,
+        "repo_type": "dataset",
+    }
+    if revision is not None:
+        tag_kwargs["revision"] = revision
+
+    try:
+        from contextlib import suppress  # noqa: PLC0415
+
+        from huggingface_hub.errors import RevisionNotFoundError  # noqa: PLC0415
+
+        with suppress(RevisionNotFoundError):
+            api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
+        api.create_tag(**tag_kwargs)
+        print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
+    except Exception as exc:  # noqa: BLE001
+        print(
+            f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
+            "Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
+            "Run: from huggingface_hub import HfApi; "
+            f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
+            flush=True,
+        )
+
+
+def main() -> None:
+    annotate()
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/scripts/lerobot_calibrate.py

@@ -54,6 +54,7 @@ from lerobot.teleoperators import (  # noqa: F401
     Teleoperator,
     TeleoperatorConfig,
     bi_openarm_leader,
+    bi_openarm_mini,
     bi_rebot_102_leader,
     bi_so_leader,
     homunculus,

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -94,6 +94,14 @@ Merge multiple datasets from a list of local dataset paths:
         --operation.repo_ids "['pusht_train', 'pusht_val']" \
         --operation.roots "['/path/to/pusht_train', '/path/to/pusht_val']"
 
+Merge multiple datasets while keeping one file per source file (no video/data stitching):
+    lerobot-edit-dataset \
+        --new_repo_id lerobot/pusht_merged \
+        --operation.type merge \
+        --operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']" \
+        --operation.concatenate_videos false \
+        --operation.concatenate_data false
+
 Remove camera feature:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht \
@@ -257,6 +265,9 @@ class SplitConfig(OperationConfig):
 class MergeConfig(OperationConfig):
     repo_ids: list[str] | None = None
     roots: list[str] | None = None
+    # When False, keep one file per source file instead of packing into shards.
+    concatenate_videos: bool = True
+    concatenate_data: bool = True
 
 
 @OperationConfig.register_subclass("remove_feature")
@@ -461,6 +472,8 @@ def handle_merge(cfg: EditDatasetConfig) -> None:
         datasets,
         output_repo_id=cfg.new_repo_id,
         output_dir=output_dir,
+        concatenate_videos=cfg.operation.concatenate_videos,
+        concatenate_data=cfg.operation.concatenate_data,
     )
 
     logging.info(f"Merged dataset saved to {output_dir}")

src/lerobot/scripts/lerobot_find_joint_limits.py

@@ -57,6 +57,7 @@ from lerobot.robots import (  # noqa: F401
 from lerobot.teleoperators import (  # noqa: F401
     TeleoperatorConfig,
     bi_openarm_leader,
+    bi_openarm_mini,
     bi_rebot_102_leader,
     bi_so_leader,
     gamepad,

src/lerobot/scripts/lerobot_record.py

@@ -137,6 +137,7 @@ from lerobot.teleoperators import (  # noqa: F401
     Teleoperator,
     TeleoperatorConfig,
     bi_openarm_leader,
+    bi_openarm_mini,
     bi_rebot_102_leader,
     bi_so_leader,
     homunculus,

src/lerobot/scripts/lerobot_rollout.py

@@ -25,6 +25,7 @@ Strategies
     --strategy.type=sentry     Continuous recording with auto-upload
     --strategy.type=highlight  Ring buffer + keystroke save
     --strategy.type=dagger     Human-in-the-loop (DAgger / RaC)
+    --strategy.type=episodic   Episode-oriented recording with reset phases
 
 Inference backends
 ------------------
@@ -111,6 +112,18 @@ Usage examples
         --display_data=true \\
         --use_torch_compile=true
 
+    # Episodic mode — episode-oriented recording with reset phases
+    lerobot-rollout \\
+        --strategy.type=episodic \\
+        --policy.path=user/my_policy \\
+        --robot.type=so100_follower \\
+        --robot.port=/dev/ttyACM0 \\
+        --teleop.type=so100_leader \\
+        --teleop.port=/dev/ttyACM1 \\
+        --dataset.repo_id=user/rollout_episodic_data \\
+        --dataset.num_episodes=20 \\
+        --dataset.single_task="Grab the cube"
+
     # Resume a previous sentry recording session
     lerobot-rollout \\
         --strategy.type=sentry \\
@@ -161,6 +174,7 @@ from lerobot.teleoperators import (  # noqa: F401
     Teleoperator,
     TeleoperatorConfig,
     bi_openarm_leader,
+    bi_openarm_mini,
     bi_rebot_102_leader,
     bi_so_leader,
     homunculus,

src/lerobot/scripts/lerobot_setup_motors.py

@@ -41,6 +41,7 @@ from lerobot.robots import (  # noqa: F401
 )
 from lerobot.teleoperators import (  # noqa: F401
     TeleoperatorConfig,
+    bi_openarm_mini,
     bi_rebot_102_leader,
     bi_so_leader,
     koch_leader,

src/lerobot/scripts/lerobot_teleoperate.py

@@ -89,6 +89,7 @@ from lerobot.teleoperators import (  # noqa: F401
     Teleoperator,
     TeleoperatorConfig,
     bi_openarm_leader,
+    bi_openarm_mini,
     bi_rebot_102_leader,
     bi_so_leader,
     gamepad,

src/lerobot/scripts/lerobot_train.py

@@ -36,6 +36,8 @@ from tqdm import tqdm
 from lerobot.common.train_utils import (
     get_step_checkpoint_dir,
     get_step_identifier,
+    load_training_batch_size,
+    load_training_num_processes,
     load_training_state,
     save_checkpoint,
     update_last_checkpoint,
@@ -43,7 +45,7 @@ from lerobot.common.train_utils import (
 from lerobot.common.wandb_utils import WandBLogger
 from lerobot.configs import parser
 from lerobot.configs.train import TrainPipelineConfig
-from lerobot.datasets import EpisodeAwareSampler, make_dataset
+from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state, make_dataset
 from lerobot.envs import close_envs, make_env, make_env_pre_post_processors
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
@@ -99,6 +101,9 @@ def update_policy(
     start_time = time.perf_counter()
     policy.train()
 
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+
     # Compute sample weights if a weighter is provided
     sample_weights = None
     weight_stats = None
@@ -158,6 +163,8 @@ def update_policy(
     train_metrics.grad_norm = grad_norm.item()
     train_metrics.lr = optimizer.param_groups[0]["lr"]
     train_metrics.update_s = time.perf_counter() - start_time
+    if torch.cuda.is_available():
+        train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
     return train_metrics, output_dict
 
 
@@ -232,14 +239,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
 
-    # Dataset loading synchronization: main process downloads first to avoid race conditions
+    # Dataset loading synchronization: the global main process downloads once to the shared
+    # dataset root, then a barrier lets every other rank read the already-populated copy.
+    # LeRobotDataset skips its snapshot_download when try_load() succeeds, so no rank re-downloads.
     if is_main_process:
         logging.info("Creating dataset")
         dataset = make_dataset(cfg)
 
     accelerator.wait_for_everyone()
 
-    # Now all other processes can safely load the dataset
+    # Other ranks read from the shared copy populated by the main process.
     if not is_main_process:
         dataset = make_dataset(cfg)
 
@@ -336,6 +345,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         preprocessor, postprocessor = make_pre_post_processors(
             policy_cfg=cfg.policy,
             pretrained_path=processor_pretrained_path,
+            pretrained_revision=getattr(cfg.policy, "pretrained_revision", None),
             **processor_kwargs,
         )
 
@@ -384,15 +394,47 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
 
     # create dataloader for offline training
-    if hasattr(active_cfg, "drop_n_last_frames"):
+    if not cfg.dataset.streaming:
+        # All non-streaming (map-style) datasets use EpisodeAwareSampler.
+        # The order is a pure function of (seed, epoch), so every rank independently produces the
+        # same permutation. accelerate then shards it disjointly across ranks via BatchSamplerShard
+        # without needing a `generator` attribute to synchronize an RNG, and resume is sample-exact.
         shuffle = False
         sampler = EpisodeAwareSampler(
             dataset.meta.episodes["dataset_from_index"],
             dataset.meta.episodes["dataset_to_index"],
             episode_indices_to_use=dataset.episodes,
-            drop_n_last_frames=active_cfg.drop_n_last_frames,
+            drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
             shuffle=True,
+            seed=cfg.seed if cfg.seed is not None else 0,
         )
+        if cfg.resume and step > 0:
+            # The resume offset depends on the (num_processes, batch_size) that produced `step`, so
+            # use the values recorded in the checkpoint (falling back to the current ones for older
+            # ckpts that did not store them).
+            saved_num_processes = load_training_num_processes(cfg.checkpoint_path)
+            saved_batch_size = load_training_batch_size(cfg.checkpoint_path)
+            ckpt_num_processes = saved_num_processes or accelerator.num_processes
+            ckpt_batch_size = saved_batch_size or cfg.batch_size
+            if is_main_process and saved_num_processes not in (None, accelerator.num_processes):
+                logging.warning(
+                    f"Resuming with num_processes={accelerator.num_processes} but the checkpoint was "
+                    f"written with num_processes={saved_num_processes}. The data order resumes at the "
+                    "right epoch/offset, but per-rank sample-exactness requires the same world size."
+                )
+            if is_main_process and saved_batch_size not in (None, cfg.batch_size):
+                logging.warning(
+                    f"Resuming with batch_size={cfg.batch_size} but the checkpoint was written with "
+                    f"batch_size={saved_batch_size}. The data order resumes at the right epoch/offset, "
+                    "but per-rank sample-exactness requires the same batch size."
+                )
+            sampler_state = compute_sampler_state(step, len(sampler), ckpt_batch_size, ckpt_num_processes)
+            sampler.load_state_dict(sampler_state)
+            if is_main_process:
+                logging.info(
+                    f"Resuming data order at epoch {sampler_state['epoch']}, "
+                    f"sample {sampler_state['start_index']}"
+                )
     else:
         shuffle = True
         sampler = None
@@ -424,12 +466,22 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     policy.train()
 
     train_metrics = {
-        "loss": AverageMeter("loss", ":.3f"),
+        # Per-rank loss reflects only one shard of the global batch; mean recovers the loss DDP
+        # is actually optimizing. grad_norm and lr are already identical on every rank (post
+        # gradient sync / deterministic scheduler) so reducing them would be a no-op collective.
+        "loss": AverageMeter("loss", ":.3f", reduction="mean"),
         "grad_norm": AverageMeter("grdn", ":.3f"),
         "lr": AverageMeter("lr", ":0.1e"),
-        "update_s": AverageMeter("updt_s", ":.3f"),
-        "dataloading_s": AverageMeter("data_s", ":.3f"),
+        # Report the slowest rank for bottleneck-style timings so multi-GPU runs surface the
+        # true straggler instead of rank 0's view.
+        "update_s": AverageMeter("updt_s", ":.3f", reduction="max"),
+        "dataloading_s": AverageMeter("data_s", ":.3f", reduction="max"),
+        # Derived from the post-reduce max step time; set once per log window on the main rank.
+        "samples_per_s": AverageMeter("smp/s", ":.0f"),
     }
+    if torch.cuda.is_available():
+        # max() because headroom is gated by the worst-case rank.
+        train_metrics["gpu_mem_gb"] = AverageMeter("mem_gb", ":.2f", reduction="max")
 
     # Keep global batch size for logging; MetricsTracker handles world size internally.
     effective_batch_size = cfg.batch_size * accelerator.num_processes
@@ -481,21 +533,29 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         if is_main_process:
             progbar.update(1)
         train_tracker.step()
-        is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 and is_main_process
+        is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0
         is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
         is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
 
         if is_log_step:
-            logging.info(train_tracker)
-            if wandb_logger:
-                wandb_log_dict = train_tracker.to_dict()
-                if output_dict:
-                    wandb_log_dict.update(output_dict)
-                # Log sample weighting statistics if enabled
-                if sample_weighter is not None:
-                    weighter_stats = sample_weighter.get_stats()
-                    wandb_log_dict.update({f"sample_weighting/{k}": v for k, v in weighter_stats.items()})
-                wandb_logger.log_dict(wandb_log_dict, step)
+            # Collective reduce must run on every rank, before the main-process gate below.
+            train_tracker.reduce_across_ranks()
+            if is_main_process:
+                # Cluster-wide throughput, derived from the already-reduced (max) step time so it
+                # reflects the slowest rank — which is what actually gates the next iteration.
+                step_time = train_tracker.update_s.avg + train_tracker.dataloading_s.avg
+                if step_time > 0:
+                    train_tracker.samples_per_s = effective_batch_size / step_time
+                logging.info(train_tracker)
+                if wandb_logger:
+                    wandb_log_dict = train_tracker.to_dict()
+                    if output_dict:
+                        wandb_log_dict.update(output_dict)
+                    # Log sample weighting statistics if enabled
+                    if sample_weighter is not None:
+                        weighter_stats = sample_weighter.get_stats()
+                        wandb_log_dict.update({f"sample_weighting/{k}": v for k, v in weighter_stats.items()})
+                    wandb_logger.log_dict(wandb_log_dict, step)
             train_tracker.reset_averages()
 
         if cfg.save_checkpoint and is_saving_step:
@@ -511,6 +571,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     scheduler=lr_scheduler,
                     preprocessor=preprocessor,
                     postprocessor=postprocessor,
+                    num_processes=accelerator.num_processes,
+                    batch_size=cfg.batch_size,
                 )
                 update_last_checkpoint(checkpoint_dir)
                 if wandb_logger:

src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py

@@ -18,7 +18,8 @@ import logging
 from functools import cached_property
 
 from lerobot.types import RobotAction
-from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
+from lerobot.utils.bimanual import BimanualMixin
+from lerobot.utils.decorators import check_if_not_connected
 
 from ..openarm_leader import OpenArmLeader, OpenArmLeaderConfig
 from ..teleoperator import Teleoperator
@@ -27,7 +28,7 @@ from .config_bi_openarm_leader import BiOpenArmLeaderConfig
 logger = logging.getLogger(__name__)
 
 
-class BiOpenArmLeader(Teleoperator):
+class BiOpenArmLeader(BimanualMixin, Teleoperator):
     """
     Bimanual OpenArm Leader Arms
     """
@@ -86,27 +87,6 @@ class BiOpenArmLeader(Teleoperator):
     def feedback_features(self) -> dict[str, type]:
         return {}
 
-    @property
-    def is_connected(self) -> bool:
-        return self.left_arm.is_connected and self.right_arm.is_connected
-
-    @check_if_already_connected
-    def connect(self, calibrate: bool = True) -> None:
-        self.left_arm.connect(calibrate)
-        self.right_arm.connect(calibrate)
-
-    @property
-    def is_calibrated(self) -> bool:
-        return self.left_arm.is_calibrated and self.right_arm.is_calibrated
-
-    def calibrate(self) -> None:
-        self.left_arm.calibrate()
-        self.right_arm.calibrate()
-
-    def configure(self) -> None:
-        self.left_arm.configure()
-        self.right_arm.configure()
-
     def setup_motors(self) -> None:
         raise NotImplementedError(
             "Motor ID configuration is typically done via manufacturer tools for CAN motors."
@@ -129,8 +109,3 @@ class BiOpenArmLeader(Teleoperator):
     def send_feedback(self, feedback: dict[str, float]) -> None:
         # TODO: Implement force feedback
         raise NotImplementedError
-
-    @check_if_not_connected
-    def disconnect(self) -> None:
-        self.left_arm.disconnect()
-        self.right_arm.disconnect()

src/lerobot/teleoperators/bi_openarm_leader/config_bi_openarm_leader.py

@@ -23,7 +23,7 @@ from ..openarm_leader import OpenArmLeaderConfigBase
 @TeleoperatorConfig.register_subclass("bi_openarm_leader")
 @dataclass
 class BiOpenArmLeaderConfig(TeleoperatorConfig):
-    """Configuration class for Bi OpenArm Follower robots."""
+    """Configuration class for Bi OpenArm Leader teleoperators."""
 
     left_arm_config: OpenArmLeaderConfigBase
     right_arm_config: OpenArmLeaderConfigBase

src/lerobot/teleoperators/bi_openarm_mini/__init__.py

@@ -0,0 +1,20 @@
+#!/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.
+
+from .bi_openarm_mini import BiOpenArmMini
+from .config_bi_openarm_mini import BiOpenArmMiniConfig
+
+__all__ = ["BiOpenArmMini", "BiOpenArmMiniConfig"]

src/lerobot/teleoperators/bi_openarm_mini/bi_openarm_mini.py

@@ -0,0 +1,101 @@
+#!/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.
+
+import logging
+from functools import cached_property
+
+from lerobot.types import RobotAction
+from lerobot.utils.bimanual import BimanualMixin
+from lerobot.utils.decorators import check_if_not_connected
+
+from ..openarm_mini import OpenArmMini, OpenArmMiniConfig
+from ..teleoperator import Teleoperator
+from .config_bi_openarm_mini import BiOpenArmMiniConfig
+
+logger = logging.getLogger(__name__)
+
+
+class BiOpenArmMini(BimanualMixin, Teleoperator):
+    """Bimanual OpenArm Mini teleoperator.
+
+    Composes two single-arm :class:`OpenArmMini` instances. Action and feedback
+    keys of each arm are namespaced with a ``left_`` / ``right_`` prefix, so a
+    bimanual leader can teleoperate a bimanual OpenArm follower.
+    """
+
+    config_class = BiOpenArmMiniConfig
+    name = "bi_openarm_mini"
+
+    def __init__(self, config: BiOpenArmMiniConfig):
+        super().__init__(config)
+        self.config = config
+
+        # `side` is forced to match left/right regardless of what the user passed
+        # on the per-arm base config — the bimanual wrapper owns the side semantics.
+        left_arm_config = OpenArmMiniConfig(
+            id=f"{config.id}_left" if config.id else None,
+            calibration_dir=config.calibration_dir,
+            port=config.left_arm_config.port,
+            side="left",
+            use_degrees=config.left_arm_config.use_degrees,
+        )
+
+        right_arm_config = OpenArmMiniConfig(
+            id=f"{config.id}_right" if config.id else None,
+            calibration_dir=config.calibration_dir,
+            port=config.right_arm_config.port,
+            side="right",
+            use_degrees=config.right_arm_config.use_degrees,
+        )
+
+        self.left_arm = OpenArmMini(left_arm_config)
+        self.right_arm = OpenArmMini(right_arm_config)
+
+    @cached_property
+    def action_features(self) -> dict[str, type]:
+        return {
+            **{f"left_{k}": v for k, v in self.left_arm.action_features.items()},
+            **{f"right_{k}": v for k, v in self.right_arm.action_features.items()},
+        }
+
+    @cached_property
+    def feedback_features(self) -> dict[str, type]:
+        return {
+            **{f"left_{k}": v for k, v in self.left_arm.feedback_features.items()},
+            **{f"right_{k}": v for k, v in self.right_arm.feedback_features.items()},
+        }
+
+    def setup_motors(self) -> None:
+        self.left_arm.setup_motors()
+        self.right_arm.setup_motors()
+
+    @check_if_not_connected
+    def get_action(self) -> RobotAction:
+        action: RobotAction = {}
+        for k, v in self.left_arm.get_action().items():
+            action[f"left_{k}"] = v
+        for k, v in self.right_arm.get_action().items():
+            action[f"right_{k}"] = v
+        return action
+
+    @check_if_not_connected
+    def send_feedback(self, feedback: dict[str, float]) -> None:
+        left_fb = {k.removeprefix("left_"): v for k, v in feedback.items() if k.startswith("left_")}
+        right_fb = {k.removeprefix("right_"): v for k, v in feedback.items() if k.startswith("right_")}
+        if left_fb:
+            self.left_arm.send_feedback(left_fb)
+        if right_fb:
+            self.right_arm.send_feedback(right_fb)

src/lerobot/teleoperators/bi_openarm_mini/config_bi_openarm_mini.py

@@ -0,0 +1,29 @@
+#!/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.
+
+from dataclasses import dataclass
+
+from ..config import TeleoperatorConfig
+from ..openarm_mini import OpenArmMiniConfigBase
+
+
+@TeleoperatorConfig.register_subclass("bi_openarm_mini")
+@dataclass
+class BiOpenArmMiniConfig(TeleoperatorConfig):
+    """Configuration class for Bi OpenArm Mini teleoperators."""
+
+    left_arm_config: OpenArmMiniConfigBase
+    right_arm_config: OpenArmMiniConfigBase

src/lerobot/teleoperators/bi_rebot_102_leader/__init__.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .bi_rebot_102_leader import BiRebotArm102Leader
-from .config_bi_rebot_102_leader import BiRebotArm102LeaderConfig
+from .bi_rebot_102_leader import BiRebot102Leader
+from .config_bi_rebot_102_leader import BiRebot102LeaderConfig
 
-__all__ = ["BiRebotArm102Leader", "BiRebotArm102LeaderConfig"]
+__all__ = ["BiRebot102Leader", "BiRebot102LeaderConfig"]

src/lerobot/teleoperators/bi_rebot_102_leader/bi_rebot_102_leader.py

@@ -18,16 +18,17 @@ import logging
 from functools import cached_property
 
 from lerobot.types import RobotAction
-from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
+from lerobot.utils.bimanual import BimanualMixin
+from lerobot.utils.decorators import check_if_not_connected
 
 from ..rebot_102_leader import RebotArm102Leader, RebotArm102LeaderTeleopConfig
 from ..teleoperator import Teleoperator
-from .config_bi_rebot_102_leader import BiRebotArm102LeaderConfig
+from .config_bi_rebot_102_leader import BiRebot102LeaderConfig
 
 logger = logging.getLogger(__name__)
 
 
-class BiRebotArm102Leader(Teleoperator):
+class BiRebot102Leader(BimanualMixin, Teleoperator):
     """Bimanual Seeed Studio StarArm102 / reBot Arm 102 leader.
 
     Composes two single-arm :class:`RebotArm102Leader` instances. Action keys of
@@ -35,10 +36,10 @@ class BiRebotArm102Leader(Teleoperator):
     leader can teleoperate a bimanual reBot B601 follower.
     """
 
-    config_class = BiRebotArm102LeaderConfig
+    config_class = BiRebot102LeaderConfig
     name = "bi_rebot_102_leader"
 
-    def __init__(self, config: BiRebotArm102LeaderConfig):
+    def __init__(self, config: BiRebot102LeaderConfig):
         super().__init__(config)
         self.config = config
 
@@ -76,27 +77,6 @@ class BiRebotArm102Leader(Teleoperator):
     def feedback_features(self) -> dict[str, type]:
         return {}
 
-    @property
-    def is_connected(self) -> bool:
-        return self.left_arm.is_connected and self.right_arm.is_connected
-
-    @check_if_already_connected
-    def connect(self, calibrate: bool = True) -> None:
-        self.left_arm.connect(calibrate)
-        self.right_arm.connect(calibrate)
-
-    @property
-    def is_calibrated(self) -> bool:
-        return self.left_arm.is_calibrated and self.right_arm.is_calibrated
-
-    def calibrate(self) -> None:
-        self.left_arm.calibrate()
-        self.right_arm.calibrate()
-
-    def configure(self) -> None:
-        self.left_arm.configure()
-        self.right_arm.configure()
-
     @check_if_not_connected
     def get_action(self) -> RobotAction:
         action_dict = {}
@@ -106,8 +86,3 @@ class BiRebotArm102Leader(Teleoperator):
 
     def send_feedback(self, feedback: dict[str, float]) -> None:
         raise NotImplementedError("Feedback is not implemented for the reBot Arm 102 leader.")
-
-    @check_if_not_connected
-    def disconnect(self) -> None:
-        self.left_arm.disconnect()
-        self.right_arm.disconnect()

src/lerobot/teleoperators/bi_rebot_102_leader/config_bi_rebot_102_leader.py

@@ -22,7 +22,7 @@ from ..rebot_102_leader import RebotArm102LeaderConfig
 
 @TeleoperatorConfig.register_subclass("bi_rebot_102_leader")
 @dataclass
-class BiRebotArm102LeaderConfig(TeleoperatorConfig):
+class BiRebot102LeaderConfig(TeleoperatorConfig):
     """Configuration class for the bimanual reBot Arm 102 leader teleoperator."""
 
     left_arm_config: RebotArm102LeaderConfig

src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py

@@ -17,7 +17,9 @@
 import logging
 from functools import cached_property
 
-from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
+from lerobot.types import RobotAction
+from lerobot.utils.bimanual import BimanualMixin
+from lerobot.utils.decorators import check_if_not_connected
 
 from ..so_leader import SOLeader, SOLeaderTeleopConfig
 from ..teleoperator import Teleoperator
@@ -26,7 +28,7 @@ from .config_bi_so_leader import BiSOLeaderConfig
 logger = logging.getLogger(__name__)
 
 
-class BiSOLeader(Teleoperator):
+class BiSOLeader(BimanualMixin, Teleoperator):
     """
     [Bimanual SO Leader Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio
     """
@@ -67,33 +69,12 @@ class BiSOLeader(Teleoperator):
     def feedback_features(self) -> dict[str, type]:
         return {}
 
-    @property
-    def is_connected(self) -> bool:
-        return self.left_arm.is_connected and self.right_arm.is_connected
-
-    @check_if_already_connected
-    def connect(self, calibrate: bool = True) -> None:
-        self.left_arm.connect(calibrate)
-        self.right_arm.connect(calibrate)
-
-    @property
-    def is_calibrated(self) -> bool:
-        return self.left_arm.is_calibrated and self.right_arm.is_calibrated
-
-    def calibrate(self) -> None:
-        self.left_arm.calibrate()
-        self.right_arm.calibrate()
-
-    def configure(self) -> None:
-        self.left_arm.configure()
-        self.right_arm.configure()
-
     def setup_motors(self) -> None:
         self.left_arm.setup_motors()
         self.right_arm.setup_motors()
 
     @check_if_not_connected
-    def get_action(self) -> dict[str, float]:
+    def get_action(self) -> RobotAction:
         action_dict = {}
 
         # Add "left_" prefix
@@ -109,8 +90,3 @@ class BiSOLeader(Teleoperator):
     def send_feedback(self, feedback: dict[str, float]) -> None:
         # TODO: Implement force feedback
         raise NotImplementedError
-
-    @check_if_not_connected
-    def disconnect(self) -> None:
-        self.left_arm.disconnect()
-        self.right_arm.disconnect()

src/lerobot/teleoperators/openarm_mini/__init__.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+# 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.
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .config_openarm_mini import OpenArmMiniConfig
+from .config_openarm_mini import OpenArmMiniConfig, OpenArmMiniConfigBase
 from .openarm_mini import OpenArmMini
 
-__all__ = ["OpenArmMini", "OpenArmMiniConfig"]
+__all__ = ["OpenArmMini", "OpenArmMiniConfig", "OpenArmMiniConfigBase"]

src/lerobot/teleoperators/openarm_mini/config_openarm_mini.py

@@ -19,12 +19,21 @@ from dataclasses import dataclass
 from ..config import TeleoperatorConfig
 
 
-@TeleoperatorConfig.register_subclass("openarm_mini")
 @dataclass
-class OpenArmMiniConfig(TeleoperatorConfig):
-    """Configuration for OpenArm Mini teleoperator with Feetech motors (dual arms)."""
+class OpenArmMiniConfigBase:
+    """Base configuration for the OpenArm Mini teleoperator (Feetech STS3215, 7DOF + gripper)."""
 
-    port_right: str = "/dev/ttyUSB0"
-    port_left: str = "/dev/ttyUSB1"
+    # Serial port for the Feetech bus (e.g., "/dev/ttyUSB0").
+    port: str
+
+    # Side of the arm: "left" or "right". Controls per-joint direction flips applied
+    # during readout. If `None`, no flipping is applied.
+    side: str | None = None
 
     use_degrees: bool = True
+
+
+@TeleoperatorConfig.register_subclass("openarm_mini")
+@dataclass
+class OpenArmMiniConfig(TeleoperatorConfig, OpenArmMiniConfigBase):
+    pass

src/lerobot/teleoperators/openarm_mini/openarm_mini.py

@@ -31,22 +31,22 @@ from .config_openarm_mini import OpenArmMiniConfig
 
 logger = logging.getLogger(__name__)
 
-# Motors whose direction is inverted during readout
-RIGHT_MOTORS_TO_FLIP = ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_7"]
-LEFT_MOTORS_TO_FLIP = ["joint_1", "joint_3", "joint_4", "joint_5", "joint_6", "joint_7"]
+# Per-side motor direction flips applied during readout.
+SIDE_MOTORS_TO_FLIP: dict[str, list[str]] = {
+    "left": ["joint_1", "joint_3", "joint_4", "joint_5", "joint_6", "joint_7"],
+    "right": ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_7"],
+}
 
-# Leader joint 6 maps to follower joint 7 and vice versa
+# Leader joint 6 ↔ follower joint 7 (symmetric — its own inverse).
 JOINT_REMAP = {"joint_6": "joint_7", "joint_7": "joint_6"}
-JOINT_REMAP_REVERSE = {"joint_7": "joint_6", "joint_6": "joint_7"}
 
 GRIPPER_TELEOP_TO_DEGREES = -0.65
 
 
 class OpenArmMini(Teleoperator):
-    """
-    OpenArm Mini Teleoperator with dual Feetech-based arms (8 motors per arm).
+    """OpenArm Mini single-arm teleoperator (Feetech STS3215, 7DOF + gripper).
 
-    Each arm has 7 joints plus a gripper, using Feetech STS3215 servos.
+    For the bimanual setup, see :class:`BiOpenArmMini` which composes two of these.
     """
 
     config_class = OpenArmMiniConfig
@@ -56,9 +56,12 @@ class OpenArmMini(Teleoperator):
         super().__init__(config)
         self.config = config
 
+        if config.side is not None and config.side not in SIDE_MOTORS_TO_FLIP:
+            raise ValueError(f"Invalid side '{config.side}'; expected 'left', 'right', or None.")
+        self._motors_to_flip: list[str] = SIDE_MOTORS_TO_FLIP.get(config.side, []) if config.side else []
+
         norm_mode_body = MotorNormMode.DEGREES
-
-        motors_right = {
+        motors = {
             "joint_1": Motor(1, "sts3215", norm_mode_body),
             "joint_2": Motor(2, "sts3215", norm_mode_body),
             "joint_3": Motor(3, "sts3215", norm_mode_body),
@@ -69,46 +72,15 @@ class OpenArmMini(Teleoperator):
             "gripper": Motor(8, "sts3215", MotorNormMode.RANGE_0_100),
         }
 
-        motors_left = {
-            "joint_1": Motor(1, "sts3215", norm_mode_body),
-            "joint_2": Motor(2, "sts3215", norm_mode_body),
-            "joint_3": Motor(3, "sts3215", norm_mode_body),
-            "joint_4": Motor(4, "sts3215", norm_mode_body),
-            "joint_5": Motor(5, "sts3215", norm_mode_body),
-            "joint_6": Motor(6, "sts3215", norm_mode_body),
-            "joint_7": Motor(7, "sts3215", norm_mode_body),
-            "gripper": Motor(8, "sts3215", MotorNormMode.RANGE_0_100),
-        }
-
-        cal_right = {
-            k.replace("right_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("right_")
-        }
-        cal_left = {
-            k.replace("left_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("left_")
-        }
-
-        self.bus_right = FeetechMotorsBus(
-            port=self.config.port_right,
-            motors=motors_right,
-            calibration=cal_right,
-        )
-
-        self.bus_left = FeetechMotorsBus(
-            port=self.config.port_left,
-            motors=motors_left,
-            calibration=cal_left,
+        self.bus = FeetechMotorsBus(
+            port=self.config.port,
+            motors=motors,
+            calibration=self.calibration,
         )
 
     @property
     def action_features(self) -> dict[str, type]:
-        # Right first, then left — matches the robot (BiOpenArmFollower) ordering
-        # and the dataset feature names recorded during data collection.
-        features: dict[str, type] = {}
-        for motor in self.bus_right.motors:
-            features[f"right_{motor}.pos"] = float
-        for motor in self.bus_left.motors:
-            features[f"left_{motor}.pos"] = float
-        return features
+        return {f"{motor}.pos": float for motor in self.bus.motors}
 
     @property
     def feedback_features(self) -> dict[str, type]:
@@ -116,14 +88,12 @@ class OpenArmMini(Teleoperator):
 
     @property
     def is_connected(self) -> bool:
-        return self.bus_right.is_connected and self.bus_left.is_connected
+        return self.bus.is_connected
 
     @check_if_already_connected
     def connect(self, calibrate: bool = True) -> None:
-        logger.info(f"Connecting right arm on {self.config.port_right}...")
-        self.bus_right.connect()
-        logger.info(f"Connecting left arm on {self.config.port_left}...")
-        self.bus_left.connect()
+        logger.info(f"Connecting arm on {self.config.port}...")
+        self.bus.connect()
 
         if calibrate:
             self.calibrate()
@@ -133,14 +103,14 @@ class OpenArmMini(Teleoperator):
 
     @property
     def is_calibrated(self) -> bool:
-        return self.bus_right.is_calibrated and self.bus_left.is_calibrated
+        return self.bus.is_calibrated
 
     def calibrate(self) -> None:
         """
-        Run calibration procedure for OpenArm Mini.
+        Run calibration procedure for a single OpenArm Mini arm.
 
         1. Disable torque
-        2. Ask user to position arms in hanging position with grippers closed
+        2. Ask user to position arm in hanging position with gripper closed
         3. Set this as zero position via half-turn homing
         4. Interactive gripper calibration (open/close positions)
         5. Save calibration
@@ -152,70 +122,51 @@ class OpenArmMini(Teleoperator):
             )
             if user_input.strip().lower() != "c":
                 logger.info(f"Using existing calibration for {self.id}")
-                cal_right = {
-                    k.replace("right_", ""): v for k, v in self.calibration.items() if k.startswith("right_")
-                }
-                cal_left = {
-                    k.replace("left_", ""): v for k, v in self.calibration.items() if k.startswith("left_")
-                }
-                self.bus_right.write_calibration(cal_right)
-                self.bus_left.write_calibration(cal_left)
+                self.bus.write_calibration(self.calibration)
                 return
 
         logger.info(f"\nRunning calibration for {self}")
 
-        self._calibrate_arm("right", self.bus_right)
-        self._calibrate_arm("left", self.bus_left)
+        self.bus.disable_torque()
 
-        self._save_calibration()
-        print(f"\nCalibration complete and saved to {self.calibration_fpath}")
+        logger.info("Setting Phase to 12 for all motors...")
+        for motor in self.bus.motors:
+            self.bus.write("Phase", motor, 12)
 
-    def _calibrate_arm(self, arm_name: str, bus: FeetechMotorsBus) -> None:
-        """Calibrate a single arm with Feetech motors."""
-        logger.info(f"\n=== Calibrating {arm_name.upper()} arm ===")
-
-        bus.disable_torque()
-
-        logger.info(f"Setting Phase to 12 for all motors in {arm_name.upper()} arm...")
-        for motor in bus.motors:
-            bus.write("Phase", motor, 12)
-
-        for motor in bus.motors:
-            bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
+        for motor in self.bus.motors:
+            self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
 
         input(
-            f"\nCalibration: Zero Position ({arm_name.upper()} arm)\n"
+            "\nCalibration: Zero Position\n"
             "Position the arm in the following configuration:\n"
             "  - Arm hanging straight down\n"
             "  - Gripper closed\n"
             "Press ENTER when ready..."
         )
 
-        homing_offsets = bus.set_half_turn_homings()
-        logger.info(f"{arm_name.capitalize()} arm zero position set.")
+        homing_offsets = self.bus.set_half_turn_homings()
+        logger.info("Arm zero position set.")
 
-        print(f"\nSetting motor ranges for {arm_name.upper()} arm\n")
+        print("\nSetting motor ranges\n")
 
         if self.calibration is None:
             self.calibration = {}
 
-        motor_resolution = bus.model_resolution_table[list(bus.motors.values())[0].model]
+        motor_resolution = self.bus.model_resolution_table[list(self.bus.motors.values())[0].model]
         max_res = motor_resolution - 1
 
-        for motor_name, motor in bus.motors.items():
-            prefixed_name = f"{arm_name}_{motor_name}"
-
+        for motor_name, motor in self.bus.motors.items():
             if motor_name == "gripper":
                 input(
-                    f"\nGripper Calibration ({arm_name.upper()} arm)\n"
-                    f"Step 1: CLOSE the gripper fully\n"
-                    f"Press ENTER when gripper is closed..."
+                    "\nGripper Calibration\n"
+                    "Step 1: CLOSE the gripper fully\n"
+                    "Press ENTER when gripper is closed..."
                 )
-                closed_pos = bus.read("Present_Position", motor_name, normalize=False)
+                closed_pos = self.bus.read("Present_Position", motor_name, normalize=False)
                 logger.info(f"  Gripper closed position recorded: {closed_pos}")
 
                 input("\nStep 2: OPEN the gripper fully\nPress ENTER when gripper is fully open...")
-                open_pos = bus.read("Present_Position", motor_name, normalize=False)
+                open_pos = self.bus.read("Present_Position", motor_name, normalize=False)
                 logger.info(f"  Gripper open position recorded: {open_pos}")
 
                 if closed_pos < open_pos:
@@ -228,16 +179,16 @@ class OpenArmMini(Teleoperator):
                     drive_mode = 1
 
                 logger.info(
-                    f"  {prefixed_name}: range set to [{range_min}, {range_max}] "
+                    f"  {motor_name}: range set to [{range_min}, {range_max}] "
                     f"(0=closed, 100=open, drive_mode={drive_mode})"
                 )
             else:
                 range_min = 0
                 range_max = max_res
                 drive_mode = 0
-                logger.info(f"  {prefixed_name}: range set to [0, {max_res}] (full motor range)")
+                logger.info(f"  {motor_name}: range set to [0, {max_res}] (full motor range)")
 
-            self.calibration[prefixed_name] = MotorCalibration(
+            self.calibration[motor_name] = MotorCalibration(
                 id=motor.id,
                 drive_mode=drive_mode,
                 homing_offset=homing_offsets[motor_name],
@@ -245,108 +196,68 @@ class OpenArmMini(Teleoperator):
                 range_max=range_max,
             )
 
-        cal_for_bus = {
-            k.replace(f"{arm_name}_", ""): v
-            for k, v in self.calibration.items()
-            if k.startswith(f"{arm_name}_")
-        }
-        bus.write_calibration(cal_for_bus)
+        self.bus.write_calibration(self.calibration)
+        self._save_calibration()
+        print(f"\nCalibration complete and saved to {self.calibration_fpath}")
 
     def configure(self) -> None:
-        self.bus_right.disable_torque()
-        self.bus_right.configure_motors()
-        for motor in self.bus_right.motors:
-            self.bus_right.write("Operating_Mode", motor, OperatingMode.POSITION.value)
-
-        self.bus_left.disable_torque()
-        self.bus_left.configure_motors()
-        for motor in self.bus_left.motors:
-            self.bus_left.write("Operating_Mode", motor, OperatingMode.POSITION.value)
+        self.bus.disable_torque()
+        self.bus.configure_motors()
+        for motor in self.bus.motors:
+            self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
 
     def setup_motors(self) -> None:
-        print("\nSetting up RIGHT arm motors...")
-        for motor in reversed(self.bus_right.motors):
-            input(f"Connect the controller board to the RIGHT '{motor}' motor only and press enter.")
-            self.bus_right.setup_motor(motor)
-            print(f"RIGHT '{motor}' motor id set to {self.bus_right.motors[motor].id}")
-
-        print("\nSetting up LEFT arm motors...")
-        for motor in reversed(self.bus_left.motors):
-            input(f"Connect the controller board to the LEFT '{motor}' motor only and press enter.")
-            self.bus_left.setup_motor(motor)
-            print(f"LEFT '{motor}' motor id set to {self.bus_left.motors[motor].id}")
+        for motor in reversed(self.bus.motors):
+            input(f"Connect the controller board to the '{motor}' motor only and press enter.")
+            self.bus.setup_motor(motor)
+            print(f"'{motor}' motor id set to {self.bus.motors[motor].id}")
 
     @check_if_not_connected
     def get_action(self) -> RobotAction:
-        """Get current action from both arms (read positions from all motors)."""
+        """Get current action (read positions from all motors)."""
         start = time.perf_counter()
 
-        right_positions = self.bus_right.sync_read("Present_Position")
-        left_positions = self.bus_left.sync_read("Present_Position")
+        positions = self.bus.sync_read("Present_Position")
 
-        # Right first, then left — matches the robot (BiOpenArmFollower) ordering
-        # and the dataset feature names recorded during data collection.
         # Joint 6↔7 remap: leader joint_6 → follower joint_7 and vice versa.
+        # Per-side direction flip is applied based on the configured `side`.
         action: dict[str, Any] = {}
-        for motor, val in right_positions.items():
+        for motor, val in positions.items():
             target = JOINT_REMAP.get(motor, motor)
             if motor == "gripper":
                 # Convert gripper from teleop 0-100 to openarms degrees: 0→0°, 100→-65°
-                action[f"right_{target}.pos"] = val * GRIPPER_TELEOP_TO_DEGREES
+                action[f"{target}.pos"] = val * GRIPPER_TELEOP_TO_DEGREES
             else:
-                action[f"right_{target}.pos"] = -val if motor in RIGHT_MOTORS_TO_FLIP else val
-        for motor, val in left_positions.items():
-            target = JOINT_REMAP.get(motor, motor)
-            if motor == "gripper":
-                action[f"left_{target}.pos"] = val * GRIPPER_TELEOP_TO_DEGREES
-            else:
-                action[f"left_{target}.pos"] = -val if motor in LEFT_MOTORS_TO_FLIP else val
+                action[f"{target}.pos"] = -val if motor in self._motors_to_flip else val
 
         dt_ms = (time.perf_counter() - start) * 1e3
         logger.debug(f"{self} read action: {dt_ms:.1f}ms")
         return action
 
     def enable_torque(self) -> None:
-        """Enable torque on both arms for position control."""
-        self.bus_right.enable_torque()
-        self.bus_left.enable_torque()
+        self.bus.enable_torque()
 
     def disable_torque(self) -> None:
-        """Disable torque on both arms for free movement."""
-        self.bus_right.disable_torque()
-        self.bus_left.disable_torque()
+        self.bus.disable_torque()
 
     def write_goal_positions(self, positions: dict[str, float]) -> None:
         """Write goal positions to motors (inverse of get_action flip/gripper/remap logic)."""
-        right_goals: dict[str, float] = {}
-        left_goals: dict[str, float] = {}
-
+        goals: dict[str, float] = {}
         for key, val in positions.items():
             if not key.endswith(".pos"):
                 continue
-            motor_name = key.removesuffix(".pos")
-            if motor_name.startswith("right_"):
-                base = motor_name.removeprefix("right_")
-                # Reverse remap: follower joint_7 → leader joint_6 and vice versa
-                target = JOINT_REMAP_REVERSE.get(base, base)
-                if base == "gripper":
-                    # Convert robot degrees to teleop 0-100: 0°→0, -65°→100
-                    right_goals[target] = val / GRIPPER_TELEOP_TO_DEGREES
-                else:
-                    # Un-flip using the ORIGINAL motor name (target = leader motor)
-                    right_goals[target] = -val if target in RIGHT_MOTORS_TO_FLIP else val
-            elif motor_name.startswith("left_"):
-                base = motor_name.removeprefix("left_")
-                target = JOINT_REMAP_REVERSE.get(base, base)
-                if base == "gripper":
-                    left_goals[target] = val / GRIPPER_TELEOP_TO_DEGREES
-                else:
-                    left_goals[target] = -val if target in LEFT_MOTORS_TO_FLIP else val
+            base = key.removesuffix(".pos")
+            # JOINT_REMAP is symmetric (its own inverse).
+            target = JOINT_REMAP.get(base, base)
+            if base == "gripper":
+                # Convert robot degrees to teleop 0-100: 0°→0, -65°→100
+                goals[target] = val / GRIPPER_TELEOP_TO_DEGREES
+            else:
+                # Un-flip using the ORIGINAL motor name (target = leader motor)
+                goals[target] = -val if target in self._motors_to_flip else val
 
-        if right_goals:
-            self.bus_right.sync_write("Goal_Position", right_goals)
-        if left_goals:
-            self.bus_left.sync_write("Goal_Position", left_goals)
+        if goals:
+            self.bus.sync_write("Goal_Position", goals)
 
     @check_if_not_connected
     def send_feedback(self, feedback: dict[str, float]) -> None:
@@ -354,6 +265,5 @@ class OpenArmMini(Teleoperator):
 
     @check_if_not_connected
     def disconnect(self) -> None:
-        self.bus_right.disconnect()
-        self.bus_left.disconnect()
+        self.bus.disconnect()
         logger.info(f"{self} disconnected.")