Merge branch 'feat/language-annotation-pipeline' into feat/smolvla-on-steerable

Resolves conflicts from 66 commits on the base branch:

* pyproject.toml — keep base's transformers>=5.4.0,<5.6.0; add the
  sentencepiece-dep entry pi052 (FAST action tokenizer) needs.
* policies/__init__.py — keep pi052 export; drop the
  RewardClassifierConfig export that base removed.
* policies/factory.py — docstring list resolution (keep pi052; drop
  reward_classifier, removed by base).
* annotations/steerable_pipeline/executor.py — adopt base's renamed
  _ensure_annotation_metadata_in_info (it already advertises the say
  tool); drop pi052's older _ensure_tools_in_info call.
* configs/train.py — keep pi052's vqa_target_fraction; adopt base's
  SampleWeightingConfig (legacy RA-BC inline params already covered
  by the migration shim base added).
* scripts/lerobot_train.py — merge pi052's per-policy processor
  rebuild + dataset_repo_id pass-through with base's active_cfg /
  is_reward_model_training tightening, and re-route vqa-weighted
  sampler to active_cfg.drop_n_last_frames.
* datasets/language_render.py — adopt base's _select_one + timestamp
  tolerance (drops pi052's stale _select_latest / per-style sort_key).
* tests — adopt base's parametrized per-camera blend + tolerance
  test; drop pi052 tests that overlap with base's tighter rewrites;
  keep pi052's flow-only / VQA-blend coverage; add a
  test_canonical_recipe_loads check on subtask_mem_vqa_speech.yaml.
* policies/pi052/processor_pi052.py — import RenderMessagesStep
  directly from render_messages_processor (base intentionally
  dropped it from lerobot.processor's re-exports).
* uv.lock — regenerated cleanly from base + pi052's pocket-tts /
  beartype.

All 67 touched tests pass (30 pi052 + 37 recipe / language-render /
pipeline / render-messages).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

tests/annotations/conftest.py

@@ -15,85 +15,24 @@
 # limitations under the License.
 """Shared fixtures for annotation-pipeline tests.
 
-Builds a minimal LeRobot-shaped dataset on disk so writer/validator tests
-can exercise real parquet reads and writes without needing a checked-in
-LFS dataset.
+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
 
-import json
 from pathlib import Path
 
-import pyarrow as pa
-import pyarrow.parquet as pq
 import pytest
 
-
-def _make_episode_table(
-    episode_index: int,
-    num_frames: int,
-    *,
-    fps: int = 10,
-    task_index: int = 0,
-) -> pa.Table:
-    timestamps = [round(i / fps, 6) for i in range(num_frames)]
-    frame_indices = list(range(num_frames))
-    return pa.Table.from_pydict(
-        {
-            "episode_index": [episode_index] * num_frames,
-            "frame_index": frame_indices,
-            "timestamp": timestamps,
-            "task_index": [task_index] * num_frames,
-            "subtask_index": [0] * num_frames,  # legacy column the writer must drop
-        }
-    )
-
-
-def _build_dataset(root: Path, episode_specs: list[tuple[int, int, str]], *, fps: int = 10) -> Path:
-    """Create a fixture dataset under ``root``.
-
-    ``episode_specs`` is a list of ``(episode_index, num_frames, task_text)``.
-    Each episode goes into its own ``data/chunk-000/file-{ep:03d}.parquet``
-    so the writer's per-shard rewrite path is exercised.
-    """
-    data_dir = root / "data" / "chunk-000"
-    data_dir.mkdir(parents=True, exist_ok=True)
-    tasks = {}
-    for episode_index, num_frames, task_text in episode_specs:
-        task_index = len(tasks)
-        if task_text not in tasks.values():
-            tasks[task_index] = task_text
-        else:
-            task_index = next(k for k, v in tasks.items() if v == task_text)
-        table = _make_episode_table(episode_index, num_frames, fps=fps, task_index=task_index)
-        path = data_dir / f"file-{episode_index:03d}.parquet"
-        pq.write_table(table, path)
-
-    meta_dir = root / "meta"
-    meta_dir.mkdir(parents=True, exist_ok=True)
-    tasks_table = pa.Table.from_pydict(
-        {
-            "task_index": list(tasks.keys()),
-            "task": list(tasks.values()),
-        }
-    )
-    pq.write_table(tasks_table, meta_dir / "tasks.parquet")
-
-    info = {
-        "codebase_version": "v3.1",
-        "fps": fps,
-        "total_episodes": len(episode_specs),
-    }
-    (meta_dir / "info.json").write_text(json.dumps(info, indent=2))
-
-    return root
+from tests.fixtures.dataset_factories import build_annotation_dataset
 
 
 @pytest.fixture
 def fixture_dataset_root(tmp_path: Path) -> Path:
     """A tiny dataset with two episodes, 12 frames each at 10 fps."""
-    return _build_dataset(
+    return build_annotation_dataset(
         tmp_path / "ds",
         episode_specs=[
             (0, 12, "Could you tidy the kitchen please?"),
@@ -105,7 +44,7 @@ def fixture_dataset_root(tmp_path: Path) -> Path:
 
 @pytest.fixture
 def single_episode_root(tmp_path: Path) -> Path:
-    return _build_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

@@ -15,22 +15,19 @@
 # limitations under the License.
 """Opt-in E2E smoke run for ``make annotation-e2e``.
 
-Builds the same fixture used by the pytest suite, 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
-without depending on conftest.py fixtures.
+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 json
 import sys
 import tempfile
 from pathlib import Path
 
-import pyarrow as pa
-import pyarrow.parquet as pq
-
 from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
 from lerobot.annotations.steerable_pipeline.executor import Executor
 from lerobot.annotations.steerable_pipeline.modules import (
@@ -41,31 +38,7 @@ from lerobot.annotations.steerable_pipeline.modules import (
 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
-
-
-def _build_dataset(root: Path) -> Path:
-    data_dir = root / "data" / "chunk-000"
-    data_dir.mkdir(parents=True, exist_ok=True)
-    n = 30
-    timestamps = [round(i / 10, 6) for i in range(n)]
-    table = pa.Table.from_pydict(
-        {
-            "episode_index": [0] * n,
-            "frame_index": list(range(n)),
-            "timestamp": timestamps,
-            "task_index": [0] * n,
-            "subtask_index": [0] * n,
-        }
-    )
-    pq.write_table(table, data_dir / "file-000.parquet")
-    meta = root / "meta"
-    meta.mkdir(parents=True, exist_ok=True)
-    pq.write_table(
-        pa.Table.from_pydict({"task_index": [0], "task": ["Pour water into the cup."]}),
-        meta / "tasks.parquet",
-    )
-    (meta / "info.json").write_text(json.dumps({"codebase_version": "v3.1", "fps": 10}))
-    return root
+from tests.fixtures.dataset_factories import build_annotation_dataset
 
 
 def _stub_responder(messages):
@@ -102,14 +75,18 @@ def _stub_responder(messages):
 
 def main() -> int:
     with tempfile.TemporaryDirectory() as tmp:
-        root = _build_dataset(Path(tmp) / "ds")
+        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,
-            module_1=PlanSubtasksMemoryModule(vlm=vlm, config=cfg.module_1),
-            module_2=InterjectionsAndSpeechModule(vlm=vlm, config=cfg.module_2, seed=cfg.seed),
-            module_3=GeneralVqaModule(vlm=vlm, config=cfg.module_3, seed=cfg.seed),
+            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(),
         )

tests/annotations/test_frames.py

@@ -0,0 +1,146 @@
+#!/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 (  # noqa: E402
+    VideoFrameProvider,
+    _decode_frames_av,
+    _decode_frames_ffmpeg,
+)
+
+
+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 test_decode_frames_av_returns_one_uint8_frame_per_timestamp(sample_video: Path) -> None:
+    """``_decode_frames_av`` decodes via PyAV directly — no torchcodec/torchvision.
+
+    This is the always-available fallback: torchcodec is unusable in some
+    containers and lerobot's ``pyav`` backend routes through the removed
+    ``torchvision.io.VideoReader``.
+    """
+    timestamps = [0.0, 1.0, 2.5]
+    frames = _decode_frames_av(sample_video, timestamps)
+
+    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_decode_frames_av_picks_nearest_frame(sample_video: Path) -> None:
+    """Repeated and out-of-order timestamps each resolve to the nearest frame."""
+    frames = _decode_frames_av(sample_video, [2.0, 0.0, 2.0])
+
+    assert len(frames) == 3
+    assert torch.equal(frames[0], frames[2])
+    assert not torch.equal(frames[0], frames[1])
+
+
+def test_decode_frames_av_raises_on_missing_file(tmp_path: Path) -> None:
+    """A missing video surfaces as an exception the caller can fall back on."""
+    with pytest.raises(Exception):  # noqa: B017, PT011
+        _decode_frames_av(tmp_path / "does_not_exist.mp4", [0.0])
+
+
+def test_decode_frames_ffmpeg_returns_one_uint8_frame_per_timestamp(sample_video: Path) -> None:
+    """``_decode_frames_ffmpeg`` shells out to the ffmpeg CLI — the always-
+    available fallback that decodes AV1 and isolates crashes to a child
+    process.
+    """
+    timestamps = [0.0, 1.0, 2.5]
+    frames = _decode_frames_ffmpeg(sample_video, timestamps)
+
+    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_decode_frames_ffmpeg_raises_on_missing_file(tmp_path: Path) -> None:
+    """A missing video raises (non-zero ffmpeg exit), never crashes the job."""
+    if shutil.which("ffmpeg") is None:
+        pytest.skip("ffmpeg not available")
+    with pytest.raises(Exception):  # noqa: B017, PT011
+        _decode_frames_ffmpeg(tmp_path / "does_not_exist.mp4", [0.0])

tests/annotations/test_modules.py

@@ -23,9 +23,9 @@ from pathlib import Path
 from typing import Any
 
 from lerobot.annotations.steerable_pipeline.config import (
-    Module1Config,
-    Module2Config,
-    Module3Config,
+    InterjectionsConfig,
+    PlanConfig,
+    VqaConfig,
 )
 from lerobot.annotations.steerable_pipeline.modules import (
     GeneralVqaModule,
@@ -80,15 +80,14 @@ def test_module1_plan_memory_subtask_smoke(fixture_dataset_root: Path, tmp_path:
                     {"text": "place the sponge into the sink", "start": 0.8, "end": 1.1},
                 ]
             },
-            "concise hierarchical PLAN": {"plan": "1. grasp\n2. wipe\n3. place"},
             "Update the memory": {"memory": "wiped the counter once"},
         },
     )
-    module = PlanSubtasksMemoryModule(vlm=vlm, config=Module1Config())
+    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("module_1")
+    rows = staging.read("plan")
 
     styles = {r["style"] for r in rows}
     assert {"subtask", "plan", "memory"}.issubset(styles)
@@ -96,10 +95,16 @@ def test_module1_plan_memory_subtask_smoke(fixture_dataset_root: Path, tmp_path:
     frame_set = set(record.frame_timestamps)
     for row in rows:
         assert row["timestamp"] in frame_set
-    # exactly one plan row at t0
-    plan_rows = [r for r in rows if r["style"] == "plan"]
-    assert len(plan_rows) == 1
+    # 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_module2_at_t0_emits_speech_only_no_interjection(fixture_dataset_root: Path, tmp_path: Path) -> None:
@@ -108,12 +113,12 @@ def test_module2_at_t0_emits_speech_only_no_interjection(fixture_dataset_root: P
     )
     module = InterjectionsAndSpeechModule(
         vlm=vlm,
-        config=Module2Config(max_interjections_per_episode=0),
+        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("module_2")
+    rows = staging.read("interjections")
     assert len(rows) == 1
     only = rows[0]
     assert only["role"] == "assistant"
@@ -126,24 +131,61 @@ def test_module2_at_t0_emits_speech_only_no_interjection(fixture_dataset_root: P
 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."},
-            "ONE realistic interruption": {
-                "interjection": "actually skip the dishes",
-                "speech": "Skipping the dishes.",
+            # Marker matches the distinctive line of
+            # ``module_2_interjection.txt``. The old marker
+            # ("ONE realistic interruption") came from a previous prompt
+            # version that asked for counterfactual interjections; the
+            # current design anchors on subtask boundaries instead, so
+            # the prompt and its marker changed.
+            "Write ONE interjection": {
+                "interjection": "now wipe the counter please",
+                "speech": "On it.",
             },
         },
     )
     module = InterjectionsAndSpeechModule(
         vlm=vlm,
-        config=Module2Config(max_interjections_per_episode=1, interjection_min_t=0.2),
+        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("module_2")
+    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"]
@@ -161,28 +203,22 @@ def test_module3_vqa_unique_per_frame_and_camera(single_episode_root: Path, tmp_
     vlm = make_canned_responder({"frame-grounded visual question": payload})
     module = GeneralVqaModule(
         vlm=vlm,
-        config=Module3Config(vqa_emission_hz=1.0, K=3),
+        config=VqaConfig(vqa_emission_hz=1.0, K=3),
         seed=1,
-        frame_provider=_StubFrameProvider(
-            cameras=("observation.images.top", "observation.images.wrist")
-        ),
+        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("module_3")
+    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"
-    ]
+    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"
+        (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))
@@ -222,17 +258,32 @@ def test_module1_attaches_video_block_to_subtask_prompt(fixture_dataset_root: Pa
     provider = _StubFrameProvider()
     module = PlanSubtasksMemoryModule(
         vlm=StubVlmClient(responder=responder),
-        config=Module1Config(max_video_frames=5, frames_per_second=10.0),
+        # 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(max_video_frames=5, frames_per_second=10.0, 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)
 
-    # the subtask call (the first VLM call) must carry exactly one video block
+    # 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"
-    first_call = captured[0]
-    content = first_call[0]["content"]
+
+    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"]
@@ -258,7 +309,7 @@ def test_module3_attaches_frame_image_block_to_prompt(single_episode_root: Path,
     provider = _StubFrameProvider()
     module = GeneralVqaModule(
         vlm=_spy_responder(captured, payload),
-        config=Module3Config(vqa_emission_hz=1.0, K=1),
+        config=VqaConfig(vqa_emission_hz=1.0, K=1),
         seed=0,
         frame_provider=provider,
     )
@@ -290,14 +341,14 @@ def test_module3_assistant_content_is_valid_json(single_episode_root: Path, tmp_
     vlm = make_canned_responder({"frame-grounded visual question": payload})
     module = GeneralVqaModule(
         vlm=vlm,
-        config=Module3Config(vqa_emission_hz=1.0, K=2),
+        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("module_3")
+    rows = staging.read("vqa")
     for row in rows:
         if row["role"] == "assistant" and row["style"] == "vqa":
             decoded = json.loads(row["content"])

tests/annotations/test_pipeline_recipe_render.py

@@ -23,9 +23,9 @@ import pyarrow.parquet as pq
 
 from lerobot.annotations.steerable_pipeline.config import (
     AnnotationPipelineConfig,
-    Module1Config,
-    Module2Config,
-    Module3Config,
+    InterjectionsConfig,
+    PlanConfig,
+    VqaConfig,
 )
 from lerobot.annotations.steerable_pipeline.executor import Executor
 from lerobot.annotations.steerable_pipeline.modules import (
@@ -35,19 +35,59 @@ from lerobot.annotations.steerable_pipeline.modules import (
 )
 from lerobot.annotations.steerable_pipeline.validator import StagingValidator
 from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
-from lerobot.configs.recipe import TrainingRecipe
+from lerobot.configs.recipe import MessageTurn, TrainingRecipe
 from lerobot.datasets.language_render import render_sample
 
 from ._helpers import make_canned_responder
 
-_RECIPE_PATH = (
-    Path(__file__).resolve().parents[2]
-    / "src"
-    / "lerobot"
-    / "configs"
-    / "recipes"
-    / "subtask_mem_vqa_speech.yaml"
-)
+def _build_pr1_style_blend_recipe() -> TrainingRecipe:
+    """Inline blend recipe that consumes every style this pipeline produces.
+
+    PR 1 used to ship ``src/lerobot/configs/recipes/pi05_hirobot.yaml`` as
+    a canonical example, but that file was dropped during PR 1 review. The
+    cross-PR 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:
@@ -74,15 +114,15 @@ def _build_executor() -> Executor:
         },
     )
     config = AnnotationPipelineConfig(
-        module_1=Module1Config(),
-        module_2=Module2Config(max_interjections_per_episode=1, interjection_min_t=0.5),
-        module_3=Module3Config(vqa_emission_hz=1.0, K=2),
+        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,
-        module_1=PlanSubtasksMemoryModule(vlm=vlm, config=config.module_1),
-        module_2=InterjectionsAndSpeechModule(vlm=vlm, config=config.module_2, seed=config.seed),
-        module_3=GeneralVqaModule(vlm=vlm, config=config.module_3, seed=config.seed),
+        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(),
     )
@@ -101,13 +141,7 @@ def test_pr1_canonical_recipe_renders_nonempty_from_pipeline_output(
     events_lists = table.column("language_events").to_pylist()
     timestamps = table.column("timestamp").to_pylist()
 
-    recipe = TrainingRecipe.from_yaml(_RECIPE_PATH) if hasattr(TrainingRecipe, "from_yaml") else None
-    if recipe is None:
-        # PR 1 may not expose from_yaml; load via PyYAML and TrainingRecipe(**...)
-        import yaml
-
-        loaded = yaml.safe_load(_RECIPE_PATH.read_text(encoding="utf-8"))
-        recipe = TrainingRecipe(**loaded)
+    recipe = _build_pr1_style_blend_recipe()
 
     rendered_any = False
     for sample_idx, (ts, persistent, events) in enumerate(

tests/annotations/test_validator.py

@@ -34,7 +34,7 @@ def _validate(root: Path, staging_dir: Path):
 def test_validator_catches_misaligned_timestamps(fixture_dataset_root: Path, tmp_path: Path) -> None:
     staging_dir = tmp_path / "stage"
     EpisodeStaging(staging_dir, 0).write(
-        "module_3",
+        "vqa",
         [
             {
                 "role": "assistant",
@@ -53,7 +53,7 @@ def test_validator_catches_misaligned_timestamps(fixture_dataset_root: Path, tmp
 def test_validator_catches_orphan_speech(fixture_dataset_root: Path, tmp_path: Path) -> None:
     staging_dir = tmp_path / "stage"
     EpisodeStaging(staging_dir, 0).write(
-        "module_2",
+        "interjections",
         [
             speech_atom(0.0, "Got it."),
             # interjection at 0.3s with NO paired speech
@@ -74,7 +74,7 @@ def test_validator_catches_orphan_speech(fixture_dataset_root: Path, tmp_path: P
 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(
-        "module_1",
+        "plan",
         [
             {
                 "role": "assistant",
@@ -93,7 +93,7 @@ def test_validator_catches_inconsistent_plan_memory(fixture_dataset_root: Path,
         ],
     )
     EpisodeStaging(staging_dir, 0).write(
-        "module_2",
+        "interjections",
         [
             speech_atom(0.0, "Got it."),
             speech_atom(0.4, "Replanning."),
@@ -115,11 +115,11 @@ def test_validator_catches_inconsistent_plan_memory(fixture_dataset_root: Path,
 def test_validator_catches_wrong_column(fixture_dataset_root: Path, tmp_path: Path) -> None:
     staging_dir = tmp_path / "stage"
     EpisodeStaging(staging_dir, 0).write(
-        "module_1",
+        "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("module_1 emitted style 'vqa'" in e or "must be persistent" in e for e in report.errors)
+    assert any("plan emitted style 'vqa'" in e or "must be persistent" in e for e in report.errors)

tests/annotations/test_vocabulary.py

@@ -0,0 +1,412 @@
+#!/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.
+"""Vocabulary-discovery phase (phase 0) tests."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from lerobot.annotations.steerable_pipeline.config import (
+    PlanConfig,
+    VocabularyConfig,
+)
+from lerobot.annotations.steerable_pipeline.modules import PlanSubtasksMemoryModule
+from lerobot.annotations.steerable_pipeline.reader import iter_episodes
+from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging
+from lerobot.annotations.steerable_pipeline.vocabulary import (
+    Vocabulary,
+    VocabularyDiscoveryModule,
+    load_vocabulary,
+    save_vocabulary,
+    vocabulary_path,
+)
+
+from ._helpers import make_canned_responder
+
+
+_CANONICAL_SUBTASKS = (
+    "grasp blue cube",
+    "place blue cube in box",
+    "retract arm",
+)
+_CANONICAL_MEMORY = (
+    "I picked up the blue cube.",
+    "I placed the blue cube in the box.",
+)
+
+
+# ---------------------------------------------------------------------------
+# Vocabulary dataclass + on-disk round-trip
+# ---------------------------------------------------------------------------
+
+
+def test_vocabulary_roundtrip(tmp_path: Path) -> None:
+    vocab = Vocabulary(
+        subtasks=_CANONICAL_SUBTASKS, memory_milestones=_CANONICAL_MEMORY
+    )
+    save_path = save_vocabulary(tmp_path, vocab)
+    assert save_path == vocabulary_path(tmp_path)
+    assert save_path.exists()
+
+    loaded = load_vocabulary(tmp_path)
+    assert loaded is not None
+    assert loaded.subtasks == _CANONICAL_SUBTASKS
+    assert loaded.memory_milestones == _CANONICAL_MEMORY
+
+
+def test_vocabulary_load_missing_returns_none(tmp_path: Path) -> None:
+    assert load_vocabulary(tmp_path) is None
+
+
+def test_vocabulary_load_malformed_returns_none(tmp_path: Path) -> None:
+    path = vocabulary_path(tmp_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text("{ not valid json", encoding="utf-8")
+    assert load_vocabulary(tmp_path) is None
+
+
+def test_vocabulary_load_empty_payload_returns_none(tmp_path: Path) -> None:
+    path = vocabulary_path(tmp_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps({"subtasks": [], "memory_milestones": []}), encoding="utf-8")
+    assert load_vocabulary(tmp_path) is None
+
+
+# ---------------------------------------------------------------------------
+# Discovery module
+# ---------------------------------------------------------------------------
+
+
+def test_vocabulary_discovery_calls_vlm_and_returns_vocab(
+    fixture_dataset_root: Path,
+) -> None:
+    vlm = make_canned_responder(
+        {
+            "canonical vocabulary": {
+                "subtasks": list(_CANONICAL_SUBTASKS),
+                "memory_milestones": list(_CANONICAL_MEMORY),
+            }
+        }
+    )
+    module = VocabularyDiscoveryModule(vlm=vlm, config=VocabularyConfig(sample_episodes=2))
+    records = list(iter_episodes(fixture_dataset_root))
+    vocab = module.discover(records)
+    assert vocab is not None
+    assert vocab.subtasks == _CANONICAL_SUBTASKS
+    assert vocab.memory_milestones == _CANONICAL_MEMORY
+
+
+def test_vocabulary_discovery_reuses_existing(fixture_dataset_root: Path) -> None:
+    """``reuse_existing=True`` short-circuits the VLM call entirely."""
+
+    def _explode(_messages):  # pragma: no cover - must not be called
+        raise AssertionError("VLM should not be invoked when reusing existing vocabulary")
+
+    from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+    vlm = StubVlmClient(responder=_explode)
+    module = VocabularyDiscoveryModule(
+        vlm=vlm, config=VocabularyConfig(reuse_existing=True)
+    )
+    records = list(iter_episodes(fixture_dataset_root))
+    existing = Vocabulary(subtasks=("a", "b"), memory_milestones=("I a.",))
+    vocab = module.discover(records, existing=existing)
+    assert vocab is existing
+
+
+def test_vocabulary_discovery_empty_payload_returns_none(
+    fixture_dataset_root: Path,
+) -> None:
+    vlm = make_canned_responder({"canonical vocabulary": {"subtasks": [], "memory_milestones": []}})
+    module = VocabularyDiscoveryModule(vlm=vlm, config=VocabularyConfig())
+    records = list(iter_episodes(fixture_dataset_root))
+    assert module.discover(records) is None
+
+
+# ---------------------------------------------------------------------------
+# PlanSubtasksMemoryModule consumes the vocabulary
+# ---------------------------------------------------------------------------
+
+
+def test_plan_module_inlines_vocab_into_subtask_prompt(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    captured: list[str] = []
+
+    def responder(messages):
+        # Find the last user text block and stash it for inspection.
+        for message in messages:
+            content = message.get("content")
+            if isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        captured.append(block.get("text", ""))
+        # Return canned subtasks; pick the first two canonical strings so
+        # the validator accepts them.
+        return {
+            "subtasks": [
+                {"text": "grasp blue cube", "start": 0.0, "end": 0.4},
+                {"text": "place blue cube in box", "start": 0.4, "end": 0.9},
+            ]
+        }
+
+    from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+    vlm = StubVlmClient(responder=responder)
+    vocab = Vocabulary(subtasks=_CANONICAL_SUBTASKS, memory_milestones=_CANONICAL_MEMORY)
+    module = PlanSubtasksMemoryModule(
+        vlm=vlm,
+        config=PlanConfig(n_task_rephrasings=0),
+        vocabulary=vocab,
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    # The subtask prompt (and the memory prompt) carries the canonical
+    # bullet list so the VLM can't paraphrase them away.
+    assert any("Canonical subtask labels:" in t for t in captured)
+    assert any("grasp blue cube" in t for t in captured)
+
+
+def test_plan_module_accepts_article_only_difference(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """Articles like 'the'/'a'/'an' are stripped during validation."""
+    from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+    def responder(_messages):
+        return {
+            "subtasks": [
+                # Same canonical phrase modulo "the" — should be accepted.
+                {"text": "grasp the blue cube", "start": 0.0, "end": 0.4},
+            ]
+        }
+
+    vlm = StubVlmClient(responder=responder)
+    vocab = Vocabulary(subtasks=_CANONICAL_SUBTASKS, memory_milestones=_CANONICAL_MEMORY)
+    module = PlanSubtasksMemoryModule(
+        vlm=vlm,
+        config=PlanConfig(n_task_rephrasings=0),
+        vocabulary=vocab,
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("plan")
+    subtask_texts = [r["content"] for r in rows if r["style"] == "subtask"]
+    assert subtask_texts == ["grasp blue cube"]
+
+
+def test_plan_module_retries_when_subtask_off_vocab(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """One-shot retry replaces an off-vocab paraphrase with the canonical form."""
+    from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+    call_count = {"n": 0}
+
+    def responder(messages):
+        call_count["n"] += 1
+        # First call: returns an off-vocab paraphrase.
+        if call_count["n"] == 1:
+            return {
+                "subtasks": [
+                    # paraphrase, not in vocab
+                    {"text": "pick up blue cube", "start": 0.0, "end": 0.4},
+                ]
+            }
+        # Second call (the retry): should contain the correction prompt;
+        # respond with the canonical phrase exactly.
+        last_user_text = ""
+        for message in messages:
+            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", "")
+        assert "NOT in the canonical vocabulary" in last_user_text
+        return {
+            "subtasks": [
+                {"text": "grasp blue cube", "start": 0.0, "end": 0.4},
+            ]
+        }
+
+    vlm = StubVlmClient(responder=responder)
+    vocab = Vocabulary(subtasks=_CANONICAL_SUBTASKS, memory_milestones=_CANONICAL_MEMORY)
+    module = PlanSubtasksMemoryModule(
+        vlm=vlm,
+        config=PlanConfig(n_task_rephrasings=0),
+        vocabulary=vocab,
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("plan")
+    subtask_texts = [r["content"] for r in rows if r["style"] == "subtask"]
+    assert subtask_texts == ["grasp blue cube"]
+    # The retry must have fired exactly once.
+    assert call_count["n"] == 2
+
+
+def test_plan_module_drops_off_vocab_subtask_after_retry(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """If the VLM stays off-vocab even after the retry, the bad span is dropped."""
+    from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+    call_count = {"n": 0}
+
+    def responder(_messages):
+        call_count["n"] += 1
+        # Both calls return the same off-vocab span — the model can't
+        # be corrected. The second call also returns one in-vocab span
+        # so the episode isn't empty; this lets us check that the
+        # off-vocab span is dropped without affecting the in-vocab one.
+        if call_count["n"] == 1:
+            return {
+                "subtasks": [
+                    {"text": "perform a fancy macarena dance", "start": 0.0, "end": 0.4},
+                    {"text": "grasp blue cube", "start": 0.4, "end": 0.9},
+                ]
+            }
+        return {
+            "subtasks": [
+                {"text": "perform a fancy macarena dance", "start": 0.0, "end": 0.4},
+                {"text": "grasp blue cube", "start": 0.4, "end": 0.9},
+            ]
+        }
+
+    vlm = StubVlmClient(responder=responder)
+    vocab = Vocabulary(subtasks=_CANONICAL_SUBTASKS, memory_milestones=_CANONICAL_MEMORY)
+    module = PlanSubtasksMemoryModule(
+        vlm=vlm,
+        config=PlanConfig(n_task_rephrasings=0),
+        vocabulary=vocab,
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("plan")
+    subtask_texts = [r["content"] for r in rows if r["style"] == "subtask"]
+    # Retry fired exactly once; bad span dropped, good span kept.
+    assert call_count["n"] == 2
+    assert subtask_texts == ["grasp blue cube"]
+
+
+def test_plan_module_bumps_collocated_subtasks_to_distinct_frames(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """Two subtasks whose starts snap to the same frame get split onto two frames.
+
+    Without this guard, both spans would emit ``style=subtask`` rows at the
+    identical persistent timestamp; the training-time renderer's
+    ``active_at(t, style=subtask)`` then raises an ambiguity error.
+    """
+    from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+    def responder(_messages):
+        # Two canonical labels with starts within one frame of each other —
+        # both snap to the same source frame, so the dedupe pass must bump
+        # the later one to the next frame.
+        return {
+            "subtasks": [
+                {"text": "grasp blue cube", "start": 0.40, "end": 0.42},
+                {"text": "place blue cube in box", "start": 0.41, "end": 0.50},
+            ]
+        }
+
+    vlm = StubVlmClient(responder=responder)
+    vocab = Vocabulary(subtasks=_CANONICAL_SUBTASKS, memory_milestones=_CANONICAL_MEMORY)
+    module = PlanSubtasksMemoryModule(
+        vlm=vlm,
+        config=PlanConfig(n_task_rephrasings=0),
+        vocabulary=vocab,
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("plan")
+    subtask_rows = [r for r in rows if r["style"] == "subtask"]
+    # Both subtasks present, both on distinct timestamps.
+    assert len(subtask_rows) == 2
+    timestamps = [r["timestamp"] for r in subtask_rows]
+    assert len(set(timestamps)) == 2, f"subtask timestamps collide: {timestamps}"
+    # Order preserved: the chronologically earlier span keeps the earlier
+    # frame, the later one was bumped onto the next available frame.
+    assert subtask_rows[0]["content"] == "grasp blue cube"
+    assert subtask_rows[1]["content"] == "place blue cube in box"
+    assert subtask_rows[1]["timestamp"] > subtask_rows[0]["timestamp"]
+
+
+def test_plan_module_empty_when_all_off_vocab_after_retry(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """All-off-vocab spans → episode comes out empty (no silent fuzzy snap)."""
+    from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+    def responder(_messages):
+        # Returns the same off-vocab spans on both attempts.
+        return {
+            "subtasks": [
+                {"text": "make a smoothie", "start": 0.0, "end": 0.4},
+                {"text": "consult the wizard", "start": 0.4, "end": 0.9},
+            ]
+        }
+
+    vlm = StubVlmClient(responder=responder)
+    vocab = Vocabulary(subtasks=_CANONICAL_SUBTASKS, memory_milestones=_CANONICAL_MEMORY)
+    module = PlanSubtasksMemoryModule(
+        vlm=vlm,
+        config=PlanConfig(n_task_rephrasings=0),
+        vocabulary=vocab,
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("plan")
+    subtask_texts = [r["content"] for r in rows if r["style"] == "subtask"]
+    # No subtask gets fabricated — better to leave the episode empty
+    # so the operator notices the vocabulary gap than to silently
+    # warp the labels.
+    assert subtask_texts == []
+
+
+def test_plan_module_without_vocab_passes_through(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """No vocabulary configured → original free-form behavior is preserved."""
+    from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+    def responder(_messages):
+        return {
+            "subtasks": [
+                {"text": "any free-form text the VLM wants", "start": 0.0, "end": 1.0},
+            ]
+        }
+
+    vlm = StubVlmClient(responder=responder)
+    module = PlanSubtasksMemoryModule(
+        vlm=vlm, config=PlanConfig(n_task_rephrasings=0)
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("plan")
+    subtask_texts = [r["content"] for r in rows if r["style"] == "subtask"]
+    assert subtask_texts == ["any free-form text the VLM wants"]

tests/annotations/test_writer.py

@@ -35,17 +35,17 @@ def _stage_episode(
     staging_dir: Path,
     episode_index: int,
     *,
-    module_1: list[dict] | None = None,
-    module_2: list[dict] | None = None,
-    module_3: list[dict] | None = None,
+    plan: list[dict] | None = None,
+    interjections: list[dict] | None = None,
+    vqa: list[dict] | None = None,
 ) -> None:
     staging = EpisodeStaging(staging_dir, episode_index)
-    if module_1 is not None:
-        staging.write("module_1", module_1)
-    if module_2 is not None:
-        staging.write("module_2", module_2)
-    if module_3 is not None:
-        staging.write("module_3", module_3)
+    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:
@@ -54,7 +54,7 @@ def test_writer_persistence_identity(fixture_dataset_root: Path, tmp_path: Path)
     _stage_episode(
         staging_dir,
         0,
-        module_1=[
+        plan=[
             {
                 "role": "assistant",
                 "content": "grasp the sponge",
@@ -94,7 +94,7 @@ def test_writer_events_exact_timestamp(fixture_dataset_root: Path, tmp_path: Pat
     _stage_episode(
         staging_dir,
         0,
-        module_2=[
+        interjections=[
             speech_atom(0.0, "Got it."),
             {
                 "role": "user",
@@ -127,7 +127,7 @@ def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> No
     _stage_episode(
         staging_dir,
         0,
-        module_1=[
+        plan=[
             {
                 "role": "assistant",
                 "content": "do X",
@@ -150,7 +150,7 @@ def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> No
                 "tool_calls": None,
             },
         ],
-        module_2=[
+        interjections=[
             speech_atom(0.0, "OK"),
             {
                 "role": "user",
@@ -161,12 +161,13 @@ def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> No
             },
             speech_atom(0.2, "Waiting"),
         ],
-        module_3=[
+        vqa=[
             {
                 "role": "user",
                 "content": "where is the cup?",
                 "style": "vqa",
                 "timestamp": 0.4,
+                "camera": "observation.images.front",
                 "tool_calls": None,
             },
             {
@@ -177,6 +178,7 @@ def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> No
                 ),
                 "style": "vqa",
                 "timestamp": 0.4,
+                "camera": "observation.images.front",
                 "tool_calls": None,
             },
         ],
@@ -199,7 +201,7 @@ def test_writer_drops_subtask_index_idempotent(fixture_dataset_root: Path, tmp_p
     _stage_episode(
         staging_dir,
         0,
-        module_1=[
+        plan=[
             {
                 "role": "assistant",
                 "content": "do X",
@@ -275,7 +277,7 @@ def test_writer_does_not_add_tools_column(fixture_dataset_root: Path, tmp_path:
     _stage_episode(
         staging_dir,
         0,
-        module_1=[
+        plan=[
             {"role": "assistant", "content": "x", "style": "subtask", "timestamp": 0.0, "tool_calls": None}
         ],
     )
@@ -285,6 +287,56 @@ def test_writer_does_not_add_tools_column(fixture_dataset_root: Path, tmp_path:
     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"