feat(streaming): random-episode admission via reshard() + multi-input-shard shuffle

Reshard parquet per row group (1 shard == 1 row group == 1 episode) and feed the
episode-pool shuffle with max_buffer_input_shards so the pool is a uniform random
sample of the corpus, independent of episodes-per-file. Add validate_row_groups
guardrails (collapsed-row-group + distributed divisibility), require datasets>=5.0.0,
make the test fixture write one row group per episode, and plumb max_buffer_input_shards
through the dataloading benchmark.

Co-authored-by: Cursor <cursoragent@cursor.com>

.pre-commit-config.yaml

@@ -65,9 +65,6 @@ 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,9 +178,3 @@ 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, reBot B601.
+**Supported Hardware:** SO100, LeKiwi, Koch, HopeJR, OMX, EarthRover, Reachy2, Gamepads, Keyboards, Phones, OpenARM, Unitree G1.
 
 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,13 +101,11 @@ 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**            | [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)                                                                                                                                                                                                                                         |
+| 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.5](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.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
 
@@ -135,7 +133,6 @@ 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
 
@@ -143,7 +140,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 Meftah, Khalil and Ellerbach, Maxime 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 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,8 +45,6 @@
     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

docs/source/annotation_pipeline.mdx

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

examples/annotations/run_hf_job.py

@@ -1,77 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6-27B VLM).
-
-Spawns one single-GPU ``h200`` job that:
-
-  1. installs ``lerobot`` from ``main`` plus the annotation extras,
-  2. boots one vllm server with Qwen3.6-27B (dense VLM),
-  3. runs the plan / interjections / vqa modules across the dataset
-     in free-form mode (each episode generates its own subtasks +
-     memory),
-  4. uploads the annotated dataset to ``--new_repo_id`` (when set)
-     or back to ``--repo_id``.
-
-Usage:
-
-    HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
-
-Adjust ``CMD`` (dataset, model, hub repo) and ``flavor`` below for your
-run. For larger datasets, scale to ``h200x4`` and raise
-``--vlm.parallel_servers`` / ``--vlm.num_gpus`` to match.
-"""
-
-import os
-
-from huggingface_hub import get_token, run_job
-
-token = os.environ.get("HF_TOKEN") or get_token()
-if not token:
-    raise RuntimeError("No HF token. Run `huggingface-cli login` or `export HF_TOKEN=hf_...`")
-
-CMD = (
-    "apt-get update -qq && apt-get install -y -qq git ffmpeg && "
-    "pip install --no-deps "
-    "'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
-    "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}")

examples/scaling/benchmark_dataloading.py

@@ -0,0 +1,547 @@
+# 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.
+
+"""Single-image dataloading benchmark across the LeRobot loaders, MADE TO RUN ON A COMPUTE CLUSTER (SLURM).
+
+This one file is both the orchestrator and the worker:
+
+  * Run it with no ``--scenario`` (from a login node) and it submits a SERIAL sbatch chain of all
+    scenarios below (no two network-bound jobs overlap, so CDN numbers stay clean).
+  * Run it with ``--scenario <name>`` and it executes that single benchmark (this is what each sbatch
+    job calls). The 2-node scenario is launched with ``srun`` and reads ``RANK``/``WORLD_SIZE`` so the
+    streaming dataset splits shards per node.
+
+Scenarios (all single-frame / non-SARM):
+  1. ``mmap_local``             map-style LeRobotDataset over a LOCAL copy (``--local_root``, no network).
+  2. ``mmap_local_maxworkers``  same, but workers scaled to saturate the node's cores (decode-bound).
+  3. ``stream_hub``             StreamingLeRobotDataset from the Hub (allenai/MolmoAct2-BimanualYAM-Dataset).
+  4. ``stream_bucket``          StreamingLeRobotDataset from a warmed storage bucket (1 node).
+  5. ``stream_bucket_2node``    same warmed bucket, 2 nodes (split_dataset_by_node, per-rank results).
+
+Reported per run: peak process-tree RSS (max memory), parallel throughput (samples/s, where a sample
+is one timestep, plus decoded_frames/s = samples/s x num_cameras),
+single-process throughput, shuffle randomness fraction (distinct episodes per batch / batch size),
+fetch vs decode split (% of single-process per-sample time), first-batch latency, and p50/p95/p99
+sample latency. Results are written as JSON + CSV under ``--out_dir``.
+
+Submit the whole chain (from a login node, inside the repo). Point the scheduler env vars at your own
+cluster's account/partition/qos, and ``--local_root`` at a local copy of the map-style dataset:
+    ACCOUNT=<account> PARTITION=<partition> QOS=<qos> \\
+        python examples/scaling/benchmark_dataloading.py --local_root /path/to/local/dataset
+"""
+
+import argparse
+import csv
+import json
+import os
+import random
+import statistics
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from lerobot.datasets import LeRobotDataset, LeRobotDatasetMetadata, StreamingLeRobotDataset
+from lerobot.datasets.partition import group_episodes_by_files, partition_episodes
+
+ROBOCASA_REPO = "pepijn223/robocasa_pretrain_human300_v4"
+MOLMO_REPO = "allenai/MolmoAct2-BimanualYAM-Dataset"
+MOLMO_BUCKET = "hf://buckets/pepijn223/MolmoAct2-BimanualYAM-Dataset-bucket"
+# MolmoAct2 is published without a codebase-version git tag, so the version-safe loader would refuse
+# it; "main" pins the branch directly and skips that check.
+MOLMO_REVISION = "main"
+
+# Per-scenario sbatch shape. mem is generous for the streaming legs (32k-episode, 3-camera, 2.35 TB
+# dataset keeps many AV1 decoders open); the local map-style leg is light. Optional ``num_workers`` /
+# ``cpus`` override the CLI defaults for that leg.
+# ``mmap_local_maxworkers``: map-style decode is CPU-bound and each worker decodes its cameras on
+# parallel threads, so the saturation point is ~num_cpus / num_cameras workers (~90 concurrent decode
+# threads). The 96-core H100 nodes here schedule at most 92 cpus/task, so we take 92 cpus / 30 workers.
+SCENARIOS = {
+    "mmap_local": {"kind": "map", "nodes": 1, "mem": "64G", "time": "01:00:00"},
+    "mmap_local_maxworkers": {
+        "kind": "map",
+        "nodes": 1,
+        "mem": "128G",
+        "time": "01:00:00",
+        "num_workers": 30,
+        "cpus": 92,
+    },
+    "stream_hub": {"kind": "stream", "nodes": 1, "mem": "250G", "time": "03:00:00"},
+    "stream_bucket": {"kind": "stream", "nodes": 1, "mem": "250G", "time": "03:00:00"},
+    "stream_bucket_2node": {"kind": "stream", "nodes": 2, "mem": "250G", "time": "03:00:00"},
+}
+
+
+def _tree_rss_bytes() -> int:
+    """Sum RSS of this process and all descendants via /proc (DataLoader workers are separate procs)."""
+    try:
+        children: dict[int, list[int]] = {}
+        for entry in os.listdir("/proc"):
+            if not entry.isdigit():
+                continue
+            try:
+                with open(f"/proc/{entry}/stat") as f:
+                    ppid = int(f.read().split(") ", 1)[1].split()[1])
+                children.setdefault(ppid, []).append(int(entry))
+            except (OSError, ValueError, IndexError):
+                pass
+        total, stack = 0, [os.getpid()]
+        while stack:
+            cur = stack.pop()
+            try:
+                with open(f"/proc/{cur}/statm") as f:
+                    total += int(f.read().split()[1]) * os.sysconf("SC_PAGE_SIZE")
+            except (OSError, ValueError, IndexError):
+                pass
+            stack.extend(children.get(cur, []))
+        return total
+    except OSError:
+        return 0
+
+
+class PeakRSSSampler:
+    """Background thread tracking peak process-tree RSS for the duration of the ``with`` block."""
+
+    def __init__(self, interval_s: float = 0.5):
+        self.interval_s = interval_s
+        self.peak_bytes = 0
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, daemon=True)
+
+    def _run(self) -> None:
+        while not self._stop.is_set():
+            self.peak_bytes = max(self.peak_bytes, _tree_rss_bytes())
+            self._stop.wait(self.interval_s)
+
+    def __enter__(self) -> "PeakRSSSampler":
+        self._thread.start()
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self._stop.set()
+        self._thread.join(timeout=2)
+
+
+def percentile(values: list[float], pct: float) -> float:
+    if not values:
+        return float("nan")
+    ordered = sorted(values)
+    k = max(0, min(len(ordered) - 1, int(round((pct / 100.0) * (len(ordered) - 1)))))
+    return ordered[k]
+
+
+class _TimedStreaming(StreamingLeRobotDataset):
+    """StreamingLeRobotDataset that times the fetch stage (parquet/network row) separately from the
+    decode stage (video decode + torch conversion in ``_finalize_sample``), so a single-process pass
+    can attribute per-sample cost to fetch vs decode. Timing lives here in the benchmark, not in the
+    library, to keep the dataset itself instrumentation-free."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.fetch_s = 0.0
+        self.decode_s = 0.0
+
+    def __iter__(self):
+        self._in_flight_epoch = self._epoch
+        self._pipeline.set_epoch(self._in_flight_epoch)
+        self._epoch += 1
+        self.video_decoder_cache = self._make_video_decoder_cache()
+        iterator = iter(self._pipeline)
+        while True:
+            t0 = time.perf_counter()
+            try:
+                row = next(iterator)
+            except StopIteration:
+                return
+            t1 = time.perf_counter()
+            sample = self._finalize_sample(row)
+            t2 = time.perf_counter()
+            self.fetch_s += t1 - t0
+            self.decode_s += t2 - t1
+            yield sample
+
+
+def select_node_episodes(
+    meta: LeRobotDatasetMetadata, num_partitions: int, index: int, cap: int
+) -> list[int]:
+    """This node's episode share, mirroring lerobot_train ``--data_partition=node``: group episodes by
+    shared video files, LPT-balance the groups by frame count, take this node's bin (capped)."""
+    episodes = list(range(meta.total_episodes))
+    from_idx = meta.episodes["dataset_from_index"]
+    to_idx = meta.episodes["dataset_to_index"]
+    lengths = [int(to_idx[ep] - from_idx[ep]) for ep in episodes]
+    if meta.video_keys:
+        file_columns = {
+            key: (meta.episodes[f"videos/{key}/chunk_index"], meta.episodes[f"videos/{key}/file_index"])
+            for key in meta.video_keys
+        }
+    else:
+        file_columns = {"data": (meta.episodes["data/chunk_index"], meta.episodes["data/file_index"])}
+    episode_file_ids = [
+        [(key, chunks[ep], files[ep]) for key, (chunks, files) in file_columns.items()] for ep in episodes
+    ]
+    groups = group_episodes_by_files(episode_file_ids)
+    if len(groups) < num_partitions:
+        groups = [[i] for i in range(len(episodes))]
+    group_lengths = [sum(lengths[i] for i in g) for g in groups]
+    bins = partition_episodes(group_lengths, num_partitions)
+    chosen = sorted(episodes[i] for g in bins[index] for i in groups[g])
+    return chosen[:cap] if cap and len(chosen) > cap else chosen
+
+
+def build_dataset(scenario: str, args: argparse.Namespace):
+    """Return (dataset, meta, is_map_style, info) for the scenario; single-frame (no delta windows)."""
+    if scenario.startswith("mmap_local"):
+        if not args.local_root:
+            raise SystemExit("mmap_local needs --local_root pointing at a local LeRobotDataset copy.")
+        meta = LeRobotDatasetMetadata(ROBOCASA_REPO, root=args.local_root)
+        episodes = select_node_episodes(meta, args.num_partitions, args.partition_index, args.max_episodes)
+        dataset = LeRobotDataset(ROBOCASA_REPO, root=args.local_root, episodes=episodes, tolerance_s=1e-3)
+        return dataset, meta, True, {"loaded_episodes": len(episodes)}
+
+    data_files_root = MOLMO_BUCKET if scenario.startswith("stream_bucket") else None
+    meta = LeRobotDatasetMetadata(MOLMO_REPO, revision=MOLMO_REVISION)
+    dataset = _TimedStreaming(
+        MOLMO_REPO,
+        revision=MOLMO_REVISION,
+        data_files_root=data_files_root,
+        episode_pool_size=args.episode_pool_size,
+        max_buffer_input_shards=args.max_buffer_input_shards,
+        video_decoder_cache_size=args.video_decoder_cache_size,
+        tolerance_s=1e-3,
+        # Throughput benchmark: don't gate on the one-row-group-per-episode invariant (a public
+        # dataset may be collapsed); reshard() still yields per-episode shards where it holds.
+        validate_row_groups=False,
+    )
+    return dataset, meta, False, {"num_shards": dataset.num_shards, "data_files_root": data_files_root}
+
+
+def _split(fetch_s: float, decode_s: float, getitem_s: float, n_probe: int) -> dict:
+    stage = fetch_s + decode_s
+    return {
+        "single_proc_samples_per_s": round(n_probe / getitem_s, 2) if getitem_s else None,
+        "fetch_pct": round(100 * fetch_s / stage, 1) if stage else None,
+        "decode_pct": round(100 * decode_s / stage, 1) if stage else None,
+    }
+
+
+def measure_fetch_decode_stream(dataset: _TimedStreaming, n_probe: int, warmup: int) -> dict:
+    """Single-process pass attributing per-sample time to fetch (parquet/network row) vs decode (video)."""
+    it = iter(dataset)
+    for _ in range(warmup):  # exclude the cold shuffle-buffer fill from the ratio
+        next(it)
+    dataset.fetch_s = dataset.decode_s = 0.0
+    t0 = time.perf_counter()
+    for _ in range(n_probe):
+        next(it)
+    return _split(dataset.fetch_s, dataset.decode_s, time.perf_counter() - t0, n_probe)
+
+
+def measure_fetch_decode_map(dataset: LeRobotDataset, n_probe: int, warmup: int) -> dict:
+    """Same split for the map-style loader: fetch = raw tabular row (``get_raw_item``), decode = the rest
+    of ``__getitem__`` (video decode + transforms). Local reads make fetch tiny and decode dominant.
+
+    Random frames are resampled past any that torchcodec fails to decode, so a single flaky frame can't
+    abort the whole benchmark (the parallel DataLoader pass draws its own fresh random frames)."""
+    rng = random.Random(0)
+    n = len(dataset)
+    fetch_s = getitem_s = 0.0
+    warmed = measured = skipped = attempts = 0
+    while measured < n_probe and attempts < (warmup + n_probe) * 10:
+        attempts += 1
+        i = rng.randrange(n)
+        try:
+            t0 = time.perf_counter()
+            dataset.get_raw_item(i)
+            t1 = time.perf_counter()
+            dataset[i]
+            t2 = time.perf_counter()
+        except Exception:
+            skipped += 1
+            continue
+        if warmed < warmup:
+            warmed += 1
+            continue
+        fetch_s += t1 - t0
+        getitem_s += t2 - t1
+        measured += 1
+    if skipped:
+        print(f"map fetch/decode probe skipped {skipped} undecodable frame(s)", flush=True)
+    return _split(fetch_s, max(0.0, getitem_s - fetch_s), getitem_s, measured)
+
+
+def run_scenario(scenario: str, args: argparse.Namespace) -> None:
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    device = torch.device(args.device)
+
+    dataset, meta, is_map_style, info = build_dataset(scenario, args)
+
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        shuffle=is_map_style,  # map-style: global random shuffle; streaming: shuffled inside the dataset
+        pin_memory=device.type == "cuda",
+        drop_last=True,
+        prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,
+        persistent_workers=args.num_workers > 0,
+    )
+
+    sample_latencies_ms: list[float] = []
+    episodes_per_batch: list[int] = []
+    samples = 0
+    first_batch_latency_s = None
+    steady_start = None
+
+    t_start = time.perf_counter()
+    t_prev = t_start
+    with PeakRSSSampler() as rss:
+        for i, batch in enumerate(loader):
+            for value in batch.values():
+                if torch.is_tensor(value):
+                    value.to(device, non_blocking=device.type == "cuda")
+            now = time.perf_counter()
+            if first_batch_latency_s is None:
+                first_batch_latency_s = now - t_start
+            if i == args.warmup_batches:
+                steady_start = now
+            elif i > args.warmup_batches:
+                sample_latencies_ms.append((now - t_prev) / args.batch_size * 1000.0)
+                samples += args.batch_size
+                ep = batch.get("episode_index")
+                if torch.is_tensor(ep):
+                    episodes_per_batch.append(int(torch.unique(ep).numel()))
+            t_prev = now
+            # Measure throughput over a fixed wall-clock window (after warmup) so every scenario is
+            # compared over the same duration regardless of its speed; num_batches is only a safety cap.
+            if steady_start is not None and (now - steady_start) >= args.duration_s:
+                break
+            if i + 1 >= args.num_batches:
+                break
+    peak_rss_gb = round(rss.peak_bytes / 1e9, 2) if rss.peak_bytes else None
+
+    now = time.perf_counter()
+    elapsed = now - t_start
+    steady_elapsed_s = (now - steady_start) if steady_start is not None else elapsed
+
+    if samples == 0:
+        raise SystemExit(
+            f"FAILED: 0 samples in {args.duration_s}s for scenario={scenario} "
+            "(inspect worker logs; try --num_workers 0 to surface the exception)."
+        )
+
+    # Single-process fetch/decode split + single-proc throughput. Run AFTER the DataLoader pass: this
+    # decodes video in the main process, which must stay decode-clean until the workers have forked
+    # (decoding before fork corrupts the workers' torchcodec state).
+    del loader
+    if is_map_style:
+        fetch_decode = measure_fetch_decode_map(dataset, args.probe_samples, args.probe_warmup)
+    else:
+        fetch_decode = measure_fetch_decode_stream(dataset, args.probe_samples, args.probe_warmup)
+
+    image_shape = list(meta.features[meta.video_keys[0]]["shape"]) if meta.video_keys else None
+    num_cameras = len(meta.video_keys)
+    results = {
+        "scenario": scenario,
+        "rank": rank,
+        "world_size": world_size,
+        "loader": "map_style" if is_map_style else "streaming",
+        "batch_size": args.batch_size,
+        "num_workers": args.num_workers,
+        "episode_pool_size": None if is_map_style else args.episode_pool_size,
+        "max_buffer_input_shards": None
+        if is_map_style
+        else (args.max_buffer_input_shards or args.episode_pool_size),
+        **info,
+        "num_cameras": num_cameras,
+        "image_shape": image_shape,
+        "fps": meta.fps,
+        "peak_rss_gb": peak_rss_gb,
+        "samples_measured": samples,
+        "steady_window_s": round(steady_elapsed_s, 2),
+        "first_batch_latency_s": round(first_batch_latency_s or float("nan"), 3),
+        # Parallel throughput over the steady window (excludes warmup + the prefetch queue it filled).
+        # A sample is one timestep (one dataset item); it decodes num_cameras video frames.
+        "samples_per_s": round(samples / steady_elapsed_s, 2) if steady_elapsed_s else 0.0,
+        "decoded_frames_per_s": round(samples / steady_elapsed_s * num_cameras, 2)
+        if steady_elapsed_s
+        else 0.0,
+        **fetch_decode,
+        # Distinct episodes per batch / batch size: ~1.0 ≈ map-style uniform, low ≈ correlated samples.
+        "shuffle_randomness_frac": round(statistics.mean(episodes_per_batch) / args.batch_size, 3)
+        if episodes_per_batch
+        else None,
+        "p50_sample_latency_ms": round(statistics.median(sample_latencies_ms), 3)
+        if sample_latencies_ms
+        else None,
+        "p95_sample_latency_ms": round(percentile(sample_latencies_ms, 95), 3),
+        "p99_sample_latency_ms": round(percentile(sample_latencies_ms, 99), 3),
+        "total_time_s": round(elapsed, 2),
+    }
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    tag = f"{scenario}_bs{args.batch_size}_w{args.num_workers}_r{rank}of{world_size}"
+    (out_dir / f"{tag}.json").write_text(json.dumps(results, indent=2))
+    flat = {k: (json.dumps(v) if isinstance(v, (dict, list)) else v) for k, v in results.items()}
+    with open(out_dir / f"{tag}.csv", "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=list(flat))
+        writer.writeheader()
+        writer.writerow(flat)
+    print(json.dumps(results, indent=2), flush=True)
+    print(f"Wrote {out_dir / tag}.json and .csv", flush=True)
+
+
+def submit_chain(args: argparse.Namespace) -> None:
+    """Submit every scenario as a serial sbatch chain (one network-bound job at a time).
+
+    Bodies are passed to ``sbatch --wrap`` as a single argv (no outer shell), so ``$SLURM_PROCID`` /
+    ``$SLURM_NTASKS`` stay literal and expand at job runtime, not at submit time.
+    """
+    this_file = Path(__file__).resolve()
+    repo_dir = str(this_file.parents[2])  # <repo>/examples/scaling/<this file>
+    logs = Path(repo_dir) / "logs"
+    logs.mkdir(exist_ok=True)
+    run = f"conda run --no-capture-output -n {args.conda_env} python"
+    common = (
+        f"--batch_size {args.batch_size} "
+        f"--prefetch_factor {args.prefetch_factor} --episode_pool_size {args.episode_pool_size} "
+        f"--video_decoder_cache_size {args.video_decoder_cache_size} --duration_s {args.duration_s} "
+        f"--num_batches {args.num_batches} --out_dir {args.out_dir}"
+    )
+    if args.max_buffer_input_shards is not None:
+        common += f" --max_buffer_input_shards {args.max_buffer_input_shards}"
+    if args.local_root:
+        common += f" --local_root {args.local_root}"
+    env_prefix = "export TOKENIZERS_PARALLELISM=false"
+    sched = []
+    for opt, env in (("--account", "ACCOUNT"), ("--partition", "PARTITION"), ("--qos", "QOS")):
+        if os.environ.get(env):
+            sched.append(f"{opt}={os.environ[env]}")
+
+    selected = args.scenarios.split(",") if args.scenarios else list(SCENARIOS)
+    prev = ""
+    for scenario in selected:
+        cfg = SCENARIOS[scenario]
+        nw = cfg.get("num_workers", args.num_workers)
+        cpus = cfg.get("cpus", nw + 4)
+        worker = f"{run} {this_file} --scenario {scenario} --num_workers {nw} {common}"
+        if cfg["nodes"] > 1:
+            # One task per node; each exports RANK/WORLD_SIZE so the stream splits shards per node.
+            inner = f"export RANK=$SLURM_PROCID WORLD_SIZE=$SLURM_NTASKS && cd {repo_dir} && {env_prefix} && {worker}"
+            body = f"srun --export=ALL bash -c '{inner}'"
+            node_flags = [f"--nodes={cfg['nodes']}", "--ntasks-per-node=1", "--gpus-per-node=1"]
+        else:
+            body = f"cd {repo_dir} && {env_prefix} && {worker}"
+            node_flags = ["--nodes=1", "--ntasks=1", "--gpus=1"]
+        cmd = [
+            "sbatch",
+            "--parsable",
+            f"--job-name=dlbench_{scenario}",
+            *node_flags,
+            f"--cpus-per-task={cpus}",
+            f"--mem={cfg['mem']}",
+            f"--time={cfg['time']}",
+            f"--output={logs}/%x-%j.out",
+            *sched,
+        ]
+        if prev:
+            cmd.append(f"--dependency=afterany:{prev}")
+        cmd += ["--wrap", body]
+        jid = subprocess.check_output(cmd, text=True).strip().split(";")[0]
+        print(f"submitted {jid}  dlbench_{scenario}{f'  (after {prev})' if prev else ''}", flush=True)
+        prev = jid
+
+    print(f"\nSubmitted {len(selected)} jobs as a serial chain. Results: {args.out_dir}/*.json", flush=True)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument(
+        "--scenario",
+        choices=list(SCENARIOS),
+        default=None,
+        help="Run ONE scenario (worker mode). Omit to submit the whole chain (orchestrator mode).",
+    )
+    p.add_argument(
+        "--scenarios",
+        type=str,
+        default=None,
+        help="Orchestrator only: comma-separated subset of scenarios to submit (default: all).",
+    )
+    p.add_argument("--local_root", type=str, default=None, help="Local LeRobotDataset copy for mmap_local.")
+    p.add_argument(
+        "--num_partitions", type=int, default=8, help="Node count for mmap_local episode partition."
+    )
+    p.add_argument("--partition_index", type=int, default=0)
+    p.add_argument(
+        "--max_episodes", type=int, default=512, help="Cap mmap_local episodes to the local share."
+    )
+    p.add_argument("--batch_size", type=int, default=64)
+    p.add_argument("--num_workers", type=int, default=8)
+    p.add_argument("--prefetch_factor", type=int, default=2)
+    p.add_argument(
+        "--episode_pool_size", type=int, default=1024, help="Streaming shuffle pool (randomness knob)."
+    )
+    p.add_argument(
+        "--max_buffer_input_shards",
+        type=int,
+        default=None,
+        help="Concurrently-live random episodes feeding the pool after reshard() "
+        "(default: episode_pool_size). The frac knob; set >= batch_size for frac->1.",
+    )
+    p.add_argument(
+        "--video_decoder_cache_size", type=int, default=32, help="Max open video decoders (bounds RAM)."
+    )
+    p.add_argument(
+        "--duration_s", type=float, default=60.0, help="Steady-state measurement window (seconds)."
+    )
+    p.add_argument(
+        "--num_batches", type=int, default=1_000_000, help="Safety cap; duration_s governs the window."
+    )
+    p.add_argument("--warmup_batches", type=int, default=5, help="Excluded from steady-state throughput.")
+    p.add_argument(
+        "--probe_samples", type=int, default=100, help="Single-process samples for fetch/decode split."
+    )
+    p.add_argument(
+        "--probe_warmup", type=int, default=10, help="Samples skipped before the fetch/decode probe."
+    )
+    p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--conda_env", type=str, default="lerobot", help="Conda env the chained jobs run in.")
+    p.add_argument("--out_dir", type=str, default="benchmarks/streaming/results_dataloading")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.scenario is None:
+        if torch.cuda.is_available():
+            print(
+                "NOTE: no --scenario given, submitting the SLURM chain. This benchmark is meant to run on a "
+                "compute cluster; run from a login node with ACCOUNT/PARTITION/QOS set.",
+                file=sys.stderr,
+            )
+        submit_chain(args)
+    else:
+        run_scenario(args.scenario, args)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -95,7 +95,7 @@ dependencies = [
 
 # ── Feature-scoped extras ──────────────────────────────────
 dataset = [
-    "datasets>=4.7.0,<5.0.0",
+    "datasets>=5.0.0,<6.0.0", # StreamingLeRobotDataset needs reshard() + shuffle(max_buffer_input_shards=...)
     "pandas>=2.0.0,<3.0.0", # NOTE: Transitive dependency of datasets
     "pyarrow>=21.0.0,<30.0.0", # NOTE: Transitive dependency of datasets
     "lerobot[av-dep]",
@@ -115,8 +115,8 @@ dataset = [
 ]
 training = [
     "lerobot[dataset]",
-    "wandb>=0.24.0,<0.28.0",
-    "lerobot[accelerate-dep]",
+    "accelerate>=1.10.0,<2.0.0",
+    "wandb>=0.24.0,<0.25.0",
 ]
 hardware = [
     "lerobot[pynput-dep]",
@@ -142,8 +142,7 @@ 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,<2.0.0", "protobuf>=6.31.1,<8.0.0"]
-accelerate-dep = ["accelerate>=1.14.0,<2.0.0"]
+grpcio-dep = ["grpcio==1.73.1", "protobuf>=6.31.1,<6.32.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"]
@@ -178,12 +177,7 @@ unitree_g1 = [
     "lerobot[matplotlib-dep]",
     "lerobot[pygame-dep]",
 ]
-# 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",
-]
+reachy2 = ["reachy2_sdk>=1.0.15,<1.1.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]"]
@@ -205,7 +199,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", "lerobot[accelerate-dep]"]
+smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "accelerate>=1.7.0,<2.0.0"]
 multi_task_dit = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]"]
 groot = [
     "lerobot[transformers-dep]",
@@ -229,23 +223,8 @@ vla_jepa = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]", "lerobot[qwen
 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,<2.0.0", "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", "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"]
@@ -338,7 +317,6 @@ 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 ----------------
@@ -355,9 +333,14 @@ explicit = true
 [tool.uv.sources]
 torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
+# Temporary: the native streaming pipeline needs batch(by_column=...) to survive shard/shuffle
+# re-creation (datasets#8259), reshard() per row group (#8193), and shuffle(max_buffer_input_shards=...)
+# (#8194) — all merged, not yet in a tagged 5.0 release. Pin to the merge commit until the next
+# datasets release ships them, then drop this and rely on the `datasets>=5.0.0` floor in `dependencies`.
+datasets = { git = "https://github.com/huggingface/datasets.git", rev = "2c45eab1bb975ac3d846f2aa6217b82adec8eba3" }
 
 [tool.setuptools.package-data]
-lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"]
+lerobot = ["envs/*.json"]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/lerobot/annotations/__init__.py

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

src/lerobot/annotations/steerable_pipeline/__init__.py

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -1,341 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""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.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
-        )
-        # Atomic replace: write to a sibling tmp path and rename so a crash
-        # mid-write can't leave a half-written shard that ``pq.read_table``
-        # would then fail to open. ``Path.replace`` is atomic on POSIX +
-        # Windows when source and target sit on the same filesystem.
-        tmp_path = path.with_suffix(path.suffix + ".tmp")
-        pq.write_table(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/train_utils.py

@@ -49,19 +49,8 @@ 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, 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 save_training_step(step: int, save_dir: Path) -> None:
+    write_json({"step": step}, save_dir / TRAINING_STEP)
 
 
 def load_training_step(save_dir: Path) -> int:
@@ -69,16 +58,6 @@ 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():
@@ -96,8 +75,6 @@ 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:
 
@@ -123,10 +100,6 @@ 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)
@@ -139,9 +112,7 @@ 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, num_processes=num_processes, batch_size=batch_size
-    )
+    save_training_state(checkpoint_dir, step, optimizer, scheduler)
 
 
 def save_training_state(
@@ -149,8 +120,6 @@ 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.
@@ -162,12 +131,10 @@ 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, num_processes=num_processes, batch_size=batch_size)
+    save_training_step(train_step, save_dir)
     save_rng_state(save_dir)
     if optimizer is not None:
         save_optimizer_state(optimizer, save_dir)

src/lerobot/configs/default.py

@@ -39,6 +39,10 @@ class DatasetConfig:
     # This reduces memory and speeds up DataLoader IPC. The training pipeline handles the conversion.
     return_uint8: bool = False
     streaming: bool = False
+    # Whole episodes each streaming consumer keeps open to shuffle across (the randomness knob).
+    # Larger mixes more episodes per batch at the cost of cold-start latency; RAM stays small because
+    # the pool holds tabular rows only. Ignored when streaming is False.
+    streaming_episode_pool_size: int = 1024
 
     def __post_init__(self) -> None:
         if self.episodes is not None:
@@ -73,8 +77,6 @@ class EvalConfig:
     # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
     # Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1.
     use_async_envs: bool = True
-    # Whether to record eval rollouts as a LeRobot v3.0 dataset on disk.
-    recording: bool = False
 
     def __post_init__(self) -> None:
         if self.batch_size == 0:

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, compute_sampler_state
+from .sampler import EpisodeAwareSampler
 from .streaming_dataset import StreamingLeRobotDataset
 from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card
 from .video_utils import VideoEncodingManager
@@ -82,7 +82,6 @@ __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

@@ -286,8 +286,6 @@ 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.
 
@@ -305,8 +303,6 @@ 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")
 
@@ -355,12 +351,8 @@ 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, concatenate_videos
-        )
-        data_idx = aggregate_data(
-            src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data
-        )
+        videos_idx = aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size)
+        data_idx = aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size)
 
         meta_idx = aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx)
 
@@ -375,9 +367,7 @@ def aggregate_datasets(
     logging.info("Aggregation complete.")
 
 
-def aggregate_videos(
-    src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
-):
+def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size):
     """Aggregates video chunks from a source dataset into the destination dataset.
 
     Handles video file concatenation and rotation based on file size limits.
@@ -389,7 +379,6 @@ def aggregate_videos(
         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.
     """
@@ -450,7 +439,7 @@ def aggregate_videos(
             src_size = get_file_size_in_mb(src_path)
             dst_size = get_file_size_in_mb(dst_path)
 
-            if not concatenate_videos or dst_size + src_size >= video_files_size_in_mb:
+            if 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)
@@ -488,7 +477,7 @@ def aggregate_videos(
     return videos_idx
 
 
-def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
+def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size):
     """Aggregates data chunks from a source dataset into the destination dataset.
 
     Reads source data files, updates indices to match the aggregated dataset,
@@ -504,7 +493,6 @@ 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.
@@ -550,7 +538,6 @@ 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,
         )
 
         # Record the mapping from source to actual destination
@@ -627,7 +614,6 @@ def append_or_create_parquet_file(
     contains_images: bool = False,
     aggr_root: Path = None,
     hf_features: datasets.Features | None = None,
-    concatenate: bool = True,
 ) -> tuple[dict[str, int], tuple[int, int]]:
     """Appends data to an existing parquet file or creates a new one based on size constraints.
 
@@ -644,7 +630,6 @@ 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.
 
     Returns:
         tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
@@ -664,7 +649,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 not concatenate or dst_size + src_size >= max_mb:
+    if 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)

src/lerobot/datasets/compute_stats.py

@@ -59,8 +59,6 @@ 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,8 +261,6 @@ 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.
 
@@ -272,8 +270,6 @@ 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")
@@ -288,8 +284,6 @@ 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(
@@ -951,8 +945,17 @@ def _write_parquet(df: pd.DataFrame, path: Path, meta: LeRobotDatasetMetadata) -
         ep_dataset = embed_images(ep_dataset)
 
     table = ep_dataset.with_format("arrow")[:]
-    writer = pq.ParquetWriter(path, schema=table.schema, compression="snappy", use_dictionary=True)
-    writer.write_table(table)
+    # Emit several row groups with a page index instead of one giant row group. A single row group forces
+    # streaming readers to materialize the whole file's columns per open shard; with random-access streaming
+    # (shuffle + delta windows) across many workers x shards that dominates RAM. Targeting ~32MB-uncompressed
+    # groups bounds per-shard memory while keeping groups large enough to scan
+    # efficiently; the page index lets readers skip to the pages they need.
+    target_row_group_bytes = 32 * 1024 * 1024
+    row_group_size = max(1, min(table.num_rows, table.num_rows * target_row_group_bytes // max(table.nbytes, 1)))
+    writer = pq.ParquetWriter(
+        path, schema=table.schema, compression="snappy", use_dictionary=True, write_page_index=True
+    )
+    writer.write_table(table, row_group_size=row_group_size)
     writer.close()
 
 

src/lerobot/datasets/factory.py

@@ -106,7 +106,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
                 delta_timestamps=delta_timestamps,
                 image_transforms=image_transforms,
                 revision=cfg.dataset.revision,
-                max_num_shards=cfg.num_workers,
+                episode_pool_size=cfg.dataset.streaming_episode_pool_size,
                 tolerance_s=cfg.tolerance_s,
                 return_uint8=True,
             )

src/lerobot/datasets/sampler.py

@@ -14,36 +14,14 @@
 # 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],
@@ -52,125 +30,63 @@ class EpisodeAwareSampler:
         drop_n_first_frames: int = 0,
         drop_n_last_frames: int = 0,
         shuffle: bool = False,
-        seed: int = 0,
+        generator: torch.Generator | None = None,
     ):
-        """
+        """Sampler that optionally incorporates episode boundary information.
+
         Args:
-            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.
+            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.
             shuffle: Whether to shuffle the indices.
-            seed: Seed the permutation is derived from (together with the epoch).
+            generator: Generator used for shuffling. Exposing this attribute (even when None) lets
+                       `accelerate` register it as the synchronized RNG in distributed training, so
+                       every rank draws the same permutation and batch shards stay disjoint. When
+                       None, shuffling falls back to the global torch RNG.
         """
         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}")
 
-        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)}"
-            )
+        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))
 
-        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():
+        if not indices:
             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._starts = starts[used]
-        self._cum_lengths = np.cumsum(lengths[used])
-        self._num_frames = int(self._cum_lengths[-1])
+        self.indices = indices
         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
+        self.generator = generator
 
     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:
-            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]))
+            for i in torch.randperm(len(self.indices), generator=self.generator):
+                yield self.indices[i]
         else:
-            for k in range(start, self._num_frames):
-                yield self._frame_index(k)
+            for i in self.indices:
+                yield i
 
     def __len__(self) -> int:
-        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}
+        return len(self.indices)

src/lerobot/datasets/streaming_dataset.py

@@ -13,16 +13,17 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from collections import deque
-from collections.abc import Callable, Generator, Iterable, Iterator
+import logging
+from collections.abc import Callable, Iterator
 from pathlib import Path
 
 import datasets
 import numpy as np
 import torch
 from datasets import load_dataset
+from datasets.distributed import split_dataset_by_node
 
-from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
+from lerobot.utils.constants import HF_LEROBOT_HOME
 
 from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
 from .feature_utils import get_delta_indices
@@ -31,207 +32,70 @@ from .utils import (
     check_version_compatibility,
     find_float_index,
     is_float_in_list,
-    safe_shard,
 )
 from .video_utils import (
     VideoDecoderCache,
     decode_video_frames_torchcodec,
 )
 
+logger = logging.getLogger(__name__)
 
-class LookBackError(Exception):
-    """
-    Exception raised when trying to look back in the history of a Backtrackable object.
-    """
-
-    pass
-
-
-class LookAheadError(Exception):
-    """
-    Exception raised when trying to look ahead in the future of a Backtrackable object.
-    """
-
-    pass
-
-
-class Backtrackable[T]:
-    """
-    Wrap any iterator/iterable so you can step back up to `history` items
-    and look ahead up to `lookahead` items.
-
-    This is useful for streaming datasets where you need to access previous and future items
-    but can't load the entire dataset into memory.
-
-    Example:
-    -------
-    ```python
-    ds = load_dataset("c4", "en", streaming=True, split="train")
-    rev = Backtrackable(ds, history=3, lookahead=2)
-
-    x0 = next(rev)  # forward
-    x1 = next(rev)
-    x2 = next(rev)
-
-    # Look ahead
-    x3_peek = rev.peek_ahead(1)  # next item without moving cursor
-    x4_peek = rev.peek_ahead(2)  # two items ahead
-
-    # Look back
-    x1_again = rev.peek_back(1)  # previous item without moving cursor
-    x0_again = rev.peek_back(2)  # two items back
-
-    # Move backward
-    x1_back = rev.prev()  # back one step
-    next(rev)  # returns x2, continues forward from where we were
-    ```
-    """
-
-    __slots__ = ("_source", "_back_buf", "_ahead_buf", "_cursor", "_history", "_lookahead")
-
-    def __init__(self, iterable: Iterable[T], *, history: int = 1, lookahead: int = 0):
-        if history < 1:
-            raise ValueError("history must be >= 1")
-        if lookahead <= 0:
-            raise ValueError("lookahead must be > 0")
-
-        self._source: Iterator[T] = iter(iterable)
-        self._back_buf: deque[T] = deque(maxlen=history)
-        self._ahead_buf: deque[T] = deque(maxlen=lookahead) if lookahead > 0 else deque()
-        self._cursor: int = 0
-        self._history = history
-        self._lookahead = lookahead
-
-    def __iter__(self) -> "Backtrackable[T]":
-        return self
-
-    def __next__(self) -> T:
-        # If we've stepped back, consume from back buffer first
-        if self._cursor < 0:  # -1 means "last item", etc.
-            self._cursor += 1
-            return self._back_buf[self._cursor]
-
-        # If we have items in the ahead buffer, use them first
-        item = self._ahead_buf.popleft() if self._ahead_buf else next(self._source)
-
-        # Add current item to back buffer and reset cursor
-        self._back_buf.append(item)
-        self._cursor = 0
-        return item
-
-    def prev(self) -> T:
-        """
-        Step one item back in history and return it.
-        Raises IndexError if already at the oldest buffered item.
-        """
-        if len(self._back_buf) + self._cursor <= 1:
-            raise LookBackError("At start of history")
-
-        self._cursor -= 1
-        return self._back_buf[self._cursor]
-
-    def peek_back(self, n: int = 1) -> T:
-        """
-        Look `n` items back (n=1 == previous item) without moving the cursor.
-        """
-        if n < 0 or n + 1 > len(self._back_buf) + self._cursor:
-            raise LookBackError("peek_back distance out of range")
-
-        return self._back_buf[self._cursor - (n + 1)]
-
-    def peek_ahead(self, n: int = 1) -> T:
-        """
-        Look `n` items ahead (n=1 == next item) without moving the cursor.
-        Fills the ahead buffer if necessary.
-        """
-        if n < 1:
-            raise LookAheadError("peek_ahead distance must be 1 or more")
-        elif n > self._lookahead:
-            raise LookAheadError("peek_ahead distance exceeds lookahead limit")
-
-        # Fill ahead buffer if we don't have enough items
-        while len(self._ahead_buf) < n:
-            try:
-                item = next(self._source)
-                self._ahead_buf.append(item)
-
-            except StopIteration as err:
-                raise LookAheadError("peek_ahead: not enough items in source") from err
-
-        return self._ahead_buf[n - 1]
-
-    def history(self) -> list[T]:
-        """
-        Return a copy of the buffered history (most recent last).
-        The list length ≤ `history` argument passed at construction.
-        """
-        if self._cursor == 0:
-            return list(self._back_buf)
-
-        # When cursor<0, slice so the order remains chronological
-        return list(self._back_buf)[: self._cursor or None]
-
-    def can_peek_back(self, steps: int = 1) -> bool:
-        """
-        Check if we can go back `steps` items without raising an IndexError.
-        """
-        return steps <= len(self._back_buf) + self._cursor
-
-    def can_peek_ahead(self, steps: int = 1) -> bool:
-        """
-        Check if we can peek ahead `steps` items.
-        This may involve trying to fill the ahead buffer.
-        """
-        if self._lookahead > 0 and steps > self._lookahead:
-            return False
-
-        # Try to fill ahead buffer to check if we can peek that far
-        try:
-            while len(self._ahead_buf) < steps:
-                if self._lookahead > 0 and len(self._ahead_buf) >= self._lookahead:
-                    return False
-                item = next(self._source)
-                self._ahead_buf.append(item)
-            return True
-        except StopIteration:
-            return False
+# Bound the default frame-level shuffle buffer: rows are tabular-only (~KB each), so this is
+# roughly a few hundred MB of host RAM per consumer at the cap.
+_MAX_DEFAULT_FRAME_BUFFER = 200_000
 
 
 class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
-    """LeRobotDataset with streaming capabilities.
+    """LeRobotDataset with streaming capabilities, built on native HF `datasets` primitives.
 
-    This class extends LeRobotDataset to add streaming functionality, allowing data to be streamed
-    rather than loaded entirely into memory. This is especially useful for large datasets that may
-    not fit in memory or when you want to quickly explore a dataset without downloading it completely.
+    The tabular side is a pure `datasets` pipeline::
 
-    The key innovation is using a Backtrackable iterator that maintains a bounded buffer of recent
-    items, allowing us to access previous frames for delta timestamps without loading the entire
-    dataset into memory.
+        load_dataset(streaming=True)                          # parquet shards from the Hub / a bucket
+          -> reshard()                                        # 1 shard == 1 row group == 1 episode
+          -> split_dataset_by_node(rank, world_size)          # disjoint shards per rank
+          -> batch(by_column="episode_index")                 # whole episodes (one per shard)
+          -> shuffle(episode_pool_size, max_buffer_input_shards)  # K random episodes, global perm
+          -> map(explode + exact delta windows)               # episode -> frames, windows are exact
+          -> shuffle(buffer_size=frame_shuffle_buffer_size)   # frame-level interleave
+
+    and this class is a thin torch ``IterableDataset`` wrapper around it that decodes video
+    per emitted sample (decode-on-exit), applies image transforms, and attaches the task
+    string. DataLoader workers are split natively by `datasets` (disjoint shards per worker),
+    and resume uses the native ``state_dict`` / ``load_state_dict``.
+
+    Random-episode admission (Plan B): the LeRobot writer stores one Parquet row group per
+    episode, so ``datasets.IterableDataset.reshard()`` makes one shard == one episode (no new
+    files; shards are (file, row_group) pairs). ``shuffle`` then permutes shard order globally and
+    fills its buffer from ``max_buffer_input_shards`` shards concurrently, so the episode pool is a
+    uniformly-random sample of the corpus regardless of how many episodes are packed per file.
+    ``max_buffer_input_shards`` is the number of concurrently-live random episodes; set it
+    ``>= batch_size`` for the per-batch distinct-episode fraction to approach 1.
+
+    Requirement: ONE ROW GROUP PER EPISODE. Recorded datasets satisfy this; bulk
+    ``df.to_parquet`` / ``push_to_hub`` / aggregate paths collapse row groups and are rejected at
+    init (see ``validate_row_groups``). Old collapsed datasets still load fine for the map-style
+    path; only this streaming random-episode path requires the invariant.
+
+    Randomness: a batch mixes up to ``episode_pool_size`` distinct episodes; delta windows are
+    exact slices of the resident episode with correct padding at episode boundaries.
+
+    Resume: ``state_dict()`` / ``load_state_dict()`` delegate to `datasets`. Samples sitting in
+    the shuffle buffers at checkpoint time are skipped on resume (documented `datasets`
+    behavior), so resume never repeats data but may drop up to roughly
+    ``episode_pool_size x episode_len + frame_shuffle_buffer_size`` frames — negligible at
+    training scale. The contract is exact with ``num_workers=0``; with DataLoader workers use
+    ``torchdata.stateful_dataloader.StatefulDataLoader``, which checkpoints each worker's
+    dataset state through this same protocol.
 
     Example:
-        Basic usage:
         ```python
-        from lerobot.common.datasets.streaming_dataset import StreamingLeRobotDataset
-
-        # Create a streaming dataset with delta timestamps
-        delta_timestamps = {
-            "observation.image": [-1.0, -0.5, 0.0],  # 1 sec ago, 0.5 sec ago, current
-            "action": [0.0, 0.1, 0.2],  # current, 0.1 sec future, 0.2 sec future
-        }
-
         dataset = StreamingLeRobotDataset(
             repo_id="your-dataset-repo-id",
-            delta_timestamps=delta_timestamps,
-            streaming=True,
-            buffer_size=1000,
+            delta_timestamps={"action": [0.0, 0.1, 0.2]},
+            episode_pool_size=1024,
         )
-
-        # Iterate over the dataset
-        for i, item in enumerate(dataset):
-            print(f"Sample {i}: Episode {item['episode_index']} Frame {item['frame_index']}")
-            # item will contain stacked frames according to delta_timestamps
-            if i >= 10:
-                break
+        for sample in dataset:
+            ...
         ```
     """
 
@@ -246,12 +110,20 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         revision: str | None = None,
         force_cache_sync: bool = False,
         streaming: bool = True,
-        buffer_size: int = 1000,
-        max_num_shards: int = 16,
+        episode_pool_size: int | None = 1024,
+        max_buffer_input_shards: int | None = None,
+        frame_shuffle_buffer_size: int | None = None,
+        buffer_size: int | None = None,
+        max_num_shards: int | None = None,
         seed: int = 42,
         rng: np.random.Generator | None = None,
         shuffle: bool = True,
         return_uint8: bool = False,
+        rank: int | None = None,
+        world_size: int | None = None,
+        video_decoder_cache_size: int | None = None,
+        data_files_root: str | None = None,
+        validate_row_groups: bool = True,
     ):
         """Initialize a StreamingLeRobotDataset.
 
@@ -267,11 +139,40 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             revision (str, optional): Git revision id (branch name, tag, or commit hash).
             force_cache_sync (bool, optional): Flag to sync and refresh local files first.
             streaming (bool, optional): Whether to stream the dataset or load it all. Defaults to True.
-            buffer_size (int, optional): Buffer size for shuffling when streaming. Defaults to 1000.
-            max_num_shards (int, optional): Number of shards to re-shard the input dataset into. Defaults to 16.
+            episode_pool_size (int, optional): Whole episodes each consumer keeps open to shuffle
+                across — the randomness knob. Larger mixes more episodes per batch (closer to
+                map-style uniform) at the cost of cold-start latency and frame-buffer RAM.
+                Defaults to 1024.
+            max_buffer_input_shards (int | None, optional): Number of shards (== episodes, after
+                ``reshard()``) the episode-pool ``shuffle`` reads from concurrently — i.e. the count
+                of concurrently-live random episodes feeding the pool from a global shard permutation.
+                Set ``>= batch_size`` for the per-batch distinct-episode fraction to approach 1.
+                Defaults to ``episode_pool_size``.
+            frame_shuffle_buffer_size (int | None, optional): Frame-level shuffle buffer after the
+                episode pool. Defaults to ``episode_pool_size x average episode length`` (capped),
+                which matches the pool's mixing radius.
+            buffer_size (int | None, optional): Deprecated; superseded by ``episode_pool_size``.
+            max_num_shards (int | None, optional): Deprecated; `datasets` handles shard-to-worker
+                assignment natively.
             seed (int, optional): Reproducibility random seed.
-            rng (np.random.Generator | None, optional): Random number generator.
-            shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
+            rng (np.random.Generator | None, optional): Deprecated; ignored.
+            shuffle (bool, optional): Whether to shuffle. False yields episodes in stream order.
+            rank (int | None, optional): This process' rank for distributed training. Each rank streams
+                a disjoint set of shards via ``split_dataset_by_node``. When omitted, resolved from
+                Accelerate (``process_index``) or the ``RANK`` env var, defaulting to 0.
+            world_size (int | None, optional): Total number of distributed processes. When omitted,
+                resolved from Accelerate or ``WORLD_SIZE``, defaulting to 1. For an even per-rank split,
+                ``num_shards % world_size == 0`` should hold (warned otherwise).
+            video_decoder_cache_size (int | None, optional): Max number of open video decoders to retain.
+                When omitted, sized to the episode pool's working set, capped at 128.
+            data_files_root (str | None, optional): fsspec root holding the bulk ``data/`` and ``videos/``
+                trees (e.g. ``hf://buckets/<owner>/<name>``). When set, parquet and video bytes are read
+                from there while metadata still loads from ``repo_id`` on the Hub.
+            validate_row_groups (bool, optional): When True (default), verify at init that the dataset
+                stores one Parquet row group per episode (sampling data-file footers) and that
+                ``num_shards`` is divisible by ``world_size`` for distributed runs, raising a clear
+                ``ValueError`` otherwise. Set False to skip the checks (e.g. single-process debugging);
+                the divisibility check then downgrades to a warning.
         """
         super().__init__()
         self.repo_id = repo_id
@@ -284,15 +185,36 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         self.tolerance_s = tolerance_s
         self.revision = revision if revision else CODEBASE_VERSION
         self.seed = seed
-        self.rng = rng if rng is not None else np.random.default_rng(seed)
+        if rng is not None:
+            logger.warning("StreamingLeRobotDataset: `rng` is deprecated and ignored; use `seed`.")
+        if buffer_size is not None:
+            logger.warning(
+                "StreamingLeRobotDataset: `buffer_size` is deprecated and ignored; "
+                "use `episode_pool_size` (whole episodes, not frames)."
+            )
+        if max_num_shards is not None:
+            logger.warning(
+                "StreamingLeRobotDataset: `max_num_shards` is deprecated and ignored; "
+                "`datasets` assigns shards to DataLoader workers natively."
+            )
         self.shuffle = shuffle
 
         self.streaming = streaming
-        self.buffer_size = buffer_size
+        self.episode_pool_size = max(1, episode_pool_size) if episode_pool_size else 1024
+        self.max_buffer_input_shards = (
+            max(1, max_buffer_input_shards) if max_buffer_input_shards else self.episode_pool_size
+        )
+        self.validate_row_groups = validate_row_groups
         self._return_uint8 = return_uint8
 
+        self.rank, self.world_size = self._resolve_distributed(rank, world_size)
+        self.video_decoder_cache_size = video_decoder_cache_size
+        self.data_files_root = data_files_root.rstrip("/") if data_files_root else None
+
         # We cache the video decoders to avoid re-initializing them at each frame (avoiding a ~10x slowdown)
         self.video_decoder_cache = None
+        self._epoch = 0
+        self._in_flight_epoch = 0
 
         if self._requested_root is not None:
             self.root.mkdir(exist_ok=True, parents=True)
@@ -314,15 +236,50 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             self.delta_timestamps = delta_timestamps
             self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
 
-        self.hf_dataset: datasets.IterableDataset = load_dataset(
-            self.repo_id if not self.streaming_from_local else str(self.root),
-            split="train",
-            streaming=self.streaming,
-            data_files="data/*/*.parquet",
-            revision=self.revision,
+        if self.data_files_root is not None:
+            # Bulk data lives in an fsspec root (e.g. an HF storage bucket); metadata stays on the Hub.
+            self.hf_dataset: datasets.IterableDataset = load_dataset(
+                "parquet",
+                split="train",
+                streaming=self.streaming,
+                data_files=f"{self.data_files_root}/data/*/*.parquet",
+            )
+        else:
+            self.hf_dataset = load_dataset(
+                self.repo_id if not self.streaming_from_local else str(self.root),
+                split="train",
+                streaming=self.streaming,
+                data_files="data/*/*.parquet",
+                revision=self.revision,
+            )
+
+        # Drop any parquet columns not declared in the dataset's feature contract. Some revisions / sources
+        # (e.g. an unversioned bucket holding `main`) carry extra, possibly variable-length annotation
+        # columns such as `language_events`; left in, they leak into the sample and break default DataLoader
+        # collation across frames of differing length. On a clean revision this is a no-op.
+        known_columns = set(self.meta.features)
+        extra_columns = [c for c in (self.hf_dataset.column_names or []) if c not in known_columns]
+        if extra_columns:
+            self.hf_dataset = self.hf_dataset.remove_columns(extra_columns)
+
+        # Reshard Parquet per row group so 1 shard == 1 row group == 1 episode (the LeRobot writer
+        # emits one row group per episode). This lets the episode-pool shuffle admit uniformly-random
+        # episodes from a global shard permutation, independent of how many episodes are packed per file.
+        if self.streaming:
+            self.hf_dataset = self.hf_dataset.reshard()
+        self.num_shards = self.hf_dataset.num_shards
+
+        if self.validate_row_groups and self.streaming:
+            self._validate_row_groups_per_episode()
+
+        avg_episode_len = max(1, round(self.meta.total_frames / max(1, self.meta.total_episodes)))
+        self.frame_shuffle_buffer_size = (
+            frame_shuffle_buffer_size
+            if frame_shuffle_buffer_size is not None
+            else min(self.episode_pool_size * avg_episode_len, _MAX_DEFAULT_FRAME_BUFFER)
         )
 
-        self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
+        self._pipeline = self._build_pipeline()
 
     @property
     def num_frames(self):
@@ -337,96 +294,270 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         return self.meta.fps
 
     @staticmethod
-    def _iter_random_indices(
-        rng: np.random.Generator, buffer_size: int, random_batch_size=100
-    ) -> Iterator[int]:
-        while True:
-            yield from (int(i) for i in rng.integers(0, buffer_size, size=random_batch_size))
+    def _resolve_distributed(rank: int | None, world_size: int | None) -> tuple[int, int]:
+        """Resolve (rank, world_size) for distributed streaming.
 
-    @staticmethod
-    def _infinite_generator_over_elements(rng: np.random.Generator, elements: list[int]) -> Iterator[int]:
-        while True:
-            yield rng.choice(elements)
+        Explicit arguments win. Otherwise prefer an already-initialized Accelerate state, then the
+        ``RANK``/``WORLD_SIZE`` env vars set by launchers, and finally fall back to single-process (0, 1).
+        """
+        import os
+
+        if rank is not None and world_size is not None:
+            return rank, world_size
+
+        try:
+            from accelerate.state import PartialState
+
+            if PartialState._shared_state:  # only read it if already initialized; never initialize here
+                state = PartialState()
+                return state.process_index, state.num_processes
+        except Exception:
+            logger.debug("Could not resolve distributed state from Accelerate; using env/defaults.")
+
+        env_rank = os.environ.get("RANK")
+        env_world = os.environ.get("WORLD_SIZE")
+        if env_rank is not None and env_world is not None:
+            return int(env_rank), int(env_world)
+
+        return 0, 1
+
+    def _resolve_data_root(self) -> str:
+        """fsspec root that holds the bulk ``data/`` parquet tree (revision-qualified for the Hub)."""
+        if self.data_files_root is not None:
+            return self.data_files_root
+        if self.streaming and not self.streaming_from_local:
+            return f"hf://datasets/{self.repo_id}@{self.revision}"
+        return str(self.root)
+
+    def _episode_files(self) -> dict[tuple[int, int], list[int]]:
+        """Map each data file ``(chunk_index, file_index)`` to the episode indices it stores."""
+        file_to_eps: dict[tuple[int, int], list[int]] = {}
+        for ep in range(self.meta.total_episodes):
+            row = self.meta.episodes[ep]
+            key = (int(row["data/chunk_index"]), int(row["data/file_index"]))
+            file_to_eps.setdefault(key, []).append(ep)
+        return file_to_eps
+
+    def _validate_row_groups_per_episode(self, sample_files: int = 32) -> None:
+        """Verify the dataset stores ONE ROW GROUP PER EPISODE so each episode is an independently
+        addressable shard after ``reshard()``. Cheap (footer-only) and sampled.
+
+        Raises:
+            ValueError: if a sampled data file collapses several episodes into fewer row groups, or
+                the whole dataset is one row group per file while holding many more episodes than files.
+        """
+        import fsspec
+        import pyarrow.parquet as pq
+
+        file_to_eps = self._episode_files()
+        num_data_files = len(file_to_eps)
+
+        # Whole-dataset extreme: reshard() could not split beyond file granularity (one row group per
+        # file) yet there are many more episodes than files -> collapsed.
+        if self.num_shards <= num_data_files and self.meta.total_episodes > self.num_shards:
+            raise ValueError(
+                f"{self.repo_id}: after reshard() the stream still has only {self.num_shards} shard(s) "
+                f"for {self.meta.total_episodes} episodes across {num_data_files} data file(s) — i.e. one "
+                "row group per file. StreamingLeRobotDataset random-episode shuffling requires ONE ROW "
+                "GROUP PER EPISODE so each episode is an independently addressable shard after reshard(). "
+                "Re-emit through the LeRobot writer (one write_table per episode) or fix the aggregate / "
+                "annotate / push_to_hub writer that collapsed the row groups, then re-upload. Recorded "
+                "datasets already satisfy this. Pass validate_row_groups=False to bypass (random-episode "
+                "quality will degrade)."
+            )
+
+        data_root = self._resolve_data_root()
+        rng = np.random.default_rng(self.seed)
+        keys = list(file_to_eps)
+        chosen = rng.choice(len(keys), size=min(sample_files, len(keys)), replace=False)
+        for i in chosen:
+            chunk_idx, file_idx = keys[int(i)]
+            n_ep = len(file_to_eps[(chunk_idx, file_idx)])
+            rel = self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
+            path = f"{data_root}/{rel}"
+            with fsspec.open(path, "rb") as f:
+                pf = pq.ParquetFile(f)
+                n_rg = pf.num_row_groups
+                num_rows = pf.metadata.num_rows
+            if n_rg < n_ep:
+                raise ValueError(
+                    f"{path}: stored as {n_rg} Parquet row group(s) ({num_rows} rows across "
+                    f"{n_ep} episodes). StreamingLeRobotDataset random-episode shuffling requires ONE ROW "
+                    "GROUP PER EPISODE so each episode becomes an independently addressable shard after "
+                    "reshard(). This file was written by a bulk df.to_parquet / push_to_hub / aggregate "
+                    "path that collapses row groups. Re-emit through the LeRobot writer (one write_table "
+                    "per episode) or fix the aggregate/annotate writer, then re-upload. Recorded datasets "
+                    "already satisfy this. Pass validate_row_groups=False to bypass (quality will degrade)."
+                )
+
+    def _build_pipeline(self) -> datasets.IterableDataset:
+        """Assemble the native tabular pipeline (everything except video decode)."""
+        ds = self.hf_dataset
+        if self.world_size > 1:
+            if ds.num_shards % self.world_size != 0:
+                msg = (
+                    f"num_shards ({ds.num_shards}) is not divisible by world_size ({self.world_size}). "
+                    "After reshard() num_shards == the episode count, and split_dataset_by_node only "
+                    "assigns shards evenly when num_shards % world_size == 0; otherwise every rank "
+                    "streams (and pays for) the full dataset and keeps only 1/world_size of it. Pin "
+                    "world_size to a divisor of the episode count, or drop/pad episodes to a divisible "
+                    "count with the dataset tools. Set validate_row_groups=False to downgrade to a warning."
+                )
+                if self.validate_row_groups:
+                    raise ValueError(msg)
+                logger.warning(msg)
+            ds = split_dataset_by_node(ds, rank=self.rank, world_size=self.world_size)
+
+        ds = ds.batch(by_column="episode_index")
+        episode_columns = list(ds.column_names or self.hf_dataset.column_names or [])
+        if self.shuffle:
+            max_input_shards = max(1, min(self.max_buffer_input_shards, ds.num_shards))
+            ds = ds.shuffle(
+                seed=self.seed,
+                buffer_size=self.episode_pool_size,
+                max_buffer_input_shards=max_input_shards,
+            )
+        # A row-count-changing batched map must drop the input columns explicitly; the exploded
+        # frames re-emit them (windowed keys replaced by their delta windows + *_is_pad masks).
+        ds = ds.map(self._explode_episodes, batched=True, remove_columns=episode_columns)
+        if self.shuffle:
+            ds = ds.shuffle(seed=self.seed + 1, buffer_size=max(2, self.frame_shuffle_buffer_size))
+        return ds
+
+    def _tabular_window_keys(self) -> list[str]:
+        if self.delta_indices is None:
+            return []
+        return [key for key in self.delta_indices if key not in self.meta.video_keys]
+
+    def _explode_episodes(self, episode_batch: dict[str, list[list]]) -> dict[str, list]:
+        """Episode batches -> per-frame rows, with exact tabular delta windows and pad masks.
+
+        Runs inside the `datasets` pipeline (plain Python values, no torch). For each windowed key
+        the original per-frame value is replaced by its delta window (list of values, clamped to
+        the episode bounds) plus a ``{key}_is_pad`` mask, mirroring the map-style dataset.
+        """
+        window_keys = set(self._tabular_window_keys())
+        out: dict[str, list] = {key: [] for key in episode_batch if key not in window_keys}
+        for key in window_keys:
+            out[key] = []
+            out[f"{key}_is_pad"] = []
+
+        num_episodes = len(episode_batch["episode_index"])
+        for e in range(num_episodes):
+            length = len(episode_batch["episode_index"][e])
+            for key, column in episode_batch.items():
+                if key in window_keys:
+                    continue
+                out[key].extend(column[e])
+            for key in window_keys:
+                episode_column = episode_batch[key][e]
+                deltas = self.delta_indices[key]
+                for t in range(length):
+                    window = []
+                    is_pad = []
+                    for delta in deltas:
+                        j = t + delta
+                        window.append(episode_column[min(max(j, 0), length - 1)])
+                        is_pad.append(not 0 <= j < length)
+                    out[key].append(window)
+                    out[f"{key}_is_pad"].append(is_pad)
+        return out
+
+    def _make_video_decoder_cache(self) -> VideoDecoderCache:
+        """Size the decoder cache to the pool's working set (pool episodes x cameras), capped at 128."""
+        if self.video_decoder_cache_size is not None:
+            return VideoDecoderCache(max_size=self.video_decoder_cache_size)
+        num_cameras = len(self.meta.video_keys)
+        if num_cameras == 0:
+            return VideoDecoderCache()
+        return VideoDecoderCache(max_size=min((self.episode_pool_size + 1) * num_cameras, 128))
 
-    # TODO(fracapuano): Implement multi-threaded prefetching to accelerate data loading.
-    # The current sequential iteration is a bottleneck. A producer-consumer pattern
-    # could be used with a ThreadPoolExecutor to run `make_frame` (especially video decoding)
-    # in parallel, feeding a queue from which this iterator will yield processed items.
     def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
-        if self.video_decoder_cache is None:
-            self.video_decoder_cache = VideoDecoderCache()
-
-        # keep the same seed across exhaustions if shuffle is False, otherwise shuffle data across exhaustions
-        rng = np.random.default_rng(self.seed) if not self.shuffle else self.rng
-
-        buffer_indices_generator = self._iter_random_indices(rng, self.buffer_size)
-
-        idx_to_backtrack_dataset = {
-            idx: self._make_backtrackable_dataset(safe_shard(self.hf_dataset, idx, self.num_shards))
-            for idx in range(self.num_shards)
-        }
-
-        # This buffer is populated while iterating on the dataset's shards
-        # the logic is to add 2 levels of randomness:
-        # (1) sample one shard at random from the ones available, and
-        # (2) sample one frame from the shard sampled at (1)
-        frames_buffer = []
-        while available_shards := list(idx_to_backtrack_dataset.keys()):
-            shard_key = next(self._infinite_generator_over_elements(rng, available_shards))
-            backtrack_dataset = idx_to_backtrack_dataset[shard_key]  # selects which shard to iterate on
-
-            try:
-                for frame in self.make_frame(backtrack_dataset):
-                    if len(frames_buffer) == self.buffer_size:
-                        i = next(buffer_indices_generator)  # samples a element from the buffer
-                        yield frames_buffer[i]
-                        frames_buffer[i] = frame
-                    else:
-                        frames_buffer.append(frame)
-                    break  # random shard sampled, switch shard
-            except (
-                RuntimeError,
-                StopIteration,
-            ):  # NOTE: StopIteration inside a generator throws a RuntimeError since python 3.7
-                del idx_to_backtrack_dataset[shard_key]  # Remove exhausted shard, onto another shard
-
-        # Once shards are all exhausted, shuffle the buffer and yield the remaining frames
-        rng.shuffle(frames_buffer)
-        yield from frames_buffer
-
-    def _get_window_steps(
-        self, delta_timestamps: dict[str, list[float]] | None = None, dynamic_bounds: bool = False
-    ) -> tuple[int, int]:
-        if delta_timestamps is None:
-            return 1, 1
-
-        if not dynamic_bounds:
-            # Fix the windows
-            lookback = LOOKBACK_BACKTRACKTABLE
-            lookahead = LOOKAHEAD_BACKTRACKTABLE
+        # `datasets` reshuffles (and re-permutes shard order) per epoch from (seed, epoch);
+        # DataLoader workers each advance their own copy's counter in lockstep. The in-flight
+        # epoch is tracked separately so a mid-iteration state_dict() records the epoch the
+        # stream position actually belongs to. Only advance when shuffling: after reshard() the
+        # stream has one shard per episode, and set_epoch(n>0) re-permutes shard order even without
+        # a shuffle op, so an unshuffled stream must pin epoch 0 to repeat the same order each pass.
+        if self.shuffle:
+            self._in_flight_epoch = self._epoch
+            self._epoch += 1
         else:
-            # Dynamically adjust the windows based on the given delta_timesteps
-            all_timestamps = sum(delta_timestamps.values(), [])
-            lookback = min(all_timestamps) * self.fps
-            lookahead = max(all_timestamps) * self.fps
+            self._in_flight_epoch = 0
+        self._pipeline.set_epoch(self._in_flight_epoch)
+        self.video_decoder_cache = self._make_video_decoder_cache()
 
-            # When lookback is >=0 it means no negative timesteps have been provided
-            lookback = 0 if lookback >= 0 else (lookback * -1)
+        iterator = iter(self._pipeline)
+        while True:
+            try:
+                row = next(iterator)
+            except StopIteration:
+                return
+            yield self._finalize_sample(row)
 
-        return lookback, lookahead
+    def _finalize_sample(self, row: dict) -> dict:
+        """Torch conversion + video decode (decode-on-exit) + transforms + task for one frame."""
+        window_keys = self._tabular_window_keys()
+        pad_masks = {f"{key}_is_pad": torch.BoolTensor(row.pop(f"{key}_is_pad")) for key in window_keys}
+        item = item_to_torch(row)
+        item.update(pad_masks)
 
-    def _make_backtrackable_dataset(self, dataset: datasets.IterableDataset) -> Backtrackable:
-        lookback, lookahead = self._get_window_steps(self.delta_timestamps)
-        return Backtrackable(dataset, history=lookback, lookahead=lookahead)
+        if len(self.meta.video_keys) > 0:
+            ep_idx = int(item["episode_index"])
+            current_ts = float(item["timestamp"])
+            # Per-camera episode-local bounds [0, duration]: out-of-episode deltas pad instead of
+            # decoding against a neighbouring episode sharing the same video file.
+            episode_boundaries_ts = {
+                key: (
+                    0.0,
+                    self.meta.episodes[ep_idx][f"videos/{key}/to_timestamp"]
+                    - self.meta.episodes[ep_idx][f"videos/{key}/from_timestamp"],
+                )
+                for key in self.meta.video_keys
+            }
+            original_timestamps = self._make_timestamps_from_indices(current_ts, self.delta_indices)
+            query_timestamps = self._get_query_timestamps(
+                current_ts, self.delta_indices, episode_boundaries_ts
+            )
+            video_frames = self._query_videos(query_timestamps, ep_idx)
+
+            if self.image_transforms is not None:
+                for cam in self.meta.camera_keys:
+                    video_frames[cam] = self.image_transforms(video_frames[cam])
+
+            item.update(video_frames)
+            if self.delta_indices is not None:
+                item.update(
+                    self._get_video_frame_padding_mask(video_frames, query_timestamps, original_timestamps)
+                )
+
+        item["task"] = self.meta.tasks.iloc[int(item["task_index"])].name
+        return item
+
+    def set_epoch(self, epoch: int) -> None:
+        """Set the epoch the next ``__iter__`` will use (reshuffles the native pipeline)."""
+        self._epoch = epoch
+
+    def state_dict(self) -> dict:
+        """Native `datasets` stream state. Exact contract with ``num_workers=0``; with DataLoader
+        workers use ``torchdata.stateful_dataloader.StatefulDataLoader`` (it checkpoints each
+        worker's copy through this protocol). Samples in the shuffle buffers are skipped on
+        resume (never repeated), bounded by the pool + frame buffer sizes.
+        """
+        return {"pipeline": self._pipeline.state_dict(), "epoch": self._in_flight_epoch}
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        # Resume continues inside the recorded epoch: the next __iter__ replays that epoch's
+        # shuffle order from the restored stream position, then advances normally.
+        self._epoch = int(state_dict.get("epoch", 0))
+        self._pipeline.load_state_dict(state_dict["pipeline"])
 
     def _make_timestamps_from_indices(
         self, start_ts: float, indices: dict[str, list[int]] | None = None
     ) -> dict[str, list[float]]:
         if indices is not None:
             return {
-                key: (
-                    start_ts + torch.tensor(indices[key]) / self.fps
-                ).tolist()  # NOTE: why not delta_timestamps directly?
+                key: (start_ts + torch.tensor(indices[key]) / self.fps).tolist()
                 for key in self.delta_timestamps
             }
         else:
@@ -463,65 +594,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
 
         return padding_mask
 
-    def make_frame(self, dataset_iterator: Backtrackable) -> Generator:
-        """Makes a frame starting from a dataset iterator"""
-        item = next(dataset_iterator)
-        item = item_to_torch(item)
-
-        updates = []  # list of "updates" to apply to the item retrieved from hf_dataset (w/o camera features)
-
-        # Get episode index from the item
-        ep_idx = item["episode_index"]
-
-        # "timestamp" restarts from 0 for each episode, whereas we need a global timestep within the single .mp4 file (given by index/fps)
-        current_ts = item["index"] / self.fps
-
-        episode_boundaries_ts = {
-            key: (
-                self.meta.episodes[ep_idx][f"videos/{key}/from_timestamp"],
-                self.meta.episodes[ep_idx][f"videos/{key}/to_timestamp"],
-            )
-            for key in self.meta.video_keys
-        }
-
-        # Apply delta querying logic if necessary
-        if self.delta_indices is not None:
-            query_result, padding = self._get_delta_frames(dataset_iterator, item)
-            updates.append(query_result)
-            updates.append(padding)
-
-        # Load video frames, when needed
-        if len(self.meta.video_keys) > 0:
-            original_timestamps = self._make_timestamps_from_indices(current_ts, self.delta_indices)
-
-            # Some timestamps might not result available considering the episode's boundaries
-            query_timestamps = self._get_query_timestamps(
-                current_ts, self.delta_indices, episode_boundaries_ts
-            )
-            video_frames = self._query_videos(query_timestamps, ep_idx)
-
-            if self.image_transforms is not None:
-                image_keys = self.meta.camera_keys
-                for cam in image_keys:
-                    video_frames[cam] = self.image_transforms(video_frames[cam])
-
-            updates.append(video_frames)
-
-            if self.delta_indices is not None:
-                # We always return the same number of frames. Unavailable frames are padded.
-                padding_mask = self._get_video_frame_padding_mask(
-                    video_frames, query_timestamps, original_timestamps
-                )
-                updates.append(padding_mask)
-
-        result = item.copy()
-        for update in updates:
-            result.update(update)
-
-        result["task"] = self.meta.tasks.iloc[item["task_index"]].name
-
-        yield result
-
     def _get_query_timestamps(
         self,
         current_ts: float,
@@ -552,11 +624,20 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
 
         item = {}
         for video_key, query_ts in query_timestamps.items():
-            root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root
-            video_path = f"{root}/{self.meta.get_video_file_path(ep_idx, video_key)}"
+            # query_ts is episode-local; shift to the absolute in-file timeline by the episode's offset.
+            from_timestamp = self.meta.episodes[ep_idx][f"videos/{video_key}/from_timestamp"]
+            shifted_query_ts = [from_timestamp + ts for ts in query_ts]
+            rel_path = str(self.meta.get_video_file_path(ep_idx, video_key))
+            if self.data_files_root is not None:
+                root = self.data_files_root
+            elif self.streaming and not self.streaming_from_local:
+                root = self.meta.url_root
+            else:
+                root = self.root
+            video_path = f"{root}/{rel_path}"
             frames = decode_video_frames_torchcodec(
                 video_path,
-                query_ts,
+                shifted_query_ts,
                 self.tolerance_s,
                 decoder_cache=self.video_decoder_cache,
                 return_uint8=self._return_uint8,
@@ -566,116 +647,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
 
         return item
 
-    def _get_delta_frames(self, dataset_iterator: Backtrackable, current_item: dict):
-        # TODO(fracapuano): Modularize this function, refactor the code
-        """Get frames with delta offsets using the backtrackable iterator.
-
-        Args:
-            current_item (dict): Current item from the iterator.
-            ep_idx (int): Episode index.
-
-        Returns:
-            tuple: (query_result, padding) - frames at delta offsets and padding info.
-        """
-        current_episode_idx = current_item["episode_index"]
-
-        # Prepare results
-        query_result = {}
-        padding = {}
-
-        for key, delta_indices in self.delta_indices.items():
-            if key in self.meta.video_keys:
-                continue  # visual frames are decoded separately
-
-            target_frames = []
-            is_pad = []
-
-            # Create a results dictionary to store frames in processing order, then reconstruct original order for stacking
-            delta_results = {}
-
-            # Separate and sort deltas by difficulty (easier operations first)
-            negative_deltas = sorted([d for d in delta_indices if d < 0], reverse=True)  # [-1, -2, -3, ...]
-            positive_deltas = sorted([d for d in delta_indices if d > 0])  # [1, 2, 3, ...]
-            zero_deltas = [d for d in delta_indices if d == 0]
-
-            # Process zero deltas (current frame)
-            for delta in zero_deltas:
-                delta_results[delta] = (
-                    current_item[key],
-                    False,
-                )
-
-            # Process negative deltas in order of increasing difficulty
-            lookback_failed = False
-
-            last_successful_frame = current_item[key]
-
-            for delta in negative_deltas:
-                if lookback_failed:
-                    delta_results[delta] = (last_successful_frame, True)
-                    continue
-
-                try:
-                    steps_back = abs(delta)
-                    if dataset_iterator.can_peek_back(steps_back):
-                        past_item = dataset_iterator.peek_back(steps_back)
-                        past_item = item_to_torch(past_item)
-
-                        if past_item["episode_index"] == current_episode_idx:
-                            delta_results[delta] = (past_item[key], False)
-                            last_successful_frame = past_item[key]
-
-                        else:
-                            raise LookBackError("Retrieved frame is from different episode!")
-                    else:
-                        raise LookBackError("Cannot go back further than the history buffer!")
-
-                except LookBackError:
-                    delta_results[delta] = (last_successful_frame, True)
-                    lookback_failed = True  # All subsequent negative deltas will also fail
-
-            # Process positive deltas in order of increasing difficulty
-            lookahead_failed = False
-            last_successful_frame = current_item[key]
-
-            for delta in positive_deltas:
-                if lookahead_failed:
-                    delta_results[delta] = (last_successful_frame, True)
-                    continue
-
-                try:
-                    if dataset_iterator.can_peek_ahead(delta):
-                        future_item = dataset_iterator.peek_ahead(delta)
-                        future_item = item_to_torch(future_item)
-
-                        if future_item["episode_index"] == current_episode_idx:
-                            delta_results[delta] = (future_item[key], False)
-                            last_successful_frame = future_item[key]
-
-                        else:
-                            raise LookAheadError("Retrieved frame is from different episode!")
-                    else:
-                        raise LookAheadError("Cannot go ahead further than the lookahead buffer!")
-
-                except LookAheadError:
-                    delta_results[delta] = (last_successful_frame, True)
-                    lookahead_failed = True  # All subsequent positive deltas will also fail
-
-            # Reconstruct original order for stacking
-            for delta in delta_indices:
-                frame, is_padded = delta_results[delta]
-
-                # add batch dimension for stacking
-                target_frames.append(frame)  # frame.unsqueeze(0))
-                is_pad.append(is_padded)
-
-            # Stack frames and add to results
-            if target_frames:
-                query_result[key] = torch.stack(target_frames)
-                padding[f"{key}_is_pad"] = torch.BoolTensor(is_pad)
-
-        return query_result, padding
-
     def _validate_delta_timestamp_keys(self, delta_timestamps: dict[list[float]]) -> None:
         """
         Validate that all keys in delta_timestamps correspond to actual features in the dataset.

src/lerobot/datasets/video_utils.py

@@ -273,7 +273,11 @@ class VideoDecoderCache:
                 self._cache.move_to_end(video_path)
                 return entry[0]
 
-            file_handle = fsspec.open(video_path).__enter__()
+            # Bound per-handle buffering: with many decoders kept open at once (one per camera per active
+            # shard, across all workers), the default fsspec read cache balloons RAM on remote backends
+            # like hf:// buckets. A small readahead cache caps each handle's footprint without hurting the
+            # mostly-sequential reads torchcodec issues.
+            file_handle = fsspec.open(video_path, cache_type="readahead", block_size=2**20).__enter__()
             try:
                 decoder = VideoDecoder(file_handle, seek_mode="approximate")
             except Exception:
@@ -481,10 +485,8 @@ 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, optionally trimming it to ``[start_time_s, end_time_s)``.
+    """Re-encode a video file using the given encoder configuration.
 
     Args:
         input_video_path: Existing video file to read.
@@ -493,17 +495,10 @@ 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:
@@ -535,10 +530,6 @@ 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",
@@ -552,14 +543,7 @@ 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/pretrained.py

@@ -29,7 +29,6 @@ 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
@@ -39,67 +38,6 @@ 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
 
@@ -290,7 +228,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=cfg
+                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
             )
             card.save(str(saved_path / "README.md"))
 
@@ -308,20 +246,9 @@ 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,
-        cfg: TrainPipelineConfig | None = None,
+        self, dataset_repo_id: str, model_type: str, license: str | None, tags: list[str] | None
     ) -> ModelCard:
-        base_model_mapping = {
-            "smolvla": "lerobot/smolvla_base",
-            "pi0": "lerobot/pi0_base",
-            "pi05": "lerobot/pi05_base",
-            "pi0_fast": "lerobot/pi0fast-base",
-            "xvla": "lerobot/xvla-base",
-        }
+        base_model = "lerobot/smolvla_base" if model_type == "smolvla" else None  # Set a base model
 
         card_data = ModelCardData(
             license=license or "apache-2.0",
@@ -330,20 +257,13 @@ 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_mapping.get(model_type),
+            base_model=base_model,
         )
 
-        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, **context)
+        card = ModelCard.from_template(card_data, template_str=template_card)
         card.validate()
         return card
 

src/lerobot/scripts/lerobot_annotate.py

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

@@ -94,14 +94,6 @@ 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 \
@@ -265,9 +257,6 @@ 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")
@@ -472,8 +461,6 @@ 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_eval.py

@@ -74,7 +74,6 @@ from tqdm import trange
 
 from lerobot.configs import parser
 from lerobot.configs.eval import EvalPipelineConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.envs import (
     check_env_attributes_and_types,
     close_envs,
@@ -85,7 +84,7 @@ from lerobot.envs import (
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.types import PolicyAction
-from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR, REWARD
+from lerobot.utils.constants import ACTION, DONE, OBS_STR, REWARD
 from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.io_utils import write_video
@@ -96,56 +95,6 @@ from lerobot.utils.utils import (
 )
 
 
-def _env_features_to_dataset_features(env_features: dict) -> dict:
-    """Convert EnvConfig.features (PolicyFeature objects) to the plain dict format for LeRobotDataset.create()."""
-    features = {}
-    for key, ft in env_features.items():
-        if ft.type.value == "visual":
-            features[key] = {
-                "dtype": "video",
-                "shape": tuple(ft.shape),
-                "names": ["channel", "height", "width"],
-            }
-        else:
-            features[key] = {"dtype": "float32", "shape": tuple(ft.shape), "names": None}
-    features["next.reward"] = {"dtype": "float32", "shape": (1,), "names": None}
-    features["next.success"] = {"dtype": "bool", "shape": (1,), "names": None}
-    features["next.done"] = {"dtype": "bool", "shape": (1,), "names": None}
-    return features
-
-
-def _build_raw_frame(
-    raw_obs: dict,
-    env_idx: int,
-    action: np.ndarray,
-    reward: float,
-    success: bool,
-    done: bool,
-    task: str,
-) -> dict:
-    """Build a dataset frame from raw env observations for one env index."""
-    frame: dict[str, Any] = {}
-    if "pixels" in raw_obs:
-        if isinstance(raw_obs["pixels"], dict):
-            for cam_name, img in raw_obs["pixels"].items():
-                frame[f"{OBS_IMAGES}.{cam_name}"] = img[env_idx]
-        else:
-            frame[OBS_IMAGE] = raw_obs["pixels"][env_idx]
-    if "agent_pos" in raw_obs:
-        frame[OBS_STATE] = raw_obs["agent_pos"][env_idx]
-    for key, val in raw_obs.items():
-        if key in ("pixels", "agent_pos"):
-            continue
-        if isinstance(val, np.ndarray):
-            frame[f"{OBS_STR}.{key}"] = val[env_idx]
-    frame[ACTION] = action
-    frame["next.reward"] = np.float32(reward)
-    frame["next.success"] = success
-    frame["next.done"] = done
-    frame["task"] = task
-    return frame
-
-
 def rollout(
     env: gym.vector.VectorEnv,
     policy: PreTrainedPolicy,
@@ -156,7 +105,6 @@ def rollout(
     seeds: list[int] | None = None,
     return_observations: bool = False,
     render_callback: Callable[[gym.vector.VectorEnv], None] | None = None,
-    recording_dataset: Any | None = None,
 ) -> dict:
     """Run a batched policy rollout once through a batch of environments.
 
@@ -197,14 +145,6 @@ def rollout(
     if render_callback is not None:
         render_callback(env)
 
-    raw_observation = deepcopy(observation) if recording_dataset is not None else None
-    task_desc = ""
-    if recording_dataset is not None:
-        try:
-            task_desc = list(env.call("task_description"))[0]
-        except (AttributeError, NotImplementedError):
-            task_desc = ""
-
     all_observations = []
     all_actions = []
     all_rewards = []
@@ -277,25 +217,6 @@ def rollout(
         else:
             successes = [False] * env.num_envs
 
-        if recording_dataset is not None and raw_observation is not None:
-            prev_done = done.copy()
-            for env_idx in range(env.num_envs):
-                if prev_done[env_idx]:
-                    continue
-                frame = _build_raw_frame(
-                    raw_observation,
-                    env_idx,
-                    action_numpy[env_idx],
-                    reward[env_idx],
-                    successes[env_idx],
-                    bool(terminated[env_idx] | truncated[env_idx]),
-                    task_desc,
-                )
-                recording_dataset.add_frame(frame)
-                if terminated[env_idx] or truncated[env_idx]:
-                    recording_dataset.save_episode()
-            raw_observation = deepcopy(observation)
-
         # Keep track of which environments are done so far.
         # Mark the episode as done if we reach the maximum step limit.
         # This ensures that the rollout always terminates cleanly at `max_steps`,
@@ -352,7 +273,6 @@ def eval_policy(
     videos_dir: Path | None = None,
     return_episode_data: bool = False,
     start_seed: int | None = None,
-    recording_dataset: Any | None = None,
 ) -> dict:
     """
     Args:
@@ -441,7 +361,6 @@ def eval_policy(
             seeds=list(seeds) if seeds else None,
             return_observations=return_episode_data,
             render_callback=render_frame if max_episodes_rendered > 0 else None,
-            recording_dataset=recording_dataset,
         )
 
         # Figure out where in each rollout sequence the first done condition was encountered (results after
@@ -644,10 +563,6 @@ def eval_main(cfg: EvalPipelineConfig):
     # Create environment-specific preprocessor and postprocessor (e.g., for LIBERO environments)
     env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env, policy_cfg=cfg.policy)
 
-    recording_dir = Path(cfg.output_dir) / "recordings" if cfg.eval.recording else None
-    max_episodes_rendered = 0 if cfg.eval.recording else 10
-    videos_dir = None if cfg.eval.recording else Path(cfg.output_dir) / "videos"
-
     with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
         info = eval_policy_all(
             envs=envs,
@@ -657,13 +572,10 @@ def eval_main(cfg: EvalPipelineConfig):
             preprocessor=preprocessor,
             postprocessor=postprocessor,
             n_episodes=cfg.eval.n_episodes,
-            max_episodes_rendered=max_episodes_rendered,
-            videos_dir=videos_dir,
-            return_episode_data=False,
+            max_episodes_rendered=10,
+            videos_dir=Path(cfg.output_dir) / "videos",
             start_seed=cfg.seed,
             max_parallel_tasks=cfg.env.max_parallel_tasks,
-            recording_dir=recording_dir,
-            env_features=cfg.env.features if cfg.eval.recording else None,
         )
         print("Overall Aggregated Metrics:")
         print(info["overall"])
@@ -706,7 +618,6 @@ def eval_one(
     videos_dir: Path | None,
     return_episode_data: bool,
     start_seed: int | None,
-    recording_dataset: Any | None = None,
 ) -> TaskMetrics:
     """Evaluates one task_id of one suite using the provided vec env."""
 
@@ -724,7 +635,6 @@ def eval_one(
         videos_dir=task_videos_dir,
         return_episode_data=return_episode_data,
         start_seed=start_seed,
-        recording_dataset=recording_dataset,
     )
 
     per_episode = task_result["per_episode"]
@@ -751,8 +661,6 @@ def run_one(
     videos_dir: Path | None,
     return_episode_data: bool,
     start_seed: int | None,
-    recording_dir: Path | None = None,
-    env_features: dict | None = None,
 ):
     """
     Run eval_one for a single (task_group, task_id, env).
@@ -764,38 +672,21 @@ def run_one(
         task_videos_dir = videos_dir / f"{task_group}_{task_id}"
         task_videos_dir.mkdir(parents=True, exist_ok=True)
 
-    recording_dataset = None
-    if recording_dir is not None and env_features is not None:
-        task_recording_dir = recording_dir / f"{task_group}_{task_id}"
-        fps = env.unwrapped.metadata.get("render_fps", 30)
-        features = _env_features_to_dataset_features(env_features)
-        recording_dataset = LeRobotDataset.create(
-            repo_id=f"eval_{task_group}_{task_id}",
-            fps=fps,
-            features=features,
-            root=str(task_recording_dir),
-            use_videos=True,
-        )
-
-    try:
-        metrics = eval_one(
-            env,
-            policy=policy,
-            env_preprocessor=env_preprocessor,
-            env_postprocessor=env_postprocessor,
-            preprocessor=preprocessor,
-            postprocessor=postprocessor,
-            n_episodes=n_episodes,
-            max_episodes_rendered=max_episodes_rendered,
-            videos_dir=task_videos_dir,
-            return_episode_data=return_episode_data,
-            start_seed=start_seed,
-            recording_dataset=recording_dataset,
-        )
-    finally:
-        if recording_dataset is not None:
-            recording_dataset.finalize()
-
+    # Call the existing eval_one (assumed to return TaskMetrics-like dict)
+    metrics = eval_one(
+        env,
+        policy=policy,
+        env_preprocessor=env_preprocessor,
+        env_postprocessor=env_postprocessor,
+        preprocessor=preprocessor,
+        postprocessor=postprocessor,
+        n_episodes=n_episodes,
+        max_episodes_rendered=max_episodes_rendered,
+        videos_dir=task_videos_dir,
+        return_episode_data=return_episode_data,
+        start_seed=start_seed,
+    )
+    # ensure we always provide video_paths key to simplify accumulation
     if max_episodes_rendered > 0:
         metrics.setdefault("video_paths", [])
     return task_group, task_id, metrics
@@ -811,8 +702,6 @@ def eval_policy_all(
     n_episodes: int,
     *,
     max_episodes_rendered: int = 0,
-    recording_dir: Path | None = None,
-    env_features: dict | None = None,
     videos_dir: Path | None = None,
     return_episode_data: bool = False,
     start_seed: int | None = None,
@@ -872,8 +761,6 @@ def eval_policy_all(
         videos_dir=videos_dir,
         return_episode_data=return_episode_data,
         start_seed=start_seed,
-        recording_dir=recording_dir,
-        env_features=env_features,
     )
 
     if max_parallel_tasks <= 1:

src/lerobot/scripts/lerobot_train.py

@@ -36,8 +36,6 @@ 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,
@@ -45,7 +43,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, compute_sampler_state, make_dataset
+from lerobot.datasets import EpisodeAwareSampler, 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
@@ -101,9 +99,6 @@ 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
@@ -163,8 +158,6 @@ 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
 
 
@@ -239,17 +232,18 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
 
-    # 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 loading synchronization: each node's local main process downloads first to avoid
+    # race conditions (the global main process only exists on node 0, so gating on it would let
+    # all ranks of the other nodes download and build the Arrow cache concurrently).
+    if accelerator.is_local_main_process:
+        if is_main_process:
+            logging.info("Creating dataset")
         dataset = make_dataset(cfg)
 
     accelerator.wait_for_everyone()
 
-    # Other ranks read from the shared copy populated by the main process.
-    if not is_main_process:
+    # Now all other processes can safely load the dataset from the local cache
+    if not accelerator.is_local_main_process:
         dataset = make_dataset(cfg)
 
     # Create environment used for evaluating checkpoints during training on simulation data.
@@ -393,47 +387,22 @@ 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 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.
+    if hasattr(active_cfg, "drop_n_last_frames") and not cfg.dataset.streaming:
         shuffle = False
+        # A dedicated generator (rather than the global torch RNG) lets accelerator.prepare
+        # synchronize the shuffle permutation across ranks, keeping batch shards disjoint even
+        # when ranks consume the global RNG asymmetrically (e.g. eval on the main process only).
+        sampler_generator = torch.Generator()
+        if cfg.seed is not None:
+            sampler_generator.manual_seed(cfg.seed)
         sampler = EpisodeAwareSampler(
             dataset.meta.episodes["dataset_from_index"],
             dataset.meta.episodes["dataset_to_index"],
             episode_indices_to_use=dataset.episodes,
-            drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
+            drop_n_last_frames=active_cfg.drop_n_last_frames,
             shuffle=True,
-            seed=cfg.seed if cfg.seed is not None else 0,
+            generator=sampler_generator,
         )
-        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
@@ -457,30 +426,27 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     # Prepare everything with accelerator
     accelerator.wait_for_everyone()
-    policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
-        policy, optimizer, dataloader, lr_scheduler
-    )
+    if cfg.dataset.streaming:
+        # The streaming IterableDataset is already rank-disjoint via split_dataset_by_node, so we must
+        # NOT hand the dataloader to accelerate: its IterableDatasetShard would keep only every
+        # world_size-th batch of each rank's already-disjoint stream (silently training on 1/N of the
+        # data while decoding all of it). Batches are moved to the device manually in the loop below.
+        policy, optimizer, lr_scheduler = accelerator.prepare(policy, optimizer, lr_scheduler)
+    else:
+        policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
+            policy, optimizer, dataloader, lr_scheduler
+        )
     dl_iter = cycle(dataloader)
 
     policy.train()
 
     train_metrics = {
-        # 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"),
+        "loss": AverageMeter("loss", ":.3f"),
         "grad_norm": AverageMeter("grdn", ":.3f"),
         "lr": AverageMeter("lr", ":0.1e"),
-        # 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"),
+        "update_s": AverageMeter("updt_s", ":.3f"),
+        "dataloading_s": AverageMeter("data_s", ":.3f"),
     }
-    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
@@ -509,6 +475,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     for _ in range(step, cfg.steps):
         start_time = time.perf_counter()
         batch = next(dl_iter)
+        if cfg.dataset.streaming:
+            # The streaming dataloader is not accelerate-prepared (see above), so move to device here.
+            batch = {k: (v.to(device, non_blocking=True) if torch.is_tensor(v) else v) for k, v in batch.items()}
         for cam_key in dataset.meta.camera_keys:
             if cam_key in batch and batch[cam_key].dtype == torch.uint8:
                 batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
@@ -532,29 +501,21 @@ 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
+        is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 and is_main_process
         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:
-            # 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)
+            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:
@@ -570,8 +531,6 @@ 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/templates/lerobot_modelcard_template.md

@@ -13,213 +13,77 @@
 [SmolVLA](https://huggingface.co/papers/2506.01844) is a compact, efficient vision-language-action model that achieves competitive performance at reduced computational costs and can be deployed on consumer-grade hardware.
 {% elif model_name == "act" %}
 [Action Chunking with Transformers (ACT)](https://huggingface.co/papers/2304.13705) is an imitation-learning method that predicts short action chunks instead of single steps. It learns from teleoperated data and often achieves high success rates.
+{% elif model_name == "tdmpc" %}
+[TD-MPC](https://huggingface.co/papers/2203.04955) combines model-free and model-based approaches to improve sample efficiency and performance in continuous control tasks by using a learned latent dynamics model and terminal value function.
 {% elif model_name == "diffusion" %}
 [Diffusion Policy](https://huggingface.co/papers/2303.04137) treats visuomotor control as a generative diffusion process, producing smooth, multi-step action trajectories that excel at contact-rich manipulation.
+{% elif model_name == "vqbet" %}
+[VQ-BET](https://huggingface.co/papers/2403.03181) combines vector-quantised action tokens with Behaviour Transformers to discretise control and achieve data-efficient imitation across diverse skills.
 {% elif model_name == "pi0" %}
-[π₀ (Pi0)](https://www.physicalintelligence.company/blog/pi0) is a general-purpose robot foundation model from Physical Intelligence: a generalist Vision-Language-Action policy that understands visual inputs, interprets natural language instructions, and controls a variety of different robots across diverse tasks. The LeRobot implementation is adapted from their open-source OpenPI repository.
+**π₀ (Pi0)**
+
+π₀ is a Vision-Language-Action model for general robot control, from Physical Intelligence. The LeRobot implementation is adapted from their open source OpenPI repository.
+
+**Model Overview**
+
+π₀ represents a breakthrough in robotics as the first general-purpose robot foundation model developed by Physical Intelligence. Unlike traditional robots that are narrow specialists programmed for repetitive motions, π₀ is designed to be a generalist policy that can understand visual inputs, interpret natural language instructions, and control a variety of different robots across diverse tasks.
+
+For more details, see the [Physical Intelligence π₀ blog post](https://www.physicalintelligence.company/blog/pi0).
 {% elif model_name == "pi05" %}
-[π₀.₅ (Pi05)](https://www.physicalintelligence.company/blog/pi05) is a Vision-Language-Action model from Physical Intelligence designed for open-world generalization: it evolves π₀ to generalize to entirely new environments and situations that were never seen during training. The LeRobot implementation is adapted from their open-source OpenPI repository.
-{% elif model_name == "molmoact2" %}
-[MolmoAct2](https://allenai.org/blog/molmoact2) is an open robotics foundation model from the Allen Institute for AI (Ai2) that maps camera images and language instructions to robot action chunks. The LeRobot implementation supports training and evaluation of the regular MolmoAct2 model.
-{% elif model_name == "vla_jepa" %}
-[VLA-JEPA](https://arxiv.org/abs/2602.10098) is a Vision-Language-Action model that combines a Qwen3-VL language backbone with a self-supervised video world model (V-JEPA2) and a flow-matching DiT action head.
+**π₀.₅ (Pi05) Policy**
+
+π₀.₅ is a Vision-Language-Action model with open-world generalization, from Physical Intelligence. The LeRobot implementation is adapted from their open source OpenPI repository.
+
+**Model Overview**
+
+π₀.₅ represents a significant evolution from π₀, developed by Physical Intelligence to address a big challenge in robotics: open-world generalization. While robots can perform impressive tasks in controlled environments, π₀.₅ is designed to generalize to entirely new environments and situations that were never seen during training.
+
+For more details, see the [Physical Intelligence π₀.₅ blog post](https://www.physicalintelligence.company/blog/pi05).
 {% elif model_name == "gaussian_actor" %}
 This is a Gaussian Actor policy (Gaussian policy with a tanh squash) — the policy-side component used by [Soft Actor-Critic (SAC)](https://huggingface.co/papers/1801.01290) and related maximum-entropy continuous-control algorithms.
-{% elif model_name == "pi0_fast" %}
-[π₀-FAST (Pi0-FAST)](https://www.physicalintelligence.company/research/fast) is a Vision-Language-Action model for general robot control, from Physical Intelligence. It models continuous robot actions with autoregressive next-token prediction using FAST (Frequency-space Action Sequence Tokenization), training up to 5x faster than diffusion-based π₀.
-{% elif model_name == "eo1" %}
-[EO-1](https://huggingface.co/papers/2508.21112) is a Vision-Language-Action model for general robot control. It pairs a Qwen2.5-VL backbone for vision-language understanding with a continuous flow-matching action head that denoises action chunks.
-{% elif model_name == "groot" %}
-[GR00T N1.5](https://github.com/NVIDIA/Isaac-GR00T) is an open, cross-embodiment foundation model from NVIDIA for generalized humanoid robot reasoning and skills. It takes language and images as input and uses a flow-matching action transformer to predict actions conditioned on vision, language, and proprioception.
-{% elif model_name == "multi_task_dit" %}
-[Multi-Task Diffusion Transformer (DiT)](https://huggingface.co/papers/2507.05331) extends Diffusion Policy with a large Diffusion Transformer and text + vision conditioning for multi-task robot learning. It supports both diffusion and flow-matching objectives and reaches high dexterity with only ~450M parameters.
-{% elif model_name == "wall_x" %}
-[WALL-OSS](https://huggingface.co/papers/2509.11766) is an open-source foundation model for embodied intelligence from XSquare Robot. Built on Qwen2.5-VL, it uses a tightly-coupled multimodal architecture with flow matching to unify semantic reasoning and high-frequency action generation for cross-embodiment control.
-{% elif model_name == "xvla" %}
-[X-VLA](https://huggingface.co/papers/2510.10274) is a soft-prompted, flow-matching Vision-Language-Action framework that treats each robot or hardware setup as a "task" encoded with a small set of learnable Soft Prompt embeddings, letting a single model reconcile diverse robot morphologies, sensors, and action spaces.
 {% else %}
-This is a **{{ model_name }}** policy trained with [LeRobot](https://github.com/huggingface/lerobot).
+_Model type not recognized — please update this template._
 {% endif %}
-{% set diagrams = {
-  "smolvla": "https://cdn-uploads.huggingface.co/production/uploads/640e21ef3c82bd463ee5a76d/aooU0a3DMtYmy_1IWMaIM.png",
-  "pi0": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-pi0%20(1).png",
-  "pi0_fast": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-pifast.png",
-  "eo1": "https://huggingface.co/datasets/HaomingSong/lerobot-documentation-images/resolve/main/lerobot/eo_pipeline.png",
-  "groot": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-groot-paper1%20(1).png",
-  "wall_x": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/walloss-lerobot-paper.png",
-  "xvla": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/xvla-architecture.png"
-} %}
-{% if diagrams.get(model_name) %}
-<p align="center">
-  <img src="{{ diagrams[model_name] }}" alt="{{ model_name }} architecture" width="85%"/>
-</p>
-{% endif %}
-
-<!-- A short demo is worth more than any description! Record a GIF/video of the policy
-running on your robot, upload it to this repo, and embed it here:
-<p align="center">
-  <img src="https://huggingface.co/<hf_user>/<policy_repo_id>/resolve/main/demo.gif" width="60%"/>
-</p>
--->
 
 This policy has been trained and pushed to the Hub using [LeRobot](https://github.com/huggingface/lerobot).
-{% set policy_docs = {
-  "act": "act",
-  "smolvla": "smolvla",
-  "pi0": "pi0",
-  "pi0_fast": "pi0fast",
-  "pi05": "pi05",
-  "molmoact2": "molmoact2",
-  "vla_jepa": "vla_jepa",
-  "eo1": "eo1",
-  "groot": "groot",
-  "xvla": "xvla",
-  "multi_task_dit": "multi_task_dit",
-  "wall_x": "walloss"
-} %}
-{% if policy_docs.get(model_name) %}Learn how to train and run it in the [LeRobot {{ model_name }} guide](https://huggingface.co/docs/lerobot/main/en/{{ policy_docs[model_name] }}), or browse the [full documentation](https://huggingface.co/docs/lerobot/index).
-{% else %}See the [full LeRobot documentation](https://huggingface.co/docs/lerobot/index).
-{% endif %}
+See the full documentation at [LeRobot Docs](https://huggingface.co/docs/lerobot/index).
+
+---
+
+## How to Get Started with the Model
+
+For a complete walkthrough, see the [training guide](https://huggingface.co/docs/lerobot/il_robots#train-a-policy).
+Below is the short version on how to train and run inference/eval:
+
+### Train from scratch
+
+```bash
+lerobot-train \
+  --dataset.repo_id=${HF_USER}/<dataset> \
+  --policy.type=act \
+  --output_dir=outputs/train/<desired_policy_repo_id> \
+  --job_name=lerobot_training \
+  --policy.device=cuda \
+  --policy.repo_id=${HF_USER}/<desired_policy_repo_id>
+  --wandb.enable=true
+```
+
+_Writes checkpoints to `outputs/train/<desired_policy_repo_id>/checkpoints/`._
+
+### Evaluate the policy/run inference
+
+```bash
+lerobot-record \
+  --robot.type=so100_follower \
+  --dataset.repo_id=<hf_user>/eval_<dataset> \
+  --policy.path=<hf_user>/<desired_policy_repo_id> \
+  --episodes=10
+```
+
+Prefix the dataset repo with **eval\_** and supply `--policy.path` pointing to a local or hub checkpoint.
 
 ---
 
 ## Model Details
 
 - **License:** {{ license | default("\[More Information Needed]", true) }}
-{% if base_model %}- **Fine-tuned from:** [{{ base_model }}](https://huggingface.co/{{ base_model }})
-{% endif %}{% if robot_type %}- **Robot type:** `{{ robot_type }}`
-{% endif %}{% if cameras %}- **Cameras:** {% for camera in cameras %}`{{ camera }}`{% if not loop.last %}, {% endif %}{% endfor %}
-{% endif %}
-{% if input_features or output_features %}
-## Inputs & Outputs
-
-The policy consumes these observation features and produces these action features.
-{% if input_features %}
-**Inputs**
-
-| Feature | Type | Shape |
-| --- | --- | --- |
-{% for name, feature in input_features.items() %}| `{{ name }}` | {{ feature.type.value }} | `{{ feature.shape }}` |
-{% endfor %}{% endif %}{% if output_features %}
-**Outputs**
-
-| Feature | Type | Shape |
-| --- | --- | --- |
-{% for name, feature in output_features.items() %}| `{{ name }}` | {{ feature.type.value }} | `{{ feature.shape }}` |
-{% endfor %}{% endif %}{% endif %}
-{% if dataset %}
-## Training Dataset
-
-- **Repository:** [{{ dataset.repo_id }}](https://huggingface.co/datasets/{{ dataset.repo_id }})
-- **Episodes:** {{ dataset.episodes }}
-- **Frames:** {{ dataset.frames }}
-- **Frame rate:** {{ dataset.fps }} FPS
-{% if dataset.tasks %}- **Task(s):** {% for task in dataset.tasks %}"{{ task }}"{% if not loop.last %}, {% endif %}{% endfor %}
-{% endif %}
-<a class="flex" href="https://huggingface.co/spaces/lerobot/visualize_dataset?path={{ dataset.repo_id }}">
-<img class="block dark:hidden" src="https://huggingface.co/datasets/huggingface/badges/resolve/main/visualize-this-dataset-xl.svg"/>
-<img class="hidden dark:block" src="https://huggingface.co/datasets/huggingface/badges/resolve/main/visualize-this-dataset-xl-dark.svg"/>
-</a>
-{% endif %}
-{% if training %}
-## Training Configuration
-
-| Setting | Value |
-| --- | --- |
-| Training steps | {{ training.steps }} |
-| Batch size | {{ training.batch_size }} |
-{% if training.optimizer %}| Optimizer | {{ training.optimizer }} |
-{% endif %}{% if training.lr %}| Learning rate | {{ training.lr }} |
-{% endif %}{% if training.seed is not none %}| Seed | {{ training.seed }} |
-{% endif %}| LeRobot version | {{ training.lerobot_version }} |
-{% endif %}
----
-
-## How to Get Started with the Model
-
-New to LeRobot? These guides cover the full workflow:
-
-- **[Install LeRobot](https://huggingface.co/docs/lerobot/main/en/installation)** — set up the `lerobot` package.
-- **[Hardware setup](https://huggingface.co/docs/lerobot/main/en/hardware_guide)** — assemble, wire, and calibrate your robot and cameras.
-- **[Record data & train a policy](https://huggingface.co/docs/lerobot/en/il_robots)** — the end-to-end imitation-learning walkthrough.
-- **[CLI cheat-sheet](https://huggingface.co/docs/lerobot/main/en/cheat-sheet)** — quick reference for the `lerobot-*` commands.
-
-The short version to run and train this policy:
-
-### Run the policy on your robot
-
-```bash
-lerobot-rollout \
-  --strategy.type=base \
-  --robot.type={{ robot_type | default("<your_robot_type>", true) }} \
-  --robot.port=<your_robot_port> \
-  --robot.cameras="{ <camera_1>: {type: opencv, index_or_path: <index_or_path>, width: 640, height: 480, fps: 30}, <camera_2>: {type: opencv, index_or_path: <index_or_path>, width: 640, height: 480, fps: 30}}" \
-  --policy.path={{ policy_repo_id | default("<hf_user>/<policy_repo_id>", true) }} \
-  --task="{% if dataset and dataset.tasks %}{{ dataset.tasks[0] }}{% else %}<your_task_description>{% endif %}" \
-  --duration=60
-```
-
-Replace the remaining `<...>` placeholders with your own values: `--robot.port` and the camera names/indices are specific to your machine, and the camera names must match the observation keys this policy was trained on.
-
-When `--strategy.type=base` is used the script doesn't record the episodes. Skipping duration will make the policy run indefinitely. For more information look at [rollout documentation](https://huggingface.co/docs/lerobot/main/en/inference).
-
-{% if base_model %}### Train your own policy
-
-This policy type is usually fine-tuned from the pretrained base model [{{ base_model }}](https://huggingface.co/{{ base_model }}):
-
-```bash
-lerobot-train \
-  --dataset.repo_id=${HF_USER}/<dataset> \
-  --policy.path={{ base_model }} \
-  --output_dir=outputs/train/<policy_repo_id> \
-  --job_name=lerobot_training \
-  --policy.device=cuda \
-  --policy.repo_id=${HF_USER}/<policy_repo_id> \
-  --wandb.enable=true
-```
-{% else %}### Train your own policy
-
-```bash
-lerobot-train \
-  --dataset.repo_id=${HF_USER}/<dataset> \
-  --policy.type={{ model_name }} \
-  --output_dir=outputs/train/<policy_repo_id> \
-  --job_name=lerobot_training \
-  --policy.device=cuda \
-  --policy.repo_id=${HF_USER}/<policy_repo_id> \
-  --wandb.enable=true
-```
-{% endif %}
-_Writes checkpoints to `outputs/train/<policy_repo_id>/checkpoints/`._
-
----
-
-## Evaluation
-
-<!-- Report real-robot results here: run the policy several times per task and count the
-successes. Delete the "No evaluation results" line and fill in this table instead:
-
-| Task | Trials | Successes | Success rate |
-| ---- | ------ | --------- | ------------ |
-| pick the lego brick | 10 | 8 | 80% |
-
-Also worth noting: anything that affects difficulty (new object positions, lighting,
-distractors, a different robot of the same type, ...).
--->
-
-_No evaluation results have been provided for this policy yet._
-
----
-
-## Citation
-
-If you use this policy, please cite the method linked in the description above, along with LeRobot:
-
-```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},
-    title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch},
-    howpublished = "\url{https://github.com/huggingface/lerobot}",
-    year = {2024}
-}
-```

src/lerobot/utils/logging_utils.py

@@ -13,39 +13,21 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from collections import defaultdict
 from collections.abc import Callable
 from typing import Any
 
-import torch
-
 from .utils import format_big_number
 
-_VALID_REDUCTIONS = ("none", "max", "mean", "sum")
-
 
 class AverageMeter:
     """
     Computes and stores the average and current value
     Adapted from https://github.com/pytorch/examples/blob/main/imagenet/main.py
-
-    Args:
-        name: Display name of the metric.
-        fmt: Format string used when rendering the metric.
-        reduction: Cross-process reduction applied by :meth:`MetricsTracker.reduce_across_ranks`
-            before logging. One of ``"none"`` (per-rank value, default), ``"max"``, ``"mean"``,
-            or ``"sum"``. Use ``"max"`` for bottleneck-style metrics (e.g. dataloading or
-            update wall time) so multi-GPU runs report the slowest rank rather than rank 0.
     """
 
-    def __init__(self, name: str, fmt: str = ":f", reduction: str = "none"):
-        if reduction not in _VALID_REDUCTIONS:
-            raise ValueError(
-                f"Invalid reduction {reduction!r} for AverageMeter; expected one of {_VALID_REDUCTIONS}."
-            )
+    def __init__(self, name: str, fmt: str = ":f"):
         self.name = name
         self.fmt = fmt
-        self.reduction = reduction
         self.reset()
 
     def reset(self) -> None:
@@ -156,37 +138,6 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
 
-    def reduce_across_ranks(self) -> None:
-        """
-        Synchronises the running averages of every metric whose ``reduction`` is not ``"none"``
-        across all distributed processes (in-place).
-
-        This is a collective operation and MUST be invoked on every rank — typically just before
-        logging. With no accelerator or in single-process runs it is a no-op. Without it, metrics
-        reported by the main process only reflect rank 0; for bottleneck-style timings
-        (``dataloading_s``, ``update_s``, ...) that means the slowest worker's stall is invisible.
-        """
-        if self.accelerator is None or self.accelerator.num_processes <= 1:
-            return
-
-        buckets: dict[str, list[str]] = defaultdict(list)
-        for name, meter in self.metrics.items():
-            if meter.reduction != "none":
-                buckets[meter.reduction].append(name)
-        if not buckets:
-            return
-
-        device = self.accelerator.device
-        for reduction, names in buckets.items():
-            tensor = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device)
-            reduced = self.accelerator.reduce(tensor, reduction=reduction)
-            for name, value in zip(names, reduced.tolist(), strict=True):
-                meter = self.metrics[name]
-                # Preserve avg == sum / count so a later .update() on this meter accumulates
-                # against the cluster view, not the stale per-rank history.
-                meter.avg = value
-                meter.sum = value * meter.count
-
     def __str__(self) -> str:
         display_list = [
             f"step:{format_big_number(self.steps)}",

tests/annotations/_helpers.py

@@ -1,58 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Helpers shared across annotation-pipeline tests."""
-
-from __future__ import annotations
-
-import json
-from typing import Any
-
-from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
-
-
-def make_canned_responder(
-    responses_by_marker: dict[str, Any],
-    default: Any = None,
-) -> StubVlmClient:
-    """Return a stub that picks a response by inspecting the user prompt.
-
-    For each call the responder examines the last user-message text and
-    returns the response keyed by the first marker substring it contains.
-    Falls back to ``default`` if no marker matches.
-    """
-
-    def responder(messages: list[dict[str, Any]]) -> Any:
-        last_user_text = ""
-        for message in messages:
-            if message.get("role") != "user":
-                continue
-            content = message.get("content")
-            if isinstance(content, str):
-                last_user_text = content
-            elif isinstance(content, list):
-                for block in content:
-                    if isinstance(block, dict) and block.get("type") == "text":
-                        last_user_text = block.get("text", "")
-        for marker, response in responses_by_marker.items():
-            if marker in last_user_text:
-                return response
-        return default
-
-    return StubVlmClient(responder=responder)
-
-
-def encode_vqa_answer(payload: dict[str, Any]) -> str:
-    return json.dumps(payload, sort_keys=True)

tests/annotations/conftest.py

@@ -1,58 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Shared fixtures for annotation-pipeline tests.
-
-The on-disk dataset builder lives with the other dataset factories in
-``tests/fixtures/dataset_factories.py`` (:func:`build_annotation_dataset`);
-these fixtures only wire it into pytest.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-# ``build_annotation_dataset`` pulls in ``lerobot.datasets`` (HF ``datasets``
-# + ``pandas``, only in the ``dataset`` extra), so it's imported lazily inside
-# each fixture — this conftest stays importable without that extra. The test
-# modules ``pytest.importorskip("datasets")`` so they skip rather than error.
-
-
-@pytest.fixture
-def fixture_dataset_root(tmp_path: Path) -> Path:
-    """A tiny dataset with two episodes, 12 frames each at 10 fps."""
-    from tests.fixtures.dataset_factories import build_annotation_dataset
-
-    return build_annotation_dataset(
-        tmp_path / "ds",
-        episode_specs=[
-            (0, 12, "Could you tidy the kitchen please?"),
-            (1, 12, "Please clean up the kitchen"),
-        ],
-        fps=10,
-    )
-
-
-@pytest.fixture
-def single_episode_root(tmp_path: Path) -> Path:
-    from tests.fixtures.dataset_factories import build_annotation_dataset
-
-    return build_annotation_dataset(
-        tmp_path / "ds_one",
-        episode_specs=[(0, 30, "Pour water from the bottle into the cup.")],
-        fps=10,
-    )

tests/annotations/run_e2e_smoke.py

@@ -1,116 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Opt-in E2E smoke run for ``make annotation-e2e``.
-
-Builds the shared annotation fixture (:func:`build_annotation_dataset`),
-runs the full annotation pipeline against it with a stub VLM, and prints a
-short report. This is intentionally not a pytest test — it exercises the
-CLI plumbing — but it reuses the same on-disk dataset builder as the pytest
-fixtures so there is no duplicated fixture code.
-"""
-
-from __future__ import annotations
-
-import sys
-import tempfile
-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.modules import (
-    GeneralVqaModule,
-    InterjectionsAndSpeechModule,
-    PlanSubtasksMemoryModule,
-)
-from lerobot.annotations.steerable_pipeline.validator import StagingValidator
-from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
-from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
-from tests.fixtures.dataset_factories import build_annotation_dataset
-
-
-def _stub_responder(messages):
-    text = ""
-    for m in messages:
-        if m.get("role") == "user":
-            content = m.get("content")
-            if isinstance(content, list):
-                for block in content:
-                    if isinstance(block, dict) and block.get("type") == "text":
-                        text = block.get("text", "")
-            elif isinstance(content, str):
-                text = content
-    if "atomic subtasks" in text:
-        return {
-            "subtasks": [
-                {"text": "grasp the bottle", "start": 0.0, "end": 1.0},
-                {"text": "pour into the cup", "start": 1.0, "end": 2.0},
-                {"text": "place the bottle down", "start": 2.0, "end": 3.0},
-            ]
-        }
-    if "compressed semantic memory" in text:
-        return {"memory": "poured once"}
-    if "acknowledgement the robot" in text:
-        return {"text": "Sure."}
-    if "compact interjection" in text:
-        return {"interjection": "use less water", "speech": "Using less water."}
-    if "frame-grounded visual question" in text:
-        return {"question": "How many cups?", "answer": {"label": "cup", "count": 1}}
-    return None
-
-
-def main() -> int:
-    with tempfile.TemporaryDirectory() as tmp:
-        root = build_annotation_dataset(
-            Path(tmp) / "ds",
-            episode_specs=[(0, 30, "Pour water into the cup.")],
-            fps=10,
-        )
-        vlm = StubVlmClient(responder=_stub_responder)
-        cfg = AnnotationPipelineConfig()
-        executor = Executor(
-            config=cfg,
-            plan=PlanSubtasksMemoryModule(vlm=vlm, config=cfg.plan),
-            interjections=InterjectionsAndSpeechModule(vlm=vlm, config=cfg.interjections, seed=cfg.seed),
-            vqa=GeneralVqaModule(vlm=vlm, config=cfg.vqa, seed=cfg.seed),
-            writer=LanguageColumnsWriter(),
-            validator=StagingValidator(),
-        )
-        summary = executor.run(root)
-        print(f"phases={[(p.name, p.episodes_processed) for p in summary.phases]}")
-        print(f"validation: {summary.validation_report.summary()}")
-        print(f"shards rewritten: {len(summary.written_paths)}")
-
-        # Assert the interjection code path actually fired — otherwise a stale
-        # canned-VLM marker would silently produce zero interjections and this
-        # smoke run would still "pass" by only printing.
-        import pyarrow.parquet as pq  # noqa: PLC0415
-
-        events = [
-            r
-            for shard in summary.written_paths
-            for ev in pq.read_table(shard).column("language_events").to_pylist()
-            for r in ev
-        ]
-        n_interjections = sum(1 for r in events if r.get("style") == "interjection")
-        n_speech = sum(1 for r in events if r.get("style") is None and r.get("role") == "assistant")
-        print(f"interjections={n_interjections} speech_atoms={n_speech}")
-        assert n_interjections > 0, "no interjection rows produced — check the interjection prompt marker"
-        assert n_speech > 0, "no speech tool-call atoms produced — check the speech prompt marker"
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())

tests/annotations/test_frames.py

@@ -1,246 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Unit tests for :class:`VideoFrameProvider` method bindings.
-
-These were prompted by a real regression: ``video_for_episode`` was once
-indented one level too deep so it ended up nested *inside* a module-level
-helper (after that function's ``return`` statement) — silently dead code
-that meant production runs with ``use_video_url=False`` would
-``AttributeError`` on ``self.frame_provider.video_for_episode(...)``. The
-existing module tests didn't catch it because they exercise stub providers.
-
-The tests below assert on the class itself (not on an instance), so a
-future reindent regression flips them to red without needing a real
-LeRobot dataset on disk.
-"""
-
-from __future__ import annotations
-
-import shutil
-import subprocess
-from pathlib import Path
-
-import pytest
-import torch
-
-pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
-
-from lerobot.annotations.steerable_pipeline.frames import VideoFrameProvider  # noqa: E402
-
-
-class _FakeMeta:
-    """Minimal metadata stub exposing ``video_keys`` / ``camera_keys``."""
-
-    def __init__(self, video_keys: list[str], image_keys: list[str], video_path: Path | None = None) -> None:
-        self.video_keys = video_keys
-        self.camera_keys = [*video_keys, *image_keys]
-        self._video_path = video_path
-        self.episodes = {0: {f"videos/{key}/from_timestamp": 0.0 for key in video_keys}}
-
-    def get_video_file_path(self, episode_index: int, camera_key: str) -> Path:
-        return self._video_path
-
-
-def test_default_camera_key_skips_image_only_cameras(tmp_path: Path, monkeypatch) -> None:
-    """The default camera must be a *video* key — image-stored cameras have no
-    ``videos/<key>/from_timestamp`` and would KeyError in the clip/decode path.
-
-    Regression: a dataset whose first ``camera_keys`` entry was an image-stored
-    camera (e.g. ``observation.images.wrist``) crashed at clip extraction.
-    """
-    fake = _FakeMeta(
-        video_keys=["observation.images.robot0_agentview_right"],
-        image_keys=["observation.images.wrist"],
-    )
-    import lerobot.datasets.dataset_metadata as meta_mod
-
-    monkeypatch.setattr(meta_mod, "LeRobotDatasetMetadata", lambda *a, **k: fake, raising=True)
-    provider = VideoFrameProvider(root=tmp_path)
-    assert provider.camera_key == "observation.images.robot0_agentview_right"
-    assert "observation.images.wrist" not in provider.camera_keys
-
-
-def test_video_for_episode_is_a_method_of_videoframeprovider():
-    """``video_for_episode`` must be a bound method, not nested dead code."""
-    assert callable(getattr(VideoFrameProvider, "video_for_episode", None))
-
-
-def test_episode_clip_path_is_a_method_of_videoframeprovider():
-    """``episode_clip_path`` is now a method (was a free function reaching
-    into ``provider._meta`` from outside the class)."""
-    assert callable(getattr(VideoFrameProvider, "episode_clip_path", None))
-
-
-def test_videoframeprovider_has_a_lock_for_concurrent_use():
-    """A ``ThreadPoolExecutor`` runs the plan / interjections / vqa phases
-    concurrently; the cache + warn-flag accesses must be guarded.
-    """
-    import threading
-
-    # Fresh-instance check via a minimal fake to avoid touching the hub.
-    # The lock is declared with ``init=False`` and has a default factory,
-    # so a constructed instance must own a real ``threading.Lock``.
-    lock_field = next(
-        (f for f in VideoFrameProvider.__dataclass_fields__.values() if f.name == "_lock"),
-        None,
-    )
-    assert lock_field is not None
-    assert lock_field.default_factory is threading.Lock
-
-
-@pytest.fixture
-def sample_video(tmp_path: Path) -> Path:
-    """A 3 s 10 fps test-pattern mp4, written with ffmpeg."""
-    if shutil.which("ffmpeg") is None:
-        pytest.skip("ffmpeg not available")
-    out = tmp_path / "sample.mp4"
-    subprocess.run(
-        [
-            "ffmpeg",
-            "-y",
-            "-f",
-            "lavfi",
-            "-i",
-            "testsrc=duration=3:size=160x120:rate=10",
-            "-pix_fmt",
-            "yuv420p",
-            str(out),
-        ],
-        check=True,
-        capture_output=True,
-    )
-    return out
-
-
-def _provider_for_video(tmp_path: Path, video: Path, monkeypatch) -> VideoFrameProvider:
-    """A provider whose single camera resolves to ``video`` via fake metadata."""
-    fake = _FakeMeta(video_keys=["observation.images.cam"], image_keys=[], video_path=video)
-    import lerobot.datasets.dataset_metadata as meta_mod
-
-    monkeypatch.setattr(meta_mod, "LeRobotDatasetMetadata", lambda *a, **k: fake, raising=True)
-    return VideoFrameProvider(root=tmp_path, tolerance_s=0.2)
-
-
-def test_decode_returns_one_uint8_frame_per_timestamp(
-    sample_video: Path, tmp_path: Path, monkeypatch
-) -> None:
-    """``_decode`` routes through ``decode_video_frames`` (torchcodec when
-    available, PyAV otherwise) — no subprocess fallback.
-    """
-    provider = _provider_for_video(tmp_path, sample_video, monkeypatch)
-    timestamps = [0.0, 1.0, 2.5]
-    frames = provider._decode(0, timestamps, "observation.images.cam")
-
-    assert len(frames) == len(timestamps)
-    for frame in frames:
-        assert isinstance(frame, torch.Tensor)
-        assert frame.dtype == torch.uint8
-        assert frame.shape == (3, 120, 160)
-
-
-def test_frames_at_snaps_mid_frame_grid_to_real_frames(
-    sample_video: Path, tmp_path: Path, monkeypatch
-) -> None:
-    """Uniform sampling grids land mid-frame; ``frames_at`` must snap them to
-    real frame timestamps before decoding.
-
-    Regression: ``decode_video_frames`` rejects queries farther than
-    ``tolerance_s`` (default 10 ms) from a decodable frame, so un-snapped
-    mid-frame queries raised ``FrameTimestampError`` wholesale and the plan
-    module silently lost its contact sheets for most episodes.
-    """
-    from types import SimpleNamespace
-
-    fake = _FakeMeta(video_keys=["observation.images.cam"], image_keys=[], video_path=sample_video)
-    import lerobot.datasets.dataset_metadata as meta_mod
-
-    monkeypatch.setattr(meta_mod, "LeRobotDatasetMetadata", lambda *a, **k: fake, raising=True)
-    provider = VideoFrameProvider(root=tmp_path)  # default 10 ms tolerance
-    # 10 fps fixture -> frames at 0.0, 0.1, ...; queries sit mid-frame.
-    record = SimpleNamespace(episode_index=0, frame_timestamps=[i / 10 for i in range(30)])
-
-    frames = provider.frames_at(record, [0.149, 1.234, 2.04], camera_key="observation.images.cam")
-
-    assert len(frames) == 3
-    for frame in frames:
-        assert isinstance(frame, torch.Tensor)
-        assert frame.shape == (3, 120, 160)
-
-
-def test_decode_returns_empty_list_on_missing_file(tmp_path: Path, monkeypatch) -> None:
-    """A missing video is a recoverable no-frames condition, never a crash."""
-    provider = _provider_for_video(tmp_path, tmp_path / "does_not_exist.mp4", monkeypatch)
-    assert provider._decode(0, [0.0], "observation.images.cam") == []
-
-
-def test_episode_clip_path_trims_via_reencode_video(tmp_path: Path, monkeypatch) -> None:
-    """Clip extraction delegates to ``video_utils.reencode_video`` with the
-    episode's ``[from_timestamp, to_timestamp)`` trim window — no subprocess.
-    """
-    from types import SimpleNamespace
-
-    import lerobot.annotations.steerable_pipeline.frames as frames_mod
-
-    src = tmp_path / "src.mp4"
-    src.write_bytes(b"src")
-    fake = _FakeMeta(video_keys=["observation.images.cam"], image_keys=[], video_path=src)
-    fake.episodes[0]["videos/observation.images.cam/from_timestamp"] = 1.5
-    fake.episodes[0]["videos/observation.images.cam/to_timestamp"] = 4.0
-    import lerobot.datasets.dataset_metadata as meta_mod
-
-    monkeypatch.setattr(meta_mod, "LeRobotDatasetMetadata", lambda *a, **k: fake, raising=True)
-
-    captured = {}
-
-    def fake_reencode(
-        input_video_path,
-        output_video_path,
-        camera_encoder=None,
-        overwrite=False,
-        start_time_s=None,
-        end_time_s=None,
-    ):
-        captured.update(
-            src=Path(input_video_path),
-            encoder=camera_encoder,
-            start_time_s=start_time_s,
-            end_time_s=end_time_s,
-        )
-        Path(output_video_path).write_bytes(b"clip")
-
-    monkeypatch.setattr(frames_mod, "reencode_video", fake_reencode, raising=True)
-    provider = VideoFrameProvider(root=tmp_path)
-    record = SimpleNamespace(episode_index=0, frame_timestamps=[0.0, 1.0])
-
-    out = provider.episode_clip_path(record, tmp_path / "clips")
-
-    assert out == tmp_path / "clips" / "ep_000000.mp4"
-    assert captured["src"] == src
-    assert captured["start_time_s"] == 1.5
-    assert captured["end_time_s"] == 4.0
-    # H.264 so the clip is decodable by vllm's libav build (sources are often AV1).
-    assert captured["encoder"].vcodec == "h264"
-
-
-def test_videoframeprovider_serializes_decodes_with_a_lock() -> None:
-    """torchcodec's cached per-file decoder is single-threaded; the provider
-    must own a dedicated lock that ``_decode`` holds around the decoder call.
-    """
-    import threading
-
-    lock_field = VideoFrameProvider.__dataclass_fields__.get("_decode_lock")
-    assert lock_field is not None
-    assert lock_field.default_factory is threading.Lock

tests/annotations/test_modules.py

@@ -1,390 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Module 1/2/3 unit tests with stubbed VLMs."""
-
-from __future__ import annotations
-
-import json
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any
-
-import PIL.Image
-import pytest
-
-# ``lerobot.annotations`` imports pull in ``lerobot.datasets`` (-> the HF
-# ``datasets`` library), which only ships under the ``dataset`` extra. Skip
-# this module in tiers without it instead of erroring at import.
-pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
-pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
-
-from lerobot.annotations.steerable_pipeline.config import (  # noqa: E402
-    InterjectionsConfig,
-    PlanConfig,
-    VqaConfig,
-)
-from lerobot.annotations.steerable_pipeline.modules import (  # noqa: E402
-    GeneralVqaModule,
-    InterjectionsAndSpeechModule,
-    PlanSubtasksMemoryModule,
-)
-from lerobot.annotations.steerable_pipeline.reader import iter_episodes  # noqa: E402
-from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging  # noqa: E402
-from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient  # noqa: E402
-
-from ._helpers import make_canned_responder  # noqa: E402
-
-
-@dataclass
-class _StubFrameProvider:
-    """Returns one sentinel object per requested timestamp."""
-
-    # A real (tiny) PIL image so the contact-sheet builder, which resizes and
-    # tiles frames, has something to draw. VQA still passes it through by
-    # identity via ``to_image_blocks``.
-    sentinel: Any = field(default_factory=lambda: PIL.Image.new("RGB", (32, 24)))
-    cameras: tuple[str, ...] = ("observation.images.top",)
-    calls: list[tuple[int, tuple[float, ...], str | None]] = field(default_factory=list)
-    video_calls: list[tuple[int, int, str | None]] = field(default_factory=list)
-
-    @property
-    def camera_keys(self) -> list[str]:
-        return list(self.cameras)
-
-    def frames_at(self, record, timestamps, camera_key=None):
-        self.calls.append((record.episode_index, tuple(timestamps), camera_key))
-        return [self.sentinel] * len(timestamps)
-
-    def video_for_episode(self, record, max_frames, camera_key=None):
-        self.video_calls.append((record.episode_index, max_frames, camera_key))
-        n = min(max_frames, len(record.frame_timestamps))
-        return [self.sentinel] * n
-
-
-def _spy_responder(captured: list[list[dict[str, Any]]], reply: Any):
-    def responder(messages):
-        captured.append(list(messages))
-        return reply
-
-    return StubVlmClient(responder=responder)
-
-
-def test_module1_plan_memory_subtask_smoke(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    vlm = make_canned_responder(
-        {
-            "atomic subtasks": {
-                "subtasks": [
-                    {"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.4},
-                    {"text": "wipe the counter from left to right", "start": 0.4, "end": 0.8},
-                    {"text": "place the sponge into the sink", "start": 0.8, "end": 1.1},
-                ]
-            },
-            "compressed semantic memory": {"memory": "wiped the counter once"},
-        },
-    )
-    module = PlanSubtasksMemoryModule(vlm=vlm, config=PlanConfig())
-    record = next(iter_episodes(fixture_dataset_root))
-    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
-    module.run_episode(record, staging)
-    rows = staging.read("plan")
-
-    styles = {r["style"] for r in rows}
-    assert {"subtask", "plan", "memory"}.issubset(styles)
-    # subtask timestamps must be exact frame timestamps
-    frame_set = set(record.frame_timestamps)
-    for row in rows:
-        assert row["timestamp"] in frame_set
-    # one plan row per subtask boundary; the first lands at t0 and each
-    # plan is the deterministic numbered list of still-todo subtasks
-    plan_rows = sorted((r for r in rows if r["style"] == "plan"), key=lambda r: r["timestamp"])
-    subtask_rows = [r for r in rows if r["style"] == "subtask"]
-    assert len(plan_rows) == len(subtask_rows)
-    assert plan_rows[0]["timestamp"] == record.frame_timestamps[0]
-    # the t0 plan enumerates all subtasks; later plans shrink
-    assert plan_rows[0]["content"].startswith("1. ")
-    assert len(plan_rows[0]["content"].splitlines()) == len(subtask_rows)
-    assert len(plan_rows[-1]["content"].splitlines()) == 1
-
-
-def test_module1_emit_memory_false_skips_memory_keeps_subtasks_and_plan(
-    fixture_dataset_root: Path, tmp_path: Path
-) -> None:
-    """``emit_memory=False`` drops ``memory`` rows (and their VLM calls) while
-    leaving subtask + plan generation intact — symmetric to ``emit_plan``."""
-    vlm = make_canned_responder(
-        {
-            "atomic subtasks": {
-                "subtasks": [
-                    {"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.4},
-                    {"text": "wipe the counter from left to right", "start": 0.4, "end": 0.8},
-                    {"text": "place the sponge into the sink", "start": 0.8, "end": 1.1},
-                ]
-            },
-            "compressed semantic memory": {"memory": "wiped the counter once"},
-        },
-    )
-    module = PlanSubtasksMemoryModule(vlm=vlm, config=PlanConfig(emit_memory=False))
-    record = next(iter_episodes(fixture_dataset_root))
-    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
-    module.run_episode(record, staging)
-    rows = staging.read("plan")
-
-    styles = {r["style"] for r in rows}
-    assert "memory" not in styles
-    assert {"subtask", "plan"}.issubset(styles)
-
-
-def test_module2_at_t0_emits_speech_only_no_interjection(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    vlm = make_canned_responder(
-        {"acknowledgement the robot": {"text": "Sure, on it."}},
-    )
-    module = InterjectionsAndSpeechModule(
-        vlm=vlm,
-        config=InterjectionsConfig(max_interjections_per_episode=0),
-    )
-    record = next(iter_episodes(fixture_dataset_root))
-    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
-    module.run_episode(record, staging)
-    rows = staging.read("interjections")
-    assert len(rows) == 1
-    only = rows[0]
-    assert only["role"] == "assistant"
-    assert only["style"] is None
-    assert only["content"] is None
-    assert only["timestamp"] == record.frame_timestamps[0]
-    assert only["tool_calls"][0]["function"]["name"] == "say"
-
-
-def test_module2_mid_episode_emits_paired_interjection_and_speech(
-    fixture_dataset_root: Path, tmp_path: Path
-) -> None:
-    """Module 2 anchors interjections on Module 1's subtask boundaries.
-
-    The executor runs Module 1 first, then Module 2 reads the subtask
-    rows back from the same staging tree (see
-    ``_mid_episode_interjections``). Reproduce that contract here by
-    seeding the staging with two subtask rows so a single ``0 → 1``
-    boundary exists for Module 2 to anchor on.
-    """
-    vlm = make_canned_responder(
-        {
-            "acknowledgement the robot": {"text": "OK."},
-            # Marker matches the distinctive line of
-            # ``interjections_interjection.txt`` ("Write ONE compact
-            # interjection ..."). Keep this in sync with that prompt's
-            # wording — the canned responder matches on substring.
-            "Write ONE compact interjection": {
-                "interjection": "now wipe the counter please",
-                "speech": "On it.",
-            },
-        },
-    )
-    module = InterjectionsAndSpeechModule(
-        vlm=vlm,
-        config=InterjectionsConfig(max_interjections_per_episode=1, interjection_min_t=0.2),
-        seed=7,
-    )
-    record = next(iter_episodes(fixture_dataset_root))
-    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
-    # Seed Module 1's subtask staging so Module 2 has a boundary to
-    # anchor on (it bails with zero rows when no spans exist — the
-    # production executor guarantees Module 1 ran first).
-    boundary_ts = float(record.frame_timestamps[len(record.frame_timestamps) // 2])
-    staging.write(
-        "plan",
-        [
-            {
-                "role": "assistant",
-                "content": "grasp the sponge",
-                "style": "subtask",
-                "timestamp": float(record.frame_timestamps[0]),
-                "tool_calls": None,
-            },
-            {
-                "role": "assistant",
-                "content": "wipe the counter",
-                "style": "subtask",
-                "timestamp": boundary_ts,
-                "tool_calls": None,
-            },
-        ],
-    )
-    module.run_episode(record, staging)
-    rows = staging.read("interjections")
-
-    interjections = [r for r in rows if r["style"] == "interjection"]
-    speeches = [r for r in rows if r["style"] is None and r["role"] == "assistant"]
-    assert len(interjections) == 1
-    assert len(speeches) >= 2  # initial t=0 + one paired with the interjection
-    inter_t = interjections[0]["timestamp"]
-    assert any(abs(s["timestamp"] - inter_t) < 1e-9 for s in speeches)
-
-
-def test_module3_vqa_unique_per_frame_and_camera(single_episode_root: Path, tmp_path: Path) -> None:
-    payload = {
-        "question": "How many cups?",
-        "answer": {"label": "cup", "count": 2, "note": "white & blue"},
-    }
-    vlm = make_canned_responder({"frame-grounded visual question": payload})
-    module = GeneralVqaModule(
-        vlm=vlm,
-        config=VqaConfig(vqa_emission_hz=1.0, K=3),
-        seed=1,
-        frame_provider=_StubFrameProvider(cameras=("observation.images.top", "observation.images.wrist")),
-    )
-    record = next(iter_episodes(single_episode_root))
-    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
-    module.run_episode(record, staging)
-    rows = staging.read("vqa")
-    # every vqa row must carry a camera tag and one of the configured cameras
-    for r in rows:
-        assert r["style"] == "vqa"
-        assert r.get("camera") in {"observation.images.top", "observation.images.wrist"}
-    # at most one (vqa, user) and one (vqa, assistant) per (timestamp, camera)
-    user_keys = [(r["timestamp"], r["camera"]) for r in rows if r["role"] == "user" and r["style"] == "vqa"]
-    assistant_keys = [
-        (r["timestamp"], r["camera"]) for r in rows if r["role"] == "assistant" and r["style"] == "vqa"
-    ]
-    assert len(user_keys) == len(set(user_keys))
-    assert len(assistant_keys) == len(set(assistant_keys))
-    # both cameras must be represented
-    assert {c for _, c in user_keys} == {"observation.images.top", "observation.images.wrist"}
-    # every emitted timestamp must be an exact source frame timestamp
-    frame_set = set(record.frame_timestamps)
-    for ts, _ in user_keys + assistant_keys:
-        assert ts in frame_set
-
-
-def test_module1_attaches_contact_sheets_to_subtask_prompt(
-    fixture_dataset_root: Path, tmp_path: Path
-) -> None:
-    """Module 1 sends timestamped contact-sheet image blocks (not a raw video block)."""
-    captured: list[list[dict[str, Any]]] = []
-    payload = {
-        "subtasks": [
-            {"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.5},
-            {"text": "wipe the counter", "start": 0.5, "end": 1.1},
-        ]
-    }
-    memory_payload = {"memory": "wiped once"}
-
-    def responder(messages):
-        captured.append(list(messages))
-        text = ""
-        for m in messages:
-            for block in m.get("content", []):
-                if isinstance(block, dict) and block.get("type") == "text":
-                    text = block.get("text", "")
-        if "compressed semantic memory" in text:
-            return memory_payload
-        return payload
-
-    provider = _StubFrameProvider()
-    module = PlanSubtasksMemoryModule(
-        vlm=StubVlmClient(responder=responder),
-        # Disable the rephrasings sub-prompt so the test's only video-bearing
-        # call is the subtask one — keeps the assertions below focused on
-        # ``_generate_subtasks`` rather than fighting the order of unrelated
-        # text-only Module-1 sub-prompts.
-        config=PlanConfig(frames_per_second=2.0, max_frames_per_prompt=60, n_task_rephrasings=0),
-        frame_provider=provider,
-    )
-    record = next(iter_episodes(fixture_dataset_root))
-    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
-    module.run_episode(record, staging)
-
-    # Find the call carrying the subtask prompt rather than blindly taking
-    # captured[0] — Module 1 issues several sub-prompts and their order is
-    # not part of the contract.
-    assert captured, "no VLM calls made"
-
-    def _prompt_text(messages):
-        for m in messages:
-            for block in m.get("content", []):
-                if isinstance(block, dict) and block.get("type") == "text":
-                    return block.get("text", "")
-        return ""
-
-    subtask_calls = [m for m in captured if "atomic subtasks" in _prompt_text(m)]
-    assert len(subtask_calls) == 1, "expected exactly one subtask-prompt VLM call"
-    content = subtask_calls[0][0]["content"]
-    video_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "video"]
-    image_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "image"]
-    text_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "text"]
-    assert video_blocks == [], "contact-sheet mode must not emit a raw video block"
-    assert len(image_blocks) >= 1, f"expected >=1 contact-sheet image block, got {content}"
-    assert all(isinstance(b["image"], PIL.Image.Image) for b in image_blocks)
-    assert len(text_blocks) == 1
-    # the prompt is prefixed with the contact-sheet reading instructions
-    assert text_blocks[0]["text"].startswith("CONTACT SHEETS")
-    # frames were decoded for this episode at episode-relative timestamps
-    assert provider.calls and provider.calls[0][0] == record.episode_index
-
-
-def test_module3_attaches_frame_image_block_to_prompt(single_episode_root: Path, tmp_path: Path) -> None:
-    """Each VQA prompt must carry a single image block at the emission frame."""
-    captured: list[list[dict[str, Any]]] = []
-    payload = {
-        "question": "How many cups?",
-        "answer": {"label": "cup", "count": 1},
-    }
-    provider = _StubFrameProvider()
-    module = GeneralVqaModule(
-        vlm=_spy_responder(captured, payload),
-        config=VqaConfig(vqa_emission_hz=1.0, K=1),
-        seed=0,
-        frame_provider=provider,
-    )
-    record = next(iter_episodes(single_episode_root))
-    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
-    module.run_episode(record, staging)
-
-    assert captured, "no VLM calls made"
-    for messages in captured:
-        content = messages[0]["content"]
-        image_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "image"]
-        text_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "text"]
-        assert len(image_blocks) == 1, f"expected 1 image block per VQA prompt, got {content}"
-        assert image_blocks[0]["image"] is provider.sentinel
-        assert len(text_blocks) == 1
-    # provider was called once per emission per camera with the exact emission timestamp
-    for ep_idx, ts_tuple, camera in provider.calls:
-        assert ep_idx == record.episode_index
-        assert len(ts_tuple) == 1
-        assert ts_tuple[0] in record.frame_timestamps
-        assert camera in provider.cameras
-
-
-def test_module3_assistant_content_is_valid_json(single_episode_root: Path, tmp_path: Path) -> None:
-    payload = {
-        "question": "Where is the cup?",
-        "answer": {"detections": [{"label": "cup", "bbox_format": "xyxy", "bbox": [10, 20, 50, 80]}]},
-    }
-    vlm = make_canned_responder({"frame-grounded visual question": payload})
-    module = GeneralVqaModule(
-        vlm=vlm,
-        config=VqaConfig(vqa_emission_hz=1.0, K=2),
-        seed=2,
-        frame_provider=_StubFrameProvider(),
-    )
-    record = next(iter_episodes(single_episode_root))
-    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
-    module.run_episode(record, staging)
-    rows = staging.read("vqa")
-    for row in rows:
-        if row["role"] == "assistant" and row["style"] == "vqa":
-            decoded = json.loads(row["content"])
-            assert "detections" in decoded

tests/annotations/test_pipeline_recipe_render.py

@@ -1,183 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""End-to-end smoke: pipeline output → canonical recipe rendering."""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-# ``pyarrow`` and the ``lerobot.datasets`` chain (-> the HF ``datasets``
-# library) only ship under the ``dataset`` extra. Skip this module in
-# tiers without it instead of erroring at import.
-pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
-pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
-
-import pyarrow.parquet as pq  # noqa: E402
-
-from lerobot.annotations.steerable_pipeline.config import (  # noqa: E402
-    AnnotationPipelineConfig,
-    InterjectionsConfig,
-    PlanConfig,
-    VqaConfig,
-)
-from lerobot.annotations.steerable_pipeline.executor import Executor  # noqa: E402
-from lerobot.annotations.steerable_pipeline.modules import (  # noqa: E402
-    GeneralVqaModule,
-    InterjectionsAndSpeechModule,
-    PlanSubtasksMemoryModule,
-)
-from lerobot.annotations.steerable_pipeline.validator import StagingValidator  # noqa: E402
-from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter  # noqa: E402
-from lerobot.configs.recipe import MessageTurn, TrainingRecipe  # noqa: E402
-from lerobot.datasets.language_render import render_sample  # noqa: E402
-
-from ._helpers import make_canned_responder  # noqa: E402
-
-
-def _build_style_blend_recipe() -> TrainingRecipe:
-    """Inline blend recipe that consumes every style this pipeline produces.
-
-    The language schema/DSL work used to ship
-    ``src/lerobot/configs/recipes/pi05_hirobot.yaml`` as a canonical
-    example, but that file was dropped during review. The contract this
-    test guards is "the recipe DSL can render non-empty messages from
-    pipeline output", which doesn't require a specific YAML — so we build
-    the equivalent blend in code.
-    """
-    return TrainingRecipe(
-        blend={
-            "low_level_execution": TrainingRecipe(
-                weight=0.35,
-                messages=[
-                    MessageTurn(
-                        role="user",
-                        content="${task}\nPlan: ${plan}\nMemory: ${memory}",
-                        stream="high_level",
-                    ),
-                    MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
-                ],
-            ),
-            "user_interjection_response": TrainingRecipe(
-                weight=0.16,
-                bindings={
-                    "speech": "emitted_at(t, role=assistant, tool_name=say)",
-                    "interjection": "emitted_at(t, style=interjection)",
-                },
-                messages=[
-                    MessageTurn(role="user", content="${task}", stream="high_level"),
-                    MessageTurn(
-                        role="user",
-                        content="${interjection}",
-                        stream="high_level",
-                        if_present="interjection",
-                    ),
-                    MessageTurn(
-                        role="assistant",
-                        content="${plan}",
-                        stream="high_level",
-                        target=True,
-                        if_present="plan",
-                        tool_calls_from="speech",
-                    ),
-                ],
-            ),
-        }
-    )
-
-
-def _build_executor() -> Executor:
-    vlm = make_canned_responder(
-        {
-            "atomic subtasks": {
-                "subtasks": [
-                    {"text": "grasp the bottle", "start": 0.0, "end": 0.5},
-                    {"text": "pour into the cup", "start": 0.5, "end": 1.0},
-                    {"text": "place the bottle down", "start": 1.0, "end": 1.5},
-                ]
-            },
-            "compressed semantic memory": {"memory": "poured once"},
-            "acknowledgement the robot": {"text": "Sure."},
-            "compact interjection": {
-                "interjection": "use less water",
-                "speech": "Using less water.",
-            },
-            "frame-grounded visual question": {
-                "question": "How many cups?",
-                "answer": {"label": "cup", "count": 1},
-            },
-        },
-    )
-    config = AnnotationPipelineConfig(
-        plan=PlanConfig(),
-        interjections=InterjectionsConfig(max_interjections_per_episode=1, interjection_min_t=0.5),
-        vqa=VqaConfig(vqa_emission_hz=1.0, K=2),
-    )
-    return Executor(
-        config=config,
-        plan=PlanSubtasksMemoryModule(vlm=vlm, config=config.plan),
-        interjections=InterjectionsAndSpeechModule(vlm=vlm, config=config.interjections, seed=config.seed),
-        vqa=GeneralVqaModule(vlm=vlm, config=config.vqa, seed=config.seed),
-        writer=LanguageColumnsWriter(),
-        validator=StagingValidator(),
-    )
-
-
-def test_canonical_recipe_renders_nonempty_from_pipeline_output(
-    single_episode_root: Path,
-) -> None:
-    executor = _build_executor()
-    summary = executor.run(single_episode_root)
-    # validator may emit warnings but no errors for the synthetic fixture
-    assert summary.validation_report.ok, summary.validation_report.summary()
-
-    table = pq.read_table(single_episode_root / "data" / "chunk-000" / "file-000.parquet")
-    persistent_lists = table.column("language_persistent").to_pylist()
-    events_lists = table.column("language_events").to_pylist()
-    timestamps = table.column("timestamp").to_pylist()
-
-    recipe = _build_style_blend_recipe()
-
-    rendered_any = False
-    for ts, persistent, events in zip(timestamps, persistent_lists, events_lists, strict=True):
-        result = render_sample(
-            recipe=recipe,
-            persistent=persistent,
-            events=events,
-            t=float(ts),
-            sample_idx=0,
-            dataset_ctx={"task": "Pour water from the bottle into the cup."},
-        )
-        if result is None:
-            continue
-        if result["messages"]:
-            rendered_any = True
-            assert result["target_message_indices"]
-            break
-    assert rendered_any, "recipe rendered no messages from pipeline output"
-
-    # Sanity: speech atom appears in events column intact
-    flat_events = [r for ev in events_lists for r in ev]
-    speech_rows = [r for r in flat_events if r.get("style") is None and r.get("role") == "assistant"]
-    assert speech_rows
-    say = speech_rows[0]["tool_calls"][0]
-    assert say["function"]["name"] == "say"
-    assert isinstance(say["function"]["arguments"]["text"], str)
-    # The pipeline does not write a ``tools`` column — the say schema lives
-    # as a constant (``SAY_TOOL_SCHEMA``) so the language row struct is the
-    # single source of truth for the v3.1 schema.
-    assert "tools" not in table.column_names

tests/annotations/test_validator.py

@@ -1,133 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Validator behavior tests."""
-
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import pytest
-
-# ``lerobot.annotations`` imports pull in ``lerobot.datasets`` (-> the HF
-# ``datasets`` library), which only ships under the ``dataset`` extra. Skip
-# this module in tiers without it instead of erroring at import.
-pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
-pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
-
-from lerobot.annotations.steerable_pipeline.reader import iter_episodes  # noqa: E402
-from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging  # noqa: E402
-from lerobot.annotations.steerable_pipeline.validator import StagingValidator  # noqa: E402
-from lerobot.annotations.steerable_pipeline.writer import speech_atom  # noqa: E402
-
-
-def _validate(root: Path, staging_dir: Path):
-    records = list(iter_episodes(root))
-    return StagingValidator().validate(records, staging_dir)
-
-
-def test_validator_catches_misaligned_timestamps(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    staging_dir = tmp_path / "stage"
-    EpisodeStaging(staging_dir, 0).write(
-        "vqa",
-        [
-            {
-                "role": "assistant",
-                "content": json.dumps({"label": "cup", "count": 2}, sort_keys=True),
-                "style": "vqa",
-                "timestamp": 9.999,  # not on any 10 fps frame
-                "tool_calls": None,
-            }
-        ],
-    )
-    report = _validate(fixture_dataset_root, staging_dir)
-    assert not report.ok
-    assert any("does not match any source frame timestamp" in e for e in report.errors)
-
-
-def test_validator_catches_orphan_speech(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    staging_dir = tmp_path / "stage"
-    EpisodeStaging(staging_dir, 0).write(
-        "interjections",
-        [
-            speech_atom(0.0, "Got it."),
-            # interjection at 0.3s with NO paired speech
-            {
-                "role": "user",
-                "content": "skip it",
-                "style": "interjection",
-                "timestamp": 0.3,
-                "tool_calls": None,
-            },
-        ],
-    )
-    report = _validate(fixture_dataset_root, staging_dir)
-    assert not report.ok
-    assert any("paired speech" in e for e in report.errors)
-
-
-def test_validator_catches_inconsistent_plan_memory(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    staging_dir = tmp_path / "stage"
-    EpisodeStaging(staging_dir, 0).write(
-        "plan",
-        [
-            {
-                "role": "assistant",
-                "content": "1. do x",
-                "style": "plan",
-                "timestamp": 0.0,
-                "tool_calls": None,
-            },
-            {
-                "role": "assistant",
-                "content": "do x",
-                "style": "subtask",
-                "timestamp": 0.0,
-                "tool_calls": None,
-            },
-        ],
-    )
-    EpisodeStaging(staging_dir, 0).write(
-        "interjections",
-        [
-            speech_atom(0.0, "Got it."),
-            speech_atom(0.4, "Replanning."),
-            {
-                "role": "user",
-                "content": "replan",
-                "style": "interjection",
-                "timestamp": 0.4,
-                "tool_calls": None,
-            },
-        ],
-    )
-    report = _validate(fixture_dataset_root, staging_dir)
-    # missing co-timestamped plan refresh at 0.4s → error
-    assert not report.ok
-    assert any("co-timestamped plan update" in e for e in report.errors)
-
-
-def test_validator_catches_wrong_column(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    staging_dir = tmp_path / "stage"
-    EpisodeStaging(staging_dir, 0).write(
-        "plan",
-        [
-            {"role": "user", "content": "where?", "style": "vqa", "timestamp": 0.0, "tool_calls": None},
-        ],
-    )
-    report = _validate(fixture_dataset_root, staging_dir)
-    assert not report.ok
-    assert any("plan emitted style 'vqa'" in e or "must be persistent" in e for e in report.errors)

tests/annotations/test_vlm_client.py

@@ -1,41 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Unit tests for ``vlm_client`` helpers."""
-
-from __future__ import annotations
-
-import pytest
-
-pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
-
-from lerobot.annotations.steerable_pipeline.vlm_client import _bind_serve_port  # noqa: E402
-
-
-def test_bind_serve_port_substitutes_placeholder() -> None:
-    # The {port} placeholder is replaced everywhere it appears, regardless of
-    # parallel vs single server — the bug was the single-server path passing
-    # it through unsubstituted.
-    cmd = "vllm serve M --max-model-len 32768 --port {port}"
-    assert _bind_serve_port(cmd, 8000) == "vllm serve M --max-model-len 32768 --port 8000"
-
-
-def test_bind_serve_port_appends_when_missing() -> None:
-    assert _bind_serve_port("vllm serve M", 8001) == "vllm serve M --port 8001"
-
-
-def test_bind_serve_port_leaves_explicit_port_untouched() -> None:
-    cmd = "vllm serve M --port 9000"
-    assert _bind_serve_port(cmd, 8000) == cmd

tests/annotations/test_writer.py

@@ -1,357 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Writer correctness tests."""
-
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import pytest
-
-# ``pyarrow`` and the ``lerobot.annotations`` -> ``lerobot.datasets`` chain
-# (-> the HF ``datasets`` library) only ship under the ``dataset`` extra.
-# Skip this module in tiers without it instead of erroring at import.
-pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
-pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
-
-import pyarrow.parquet as pq  # noqa: E402
-
-from lerobot.annotations.steerable_pipeline.reader import iter_episodes  # noqa: E402
-from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging  # noqa: E402
-from lerobot.annotations.steerable_pipeline.writer import (  # noqa: E402
-    LanguageColumnsWriter,
-    speech_atom,
-)
-
-
-def _stage_episode(
-    staging_dir: Path,
-    episode_index: int,
-    *,
-    plan: list[dict] | None = None,
-    interjections: list[dict] | None = None,
-    vqa: list[dict] | None = None,
-) -> None:
-    staging = EpisodeStaging(staging_dir, episode_index)
-    if plan is not None:
-        staging.write("plan", plan)
-    if interjections is not None:
-        staging.write("interjections", interjections)
-    if vqa is not None:
-        staging.write("vqa", vqa)
-
-
-def test_writer_persistence_identity(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    """Every frame in an episode has a byte-identical persistent list."""
-    staging_dir = tmp_path / "stage"
-    _stage_episode(
-        staging_dir,
-        0,
-        plan=[
-            {
-                "role": "assistant",
-                "content": "grasp the sponge",
-                "style": "subtask",
-                "timestamp": 0.0,
-                "tool_calls": None,
-            },
-            {
-                "role": "assistant",
-                "content": "1. wipe\n2. dry",
-                "style": "plan",
-                "timestamp": 0.0,
-                "tool_calls": None,
-            },
-            {
-                "role": "assistant",
-                "content": "wiped the counter",
-                "style": "memory",
-                "timestamp": 0.5,
-                "tool_calls": None,
-            },
-        ],
-    )
-    records = list(iter_episodes(fixture_dataset_root))
-    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
-
-    table = pq.read_table(fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet")
-    persistent = table.column("language_persistent").to_pylist()
-    first = persistent[0]
-    assert first  # non-empty
-    for row in persistent:
-        assert row == first, "persistent slice must be byte-identical across all frames"
-
-
-def test_writer_events_exact_timestamp(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    staging_dir = tmp_path / "stage"
-    _stage_episode(
-        staging_dir,
-        0,
-        interjections=[
-            speech_atom(0.0, "Got it."),
-            {
-                "role": "user",
-                "content": "skip the dishes",
-                "style": "interjection",
-                "timestamp": 0.5,
-                "tool_calls": None,
-            },
-            speech_atom(0.5, "Skipping the dishes."),
-        ],
-    )
-    records = list(iter_episodes(fixture_dataset_root))
-    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
-
-    table = pq.read_table(fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet")
-    timestamps = table.column("timestamp").to_pylist()
-    events = table.column("language_events").to_pylist()
-    for ts, ev in zip(timestamps, events, strict=True):
-        if abs(ts - 0.0) < 1e-9:
-            assert any(r["role"] == "assistant" and r.get("style") is None for r in ev), ev
-        elif abs(ts - 0.5) < 1e-9:
-            assert any(r.get("style") == "interjection" for r in ev), ev
-            assert any(r.get("style") is None for r in ev), ev
-        else:
-            assert ev == []
-
-
-def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    staging_dir = tmp_path / "stage"
-    _stage_episode(
-        staging_dir,
-        0,
-        plan=[
-            {
-                "role": "assistant",
-                "content": "do X",
-                "style": "subtask",
-                "timestamp": 0.0,
-                "tool_calls": None,
-            },
-            {
-                "role": "assistant",
-                "content": "1. do X",
-                "style": "plan",
-                "timestamp": 0.0,
-                "tool_calls": None,
-            },
-            {
-                "role": "assistant",
-                "content": "did X",
-                "style": "memory",
-                "timestamp": 0.3,
-                "tool_calls": None,
-            },
-        ],
-        interjections=[
-            speech_atom(0.0, "OK"),
-            {
-                "role": "user",
-                "content": "wait",
-                "style": "interjection",
-                "timestamp": 0.2,
-                "tool_calls": None,
-            },
-            speech_atom(0.2, "Waiting"),
-        ],
-        vqa=[
-            {
-                "role": "user",
-                "content": "where is the cup?",
-                "style": "vqa",
-                "timestamp": 0.4,
-                "camera": "observation.images.front",
-                "tool_calls": None,
-            },
-            {
-                "role": "assistant",
-                "content": json.dumps(
-                    {"detections": [{"label": "cup", "bbox_format": "xyxy", "bbox": [1, 2, 3, 4]}]},
-                    sort_keys=True,
-                ),
-                "style": "vqa",
-                "timestamp": 0.4,
-                "camera": "observation.images.front",
-                "tool_calls": None,
-            },
-        ],
-    )
-    records = list(iter_episodes(fixture_dataset_root))
-    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
-    table = pq.read_table(fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet")
-
-    persistent = table.column("language_persistent").to_pylist()[0]
-    persistent_styles = {r["style"] for r in persistent}
-    assert persistent_styles == {"subtask", "plan", "memory"}
-
-    all_events = [r for ev in table.column("language_events").to_pylist() for r in ev]
-    event_styles = {r.get("style") for r in all_events}
-    assert event_styles == {None, "interjection", "vqa"}
-
-
-def test_writer_drops_subtask_index_idempotent(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    staging_dir = tmp_path / "stage"
-    _stage_episode(
-        staging_dir,
-        0,
-        plan=[
-            {
-                "role": "assistant",
-                "content": "do X",
-                "style": "subtask",
-                "timestamp": 0.0,
-                "tool_calls": None,
-            },
-        ],
-    )
-    records = list(iter_episodes(fixture_dataset_root))
-    writer = LanguageColumnsWriter()
-    writer.write_all(records, staging_dir, fixture_dataset_root)
-
-    path = fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet"
-    table_a = pq.read_table(path)
-    assert "subtask_index" not in table_a.column_names
-    assert "language_persistent" in table_a.column_names
-    assert "language_events" in table_a.column_names
-    # The writer no longer emits a dataset-level ``tools`` column; the
-    # ``say`` tool schema lives as a code constant (``SAY_TOOL_SCHEMA``)
-    # so the parquet stays small and the pipeline doesn't extend the schema.
-    assert "tools" not in table_a.column_names
-
-    # second pass — must produce identical bytes for the language columns
-    records_again = list(iter_episodes(fixture_dataset_root))
-    writer.write_all(records_again, staging_dir, fixture_dataset_root)
-    table_b = pq.read_table(path)
-    assert (
-        table_a.column("language_persistent").to_pylist() == table_b.column("language_persistent").to_pylist()
-    )
-    assert table_a.column("language_events").to_pylist() == table_b.column("language_events").to_pylist()
-
-
-def test_writer_normalize_rejects_misrouted_persistent_style() -> None:
-    """``_normalize_persistent_row`` must reject any non-persistent style."""
-    from lerobot.annotations.steerable_pipeline.writer import _normalize_persistent_row
-
-    with pytest.raises(ValueError, match="non-persistent style"):
-        _normalize_persistent_row(
-            {"role": "assistant", "content": "oops", "style": "vqa", "timestamp": 0.0, "tool_calls": None}
-        )
-
-
-def test_writer_normalize_rejects_misrouted_event_style() -> None:
-    """``_normalize_event_row`` must reject any persistent style."""
-    from lerobot.annotations.steerable_pipeline.writer import _normalize_event_row
-
-    with pytest.raises(ValueError):
-        _normalize_event_row({"role": "assistant", "content": "oops", "style": "subtask", "tool_calls": None})
-
-
-def test_say_tool_schema_constant_is_well_formed() -> None:
-    """``SAY_TOOL_SCHEMA`` (and ``DEFAULT_TOOLS``) replace the parquet
-    ``tools`` column — chat-template consumers import them directly.
-    """
-    from lerobot.annotations.steerable_pipeline.writer import (
-        DEFAULT_TOOLS,
-        SAY_TOOL_SCHEMA,
-    )
-
-    assert DEFAULT_TOOLS == [SAY_TOOL_SCHEMA]
-    assert SAY_TOOL_SCHEMA["function"]["name"] == "say"
-    params = SAY_TOOL_SCHEMA["function"]["parameters"]
-    assert params["properties"]["text"]["type"] == "string"
-    assert params["required"] == ["text"]
-
-
-def test_writer_does_not_add_tools_column(fixture_dataset_root: Path, tmp_path: Path) -> None:
-    """Re-running on a parquet that already has a legacy ``tools`` column
-    must drop it cleanly so reruns converge to the v3.1 schema.
-    """
-    staging_dir = tmp_path / "stage"
-    _stage_episode(
-        staging_dir,
-        0,
-        plan=[
-            {"role": "assistant", "content": "x", "style": "subtask", "timestamp": 0.0, "tool_calls": None}
-        ],
-    )
-    records = list(iter_episodes(fixture_dataset_root))
-    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
-    table = pq.read_table(fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet")
-    assert "tools" not in table.column_names
-
-
-def test_annotation_metadata_sync_allows_non_streaming_load(
-    fixture_dataset_root: Path, tmp_path: Path
-) -> None:
-    """Annotated parquet columns must be declared in ``meta/info.json``.
-
-    ``LeRobotDataset`` loads non-streaming datasets by casting parquet
-    against metadata-derived HF features. If the annotation writer adds
-    language columns but metadata stays stale, that cast fails with a column
-    mismatch.
-    """
-    from lerobot.annotations.steerable_pipeline.executor import Executor
-    from lerobot.datasets.feature_utils import get_hf_features_from_features
-    from lerobot.datasets.io_utils import load_info, load_nested_dataset
-    from lerobot.datasets.language import LANGUAGE_EVENTS, LANGUAGE_PERSISTENT, language_feature_info
-
-    info_path = fixture_dataset_root / "meta" / "info.json"
-    info = json.loads(info_path.read_text())
-    info["features"] = {
-        "episode_index": {"dtype": "int64", "shape": (1,), "names": None},
-        "frame_index": {"dtype": "int64", "shape": (1,), "names": None},
-        "timestamp": {"dtype": "float32", "shape": (1,), "names": None},
-        "task_index": {"dtype": "int64", "shape": (1,), "names": None},
-    }
-    info_path.write_text(json.dumps(info, indent=2))
-
-    staging_dir = tmp_path / "stage"
-    _stage_episode(
-        staging_dir,
-        0,
-        plan=[
-            {"role": "assistant", "content": "do X", "style": "subtask", "timestamp": 0.0, "tool_calls": None}
-        ],
-    )
-    records = list(iter_episodes(fixture_dataset_root))
-    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
-
-    Executor._ensure_annotation_metadata_in_info(fixture_dataset_root)
-
-    synced = load_info(fixture_dataset_root)
-    for key, feature in language_feature_info().items():
-        assert synced["features"][key] == feature
-
-    hf_features = get_hf_features_from_features(synced["features"])
-    dataset = load_nested_dataset(fixture_dataset_root / "data", features=hf_features)
-
-    assert LANGUAGE_PERSISTENT in dataset.column_names
-    assert LANGUAGE_EVENTS in dataset.column_names
-    assert len(dataset) == 24
-
-
-def test_speech_atom_shape_matches_plan_spec() -> None:
-    atom = speech_atom(2.5, "I'm cleaning up!")
-    assert atom["role"] == "assistant"
-    assert atom["style"] is None
-    assert atom["content"] is None
-    assert atom["timestamp"] == 2.5
-    assert isinstance(atom["tool_calls"], list)
-    call = atom["tool_calls"][0]
-    assert call["type"] == "function"
-    assert call["function"]["name"] == "say"
-    assert call["function"]["arguments"]["text"] == "I'm cleaning up!"

tests/datasets/test_aggregate.py

@@ -289,52 +289,6 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
     assert_dataset_iteration_works(aggr_ds)
 
 
-def test_aggregate_datasets_without_concatenation(tmp_path, lerobot_dataset_factory):
-    """With concatenation disabled, each source file is kept as its own destination file."""
-    ds_0 = lerobot_dataset_factory(
-        root=tmp_path / "no_stitch_0",
-        repo_id=f"{DUMMY_REPO_ID}_no_stitch_0",
-        total_episodes=3,
-        total_frames=60,
-    )
-    ds_1 = lerobot_dataset_factory(
-        root=tmp_path / "no_stitch_1",
-        repo_id=f"{DUMMY_REPO_ID}_no_stitch_1",
-        total_episodes=4,
-        total_frames=80,
-    )
-
-    aggr_root = tmp_path / "no_stitch_aggr"
-    aggregate_datasets(
-        repo_ids=[ds_0.repo_id, ds_1.repo_id],
-        roots=[ds_0.root, ds_1.root],
-        aggr_repo_id=f"{DUMMY_REPO_ID}_no_stitch_aggr",
-        aggr_root=aggr_root,
-        concatenate_videos=False,
-        concatenate_data=False,
-    )
-
-    with (
-        patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
-        patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
-    ):
-        mock_get_safe_version.return_value = "v3.0"
-        mock_snapshot_download.return_value = str(aggr_root)
-        aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_no_stitch_aggr", root=aggr_root)
-
-    assert_episode_and_frame_counts(
-        aggr_ds, ds_0.num_episodes + ds_1.num_episodes, ds_0.num_frames + ds_1.num_frames
-    )
-    assert_dataset_iteration_works(aggr_ds)
-    assert_video_timestamps_within_bounds(aggr_ds)
-
-    # Two single-file sources stay as two files each, instead of being packed together.
-    assert len(list((aggr_root / "data").rglob("*.parquet"))) == 2
-    assert aggr_ds.meta.video_keys, "Test fixture should produce at least one video feature"
-    for key in aggr_ds.meta.video_keys:
-        assert len(list((aggr_root / "videos" / key).rglob("*.mp4"))) == 2
-
-
 @pytest.mark.parametrize("mutation", ["mismatched_value", "missing_key"])
 def test_aggregate_incomplete_video_encoder_info_warns_and_nuls_encoders(
     tmp_path, lerobot_dataset_factory, caplog, mutation

tests/datasets/test_compute_stats.py

@@ -83,29 +83,6 @@ def test_get_feature_stats_images():
     assert stats["min"].shape == stats["max"].shape == stats["mean"].shape == stats["std"].shape
 
 
-def test_get_feature_stats_uint8_images_preserves_std():
-    data = np.array(
-        [
-            [
-                [[0, 64], [128, 255]],
-                [[255, 128], [64, 0]],
-                [[32, 96], [160, 224]],
-            ],
-            [
-                [[16, 80], [144, 240]],
-                [[240, 144], [80, 16]],
-                [[48, 112], [176, 208]],
-            ],
-        ],
-        dtype=np.uint8,
-    )
-
-    stats = get_feature_stats(data, axis=(0, 2, 3), keepdims=True)
-
-    expected_std = data.transpose(0, 2, 3, 1).reshape(-1, 3).std(axis=0).reshape(1, 3, 1, 1)
-    np.testing.assert_allclose(stats["std"], expected_std)
-
-
 def test_get_feature_stats_axis_0_keepdims(sample_array):
     expected = {
         "min": np.array([[1, 2, 3]]),

tests/datasets/test_sampler.py

@@ -114,17 +114,28 @@ def test_shuffle():
     assert set(sampler) == {0, 1, 2, 3, 4, 5}
 
 
-def test_shuffle_is_reproducible_across_instances():
-    # The order is a pure function of (seed, epoch), so two fresh samplers (e.g. two ranks)
-    # produce the same permutation without any generator synchronization.
-    sampler_a = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
-    sampler_b = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
-    epoch_0 = list(sampler_a)
-    assert list(sampler_b) == epoch_0
+def test_shuffle_with_generator_is_deterministic():
+    # Two samplers shuffling with same-seed generators must yield identical permutations.
+    # This is what keeps batch shards disjoint across ranks in distributed training, where
+    # accelerate synchronizes the sampler's generator state instead of the global torch RNG.
+    sampler_a = EpisodeAwareSampler([0], [6], shuffle=True, generator=torch.Generator().manual_seed(42))
+    sampler_b = EpisodeAwareSampler([0], [6], shuffle=True, generator=torch.Generator().manual_seed(42))
+    assert list(sampler_a) == list(sampler_b)
+
     # Desyncing the global RNG must not affect the permutation.
-    sampler_c = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
+    sampler_c = EpisodeAwareSampler([0], [6], shuffle=True, generator=torch.Generator().manual_seed(42))
+    order_before = list(sampler_c)
+    sampler_c.generator.manual_seed(42)
     torch.randperm(1000)  # consume global RNG, as rank-asymmetric code (e.g. eval) would
-    assert list(sampler_c) == epoch_0
+    assert list(sampler_c) == order_before
+
+
+def test_generator_attribute_defaults_to_none():
+    # accelerate detects synchronizable samplers via `hasattr(sampler, "generator")`,
+    # so the attribute must exist even when no generator is passed.
+    sampler = EpisodeAwareSampler([0], [6], shuffle=True)
+    assert sampler.generator is None
+    assert set(sampler) == {0, 1, 2, 3, 4, 5}
 
 
 def test_negative_drop_first_frames_raises():
@@ -150,87 +161,3 @@ def test_partial_episode_drop_warns(caplog):
     # Episode 0 is skipped (1 frame, drop 1), Episode 1 keeps frames 2-5
     assert sampler.indices == [2, 3, 4, 5]
     assert "Episode 0" in caplog.text
-
-
-# --- seeded (seed, epoch) shuffling, resume, and state ---
-
-from lerobot.datasets.sampler import compute_sampler_state  # noqa: E402
-
-EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6])  # episodes of 2, 1 and 3 frames
-
-
-@pytest.mark.parametrize("num_frames", [1, 2, 3, 37, 64, 100])
-def test_deterministic_sampler_shuffle_is_permutation(num_frames):
-    for seed in (0, 1, 1234):
-        sampler = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=seed)
-        assert sorted(sampler) == list(range(num_frames))
-
-
-def test_deterministic_sampler_epochs_reproduce_and_differ():
-    sampler_a = EpisodeAwareSampler([0], [100], shuffle=True, seed=42)
-    sampler_b = EpisodeAwareSampler([0], [100], shuffle=True, seed=42)
-    epoch_0 = list(sampler_a)
-    assert list(sampler_b) == epoch_0  # same (seed, epoch) -> same order on any process
-    epoch_1 = list(sampler_a)  # __iter__ auto-advances the epoch
-    assert epoch_1 != epoch_0
-    assert sorted(epoch_1) == sorted(epoch_0)
-    sampler_a.set_epoch(0)
-    assert list(sampler_a) == epoch_0
-    assert list(EpisodeAwareSampler([0], [100], shuffle=True, seed=7)) != epoch_0
-
-
-def test_deterministic_sampler_resume_mid_epoch():
-    reference = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, seed=42)
-    epoch_0 = list(reference)
-    epoch_1 = list(reference)
-    for start in (0, 1, 4, len(epoch_0)):
-        resumed = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, seed=42)
-        resumed.load_state_dict({"epoch": 0, "start_index": start})
-        assert list(resumed) == epoch_0[start:]
-        # the resumed sampler continues into the same epoch 1 as the uninterrupted one
-        assert list(resumed) == epoch_1
-
-
-def test_deterministic_sampler_construction_stores_only_boundaries():
-    # Construction is O(num_episodes), not O(num_frames): a million-frame single episode
-    # instantiates from just its boundaries without materializing a per-frame index list.
-    num_frames = 1_000_000
-    sampler = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0)
-    assert len(sampler) == num_frames
-    assert sampler._starts.shape == (1,) and sampler._cum_lengths.shape == (1,)
-
-
-def test_deterministic_sampler_resume_is_exact_at_scale():
-    # Seeded randperm makes resume sample-exact at non-trivial sizes: regenerating the epoch's
-    # permutation and slicing from the saved offset reproduces the remaining order exactly.
-    num_frames = 100_000
-    reference = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0)
-    epoch_0 = list(reference)
-    assert sorted(epoch_0) == list(range(num_frames))
-    start = num_frames - 5
-    resumed = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0)
-    resumed.load_state_dict({"epoch": 0, "start_index": start})
-    assert list(resumed) == epoch_0[start:]
-
-
-def test_compute_sampler_state():
-    # 100 frames, batch 10, 2 ranks -> 10 underlying batches, 5 per rank per epoch.
-    assert compute_sampler_state(step=0, num_frames=100, batch_size=10, num_processes=2) == {
-        "epoch": 0,
-        "start_index": 0,
-    }
-    # step 7 -> epoch 1, 2 per-rank batches in = 2 * 10 * 2 = 40 samples in
-    assert compute_sampler_state(step=7, num_frames=100, batch_size=10, num_processes=2) == {
-        "epoch": 1,
-        "start_index": 40,
-    }
-    # uneven epoch: 95 frames -> 10 underlying batches (last short), still 5 per rank
-    assert compute_sampler_state(step=12, num_frames=95, batch_size=10, num_processes=2) == {
-        "epoch": 2,
-        "start_index": 40,
-    }
-    # uneven sharding: 105 frames -> 11 underlying batches, 6 per rank (even_batches pads)
-    assert compute_sampler_state(step=11, num_frames=105, batch_size=10, num_processes=2) == {
-        "epoch": 1,
-        "start_index": 100,
-    }

tests/datasets/test_streaming.py

@@ -13,7 +13,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import numpy as np
 import pytest
 import torch
 
@@ -25,52 +24,6 @@ from lerobot.utils.constants import ACTION
 from tests.fixtures.constants import DUMMY_REPO_ID
 
 
-def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int]:
-    """Replicates the shuffling logic of StreamingLeRobotDataset to get the expected order of indices."""
-    rng = np.random.default_rng(streaming_ds.seed)
-    buffer_size = streaming_ds.buffer_size
-    num_shards = streaming_ds.num_shards
-
-    shards_indices = []
-    for shard_idx in range(num_shards):
-        shard = streaming_ds.hf_dataset.shard(num_shards, index=shard_idx)
-        shard_indices = [item["index"] for item in shard]
-        shards_indices.append(shard_indices)
-
-    shard_iterators = {i: iter(s) for i, s in enumerate(shards_indices)}
-
-    buffer_indices_generator = streaming_ds._iter_random_indices(rng, buffer_size)
-
-    frames_buffer = []
-    expected_indices = []
-
-    while shard_iterators:  # While there are still available shards
-        available_shard_keys = list(shard_iterators.keys())
-        if not available_shard_keys:
-            break
-
-        # Call _infinite_generator_over_elements with current available shards (key difference!)
-        shard_key = next(streaming_ds._infinite_generator_over_elements(rng, available_shard_keys))
-
-        try:
-            frame_index = next(shard_iterators[shard_key])
-
-            if len(frames_buffer) == buffer_size:
-                i = next(buffer_indices_generator)
-                expected_indices.append(frames_buffer[i])
-                frames_buffer[i] = frame_index
-            else:
-                frames_buffer.append(frame_index)
-
-        except StopIteration:
-            del shard_iterators[shard_key]  # Remove exhausted shard
-
-    rng.shuffle(frames_buffer)
-    expected_indices.extend(frames_buffer)
-
-    return expected_indices
-
-
 def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
     """Test if are correctly accessed"""
     ds_num_frames = 400
@@ -120,10 +73,9 @@ def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
     [False, True],
 )
 def test_frames_order_over_epochs(tmp_path, lerobot_dataset_factory, shuffle):
-    """Test if streamed frames correspond to shuffling operations over in-memory dataset."""
+    """Each epoch covers every frame exactly once; shuffle reshuffles across epochs."""
     ds_num_frames = 400
     ds_num_episodes = 10
-    buffer_size = 100
     seed = 42
     n_epochs = 3
 
@@ -138,25 +90,17 @@ def test_frames_order_over_epochs(tmp_path, lerobot_dataset_factory, shuffle):
     )
 
     streaming_ds = StreamingLeRobotDataset(
-        repo_id=repo_id, root=local_path, buffer_size=buffer_size, seed=seed, shuffle=shuffle
+        repo_id=repo_id, root=local_path, episode_pool_size=4, seed=seed, shuffle=shuffle
     )
 
-    first_epoch_indices = [frame["index"] for frame in streaming_ds]
-    expected_indices = get_frames_expected_order(streaming_ds)
-
-    assert first_epoch_indices == expected_indices, "First epoch indices do not match expected indices"
-
-    expected_indices = get_frames_expected_order(streaming_ds)
-    for _ in range(n_epochs):
-        streaming_indices = [frame["index"] for frame in streaming_ds]
-        frames_match = all(
-            s_index == e_index for s_index, e_index in zip(streaming_indices, expected_indices, strict=True)
-        )
-
-        if shuffle:
-            assert not frames_match
-        else:
-            assert frames_match
+    epochs = [[int(frame["index"]) for frame in streaming_ds] for _ in range(n_epochs)]
+    for epoch_indices in epochs:
+        assert sorted(epoch_indices) == list(range(ds_num_frames)), "epoch did not cover every frame once"
+    if shuffle:
+        assert epochs[0] != epochs[1], "shuffle did not reshuffle across epochs"
+        assert epochs[0] != list(range(ds_num_frames)), "shuffle left the stream in sequential order"
+    else:
+        assert epochs[0] == epochs[1] == epochs[2], "unshuffled epochs must repeat the same order"
 
 
 @pytest.mark.parametrize(
@@ -164,15 +108,11 @@ def test_frames_order_over_epochs(tmp_path, lerobot_dataset_factory, shuffle):
     [False, True],
 )
 def test_frames_order_with_shards(tmp_path, lerobot_dataset_factory, shuffle):
-    """Test if streamed frames correspond to shuffling operations over in-memory dataset with multiple shards."""
+    """Multi-shard streams keep exactly-once coverage and deterministic per-seed order."""
     ds_num_frames = 100
     ds_num_episodes = 10
-    buffer_size = 10
-
     seed = 42
-    n_epochs = 3
     data_file_size_mb = 0.001
-
     chunks_size = 1
 
     local_path = tmp_path / "test"
@@ -187,31 +127,21 @@ def test_frames_order_with_shards(tmp_path, lerobot_dataset_factory, shuffle):
         chunks_size=chunks_size,
     )
 
-    streaming_ds = StreamingLeRobotDataset(
-        repo_id=repo_id,
-        root=local_path,
-        buffer_size=buffer_size,
-        seed=seed,
-        shuffle=shuffle,
-        max_num_shards=4,
-    )
-
-    first_epoch_indices = [frame["index"] for frame in streaming_ds]
-    expected_indices = get_frames_expected_order(streaming_ds)
-
-    assert first_epoch_indices == expected_indices, "First epoch indices do not match expected indices"
-
-    for _ in range(n_epochs):
-        streaming_indices = [
-            frame["index"] for frame in streaming_ds
-        ]  # NOTE: this is the same as first_epoch_indices
-        frames_match = all(
-            s_index == e_index for s_index, e_index in zip(streaming_indices, expected_indices, strict=True)
+    def make_ds():
+        return StreamingLeRobotDataset(
+            repo_id=repo_id,
+            root=local_path,
+            episode_pool_size=3,
+            seed=seed,
+            shuffle=shuffle,
+            max_num_shards=4,
         )
-        if shuffle:
-            assert not frames_match
-        else:
-            assert frames_match
+
+    first = [int(frame["index"]) for frame in make_ds()]
+    again = [int(frame["index"]) for frame in make_ds()]
+
+    assert sorted(first) == list(range(ds_num_frames)), "epoch did not cover every frame once"
+    assert first == again, "same seed must reproduce the same order"
 
 
 @pytest.mark.parametrize(
@@ -288,6 +218,11 @@ def test_frames_with_delta_consistency(tmp_path, lerobot_dataset_factory, state_
 
                 check = torch.allclose(left, right) and left.shape == right.shape
 
+            else:
+                # Scalar numerics: streaming yields python floats/ints where map-style yields
+                # 0-dim tensors (long-standing accepted difference). Compare by value.
+                check = float(left) == float(right)
+
             key_checks.append((key, check))
 
         assert all(t[1] for t in key_checks), (

tests/datasets/test_streaming_distributed.py

@@ -0,0 +1,100 @@
+# 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.
+
+"""End-to-end distributed streaming smoke test under a real `accelerate launch`.
+
+Mirrors tests/training/test_multi_gpu.py but runs on CPU and only checks the dataloading contract: with
+two processes, `split_dataset_by_node` (auto-resolved from the Accelerate state) must give each rank a
+disjoint set of frames that together cover the dataset. Skips if the environment can't actually spawn
+>= 2 processes (e.g. local macOS multi-CPU), so it never silently passes as a single process.
+"""
+
+import json
+import shutil
+import subprocess
+import sys
+
+import pytest
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+from tests.fixtures.constants import DUMMY_REPO_ID
+
+WORKER = """
+import json, sys
+from accelerate import PartialState
+from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
+
+root, repo_id, out_dir = sys.argv[1], sys.argv[2], sys.argv[3]
+state = PartialState()
+ds = StreamingLeRobotDataset(
+    repo_id=repo_id, root=root, shuffle=False, episode_pool_size=8, max_num_shards=8
+)
+indices = [int(frame["index"]) for frame in ds]
+payload = {"rank": state.process_index, "world": state.num_processes, "indices": indices}
+with open(f"{out_dir}/rank_{state.process_index}.json", "w") as f:
+    json.dump(payload, f)
+"""
+
+
+@pytest.mark.skipif(shutil.which("accelerate") is None, reason="accelerate CLI not available")
+def test_accelerate_launch_ranks_are_disjoint(tmp_path, lerobot_dataset_factory):
+    total_frames = 160
+    repo_id = f"{DUMMY_REPO_ID}-acc"
+    root = tmp_path / "ds"
+    lerobot_dataset_factory(
+        root=root,
+        repo_id=repo_id,
+        total_episodes=8,
+        total_frames=total_frames,
+        use_videos=False,
+        data_files_size_in_mb=0.001,
+        chunks_size=1,
+    )
+
+    worker = tmp_path / "worker.py"
+    worker.write_text(WORKER)
+    out_dir = tmp_path / "out"
+    out_dir.mkdir()
+
+    cmd = [
+        "accelerate",
+        "launch",
+        "--num_processes=2",
+        "--num_machines=1",
+        "--mixed_precision=no",
+        "--dynamo_backend=no",
+        "--cpu",
+        str(worker),
+        str(root),
+        repo_id,
+        str(out_dir),
+    ]
+    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
+    assert result.returncode == 0, (
+        f"accelerate launch failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+    )
+
+    payloads = [json.loads(p.read_text()) for p in sorted(out_dir.glob("rank_*.json"))]
+    if len(payloads) < 2 or any(p["world"] < 2 for p in payloads):
+        pytest.skip("environment did not spawn >= 2 distributed processes (e.g. local macOS multi-CPU)")
+
+    rank_sets = [set(p["indices"]) for p in payloads]
+    assert rank_sets[0].isdisjoint(rank_sets[1]), "ranks streamed overlapping frames under accelerate launch"
+    assert set().union(*rank_sets) == set(range(total_frames)), "ranks did not jointly cover all frames"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

tests/datasets/test_streaming_native.py

@@ -0,0 +1,430 @@
+# 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.
+
+"""Tests for the HF-native large-scale streaming additions: distributed (per-rank) sharding,
+DataLoader worker splitting, the episode pool (randomness, coverage, exact deltas), video
+prefetching, deterministic fast-forward resume, and schema parity."""
+
+import pytest
+import torch
+from torch.utils.data import DataLoader
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
+from lerobot.utils.constants import ACTION
+from tests.fixtures.constants import DUMMY_REPO_ID
+
+
+def _make_local_dataset(factory, root, repo_id, *, total_episodes, total_frames, use_videos=False, **kw):
+    factory(
+        root=root,
+        repo_id=repo_id,
+        total_episodes=total_episodes,
+        total_frames=total_frames,
+        use_videos=use_videos,
+        data_files_size_in_mb=0.001,
+        chunks_size=1,
+        **kw,
+    )
+
+
+def _stream_indices(ds: StreamingLeRobotDataset) -> list[int]:
+    return [int(frame["index"]) for frame in ds]
+
+
+def test_resolve_distributed_prefers_explicit_then_env(monkeypatch):
+    assert StreamingLeRobotDataset._resolve_distributed(2, 8) == (2, 8)
+
+    monkeypatch.delenv("RANK", raising=False)
+    monkeypatch.delenv("WORLD_SIZE", raising=False)
+    # No accelerate state, no env -> single process.
+    assert StreamingLeRobotDataset._resolve_distributed(None, None) == (0, 1)
+
+    monkeypatch.setenv("RANK", "3")
+    monkeypatch.setenv("WORLD_SIZE", "4")
+    assert StreamingLeRobotDataset._resolve_distributed(None, None) == (3, 4)
+
+
+def test_split_by_node_disjoint_across_ranks(tmp_path, lerobot_dataset_factory):
+    """Each rank must stream a disjoint set of frames, and the ranks together must cover every frame."""
+    repo_id = f"{DUMMY_REPO_ID}-ranks"
+    total_frames, total_episodes = 200, 8
+    _make_local_dataset(
+        lerobot_dataset_factory,
+        tmp_path / "ds",
+        repo_id,
+        total_episodes=total_episodes,
+        total_frames=total_frames,
+    )
+
+    world_size = 2
+    per_rank = []
+    for rank in range(world_size):
+        ds = StreamingLeRobotDataset(
+            repo_id=repo_id,
+            root=tmp_path / "ds",
+            shuffle=False,
+            episode_pool_size=8,
+            max_num_shards=8,
+            rank=rank,
+            world_size=world_size,
+        )
+        per_rank.append(set(_stream_indices(ds)))
+
+    assert per_rank[0].isdisjoint(per_rank[1]), (
+        "ranks streamed overlapping frames (duplicate data across GPUs)"
+    )
+    assert per_rank[0] | per_rank[1] == set(range(total_frames)), "ranks did not jointly cover all frames"
+
+
+def test_dataloader_workers_no_duplicates_within_rank(tmp_path, lerobot_dataset_factory):
+    """DataLoader workers within a rank must split shards so no frame is yielded twice."""
+    repo_id = f"{DUMMY_REPO_ID}-workers"
+    total_frames, total_episodes = 120, 8
+    _make_local_dataset(
+        lerobot_dataset_factory,
+        tmp_path / "ds",
+        repo_id,
+        total_episodes=total_episodes,
+        total_frames=total_frames,
+    )
+
+    ds = StreamingLeRobotDataset(
+        repo_id=repo_id, root=tmp_path / "ds", shuffle=False, episode_pool_size=4, max_num_shards=4
+    )
+    loader = DataLoader(ds, batch_size=None, num_workers=2)
+    indices = [int(batch["index"]) for batch in loader]
+
+    assert len(indices) == len(set(indices)), "DataLoader workers yielded duplicate frames within a rank"
+
+
+def test_sarm_window_covers_long_horizon_without_padding(tmp_path, lerobot_dataset_factory):
+    """A delta window longer than the old 100-frame ceiling must fetch real frames, not pad them.
+
+    SARM uses a window of 8 steps spaced 1s (~160 frames @ fps20). Here fps=30, so +5s = 150 frames > 100.
+    """
+    repo_id = f"{DUMMY_REPO_ID}-sarm"
+    # A single long episode so a +150-frame lookahead is unambiguously inside the episode (the fixture
+    # gives episodes variable lengths, so multi-episode boundaries can't be assumed).
+    episode_frames = 300
+    _make_local_dataset(
+        lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=1, total_frames=episode_frames
+    )
+
+    horizon_s = 5.0  # 150 frames @ fps30, well beyond LOOKAHEAD_BACKTRACKTABLE=100
+    delta_timestamps = {ACTION: [0.0, horizon_s]}
+    ds = StreamingLeRobotDataset(
+        repo_id=repo_id,
+        root=tmp_path / "ds",
+        shuffle=False,
+        episode_pool_size=1,
+        max_num_shards=1,
+        delta_timestamps=delta_timestamps,
+    )
+
+    horizon_frames = int(round(horizon_s * ds.fps))
+    assert horizon_frames > 100, "test must exceed the old LOOKAHEAD_BACKTRACKTABLE ceiling"
+    checked = 0
+    for frame in ds:
+        idx = int(frame["index"])
+        # The +horizon target is inside the single episode -> it must be a real frame, not padding.
+        if idx + horizon_frames < episode_frames:
+            assert not bool(frame[f"{ACTION}_is_pad"][-1]), (
+                f"frame {idx}: +{horizon_frames} target was padded; long delta window did not reach it"
+            )
+            checked += 1
+    assert checked > 0, "test did not exercise any in-episode long-horizon frame"
+
+
+def test_pool_order_is_deterministic_per_seed(tmp_path, lerobot_dataset_factory):
+    repo_id = f"{DUMMY_REPO_ID}-seeds"
+    _make_local_dataset(lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=6, total_frames=120)
+
+    def order(seed):
+        return _stream_indices(
+            StreamingLeRobotDataset(
+                repo_id=repo_id,
+                root=tmp_path / "ds",
+                shuffle=True,
+                seed=seed,
+                episode_pool_size=4,
+                max_num_shards=2,
+            )
+        )
+
+    assert order(0) == order(0), "same seed must reproduce the same order"
+    assert order(0) != order(1), "different seeds should give different orders"
+
+
+def test_pool_epochs_reshuffle_and_cover(tmp_path, lerobot_dataset_factory):
+    """Consecutive passes over the same dataset object reshuffle (epoch advances) but keep coverage."""
+    repo_id = f"{DUMMY_REPO_ID}-epochs"
+    total_frames = 120
+    _make_local_dataset(
+        lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=6, total_frames=total_frames
+    )
+    ds = StreamingLeRobotDataset(
+        repo_id=repo_id, root=tmp_path / "ds", shuffle=True, seed=3, episode_pool_size=4, max_num_shards=2
+    )
+    epoch_0 = _stream_indices(ds)
+    epoch_1 = _stream_indices(ds)
+    assert sorted(epoch_0) == sorted(epoch_1) == list(range(total_frames))
+    assert epoch_0 != epoch_1, "epoch did not reshuffle"
+
+
+def test_pool_mixes_episodes(tmp_path, lerobot_dataset_factory):
+    """Early samples should already come from several distinct episodes (the pool's purpose)."""
+    repo_id = f"{DUMMY_REPO_ID}-mix"
+    _make_local_dataset(lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=8, total_frames=200)
+    ds = StreamingLeRobotDataset(
+        repo_id=repo_id, root=tmp_path / "ds", shuffle=True, seed=0, episode_pool_size=8, max_num_shards=4
+    )
+    episodes_in_head = {int(frame["episode_index"]) for _, frame in zip(range(20), ds, strict=False)}
+    assert len(episodes_in_head) >= 3, f"pool did not mix episodes: {episodes_in_head}"
+
+
+def test_schema_parity_with_map_style(tmp_path, lerobot_dataset_factory):
+    """Streamed samples must have the same keys / shapes / dtypes as map-style LeRobotDataset."""
+    repo_id = f"{DUMMY_REPO_ID}-parity"
+    map_ds = lerobot_dataset_factory(
+        root=tmp_path / "ds", repo_id=repo_id, total_episodes=4, total_frames=80, use_videos=True
+    )
+    stream_ds = StreamingLeRobotDataset(
+        repo_id=repo_id, root=tmp_path / "ds", shuffle=False, episode_pool_size=4, max_num_shards=2
+    )
+
+    map_frame = map_ds[0]
+    stream_frame = next(iter(stream_ds))
+
+    assert set(stream_frame) == set(map_frame), set(stream_frame) ^ set(map_frame)
+    for key, value in stream_frame.items():
+        ref = map_frame[key]
+        if isinstance(value, torch.Tensor):
+            assert isinstance(ref, torch.Tensor) and value.shape == ref.shape and value.dtype == ref.dtype, (
+                f"{key}: stream {tuple(value.shape)}/{value.dtype} vs map {tuple(ref.shape)}/{ref.dtype}"
+            )
+        elif isinstance(value, str):
+            assert isinstance(ref, str), f"{key}: {type(value)} vs {type(ref)}"
+        else:
+            # Scalar numerics: streaming yields python floats where map-style yields 0-dim tensors
+            # (a long-standing, accepted difference). Compare by value rather than exact type.
+            assert float(value) == float(ref), f"{key}: {value} vs {ref}"
+
+
+def test_video_path_resolution_local(tmp_path, lerobot_dataset_factory, monkeypatch):
+    """For a local (prewarmed) root, video decode must be issued against the local path, not hf://."""
+    import lerobot.datasets.streaming_dataset as sd
+
+    repo_id = f"{DUMMY_REPO_ID}-vpath"
+    lerobot_dataset_factory(
+        root=tmp_path / "ds", repo_id=repo_id, total_episodes=2, total_frames=40, use_videos=True
+    )
+    ds = StreamingLeRobotDataset(
+        repo_id=repo_id, root=tmp_path / "ds", shuffle=False, episode_pool_size=1, max_num_shards=1
+    )
+
+    seen_paths = []
+
+    def fake_decode(video_path, query_ts, *args, **kwargs):
+        seen_paths.append(str(video_path))
+        return torch.zeros(len(query_ts), 3, 64, 96)
+
+    monkeypatch.setattr(sd, "decode_video_frames_torchcodec", fake_decode)
+    next(iter(ds))
+
+    assert seen_paths, "no video decode was issued"
+    assert all(str(ds.root) in p and not p.startswith("hf://") for p in seen_paths), seen_paths
+
+
+def test_shuffle_decorrelates_output_order(tmp_path, lerobot_dataset_factory):
+    """With shuffle on, streamed frame order must differ from the underlying sequential order."""
+    repo_id = f"{DUMMY_REPO_ID}-shuf"
+    _make_local_dataset(lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=8, total_frames=200)
+    ordered = _stream_indices(
+        StreamingLeRobotDataset(
+            repo_id=repo_id, root=tmp_path / "ds", shuffle=False, episode_pool_size=1, max_num_shards=1
+        )
+    )
+    shuffled = _stream_indices(
+        StreamingLeRobotDataset(
+            repo_id=repo_id, root=tmp_path / "ds", shuffle=True, episode_pool_size=8, max_num_shards=4, seed=0
+        )
+    )
+    assert sorted(shuffled) == sorted(ordered), "shuffling changed the set of frames"
+    assert shuffled != ordered, "shuffle did not decorrelate output order"
+
+
+def test_native_resume_never_repeats_and_loss_is_bounded(tmp_path, lerobot_dataset_factory):
+    """Native state_dict resume: no sample is re-yielded; loss is bounded by the shuffle buffers."""
+    repo_id = f"{DUMMY_REPO_ID}-native-resume"
+    total_frames = 100
+    _make_local_dataset(
+        lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=5, total_frames=total_frames
+    )
+
+    def fresh_ds():
+        return StreamingLeRobotDataset(
+            repo_id=repo_id,
+            root=tmp_path / "ds",
+            shuffle=True,
+            seed=7,
+            episode_pool_size=2,
+            frame_shuffle_buffer_size=8,
+        )
+
+    ds = fresh_ds()
+    it = iter(ds)
+    consumed = [int(next(it)["index"]) for _ in range(30)]
+    state = ds.state_dict()
+
+    resumed_ds = fresh_ds()
+    resumed_ds.load_state_dict(state)
+    rest = [int(frame["index"]) for frame in resumed_ds]
+
+    assert not set(consumed) & set(rest), "resume re-yielded already-seen frames"
+    # in-flight buffer contents are skipped on resume (documented datasets behavior):
+    # bounded by the episode pool (2 episodes of <= ~30 frames here) + frame buffer (8)
+    covered = len(set(consumed) | set(rest))
+    max_in_flight = 2 * 30 + 8
+    assert covered >= total_frames - max_in_flight
+    assert covered + len(consumed) >= total_frames - max_in_flight
+
+
+def test_pipeline_uses_native_primitives(tmp_path, lerobot_dataset_factory):
+    """The tabular pipeline is pure datasets: batch(by_column) + shuffle + map + shuffle."""
+    repo_id = f"{DUMMY_REPO_ID}-native-pipe"
+    _make_local_dataset(lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=4, total_frames=80)
+    ds = StreamingLeRobotDataset(repo_id=repo_id, root=tmp_path / "ds", shuffle=True, episode_pool_size=2)
+    import datasets as hf_datasets
+
+    assert isinstance(ds._pipeline, hf_datasets.IterableDataset)
+    state = ds._pipeline.state_dict()  # the native resume protocol is available end-to-end
+    assert state is not None
+
+
+# --- Plan B: random-episode admission via reshard() + multi-input-shard shuffle ---
+
+
+def test_reshard_makes_one_shard_per_episode(tmp_path, lerobot_dataset_factory):
+    """With one row group per episode (the writer's invariant), reshard() turns each episode into its
+    own shard, so num_shards == total_episodes even when many episodes share a single data file."""
+    import pyarrow.parquet as pq
+
+    repo_id = f"{DUMMY_REPO_ID}-reshard"
+    total_episodes = 3
+    # Default (large) data-file size packs all (unequal-length) episodes into one file, so the only way
+    # num_shards can reach total_episodes is per-row-group resharding.
+    lerobot_dataset_factory(
+        root=tmp_path / "ds",
+        repo_id=repo_id,
+        total_episodes=total_episodes,
+        total_frames=90,
+        use_videos=False,
+    )
+    ds = StreamingLeRobotDataset(repo_id=repo_id, root=tmp_path / "ds", shuffle=False, episode_pool_size=3)
+
+    file_to_eps = ds._episode_files()
+    assert len(file_to_eps) == 1, "test expects all episodes packed into a single data file"
+    for (chunk_idx, file_idx), eps in file_to_eps.items():
+        rel = ds.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
+        assert pq.ParquetFile(str(ds.root / rel)).num_row_groups == len(eps)
+
+    assert ds.num_shards == total_episodes
+
+
+def test_max_buffer_input_shards_admits_random_episodes(tmp_path, lerobot_dataset_factory):
+    """max_buffer_input_shards (== concurrently-live random episodes) drives the per-batch episode mix:
+    a single batch should already span most of the live episodes."""
+    repo_id = f"{DUMMY_REPO_ID}-frac"
+    total_episodes = 8
+    lerobot_dataset_factory(
+        root=tmp_path / "ds",
+        repo_id=repo_id,
+        total_episodes=total_episodes,
+        total_frames=240,
+        use_videos=False,
+    )
+    ds = StreamingLeRobotDataset(
+        repo_id=repo_id,
+        root=tmp_path / "ds",
+        shuffle=True,
+        seed=0,
+        episode_pool_size=total_episodes,
+        max_buffer_input_shards=total_episodes,
+    )
+    assert ds.max_buffer_input_shards == total_episodes
+
+    batch = 32
+    head = {int(frame["episode_index"]) for _, frame in zip(range(batch), ds, strict=False)}
+    assert len(head) >= min(total_episodes, batch) - 2, f"batch did not mix random episodes: {head}"
+
+
+def test_collapsed_row_groups_raise(tmp_path, lerobot_dataset_factory):
+    """A data file that collapses several episodes into a single row group (bulk df.to_parquet /
+    push_to_hub) must be rejected with an actionable error: reshard() cannot address its episodes."""
+    import pyarrow.parquet as pq
+
+    repo_id = f"{DUMMY_REPO_ID}-collapsed"
+    lerobot_dataset_factory(
+        root=tmp_path / "ds", repo_id=repo_id, total_episodes=3, total_frames=90, use_videos=False
+    )
+    # Rewrite every data file as a single row group (simulating the aggregate/push_to_hub collapse).
+    for parquet_path in (tmp_path / "ds" / "data").rglob("*.parquet"):
+        pq.write_table(pq.read_table(parquet_path), parquet_path)
+
+    with pytest.raises(ValueError, match="ONE ROW GROUP PER EPISODE"):
+        StreamingLeRobotDataset(repo_id=repo_id, root=tmp_path / "ds", shuffle=False, episode_pool_size=3)
+
+
+def test_collapsed_row_groups_can_be_bypassed(tmp_path, lerobot_dataset_factory):
+    """validate_row_groups=False skips the row-group check (collapsed datasets still load, degraded)."""
+    import pyarrow.parquet as pq
+
+    repo_id = f"{DUMMY_REPO_ID}-collapsed-bypass"
+    lerobot_dataset_factory(
+        root=tmp_path / "ds", repo_id=repo_id, total_episodes=3, total_frames=90, use_videos=False
+    )
+    for parquet_path in (tmp_path / "ds" / "data").rglob("*.parquet"):
+        pq.write_table(pq.read_table(parquet_path), parquet_path)
+
+    ds = StreamingLeRobotDataset(
+        repo_id=repo_id, root=tmp_path / "ds", shuffle=False, episode_pool_size=3, validate_row_groups=False
+    )
+    assert sorted(int(frame["index"]) for frame in ds) == list(range(90))
+
+
+def test_distributed_divisibility_guard_raises(tmp_path, lerobot_dataset_factory):
+    """When num_shards (== episodes after reshard) is not divisible by world_size, every rank would
+    stream the whole dataset; the guard must raise instead of silently degrading."""
+    repo_id = f"{DUMMY_REPO_ID}-divis"
+    lerobot_dataset_factory(
+        root=tmp_path / "ds", repo_id=repo_id, total_episodes=3, total_frames=90, use_videos=False
+    )
+    with pytest.raises(ValueError, match="not divisible by world_size"):
+        StreamingLeRobotDataset(
+            repo_id=repo_id, root=tmp_path / "ds", shuffle=False, episode_pool_size=3, rank=0, world_size=2
+        )
+
+    # Bypassing the guard downgrades it to a warning (no raise).
+    ds = StreamingLeRobotDataset(
+        repo_id=repo_id,
+        root=tmp_path / "ds",
+        shuffle=False,
+        episode_pool_size=3,
+        rank=0,
+        world_size=2,
+        validate_row_groups=False,
+    )
+    assert ds.num_shards == 3

tests/datasets/test_video_encoding.py

@@ -504,19 +504,6 @@ class TestReencodeVideo:
         assert info["video.g"] == 6
         assert info["video.crf"] == 23
 
-    @require_h264
-    def test_reencode_video_trim_window(self, tmp_path):
-        src = TEST_ARTIFACTS_DIR / "clip_6frames.mp4"
-        out = tmp_path / "trim_window.mp4"
-        cfg = VideoEncoderConfig(vcodec="h264")
-        reencode_video(src, out, camera_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True)
-
-        with av.open(str(out)) as container:
-            frames = list(container.decode(video=0))
-        # Only the frames at 0.067 and 0.1 s fall inside [0.05, 0.12).
-        assert len(frames) == 2
-        assert frames[0].time == pytest.approx(0.0, abs=1e-3)
-
 
 class TestConcatenateVideoFiles:
     def test_two_clips_frame_count(self, tmp_path):

tests/fixtures/dataset_factories.py

@@ -552,64 +552,3 @@ def lerobot_dataset_factory(
 @pytest.fixture(scope="session")
 def empty_lerobot_dataset_factory() -> LeRobotDatasetFactory:
     return partial(LeRobotDataset.create, repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS)
-
-
-def build_annotation_dataset(
-    root: Path,
-    episode_specs: list[tuple[int, int, str]],
-    *,
-    fps: int = 10,
-) -> Path:
-    """Build a minimal LeRobot-shaped dataset on disk for annotation tests.
-
-    ``episode_specs`` is a list of ``(episode_index, num_frames, task_text)``.
-    Each episode is written to its own
-    ``data/chunk-000/file-{ep:03d}.parquet`` so the writer's per-shard
-    rewrite path is exercised. The dataset carries the minimum
-    ``meta/tasks.parquet`` + ``meta/info.json`` the reader / executor need;
-    it has no videos, so the modules fall back to text-only prompts.
-
-    Shared by the annotation-pipeline pytest fixtures (``tests/annotations/
-    conftest.py``) and the opt-in E2E smoke run so the fixture shape lives
-    in exactly one place.
-    """
-    from lerobot.datasets.io_utils import write_tasks
-    from lerobot.utils.io_utils import write_json
-
-    data_dir = root / "data" / "chunk-000"
-    data_dir.mkdir(parents=True, exist_ok=True)
-
-    tasks: dict[int, str] = {}
-    for episode_index, num_frames, task_text in episode_specs:
-        if task_text not in tasks.values():
-            tasks[len(tasks)] = task_text
-        task_index = next(k for k, v in tasks.items() if v == task_text)
-        frame = pd.DataFrame(
-            {
-                "episode_index": [episode_index] * num_frames,
-                "frame_index": list(range(num_frames)),
-                "timestamp": [round(i / fps, 6) for i in range(num_frames)],
-                "task_index": [task_index] * num_frames,
-                "subtask_index": [0] * num_frames,  # legacy column the writer must drop
-            }
-        )
-        frame.to_parquet(data_dir / f"file-{episode_index:03d}.parquet", index=False)
-
-    # Canonical tasks frame: indexed by task string with a ``task_index``
-    # column, matching what ``lerobot.datasets.io_utils.load_tasks`` expects.
-    tasks_df = pd.DataFrame(
-        {"task_index": list(tasks.keys())},
-        index=pd.Index(list(tasks.values()), name="task"),
-    )
-    write_tasks(tasks_df, root)
-
-    write_json(
-        {
-            "codebase_version": "v3.1",
-            "fps": fps,
-            "features": {},
-            "total_episodes": len(episode_specs),
-        },
-        root / "meta" / "info.json",
-    )
-    return root

tests/fixtures/files.py

@@ -17,6 +17,7 @@ from pathlib import Path
 import datasets
 import numpy as np
 import pandas as pd
+import pyarrow.parquet as pq
 import pytest
 from datasets import Dataset
 
@@ -35,6 +36,24 @@ from lerobot.datasets.utils import (
 )
 
 
+def _to_parquet_one_row_group_per_episode(hf_dataset: Dataset, path: Path) -> None:
+    """Write ``hf_dataset`` to ``path`` with one Parquet row group per episode.
+
+    Mirrors the LeRobot recording writer (one ``write_table`` per episode) so each episode stays an
+    independently addressable shard after ``datasets.IterableDataset.reshard()``, which
+    ``StreamingLeRobotDataset`` relies on. ``Dataset.to_parquet`` would collapse the file into a
+    single row group instead.
+    """
+    table = hf_dataset.with_format("arrow")[:]
+    episode_index = np.asarray(hf_dataset["episode_index"])
+    boundaries = np.where(np.diff(episode_index) != 0)[0] + 1
+    starts = [0, *boundaries.tolist()]
+    ends = [*boundaries.tolist(), len(episode_index)]
+    with pq.ParquetWriter(str(path), table.schema) as writer:
+        for start, end in zip(starts, ends, strict=True):
+            writer.write_table(table.slice(start, end - start))
+
+
 def write_hf_dataset(
     hf_dataset: Dataset,
     local_dir: Path,
@@ -67,7 +86,7 @@ def write_hf_dataset(
         # If the dataset is small enough, write it to a single file.
         path = local_dir / DEFAULT_DATA_PATH.format(chunk_index=0, file_index=0)
         path.parent.mkdir(parents=True, exist_ok=True)
-        hf_dataset.to_parquet(path)
+        _to_parquet_one_row_group_per_episode(hf_dataset, path)
         return
 
     # If the dataset is too large, split it into smaller chunks, keeping episodes whole.
@@ -114,8 +133,8 @@ def write_hf_dataset(
         path = local_dir / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
         path.parent.mkdir(parents=True, exist_ok=True)
 
-        # Write the shard to a Parquet file.
-        dataset_shard.to_parquet(path)
+        # Write the shard to a Parquet file (one row group per episode).
+        _to_parquet_one_row_group_per_episode(dataset_shard, path)
 
         # Update chunk and file indices for the next iteration.
         chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, chunk_size)

tests/scripts/test_edit_dataset_parsing.py

@@ -66,20 +66,6 @@ class TestOperationTypeParsing:
         with pytest.raises(ValueError, match="--new_repo_id is required for merge"):
             _validate_config(cfg)
 
-    @pytest.mark.parametrize("flag", ["concatenate_videos", "concatenate_data"])
-    def test_merge_concatenate_flag_defaults_true(self, flag):
-        cfg = parse_cfg(["--new_repo_id", "test/merged", "--operation.type", "merge"])
-        assert isinstance(cfg.operation, MergeConfig)
-        assert getattr(cfg.operation, flag) is True
-
-    @pytest.mark.parametrize("flag", ["concatenate_videos", "concatenate_data"])
-    def test_merge_concatenate_flag_can_be_disabled(self, flag):
-        cfg = parse_cfg(
-            ["--new_repo_id", "test/merged", "--operation.type", "merge", f"--operation.{flag}", "false"]
-        )
-        assert isinstance(cfg.operation, MergeConfig)
-        assert getattr(cfg.operation, flag) is False
-
     def test_non_merge_requires_repo_id(self):
         cfg = parse_cfg(["--operation.type", "delete_episodes"])
         with pytest.raises(ValueError, match="--repo_id is required for delete_episodes"):

tests/scripts/test_lerobot_annotate.py

@@ -1,86 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-from types import SimpleNamespace
-
-import pytest
-
-# ``lerobot.scripts.lerobot_annotate`` (and the ``_push_to_hub`` path it
-# exercises) imports ``lerobot.datasets``, which only ships under the
-# ``dataset`` extra. Skip in tiers without it instead of erroring.
-pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
-
-
-def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
-    from lerobot.scripts.lerobot_annotate import _push_to_hub
-
-    root = tmp_path / "dataset"
-    (root / "meta").mkdir(parents=True)
-    (root / "meta" / "info.json").write_text(json.dumps({"codebase_version": "v3.0"}))
-
-    calls = {}
-
-    class FakeHfApi:
-        def create_repo(self, **kwargs):
-            calls["create_repo"] = kwargs
-
-        def upload_folder(self, **kwargs):
-            calls["upload_folder"] = kwargs
-            return SimpleNamespace(oid="abc123")
-
-        def delete_tag(self, repo_id, **kwargs):
-            import requests
-            from huggingface_hub.errors import RevisionNotFoundError
-
-            calls["delete_tag"] = {"repo_id": repo_id, **kwargs}
-            # Simulate the common case: no stale tag to delete.
-            raise RevisionNotFoundError("no such tag", response=requests.Response())
-
-        def create_tag(self, **kwargs):
-            calls["create_tag"] = kwargs
-
-    monkeypatch.setattr("huggingface_hub.HfApi", FakeHfApi)
-
-    cfg = SimpleNamespace(
-        repo_id="source/dataset",
-        new_repo_id="annotated/dataset",
-        push_private=True,
-        push_commit_message=None,
-    )
-
-    _push_to_hub(root, cfg)
-
-    assert calls["create_repo"] == {
-        "repo_id": "annotated/dataset",
-        "repo_type": "dataset",
-        "private": True,
-        "exist_ok": True,
-    }
-    assert calls["upload_folder"]["repo_id"] == "annotated/dataset"
-    # A stale tag (e.g. from a previous annotation run) is deleted first so
-    # the new tag always points at the upload we just made.
-    assert calls["delete_tag"] == {
-        "repo_id": "annotated/dataset",
-        "tag": "v3.0",
-        "repo_type": "dataset",
-    }
-    assert calls["create_tag"] == {
-        "repo_id": "annotated/dataset",
-        "tag": "v3.0",
-        "repo_type": "dataset",
-        "revision": "abc123",
-    }

tests/utils/test_logging_utils.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 
 import pytest
-import torch
 
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 
@@ -26,16 +25,8 @@ def mock_metrics():
 
 
 class MockAccelerator:
-    def __init__(self, num_processes: int, reduce_fn=None):
+    def __init__(self, num_processes: int):
         self.num_processes = num_processes
-        self.device = torch.device("cpu")
-        self._reduce_fn = reduce_fn
-
-    def reduce(self, tensor, reduction="mean"):
-        # In single-process tests we just want a deterministic stand-in for accelerate's reduce.
-        if self._reduce_fn is not None:
-            return self._reduce_fn(tensor, reduction)
-        return tensor
 
 
 def test_average_meter_initialization():
@@ -166,70 +157,3 @@ def test_metrics_tracker_reset_averages(mock_metrics):
     tracker.reset_averages()
     assert tracker.loss.avg == 0.0
     assert tracker.accuracy.avg == 0.0
-
-
-def test_average_meter_invalid_reduction():
-    with pytest.raises(ValueError):
-        AverageMeter("loss", reduction="median")
-
-
-def test_average_meter_reduction_stored():
-    meter = AverageMeter("updt_s", reduction="max")
-    assert meter.reduction == "max"
-
-
-def test_metrics_tracker_reduce_across_ranks_no_accelerator():
-    metrics = {"update_s": AverageMeter("update_s", reduction="max")}
-    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
-    tracker.update_s = 0.5
-    tracker.reduce_across_ranks()  # no-op without accelerator
-    assert tracker.update_s.avg == 0.5
-
-
-def test_metrics_tracker_reduce_across_ranks_single_process():
-    metrics = {"update_s": AverageMeter("update_s", reduction="max")}
-    tracker = MetricsTracker(
-        batch_size=32,
-        num_frames=1000,
-        num_episodes=50,
-        metrics=metrics,
-        accelerator=MockAccelerator(num_processes=1),
-    )
-    tracker.update_s = 0.5
-    tracker.reduce_across_ranks()  # no-op when world size is 1
-    assert tracker.update_s.avg == 0.5
-
-
-def test_metrics_tracker_reduce_across_ranks_invokes_reduce():
-    captured = {}
-
-    def fake_reduce(tensor, reduction):
-        captured["reduction"] = reduction
-        captured["values"] = tensor.clone()
-        # Pretend the slowest rank reported 0.9 instead of this rank's 0.4.
-        return torch.tensor([0.9], dtype=tensor.dtype, device=tensor.device)
-
-    metrics = {
-        "loss": AverageMeter("loss"),  # reduction="none" -> not touched
-        "update_s": AverageMeter("update_s", reduction="max"),
-    }
-    tracker = MetricsTracker(
-        batch_size=32,
-        num_frames=1000,
-        num_episodes=50,
-        metrics=metrics,
-        accelerator=MockAccelerator(num_processes=4, reduce_fn=fake_reduce),
-    )
-    tracker.loss = 1.0
-    tracker.update_s = 0.4
-    tracker.reduce_across_ranks()
-
-    assert captured["reduction"] == "max"
-    assert torch.allclose(captured["values"], torch.tensor([0.4]))
-    assert tracker.update_s.avg == pytest.approx(0.9)
-    # Metrics without a reduction stay untouched.
-    assert tracker.loss.avg == 1.0
-    # Invariant: avg == sum / count must hold after reduce, so subsequent .update() calls
-    # accumulate against the cluster view rather than the stale per-rank sum.
-    meter = tracker.update_s
-    assert meter.sum / meter.count == pytest.approx(meter.avg)

tests/utils/test_train_utils.py

@@ -20,8 +20,6 @@ from unittest.mock import Mock, patch
 from lerobot.common.train_utils import (
     get_step_checkpoint_dir,
     get_step_identifier,
-    load_training_batch_size,
-    load_training_num_processes,
     load_training_state,
     load_training_step,
     save_checkpoint,
@@ -65,28 +63,6 @@ def test_load_training_step(tmp_path):
     assert loaded_step == step
 
 
-def test_save_training_state_records_num_processes(tmp_path, optimizer, scheduler):
-    save_training_state(tmp_path, 10, optimizer, scheduler, num_processes=4)
-    assert load_training_num_processes(tmp_path) == 4
-
-
-def test_load_training_num_processes_absent_returns_none(tmp_path, optimizer, scheduler):
-    # Checkpoints written before the world size was recorded must still load (back-compat).
-    save_training_state(tmp_path, 10, optimizer, scheduler)
-    assert load_training_num_processes(tmp_path) is None
-
-
-def test_save_training_state_records_batch_size(tmp_path, optimizer, scheduler):
-    save_training_state(tmp_path, 10, optimizer, scheduler, batch_size=32)
-    assert load_training_batch_size(tmp_path) == 32
-
-
-def test_load_training_batch_size_absent_returns_none(tmp_path, optimizer, scheduler):
-    # Checkpoints written before the batch size was recorded must still load (back-compat).
-    save_training_state(tmp_path, 10, optimizer, scheduler)
-    assert load_training_batch_size(tmp_path) is None
-
-
 def test_update_last_checkpoint(tmp_path):
     checkpoint = tmp_path / "0005"
     checkpoint.mkdir()

uv.lock

@@ -59,7 +59,7 @@ wheels = [
 
 [[package]]
 name = "accelerate"
-version = "1.14.0"
+version = "1.13.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "huggingface-hub" },
@@ -71,9 +71,9 @@ dependencies = [
     { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
     { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/8d/75/94cd5d389649578aca399e5aa822637eec18319a1dadc400ffe2f9a7493f/accelerate-1.14.0.tar.gz", hash = "sha256:41b9c4377a54e0b460a959b0defa1b736e4ca0a2373252d9a539964c2afe3c8d", size = 412167, upload-time = "2026-06-11T13:45:52.326Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/ca/14/787e5498cd062640f0f3d92ef4ae4063174f76f9afd29d13fc52a319daae/accelerate-1.13.0.tar.gz", hash = "sha256:d631b4e0f5b3de4aff2d7e9e6857d164810dfc3237d54d017f075122d057b236", size = 402835, upload-time = "2026-03-04T19:34:12.359Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/a8/db/253133d7e7cb40d3af384bb2f5c0b4a2b7fdcffbc95c688cc67a20a3c103/accelerate-1.14.0-py3-none-any.whl", hash = "sha256:e94390c2863b873be18f623f9df48a0d8fe5eff13ea7f1a00092b0a7904888c6", size = 389246, upload-time = "2026-06-11T13:45:50.477Z" },
+    { url = "https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl", hash = "sha256:cf1a3efb96c18f7b152eb0fa7490f3710b19c3f395699358f08decca2b8b62e0", size = 383744, upload-time = "2026-03-04T19:34:10.313Z" },
 ]
 
 [[package]]
@@ -1084,8 +1084,8 @@ wheels = [
 
 [[package]]
 name = "datasets"
-version = "4.8.5"
-source = { registry = "https://pypi.org/simple" }
+version = "5.0.1.dev0"
+source = { git = "https://github.com/huggingface/datasets.git?rev=2c45eab1bb975ac3d846f2aa6217b82adec8eba3#2c45eab1bb975ac3d846f2aa6217b82adec8eba3" }
 dependencies = [
     { name = "dill" },
     { name = "filelock" },
@@ -1102,10 +1102,6 @@ dependencies = [
     { name = "tqdm" },
     { name = "xxhash" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/66/34/14cd8e76f907f7d4dca2334cfeec9f81d30fd15c25a015f99aaea694eaed/datasets-4.8.5.tar.gz", hash = "sha256:0f0c1c3d56ffff2c93b2f4c63c95bac94f3d7e8621aea2a2a576275233bba772", size = 605649, upload-time = "2026-04-27T15:43:57.384Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/65/99/00f3196036501b53032c4b1ab8337a0b978dee832ed276dae3815df4e8b5/datasets-4.8.5-py3-none-any.whl", hash = "sha256:5079900781719c0e063a8efdd2cd95a31ad0c63209178669cd23cf1b926149ff", size = 528973, upload-time = "2026-04-27T15:43:53.702Z" },
-]
 
 [[package]]
 name = "debugpy"
@@ -1207,15 +1203,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047, upload-time = "2025-07-17T16:51:58.613Z" },
 ]
 
-[[package]]
-name = "distro"
-version = "1.9.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/fc/f8/98eea607f65de6527f8a2e8885fc8015d3e6f5775df186e443e0964a11c3/distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed", size = 60722, upload-time = "2023-12-24T09:54:32.31Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" },
-]
-
 [[package]]
 name = "dm-control"
 version = "1.0.41"
@@ -2261,78 +2248,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" },
 ]
 
-[[package]]
-name = "jiter"
-version = "0.15.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/66/b5/55f06bb281d92fb3cc86d14e1def2bd908bb77693183e7cb1f5a3c388b0c/jiter-0.15.0.tar.gz", hash = "sha256:4251acc80e2b7c9b7b8823456ea0fceeb0734dac2df7636d3c711b38476b5a76", size = 166640, upload-time = "2026-05-19T10:09:48.361Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/44/53/4f6bddbcde3c71e56d0aa1337ec95950f3d27dd4153e25aadf0feac71751/jiter-0.15.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:0e90a1c315a0226ec822d973817967f9223b7701546c8c2a7913e7ab0926294d", size = 308793, upload-time = "2026-05-19T10:07:35.25Z" },
-    { url = "https://files.pythonhosted.org/packages/01/84/c01099b59a285a1ebba64ae93f62bfa036675340fd1b0045ae65890a0442/jiter-0.15.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8c9004af7c8d67cce7f1aae1026fb55607f4aa600710d08ede3a3ce4aeefe7e0", size = 309570, upload-time = "2026-05-19T10:07:36.919Z" },
-    { url = "https://files.pythonhosted.org/packages/58/64/8fb7f9d45bb98190355454cd04dad8d8f27223d6bd52f83af07f637168a6/jiter-0.15.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c210f8b35dc6f30aafd4b4365ca89b9d1189f21ab49b8e68fa6322a847aef138", size = 336783, upload-time = "2026-05-19T10:07:38.694Z" },
-    { url = "https://files.pythonhosted.org/packages/c3/b6/f5739011d009b3a30f6a53c5240979030ba29ae46a8c67e3a15759f7c37d/jiter-0.15.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5f30bae8bc1c2d613e28e5af3e8cceb09b742f1c8a8a5f839fb67afaffc03b61", size = 363555, upload-time = "2026-05-19T10:07:40.832Z" },
-    { url = "https://files.pythonhosted.org/packages/e5/12/98a9d9f766665e8a3b6252454e17cb0c464606a28cf2fa09399b003345fa/jiter-0.15.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c60e71b6d10cfc284c9bf36bd885e8d44c46f688ce50aa91b5edd90181dea687", size = 452255, upload-time = "2026-05-19T10:07:42.62Z" },
-    { url = "https://files.pythonhosted.org/packages/e8/d5/60f972840f79c5e7544fce567c56f1e4e50468f996baba3e78d823dd62a6/jiter-0.15.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0ab068bce62a45aa3e7367eceaffb5dde60b7eb853be8dece45132e3d0ff4879", size = 373559, upload-time = "2026-05-19T10:07:44.201Z" },
-    { url = "https://files.pythonhosted.org/packages/ee/cf/d46ef1234ba335aabc2f013210db8e0821a22f5e644a2e9449df199ecc23/jiter-0.15.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa248c9eb220197d363f688818dac2fd4b2f0cd7d843ca7105d652034823427d", size = 346055, upload-time = "2026-05-19T10:07:46.005Z" },
-    { url = "https://files.pythonhosted.org/packages/f0/63/4d2749d8d54d230bad9b3a6b0d00cc28c6ff6b2fdffc26a8ccf76cc5a974/jiter-0.15.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:2a77aadd57cac1682e4401a72724d2796d89a4ba129b1a5812aa94ee480826eb", size = 351406, upload-time = "2026-05-19T10:07:47.855Z" },
-    { url = "https://files.pythonhosted.org/packages/d9/b9/9965b990035d8773328e0a8c8b457a87bf2b19f6c4126d9d99296be5d16a/jiter-0.15.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2ae901f3a55bfafdde31d289590fa25e3245735a2b1e8c7cc15871710a002871", size = 389357, upload-time = "2026-05-19T10:07:49.665Z" },
-    { url = "https://files.pythonhosted.org/packages/2d/55/9ddf903deda1413e87fed792f416b7123daee5b8efbad6a202a7421c36a5/jiter-0.15.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:f0b271b462769543716f92d3a4f90527df6ef5ed05ee95ec4137f513e21e1b77", size = 517263, upload-time = "2026-05-19T10:07:51.537Z" },
-    { url = "https://files.pythonhosted.org/packages/e8/76/a0c40ad064d3a20a4fde231e35d56e9a01ce82164278180e82d5daf85469/jiter-0.15.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:2fb6a5d26af81fc0f00f9360a891e05cf755e149bba391c4d563adc54812973d", size = 548646, upload-time = "2026-05-19T10:07:53.196Z" },
-    { url = "https://files.pythonhosted.org/packages/23/4f/eca9b954942916ba2f453891b8593ab444cd872396fe66a3936616f236f3/jiter-0.15.0-cp312-cp312-win32.whl", hash = "sha256:c2f6bb8b5216ab9e7873bc08b5d7bef2b8abbb578a3069bf1cd14a45d71d771d", size = 206427, upload-time = "2026-05-19T10:07:55.307Z" },
-    { url = "https://files.pythonhosted.org/packages/95/bf/8ead82a87495149542748e828d153fd232a512a22c83b02c4815c1a9c7d8/jiter-0.15.0-cp312-cp312-win_amd64.whl", hash = "sha256:40b2c7e92c44a84d748d21706c68dc6ff8161d80b59c99d774721a0d2317d7c7", size = 197300, upload-time = "2026-05-19T10:07:56.651Z" },
-    { url = "https://files.pythonhosted.org/packages/f4/e4/9b8a78fb2d894471bc344e37f1949bdd784bd914d031dba0ba3a40c71dd7/jiter-0.15.0-cp312-cp312-win_arm64.whl", hash = "sha256:cc0bc345cf2df9d1c00ac443f50d543c1ccfa8b0422cb85b1ab70d681c0b255b", size = 192702, upload-time = "2026-05-19T10:07:58.307Z" },
-    { url = "https://files.pythonhosted.org/packages/e5/f4/f708c900ecee41b2025ef8413d5351e5649eb2125c506f6720cc69b06f5c/jiter-0.15.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:1c11465f97e2abf45a014b83b730222f8f1c5335e802c7055a67d50de6f1f4e3", size = 307829, upload-time = "2026-05-19T10:07:59.704Z" },
-    { url = "https://files.pythonhosted.org/packages/86/59/db537c0949e83668c38481d426b9f2fd5ab758c4ee53a811dd0a510626a0/jiter-0.15.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d1e7b1776f0797956c509e123d0952d10d293a9492dea9f288ab9570ec01d1a5", size = 308445, upload-time = "2026-05-19T10:08:01.184Z" },
-    { url = "https://files.pythonhosted.org/packages/37/38/ea0e13b18c30ef951da0d47d39e7fa9edb82a93a62990ffbd7cea9b622d4/jiter-0.15.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:351a341c2105aa430b7047e30f1bf7975f6313b00165d3fc07be2edaf741f279", size = 336181, upload-time = "2026-05-19T10:08:02.688Z" },
-    { url = "https://files.pythonhosted.org/packages/58/fc/2303901b16c4ba05865588990a420c0b4156270b44379c20931544a1d962/jiter-0.15.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4ab395feec8d249ec4044e228e98a7033f043426a265df439dc3698823f0a4e4", size = 362985, upload-time = "2026-05-19T10:08:04.394Z" },
-    { url = "https://files.pythonhosted.org/packages/5b/6f/11bace093c52e7d4d26c8e606ccd7ae8c972189622469ec0d9e28161e28b/jiter-0.15.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a2a438005b6f22d0273413484d6094d7c2c5d10ec1b3a3bf128e0d1d3ba53258", size = 453292, upload-time = "2026-05-19T10:08:05.967Z" },
-    { url = "https://files.pythonhosted.org/packages/22/db/987f2f086ca4d7a6582eb4ccd513f9b26b42d9e4243a087609a3137a8fc7/jiter-0.15.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f18f85e4218d1b40f000f42a92239a7a61a902cd42c65e6c360dbd17dcb20894", size = 373501, upload-time = "2026-05-19T10:08:07.857Z" },
-    { url = "https://files.pythonhosted.org/packages/8f/7c/89fbcabb2739b7a5b8dc959a1b6c5761f6484f5fed3486854b3c789bb1de/jiter-0.15.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d1aa62e277fc1cbd80e6deacae6f4d983b41b3d7728e0645c5d741a6149bba45", size = 344683, upload-time = "2026-05-19T10:08:09.431Z" },
-    { url = "https://files.pythonhosted.org/packages/30/6f/6cca7692e7dddfec6d8d76c54dc97f2af2a41df4ac0674b999df1f09a5f3/jiter-0.15.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:6550fa135c7deb8ead6af49ed7ff648532ea8334a1447fe34a36315ef79c5c29", size = 350892, upload-time = "2026-05-19T10:08:11.352Z" },
-    { url = "https://files.pythonhosted.org/packages/39/14/0338d6190cb8e6d22e677ab1d4eabd4117f67cca70c54cd04b82ff64e068/jiter-0.15.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:066f8f33f18b2419cd8213b2436fa7fbc9c499f315971cfa3ce1f9820c001b1b", size = 388723, upload-time = "2026-05-19T10:08:12.912Z" },
-    { url = "https://files.pythonhosted.org/packages/90/31/cc19f4a1bdb6afb09ce6a2f2615aa8d44d994eba0d8e6105ed1af920e736/jiter-0.15.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:75e8a04e91432dde9f1838373cf93d23726c79d3e908d319acf0e796f85592e7", size = 516648, upload-time = "2026-05-19T10:08:14.808Z" },
-    { url = "https://files.pythonhosted.org/packages/49/9f/833c541512cd091b63c10c0381973dfe11bc7a503a818c16384417e0c81e/jiter-0.15.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:a97261f1fccb8e50ecd2890a96e46efdc3f57c80a197324c6777827231eca712", size = 547382, upload-time = "2026-05-19T10:08:16.927Z" },
-    { url = "https://files.pythonhosted.org/packages/d2/11/e7b70e91f90bc4477e8eee9e8a5f7cf3cb41b4525d6394dc98a714eb8f7f/jiter-0.15.0-cp313-cp313-win32.whl", hash = "sha256:c77496cb10bd7549690fbbab3e5ec05857b83e49276f4a9423a766ddd2afcd4c", size = 205845, upload-time = "2026-05-19T10:08:18.401Z" },
-    { url = "https://files.pythonhosted.org/packages/4b/23/5c20d9ad6f02c493e4023e5d2d09e1c1f15fe2753c9102c544aff068a88e/jiter-0.15.0-cp313-cp313-win_amd64.whl", hash = "sha256:b15741f501469009ae0ae90b7147958a664a7dede40aa7ff174a8a4645f546d0", size = 196842, upload-time = "2026-05-19T10:08:20.131Z" },
-    { url = "https://files.pythonhosted.org/packages/6b/11/1eb400ef248e8c925fd883fbe325daf5e42cd1b0d308539dd332bd4f7ffc/jiter-0.15.0-cp313-cp313-win_arm64.whl", hash = "sha256:5d6a60072b44c3c2b797a7ddcbcbbf2b34ea3cfd4721580fbfd2a09d9d9b84ba", size = 192212, upload-time = "2026-05-19T10:08:21.807Z" },
-    { url = "https://files.pythonhosted.org/packages/8a/60/2fd8d7c79da8acf9b7b277c7616847773779356b92acfc9bb158452174da/jiter-0.15.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ef1fd24d9413f6209e00d3d5a453e67acfe004a25cc6c8e8484faed4311ab9e8", size = 315065, upload-time = "2026-05-19T10:08:23.218Z" },
-    { url = "https://files.pythonhosted.org/packages/46/f4/008fb7d65e8ac2abf00811651a661e025c4ba80bbc6f378450384ddd3aed/jiter-0.15.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:144f8e72cb53dab146347b91cceac01f5481237f2b93b4a339a1ee8f8878b67c", size = 339444, upload-time = "2026-05-19T10:08:24.701Z" },
-    { url = "https://files.pythonhosted.org/packages/00/55/90b0c7b9c6896c0f2a591dd36d36b71d22e09674bfef178fa03ba3f81499/jiter-0.15.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:553fcac2ef2cb990877f9fc0833b8b629a3e6a5670b6b5fd58219b41a653ddc4", size = 347779, upload-time = "2026-05-19T10:08:26.408Z" },
-    { url = "https://files.pythonhosted.org/packages/51/6b/69666cec5000fd57734c118437394516c749ae8dbeea9fb66d6fef9c4775/jiter-0.15.0-cp313-cp313t-win_amd64.whl", hash = "sha256:774f93f65031856bf14ad9f59bdcab8b8cad501e5ceabd51ba3525f76937a25b", size = 200395, upload-time = "2026-05-19T10:08:28.055Z" },
-    { url = "https://files.pythonhosted.org/packages/39/04/a6aa62cd27e8149b0d28df5561f10f6cceaf7935a9ccf3f1c5a05f9a0cd8/jiter-0.15.0-cp313-cp313t-win_arm64.whl", hash = "sha256:f1e1754960f38ec40613a07e5e372df67acb3b890fb383b6fb3de3e49ddbf3c7", size = 190516, upload-time = "2026-05-19T10:08:29.35Z" },
-    { url = "https://files.pythonhosted.org/packages/eb/d2/079f350ebf7859d081de30aa890f9e3be68516f754f3ba32366ffff4dcee/jiter-0.15.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:ac0d9ddea4350974be7a221fc25895f251a8fee748c889bdced2141c0fec1a49", size = 308884, upload-time = "2026-05-19T10:08:31.667Z" },
-    { url = "https://files.pythonhosted.org/packages/04/4e/a2c30a7f69b48c03b20935d647479106fe932f6e63f75faf53937197e05d/jiter-0.15.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:01a8222cf05ab1128e239421156c207949808acaaea2bdfd33130ae666786e86", size = 310028, upload-time = "2026-05-19T10:08:33.304Z" },
-    { url = "https://files.pythonhosted.org/packages/40/90/2e7cdfd3cf8ca967be38c48f5cf474d79f089efaf559a40f15984a77ae69/jiter-0.15.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:182226cbc930c9fab81bc2e41a4da672f89539906dadb05e75670ac07b94f71f", size = 337485, upload-time = "2026-05-19T10:08:35.259Z" },
-    { url = "https://files.pythonhosted.org/packages/9b/11/15a1aa28b120b8ee5b4f1fb894c125046225f09847738bd64233d3b84883/jiter-0.15.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:71683c38c825452999b5717fcae07ea708e8c93003e808be4319c1b02e3d176e", size = 364223, upload-time = "2026-05-19T10:08:36.694Z" },
-    { url = "https://files.pythonhosted.org/packages/b7/25/f442e8af5f3d0dcf47b39e83a0efd9ee45ea946aa6d04625dc3181eae3b6/jiter-0.15.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:30f2218e6a9e5c18bc10fe6d41ac189c442c88eacf11bad9f28ef95a9bef00e6", size = 456387, upload-time = "2026-05-19T10:08:38.143Z" },
-    { url = "https://files.pythonhosted.org/packages/da/f4/37f2d2c9f64f49af7da652ed7532bb5a2372e588e6927c3fdd76f911db65/jiter-0.15.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5157de9f76eb4bc5ea74a1219366a25f945ad305641d74e04f59c54087091aa9", size = 374461, upload-time = "2026-05-19T10:08:39.869Z" },
-    { url = "https://files.pythonhosted.org/packages/60/28/edcfbbbf0cb15436f36664a8908a0df47ab9006298d4cd937dc08ea932d6/jiter-0.15.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:90c5db5527c221249a876160663ab891ace358c17f7b9c93ec1478b7f0550e5c", size = 345924, upload-time = "2026-05-19T10:08:41.668Z" },
-    { url = "https://files.pythonhosted.org/packages/47/13/89fba6398dab7f202b7278c4b4aac122399d2c0183971c4a57a3b7088df5/jiter-0.15.0-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:3e4540b8e74e4268811ac05db226a6a128ff572e7e0ce3f1163b693cadb184cd", size = 352283, upload-time = "2026-05-19T10:08:43.091Z" },
-    { url = "https://files.pythonhosted.org/packages/1b/da/0f6af8cef2c565a1ab44d970f268c43ccaa72707386ea6388e6fe2b6cd26/jiter-0.15.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:62ebd14e47e9aed9df4472afcb2663668ce4d74891cd54f86bf6e44029d6dc89", size = 389985, upload-time = "2026-05-19T10:08:44.915Z" },
-    { url = "https://files.pythonhosted.org/packages/a1/ec/b9cb7d6d29e24ee14910266157d2a279d7a8f60ee0df7fa840882976ba64/jiter-0.15.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:0be6f5ad41a809f303f416d17cec92a7a725902fb9b4f3de3d19362ac0ef8554", size = 517695, upload-time = "2026-05-19T10:08:46.486Z" },
-    { url = "https://files.pythonhosted.org/packages/64/5e/6d1bda880723aae0ad86b4b763f044362448efe31e3e819635d41cb03451/jiter-0.15.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:813dfbb17d65328bf86e5f0905dd277ba2265d3ca20556e86c0c7035b7182e5a", size = 548868, upload-time = "2026-05-19T10:08:48.026Z" },
-    { url = "https://files.pythonhosted.org/packages/0c/72/7de501cf38dcacaf35098796f3a50e0f2e338baba18a58946c618544b809/jiter-0.15.0-cp314-cp314-win32.whl", hash = "sha256:50e51156192722a9c58db112837d3f8ef96fb3c5ecc14e95f409134b08b158ec", size = 206380, upload-time = "2026-05-19T10:08:49.738Z" },
-    { url = "https://files.pythonhosted.org/packages/1e/a9/e19addf4b0c1bdce52c6da12351e6bc42c340c45e7c09e2158e46d293ccc/jiter-0.15.0-cp314-cp314-win_amd64.whl", hash = "sha256:30ce1a5d16b5641dc935d50ef775af6a0871e3d14ab05d6fc54dff371b78e558", size = 197687, upload-time = "2026-05-19T10:08:51.088Z" },
-    { url = "https://files.pythonhosted.org/packages/f2/c9/776b1db01db25fc6c1d58d1979a37b0a9fe787e5f5b1d062d2eaacb77923/jiter-0.15.0-cp314-cp314-win_arm64.whl", hash = "sha256:510c8b3c17a0ed9ac69850c0438dada3c9b82d9c4d589fcb62002a5a9cf3a866", size = 192571, upload-time = "2026-05-19T10:08:52.451Z" },
-    { url = "https://files.pythonhosted.org/packages/a0/f6/45bb4670bacf300fd2c7abadbfb3af376e5f1b6ae75fd9bc069891d15870/jiter-0.15.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7553333dd0930c104a5a0db8df72bf7219fe663d731383b576bb6ed6351c984d", size = 317151, upload-time = "2026-05-19T10:08:53.867Z" },
-    { url = "https://files.pythonhosted.org/packages/d7/68/ed635ad5acd7b73e454283083bbb7c8205ad10e88b0d9d7d793b09fe8226/jiter-0.15.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2143ab06181d2b029eedcb6af3cebe95f11bbac62441781860f98ee9330a6a6", size = 341243, upload-time = "2026-05-19T10:08:55.383Z" },
-    { url = "https://files.pythonhosted.org/packages/5d/db/3ff4176b817b8ea33879e71e13d8bc2b0d481a7ed3fe9e080f333d415c16/jiter-0.15.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6eac374c5c975709b69c10f09afd199df74150172156ad10c8d4fd785b7da995", size = 363629, upload-time = "2026-05-19T10:08:56.928Z" },
-    { url = "https://files.pythonhosted.org/packages/ab/24/5f8270e0ba9c883582f96f722f8a0b58015c7ce1f8c6d4571cf394e99b6b/jiter-0.15.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b3b3b775e33d3bfaec9899edc526ae97b0da0bf9d071a46124ba419149a414f8", size = 456198, upload-time = "2026-05-19T10:08:58.618Z" },
-    { url = "https://files.pythonhosted.org/packages/45/5b/76fc02b0b5c54c3d18c60653156e2f76fde1816f9b4722db68d6ee2c897e/jiter-0.15.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:eda3071db3346334beae1360b46da4606da57bf3528c167b3c38533afaf9f2c5", size = 373710, upload-time = "2026-05-19T10:09:00.151Z" },
-    { url = "https://files.pythonhosted.org/packages/c4/52/4310821b0ea9277994d3e1f49fc6a4b34e4800caebacb2c0af81da59a454/jiter-0.15.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c6694a173ecabc12eb60efbc0b474464ead1951ff65cd8b1e72100715c64512b", size = 349901, upload-time = "2026-05-19T10:09:01.621Z" },
-    { url = "https://files.pythonhosted.org/packages/93/fe/67648c35b3594fba8854ac64cc8a826d8bcd18324bbdb53d77697c60b6ef/jiter-0.15.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:a254e10b593624d230c365b6d616b22ca0ad65e63a16e6631c2b3466022e6ba8", size = 352438, upload-time = "2026-05-19T10:09:03.216Z" },
-    { url = "https://files.pythonhosted.org/packages/cb/28/0a1879d07ad6b3e025a2750027363452ced93c2d16d1c9d4b153ffd51c91/jiter-0.15.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d8d2955167274e15d79a7a020afdd9b39c990eb80b2d89fca695d92dcfdd38ec", size = 388152, upload-time = "2026-05-19T10:09:04.741Z" },
-    { url = "https://files.pythonhosted.org/packages/c1/78/46c6f6b56ba85c90021f4afd72ed42f691f8f84daacb5fe27277070e3858/jiter-0.15.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:acf4ee4d1fc55917239fe72972fb292dd773055d05eb040d36f4326e02cc2c0e", size = 517707, upload-time = "2026-05-19T10:09:06.231Z" },
-    { url = "https://files.pythonhosted.org/packages/ca/cb/720662d4c88fcad606e826fef5424365527ba43ce4868a479aed8f8c507e/jiter-0.15.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:e7196e56f1cd69af1dbb07dff02dcfb260a50b45a82d409d92a06fedb32473b5", size = 548241, upload-time = "2026-05-19T10:09:08.093Z" },
-    { url = "https://files.pythonhosted.org/packages/60/e3/935b8034fd143f21125c87d51404a9e0e1449186a494405721ff5d1d695e/jiter-0.15.0-cp314-cp314t-win32.whl", hash = "sha256:7f6163c0f10b055245f814dcc59f4818da60dfe72f3e72ab89fc24b6bd5e9c52", size = 207950, upload-time = "2026-05-19T10:09:09.616Z" },
-    { url = "https://files.pythonhosted.org/packages/93/59/984fd9ece895953dad3e0880a650e766f5a2da2c5514f0eafdaaabbeb5f9/jiter-0.15.0-cp314-cp314t-win_amd64.whl", hash = "sha256:980c256edb05b78a111b99c4de3b1d32e31634b867fd1fc2cf726e7b7bba9854", size = 200055, upload-time = "2026-05-19T10:09:11.367Z" },
-    { url = "https://files.pythonhosted.org/packages/0e/a4/cf8d779feb133a27a2e3bc833bccb9e13aa332cdf820497ebf72c10ce8c3/jiter-0.15.0-cp314-cp314t-win_arm64.whl", hash = "sha256:66b1880df2d01e206e8339769d1c7c1753bcb653efd6289e203f6f24ebada0c0", size = 191244, upload-time = "2026-05-19T10:09:12.74Z" },
-    { url = "https://files.pythonhosted.org/packages/73/38/505941b2b092fd5bbbd60a52a880db1173f1690ae6751bed3af1c9ddcb4e/jiter-0.15.0-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:631f13a3d04e97d4e083993b10f4b99530e3a10d953e2eb5e196b7dc7f812ce0", size = 303769, upload-time = "2026-05-19T10:09:42.203Z" },
-    { url = "https://files.pythonhosted.org/packages/e7/95/a06692b29e77473f286e1ec1f426d3ca44d7b5843be8ad21d7a5f3fcdcc0/jiter-0.15.0-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:b6c0ffae686c39bf3737be60793783267628783ea42545632c10b291105aee45", size = 305128, upload-time = "2026-05-19T10:09:43.657Z" },
-    { url = "https://files.pythonhosted.org/packages/23/85/7270d7ad41d6061a25b950c6bf91d638bd9aacb113200a8c8d57a055fd67/jiter-0.15.0-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1d54fb5b31dea401a41af3f8a7d2512e9b6a6a005491e6166c7e4ffab9639a9c", size = 340459, upload-time = "2026-05-19T10:09:45.452Z" },
-    { url = "https://files.pythonhosted.org/packages/c8/8d/302cb2057b7513327b4d575cff6b1d066ee6431a5357fc3f8867cd684406/jiter-0.15.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:54d5d6090cdc1b7c9e780dfb04949a990adb1e301a2fc0bbcee7de4638d33f9a", size = 344469, upload-time = "2026-05-19T10:09:46.864Z" },
-]
-
 [[package]]
 name = "json5"
 version = "0.14.0"
@@ -2768,9 +2683,6 @@ dependencies = [
 ]
 
 [package.optional-dependencies]
-accelerate-dep = [
-    { name = "accelerate" },
-]
 all = [
     { name = "accelerate" },
     { name = "av" },
@@ -2841,16 +2753,6 @@ aloha = [
     { name = "scipy" },
     { name = "torchcodec", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'arm64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" },
 ]
-annotations = [
-    { name = "av" },
-    { name = "datasets" },
-    { name = "jsonlines" },
-    { name = "openai" },
-    { name = "pandas" },
-    { name = "pyarrow" },
-    { name = "torchcodec", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'arm64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" },
-    { name = "transformers" },
-]
 async = [
     { name = "contourpy" },
     { name = "grpcio" },
@@ -3080,8 +2982,6 @@ qwen-vl-utils-dep = [
     { name = "qwen-vl-utils" },
 ]
 reachy2 = [
-    { name = "grpcio" },
-    { name = "protobuf" },
     { name = "reachy2-sdk" },
 ]
 rebot = [
@@ -3169,11 +3069,12 @@ xvla = [
 
 [package.metadata]
 requires-dist = [
-    { name = "accelerate", marker = "extra == 'accelerate-dep'", specifier = ">=1.14.0,<2.0.0" },
+    { name = "accelerate", marker = "extra == 'smolvla'", specifier = ">=1.7.0,<2.0.0" },
+    { name = "accelerate", marker = "extra == 'training'", specifier = ">=1.10.0,<2.0.0" },
     { name = "av", marker = "extra == 'av-dep'", specifier = ">=15.0.0,<16.0.0" },
     { name = "cmake", specifier = ">=3.29.0.1,<4.2.0" },
     { name = "contourpy", marker = "extra == 'matplotlib-dep'", specifier = ">=1.3.0,<2.0.0" },
-    { name = "datasets", marker = "extra == 'dataset'", specifier = ">=4.7.0,<5.0.0" },
+    { name = "datasets", marker = "extra == 'dataset'", git = "https://github.com/huggingface/datasets.git?rev=2c45eab1bb975ac3d846f2aa6217b82adec8eba3" },
     { name = "debugpy", marker = "extra == 'dev'", specifier = ">=1.8.1,<1.9.0" },
     { name = "decord", marker = "(platform_machine == 'AMD64' and extra == 'groot') or (platform_machine == 'x86_64' and extra == 'groot')", specifier = ">=0.6.0,<1.0.0" },
     { name = "deepdiff", marker = "extra == 'deepdiff-dep'", specifier = ">=7.0.1,<9.0.0" },
@@ -3186,9 +3087,8 @@ requires-dist = [
     { name = "fastapi", marker = "extra == 'phone'", specifier = "<1.0" },
     { name = "feetech-servo-sdk", marker = "extra == 'feetech'", specifier = ">=1.0.0,<2.0.0" },
     { name = "flash-attn", marker = "sys_platform != 'darwin' and extra == 'groot'", specifier = ">=2.5.9,<3.0.0" },
-    { name = "grpcio", marker = "extra == 'grpcio-dep'", specifier = ">=1.73.1,<2.0.0" },
-    { name = "grpcio", marker = "extra == 'reachy2'", specifier = "<=1.73.1" },
-    { name = "grpcio-tools", marker = "extra == 'dev'", specifier = ">=1.73.1,<2.0.0" },
+    { name = "grpcio", marker = "extra == 'grpcio-dep'", specifier = "==1.73.1" },
+    { name = "grpcio-tools", marker = "extra == 'dev'", specifier = "==1.73.1" },
     { name = "gym-aloha", marker = "extra == 'aloha'", specifier = ">=0.1.4,<0.2.0" },
     { name = "gym-hil", marker = "extra == 'hilserl'", specifier = ">=0.1.14,<0.2.0" },
     { name = "gym-pusht", marker = "extra == 'pusht'", specifier = ">=0.1.5,<0.2.0" },
@@ -3200,8 +3100,6 @@ requires-dist = [
     { name = "ipykernel", marker = "extra == 'notebook'", specifier = ">=6.0.0,<7.0.0" },
     { name = "jsonlines", marker = "extra == 'dataset'", specifier = ">=4.0.0,<5.0.0" },
     { name = "jupyter", marker = "extra == 'notebook'", specifier = ">=1.0.0,<2.0.0" },
-    { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'smolvla'" },
-    { name = "lerobot", extras = ["accelerate-dep"], marker = "extra == 'training'" },
     { name = "lerobot", extras = ["aloha"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["async"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["av-dep"], marker = "extra == 'dataset'" },
@@ -3212,7 +3110,6 @@ requires-dist = [
     { name = "lerobot", extras = ["damiao"], marker = "extra == 'openarms'" },
     { name = "lerobot", extras = ["dataset"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["dataset"], marker = "extra == 'aloha'" },
-    { name = "lerobot", extras = ["dataset"], marker = "extra == 'annotations'" },
     { name = "lerobot", extras = ["dataset"], marker = "extra == 'core-scripts'" },
     { name = "lerobot", extras = ["dataset"], marker = "extra == 'dataset-viz'" },
     { name = "lerobot", extras = ["dataset"], marker = "extra == 'hilserl'" },
@@ -3297,7 +3194,6 @@ requires-dist = [
     { name = "lerobot", extras = ["test"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["topreward"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["training"], marker = "extra == 'all'" },
-    { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'annotations'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'eo1'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'groot'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'hilserl'" },
@@ -3332,7 +3228,6 @@ requires-dist = [
     { name = "numpy", specifier = ">=2.0.0,<2.3.0" },
     { name = "onnx", marker = "extra == 'unitree-g1'", specifier = ">=1.16.0,<2.0.0" },
     { name = "onnxruntime", marker = "extra == 'unitree-g1'", specifier = ">=1.16.0,<2.0.0" },
-    { name = "openai", marker = "extra == 'annotations'", specifier = ">=1.40,<2.0" },
     { name = "opencv-python-headless", specifier = ">=4.9.0,<4.14.0" },
     { name = "packaging", specifier = ">=24.2,<26.0" },
     { name = "pandas", marker = "extra == 'dataset'", specifier = ">=2.0.0,<3.0.0" },
@@ -3341,8 +3236,7 @@ requires-dist = [
     { name = "pillow", specifier = ">=10.0.0,<13.0.0" },
     { name = "placo", marker = "extra == 'placo-dep'", specifier = ">=0.9.6,<0.9.16" },
     { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.7.0,<5.0.0" },
-    { name = "protobuf", marker = "extra == 'grpcio-dep'", specifier = ">=6.31.1,<8.0.0" },
-    { name = "protobuf", marker = "extra == 'reachy2'", specifier = "<=6.32.0" },
+    { name = "protobuf", marker = "extra == 'grpcio-dep'", specifier = ">=6.31.1,<6.32.0" },
     { name = "pyarrow", marker = "extra == 'dataset'", specifier = ">=21.0.0,<30.0.0" },
     { name = "pydantic", marker = "extra == 'sarm'", specifier = ">=2.0.0,<3.0.0" },
     { name = "pygame", marker = "extra == 'pygame-dep'", specifier = ">=2.5.1,<2.7.0" },
@@ -3379,9 +3273,9 @@ requires-dist = [
     { name = "torchvision", marker = "sys_platform == 'linux'", specifier = ">=0.22.0,<0.27.0", index = "https://download.pytorch.org/whl/cu128" },
     { name = "tqdm", specifier = ">=4.66.0,<5.0.0" },
     { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
-    { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" },
+    { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.25.0" },
 ]
-provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
+provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "hilserl", "vla-jepa", "async", "peft", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
 
 [[package]]
 name = "librt"
@@ -4474,25 +4368,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/93/99/f2be40a31b908d96b861ae0ce98582fa376c18a7f816b9d5eb4cd6aa0a4c/onnxruntime-1.26.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4eefd386a45202aefb7a5132b94f32df9d506c9edcc7faf2fc60d65183f4b183", size = 18197382, upload-time = "2026-05-08T19:07:46.965Z" },
 ]
 
-[[package]]
-name = "openai"
-version = "1.109.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "anyio" },
-    { name = "distro" },
-    { name = "httpx" },
-    { name = "jiter" },
-    { name = "pydantic" },
-    { name = "sniffio" },
-    { name = "tqdm" },
-    { name = "typing-extensions" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/c6/a1/a303104dc55fc546a3f6914c842d3da471c64eec92043aef8f652eb6c524/openai-1.109.1.tar.gz", hash = "sha256:d173ed8dbca665892a6db099b4a2dfac624f94d20a93f46eb0b56aae940ed869", size = 564133, upload-time = "2025-09-24T13:00:53.075Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/1d/2a/7dd3d207ec669cacc1f186fd856a0f61dbc255d24f6fdc1a6715d6051b0f/openai-1.109.1-py3-none-any.whl", hash = "sha256:6bcaf57086cf59159b8e27447e4e7dd019db5d29a438072fbd49c290c7e65315", size = 948627, upload-time = "2025-09-24T13:00:50.754Z" },
-]
-
 [[package]]
 name = "opencv-python"
 version = "4.13.0.92"
@@ -6211,15 +6086,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" },
 ]
 
-[[package]]
-name = "sniffio"
-version = "1.3.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372, upload-time = "2024-02-25T23:20:04.057Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" },
-]
-
 [[package]]
 name = "soupsieve"
 version = "2.8.3"