admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 17:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

review: address CarolinePascal feedback

Browse Source 
- name the three modules everywhere (plan / interjections / vqa) instead
  of module_1/2/3 — config classes, config fields, executor params,
  staging keys and phase names now carry the module name
- rename examples/annotation -> examples/annotations; add the Apache
  header to run_hf_job.py
- drop the unused GeneralVqaModule._generate_one
- remove "PR 1" references from comments/docstrings
- frames.py: rely on the always-defined LeRobotDatasetMetadata.camera_keys
- executor.py: read/write meta/info.json via load_info / write_info
- reader.py: load meta/tasks.parquet via io_utils.load_tasks
- make --push_to_hub a bool; push the annotated dataset back to --repo_id
- move the on-disk test dataset builder into tests/fixtures
  (build_annotation_dataset); run_e2e_smoke reuses it
- clarify in the docs that the vqa module grounds each pair on a single
  frame (K = per-tick anchor count)
- hoist stdlib dynamic imports to module scope

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Pepijn Kooijmans
2026-05-18 12:03:25 +02:00
parent 965d42825f
commit fd18beb3a1
23 changed files with 383 additions and 412 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/annotation_pipeline.mdx
+36 -30
View File
@@ -10,15 +10,15 @@
Three modules write into a per-episode staging tree, then a single writer Three modules write into a per-episode staging tree, then a single writer
rewrites the data shards in place: rewrites the data shards in place:
| Style / atom | Column | Module | | Style / atom | Column | Module |
| ------------------------------------------- | --------------------- | -------- | | ------------------------------------------- | --------------------- | -------------- |
| `subtask` (Pi0.7-style "how, not what") | `language_persistent` | Module 1 | | `subtask` (Pi0.7-style "how, not what") | `language_persistent` | `plan` |
| `plan` (initial + refresh on interjection) | `language_persistent` | Module 1 | | `plan` (initial + refresh on interjection) | `language_persistent` | `plan` |
| `memory` (MEM-style compression) | `language_persistent` | Module 1 | | `memory` (MEM-style compression) | `language_persistent` | `plan` |
| `task_aug` (rephrasings of canonical task) | `language_persistent` | Module 1 | | `task_aug` (rephrasings of canonical task) | `language_persistent` | `plan` |
| `interjection` | `language_events` | Module 2 | | `interjection` | `language_events` | `interjections`|
| speech tool-call atom (`style=null`, `say`) | `language_events` | Module 2 | | speech tool-call atom (`style=null`, `say`) | `language_events` | `interjections`|
| `vqa` (user / assistant pair) | `language_events` | Module 3 | | `vqa` (user / assistant pair) | `language_events` | `vqa` |
The writer does **not** add a `tools` column to the parquet — the tool The writer does **not** add a `tools` column to the parquet — the tool
catalog lives at `meta/info.json["tools"]` instead (see catalog lives at `meta/info.json["tools"]` instead (see
@@ -45,20 +45,24 @@ uv run lerobot-annotate \
--vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct --vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct
``` ```
The pipeline attaches actual camera footage to every Module 1/2/3 prompt The pipeline attaches actual camera footage to every `plan` /
by default, decoded from the dataset's first `observation.images.*` `interjections` / `vqa` prompt by default, decoded from the dataset's
stream. Override with `--vlm.camera_key=observation.images.<name>` to first `observation.images.*` stream. Override with
pin a specific viewpoint. Datasets with no video tracks fall back to `--vlm.camera_key=observation.images.<name>` to pin a specific
text-only prompts automatically. viewpoint. Datasets with no video tracks fall back to text-only prompts
automatically.
**Module 1 sees the whole episode as one video block.** Subtask **The `plan` module sees the whole episode as one video block.** Subtask
decomposition gets a `{"type":"video", "video":[<frames>]}` block decomposition gets a `{"type":"video", "video":[<frames>]}` block
covering the entire demonstration; Qwen-VL pools temporally on its own covering the entire demonstration; Qwen-VL pools temporally on its own
and decides where to cut. There is no keyframe stride or count knob — and decides where to cut. There is no keyframe stride or count knob —
`--module_1.max_video_frames` (default 128) only caps the frames packed `--plan.max_video_frames` (default 128) only caps the frames packed
into the video block as a model-capacity bound. Module 2 attaches a into the video block as a model-capacity bound. The `interjections`
short window of frames around the interjection timestamp; Module 3 module attaches a short window of frames straddling the interjection
attaches the exact emission frame to each VQA pair. timestamp. The `vqa` module grounds each VQA pair on a single frame —
its `--vqa.K` knob sets how many consecutive frames each emission tick
anchors, and every anchored frame gets its own VQA pair on that one
frame (there is no per-pair frame window).
## Running on Hugging Face Jobs ## Running on Hugging Face Jobs
@@ -67,15 +71,16 @@ Distributed annotation is delegated to
ships a launcher script you copy and edit for your dataset: ships a launcher script you copy and edit for your dataset:
```bash ```bash
HF_TOKEN=hf_... uv run python examples/annotation/run_hf_job.py HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
``` ```
[`examples/annotation/run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotation/run_hf_job.py) [`examples/annotations/run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
spawns one `h200x2` job that: spawns one `h200x2` job that:
1. installs the branch under test plus the annotation extras, 1. installs the branch under test plus the annotation extras,
2. boots two vllm servers (one per GPU) for the chosen model, 2. boots two vllm servers (one per GPU) for the chosen model,
3. runs Modules 1 / 2 / 3 across the dataset via `lerobot-annotate`, 3. runs the `plan` / `interjections` / `vqa` modules across the dataset
via `lerobot-annotate`,
4. uploads the annotated dataset to `--push_to_hub`. 4. uploads the annotated dataset to `--push_to_hub`.
To target a different dataset, model, or hub repo, edit the `CMD` block To target a different dataset, model, or hub repo, edit the `CMD` block
@@ -115,7 +120,8 @@ Each module writes its raw output to
`<root>/.annotate_staging/episode_{N:06d}/<module>.jsonl`. That makes `<root>/.annotate_staging/episode_{N:06d}/<module>.jsonl`. That makes
prompt iteration cheap — re-running one module overwrites only its own prompt iteration cheap — re-running one module overwrites only its own
JSONL file before the writer composes the final parquet. Modules can be JSONL file before the writer composes the final parquet. Modules can be
disabled via `--module_1.enabled=false` (and similarly for 2 and 3) to disabled via `--plan.enabled=false` (and likewise `--interjections.enabled`
/ `--vqa.enabled`) to
test them in isolation. test them in isolation.
## Validation/report checks before final write ## Validation/report checks before final write
@@ -134,18 +140,18 @@ Errors abort the writer (`--skip_validation=true` overrides for debugging).
## Paper inspirations per module ## Paper inspirations per module
- **Module 1 — subtasks.** Hi Robot ([Shi 2025](https://arxiv.org/abs/2502.19417)) - **`plan` module — subtasks.** Hi Robot ([Shi 2025](https://arxiv.org/abs/2502.19417))
atom granularity ("pick up one piece of lettuce", "place bowl to box"); atom granularity ("pick up one piece of lettuce", "place bowl to box");
Pi0.7 ([Physical Intelligence 2025](https://pi.website/pi07)) "how, not Pi0.7 ([Physical Intelligence 2025](https://pi.website/pi07)) "how, not
what" detail. what" detail.
- **Module 1 — memory.** MEM ([Torne 2026](https://arxiv.org/abs/2603.03596)) - **`plan` module — memory.** MEM ([Torne 2026](https://arxiv.org/abs/2603.03596))
compression directive: keep only minimal relevant information; functional compression directive: keep only minimal relevant information; functional
outcomes preserved, specific attributes dropped. outcomes preserved, specific attributes dropped.
- **Module 2 — interjections.** Hi Robot scenario taxonomy: negative task, - **`interjections` module.** Hi Robot scenario taxonomy: negative task,
situated correction, specific constraint, preference. Speech is a situated correction, specific constraint, preference. Speech is a
tool-call-only atom (`tool_calls=[{type:function, function:{name:"say", tool-call-only atom (`tool_calls=[{type:function, function:{name:"say",
arguments:{text:...}}}]`). arguments:{text:...}}}]`).
- **Module 3 — VQA.** ECoT ([Zawalski 2024](https://arxiv.org/abs/2407.08693)) - **`vqa` module.** ECoT ([Zawalski 2024](https://arxiv.org/abs/2407.08693))
grounded features (bounding boxes in pixel `[x_min, y_min, x_max, y_max]`, grounded features (bounding boxes in pixel `[x_min, y_min, x_max, y_max]`,
keypoints) and Steerable VLA Policies ([Zhao 2025](https://arxiv.org/abs/2509.07626)) keypoints) and Steerable VLA Policies ([Zhao 2025](https://arxiv.org/abs/2509.07626))
multi-abstraction grounding. Pi0.7 also grounds answers across multi-abstraction grounding. Pi0.7 also grounds answers across
@@ -157,9 +163,9 @@ references rather than rewriting from scratch.
## Compute and list-size estimates ## Compute and list-size estimates
Per episode, the pipeline issues O(`max_steps`) Module 1 calls, Per episode, the pipeline issues O(`max_steps`) `plan`-module calls,
O(`max_interjections_per_episode`) Module 2 calls, and O(`max_interjections_per_episode`) `interjections`-module calls, and
O(`vqa_emission_hz × episode_seconds`) Module 3 calls. With defaults O(`vqa_emission_hz × episode_seconds`) `vqa`-module calls. With defaults
(8 subtasks, 1 interjection, 1 Hz × 3 pairs) and 30-second episodes, that (8 subtasks, 1 interjection, 1 Hz × 3 pairs) and 30-second episodes, that
is ~50 VLM calls per episode. `language_persistent` per episode is ~10s of is ~50 VLM calls per episode. `language_persistent` per episode is ~10s of
KB at most (parquet dictionary-encodes one entry per episode); KB at most (parquet dictionary-encodes one entry per episode);
examples/annotation/run_hf_job.py → examples/annotations/run_hf_job.py
+28 -10
View File
@@ -1,18 +1,35 @@
#!/usr/bin/env python #!/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 MoE). """Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6 MoE).
Spawns one ``h200x2`` job that: Spawns one ``h200x2`` job that:
1. installs this branch of ``lerobot`` plus the annotation extras, 1. installs this branch of ``lerobot`` plus the annotation extras,
2. boots two vllm servers (one per GPU) with Qwen3.6-35B-A3B-FP8, 2. boots two vllm servers (one per GPU) with Qwen3.6-35B-A3B-FP8,
3. runs Module 1/2/3 across the dataset (per-camera VQA via PR 3471), 3. runs the plan / interjections / vqa modules across the dataset,
4. uploads the annotated dataset to ``--push_to_hub``. 4. uploads the annotated dataset back to ``--repo_id``.
``--repo_id`` is both the download source and, with ``--push_to_hub=true``,
the upload destination — the job annotates the dataset in place.
Usage: Usage:
HF_TOKEN=hf_... uv run python examples/annotation/run_hf_job.py HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
Adjust ``CMD`` below to point at your own dataset / target hub repo. Adjust ``CMD`` below to point at your own dataset.
""" """
import os import os
@@ -36,7 +53,9 @@ CMD = (
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && " "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
"export VLLM_VIDEO_BACKEND=pyav && " "export VLLM_VIDEO_BACKEND=pyav && "
"lerobot-annotate " "lerobot-annotate "
"--repo_id=imstevenpmwork/super_poulain_draft " # The dataset to annotate; also the push destination (annotate in place).
"--repo_id=<your-org>/<your-dataset> "
"--push_to_hub=true "
"--vlm.backend=openai " "--vlm.backend=openai "
"--vlm.model_id=Qwen/Qwen3.6-35B-A3B-FP8 " "--vlm.model_id=Qwen/Qwen3.6-35B-A3B-FP8 "
"--vlm.parallel_servers=2 " "--vlm.parallel_servers=2 "
@@ -50,11 +69,10 @@ CMD = (
"--executor.episode_parallelism=32 " "--executor.episode_parallelism=32 "
"--vlm.chat_template_kwargs='{enable_thinking: false}' " "--vlm.chat_template_kwargs='{enable_thinking: false}' "
"--vlm.camera_key=observation.images.wrist " "--vlm.camera_key=observation.images.wrist "
"--module_1.frames_per_second=1.0 " "--plan.frames_per_second=1.0 "
"--module_1.use_video_url=true " "--plan.use_video_url=true "
"--module_1.use_video_url_fps=1.0 " "--plan.use_video_url_fps=1.0 "
"--module_3.K=1 --module_3.vqa_emission_hz=0.2 " "--vqa.K=1 --vqa.vqa_emission_hz=0.2"
"--push_to_hub=pepijn223/super_poulain_qwen36moe-3"
) )
job = run_job( job = run_job(
pyproject.toml
+1 -1
View File
@@ -205,7 +205,7 @@ peft = ["lerobot[transformers-dep]", "lerobot[peft-dep]"]
# on Linux, with a transformers fallback elsewhere; openai is the default # on Linux, with a transformers fallback elsewhere; openai is the default
# backend and talks to any OpenAI-compatible server (``vllm serve`` / # backend and talks to any OpenAI-compatible server (``vllm serve`` /
# ``transformers serve`` / hosted endpoints). Distributed execution is # ``transformers serve`` / hosted endpoints). Distributed execution is
# delegated to Hugging Face Jobs (see examples/annotation/run_hf_job.py). # delegated to Hugging Face Jobs (see examples/annotations/run_hf_job.py).
annotations = [ annotations = [
"lerobot[dataset]", "lerobot[dataset]",
"lerobot[transformers-dep]", "lerobot[transformers-dep]",
src/lerobot/annotations/steerable_pipeline/__init__.py
+3 -3
View File
@@ -19,9 +19,9 @@
The pipeline is decomposed into three independently runnable modules whose The pipeline is decomposed into three independently runnable modules whose
outputs are staged per-episode before a final parquet rewrite: outputs are staged per-episode before a final parquet rewrite:
- :mod:`.modules.plan_subtasks_memory` (Module 1) — persistent styles - :mod:`.modules.plan_subtasks_memory` (the ``plan`` module) — persistent styles
- :mod:`.modules.interjections_and_speech` (Module 2) — event styles + speech - :mod:`.modules.interjections_and_speech` (the ``interjections`` module) — event styles + speech
- :mod:`.modules.general_vqa` (Module 3) — event-style VQA pairs - :mod:`.modules.general_vqa` (the ``vqa`` module) — event-style VQA pairs
""" """
from .config import AnnotationPipelineConfig from .config import AnnotationPipelineConfig
src/lerobot/annotations/steerable_pipeline/config.py
+24 -20
View File
@@ -22,12 +22,12 @@ from typing import Any
@dataclass @dataclass
class Module1Config: class PlanConfig:
"""Module 1: plan + subtasks + memory + task augmentation. """``plan`` module: plan + subtasks + memory + task augmentation.
Module 1 attaches the whole episode as one Qwen-VL video block; The ``plan`` module attaches the whole episode as one Qwen-VL video
``max_video_frames`` only caps the frames packed in (a model-capacity block; ``max_video_frames`` only caps the frames packed in (a
bound, not an annotation-logic knob). model-capacity bound, not an annotation-logic knob).
""" """
enabled: bool = True enabled: bool = True
@@ -39,8 +39,8 @@ class Module1Config:
# When to derive the task from the video instead of using # When to derive the task from the video instead of using
# ``record.episode_task``: ``off``, ``if_short`` (short / placeholder / # ``record.episode_task``: ``off``, ``if_short`` (short / placeholder /
# missing canonical task), or ``always``. The derived task replaces the # missing canonical task), or ``always``. The derived task replaces the
# canonical one for every Module-1 prompt; ``meta/tasks.parquet`` is # canonical one for every ``plan``-module prompt; ``meta/tasks.parquet``
# never modified. # is never modified.
derive_task_from_video: str = "if_short" derive_task_from_video: str = "if_short"
derive_task_min_words: int = 3 derive_task_min_words: int = 3
@@ -51,21 +51,22 @@ class Module1Config:
min_subtask_seconds: float = 1.5 min_subtask_seconds: float = 1.5
plan_max_steps: int = 8 plan_max_steps: int = 8
# When True (and backend supports it, e.g. ``openai``), Module 1 sends a # When True (and backend supports it, e.g. ``openai``), the ``plan``
# ``video_url`` block pointing at a per-episode mp4 subclip and lets the # module sends a ``video_url`` block pointing at a per-episode mp4
# server sample frames at ``use_video_url_fps``. # subclip and lets the server sample frames at ``use_video_url_fps``.
use_video_url: bool = False use_video_url: bool = False
use_video_url_fps: float = 1.0 use_video_url_fps: float = 1.0
@dataclass @dataclass
class Module2Config: class InterjectionsConfig:
"""Module 2: interjections + paired speech.""" """``interjections`` module: interjections + paired speech."""
enabled: bool = True enabled: bool = True
# Each interjection emits a paired ``(interjection, speech)`` event row # Each interjection emits a paired ``(interjection, speech)`` event row
# and triggers a ``plan`` refresh at the same timestamp via Module 1. # and triggers a ``plan`` refresh at the same timestamp via the
# ``plan`` module.
max_interjections_per_episode: int = 3 max_interjections_per_episode: int = 3
interjection_min_t: float = 2.0 interjection_min_t: float = 2.0
@@ -77,8 +78,8 @@ class Module2Config:
@dataclass @dataclass
class Module3Config: class VqaConfig:
"""Module 3: general VQA.""" """``vqa`` module: general VQA."""
enabled: bool = True enabled: bool = True
vqa_emission_hz: float = 1.0 vqa_emission_hz: float = 1.0
@@ -161,6 +162,8 @@ class AnnotationPipelineConfig:
revisions of the same dataset live in separate copies. revisions of the same dataset live in separate copies.
""" """
# Hub dataset id. Used as the download source when ``root`` is unset,
# and as the destination repo when ``push_to_hub`` is enabled.
repo_id: str | None = None repo_id: str | None = None
root: Path | None = None root: Path | None = None
@@ -169,9 +172,9 @@ class AnnotationPipelineConfig:
seed: int = 1729 seed: int = 1729
module_1: Module1Config = field(default_factory=Module1Config) plan: PlanConfig = field(default_factory=PlanConfig)
module_2: Module2Config = field(default_factory=Module2Config) interjections: InterjectionsConfig = field(default_factory=InterjectionsConfig)
module_3: Module3Config = field(default_factory=Module3Config) vqa: VqaConfig = field(default_factory=VqaConfig)
vlm: VlmConfig = field(default_factory=VlmConfig) vlm: VlmConfig = field(default_factory=VlmConfig)
executor: ExecutorConfig = field(default_factory=ExecutorConfig) executor: ExecutorConfig = field(default_factory=ExecutorConfig)
@@ -179,8 +182,9 @@ class AnnotationPipelineConfig:
skip_validation: bool = False skip_validation: bool = False
only_episodes: tuple[int, ...] | None = None only_episodes: tuple[int, ...] | None = None
# Upload the annotated dataset to the Hugging Face Hub when set. # When True, upload the annotated dataset back to ``repo_id`` on the
push_to_hub: str | None = None # Hugging Face Hub. ``repo_id`` must be set for this to take effect.
push_to_hub: bool = False
push_private: bool = False push_private: bool = False
push_commit_message: str | None = None push_commit_message: str | None = None
src/lerobot/annotations/steerable_pipeline/executor.py
+46 -58
View File
@@ -17,19 +17,20 @@
The executor plans **six phases** in the dependency order from the plan: The executor plans **six phases** in the dependency order from the plan:
phase 1: Module 1 (plan + subtasks + memory) phase 1: ``plan`` module (plan + subtasks + memory)
phase 2: Module 2 (interjections + speech) phase 2: ``interjections`` module (interjections + speech)
phase 3: Module 1 plan-update pass — re-runs plan emission at every phase 3: ``plan`` plan-update pass — re-runs plan emission at every
interjection timestamp produced by phase 2 interjection timestamp produced by phase 2
phase 4: Module 3 (VQA) phase 4: ``vqa`` module (VQA)
phase 5: validator phase 5: validator
phase 6: writer phase 6: writer
Phase 3 is why Module 1 must be re-entered after Module 2 — to refresh Phase 3 is why the ``plan`` module must be re-entered after the
``plan`` rows at interjection timestamps. ``interjections`` module — to refresh ``plan`` rows at interjection
timestamps.
Distributed execution is provided by Hugging Face Jobs (see Distributed execution is provided by Hugging Face Jobs (see
``examples/annotation/run_hf_job.py``); the runner inside the job ``examples/annotations/run_hf_job.py``); the runner inside the job
invokes ``lerobot-annotate`` which uses this in-process executor. invokes ``lerobot-annotate`` which uses this in-process executor.
Episode-level concurrency is controlled by Episode-level concurrency is controlled by
``ExecutorConfig.episode_parallelism``. ``ExecutorConfig.episode_parallelism``.
@@ -38,6 +39,8 @@ Episode-level concurrency is controlled by
from __future__ import annotations from __future__ import annotations
import logging import logging
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -71,7 +74,7 @@ class PipelineRunSummary:
@dataclass @dataclass
class Executor: class Executor:
"""Run all four phases over a dataset root in-process. """Run all six phases over a dataset root in-process.
Episode-level concurrency comes from ``ExecutorConfig.episode_parallelism`` Episode-level concurrency comes from ``ExecutorConfig.episode_parallelism``
(a thread pool); cluster-level concurrency comes from running this (a thread pool); cluster-level concurrency comes from running this
@@ -80,9 +83,9 @@ class Executor:
""" """
config: AnnotationPipelineConfig config: AnnotationPipelineConfig
module_1: Any # PlanSubtasksMemoryModule plan: Any # PlanSubtasksMemoryModule
module_2: Any # InterjectionsAndSpeechModule interjections: Any # InterjectionsAndSpeechModule
module_3: Any # GeneralVqaModule vqa: Any # GeneralVqaModule
writer: LanguageColumnsWriter writer: LanguageColumnsWriter
validator: StagingValidator validator: StagingValidator
@@ -99,16 +102,16 @@ class Executor:
phases: list[PhaseResult] = [] phases: list[PhaseResult] = []
# Phase 1: Module 1 (plan + subtasks + memory) # Phase 1: ``plan`` module (plan + subtasks + memory)
phases.append(self._run_module_phase("module_1", records, staging_dir, self.module_1)) phases.append(self._run_module_phase("plan", records, staging_dir, self.plan))
# Phase 2: Module 2 (interjections + speech). Module 2 reads # Phase 2: ``interjections`` module (interjections + speech). It
# Module 1's subtask rows from the same staging tree to ground # reads the ``plan`` module's subtask rows from the same staging
# the interjection prompt in the correct local subtask. # tree to ground the interjection prompt in the correct local subtask.
phases.append(self._run_module_phase("module_2", records, staging_dir, self.module_2)) phases.append(self._run_module_phase("interjections", records, staging_dir, self.interjections))
# Phase 3: Module 1 plan-update pass at interjection timestamps. # Phase 3: ``plan`` plan-update pass at interjection timestamps.
phases.append(self._run_plan_update_phase(records, staging_dir)) phases.append(self._run_plan_update_phase(records, staging_dir))
# Phase 4: Module 3 (VQA) # Phase 4: ``vqa`` module (VQA)
phases.append(self._run_module_phase("module_3", records, staging_dir, self.module_3)) phases.append(self._run_module_phase("vqa", records, staging_dir, self.vqa))
print("[annotate] running validator...", flush=True) print("[annotate] running validator...", flush=True)
report = self.validator.validate(records, staging_dir) report = self.validator.validate(records, staging_dir)
@@ -135,50 +138,37 @@ class Executor:
those columns too, otherwise non-streaming ``LeRobotDataset`` loads those columns too, otherwise non-streaming ``LeRobotDataset`` loads
cast against the old schema and fail on the extra parquet columns. cast against the old schema and fail on the extra parquet columns.
""" """
import json # noqa: PLC0415 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 from lerobot.datasets.language import SAY_TOOL_SCHEMA, language_feature_info # noqa: PLC0415
info_path = root / "meta" / "info.json" info_path = root / "meta" / "info.json"
if not info_path.exists(): if not info_path.exists():
return return
try: try:
info = json.loads(info_path.read_text()) info = load_info(root)
except Exception as exc: # noqa: BLE001 except Exception as exc: # noqa: BLE001
print(f"[annotate] could not read {info_path}: {exc}", flush=True) print(f"[annotate] could not read {info_path}: {exc}", flush=True)
return return
changed = False changed = False
features = info.get("features") merged_features = {**info.features, **language_feature_info()}
if not isinstance(features, dict): if merged_features != info.features:
features = {} info.features = merged_features
merged_features = {**features, **language_feature_info()}
if merged_features != features:
info["features"] = merged_features
changed = True changed = True
existing = info.get("tools") existing = info.tools or []
if not isinstance(existing, list):
existing = []
names = {(t.get("function") or {}).get("name") for t in existing if isinstance(t, dict)} names = {(t.get("function") or {}).get("name") for t in existing if isinstance(t, dict)}
merged = list(existing)
if SAY_TOOL_SCHEMA["function"]["name"] not in names: if SAY_TOOL_SCHEMA["function"]["name"] not in names:
merged.append(SAY_TOOL_SCHEMA) info.tools = [*existing, SAY_TOOL_SCHEMA]
if merged != existing:
info["tools"] = merged
changed = True changed = True
if changed: if changed:
# Atomic replace — info.json is load-bearing for dataset write_info(info, root)
# metadata, so a crash mid-write would brick the dataset.
tmp_info = info_path.with_suffix(info_path.suffix + ".tmp")
tmp_info.write_text(json.dumps(info, indent=2))
tmp_info.replace(info_path)
print( print(
"[annotate] meta/info.json: " "[annotate] meta/info.json: "
f"language_features={list(language_feature_info())}, " f"language_features={list(language_feature_info())}, "
f"tools={[t['function']['name'] for t in merged]}", f"tools={[t['function']['name'] for t in (info.tools or [])]}",
flush=True, flush=True,
) )
@@ -189,9 +179,6 @@ class Executor:
staging_dir: Path, staging_dir: Path,
module: Any, module: Any,
) -> PhaseResult: ) -> PhaseResult:
import time as _time # noqa: PLC0415
from concurrent.futures import ThreadPoolExecutor, as_completed # noqa: PLC0415
if not module.enabled: if not module.enabled:
print(f"[annotate] phase={name} skipped (module disabled)", flush=True) print(f"[annotate] phase={name} skipped (module disabled)", flush=True)
return PhaseResult(name=name, episodes_processed=0, episodes_skipped=len(records)) return PhaseResult(name=name, episodes_processed=0, episodes_skipped=len(records))
@@ -201,14 +188,14 @@ class Executor:
f"[annotate] phase={name} starting on {n} episode(s) (parallelism={parallelism})", f"[annotate] phase={name} starting on {n} episode(s) (parallelism={parallelism})",
flush=True, flush=True,
) )
t0 = _time.time() t0 = time.time()
def _do(idx_record: tuple[int, EpisodeRecord]) -> tuple[int, int, float]: def _do(idx_record: tuple[int, EpisodeRecord]) -> tuple[int, int, float]:
i, record = idx_record i, record = idx_record
ep_start = _time.time() ep_start = time.time()
staging = EpisodeStaging(staging_dir, record.episode_index) staging = EpisodeStaging(staging_dir, record.episode_index)
module.run_episode(record, staging) module.run_episode(record, staging)
return i, record.episode_index, _time.time() - ep_start return i, record.episode_index, time.time() - ep_start
processed = 0 processed = 0
if parallelism == 1: if parallelism == 1:
@@ -230,38 +217,39 @@ class Executor:
f"(idx={ep_idx}, submit_order={i}) done in {elapsed:.1f}s", f"(idx={ep_idx}, submit_order={i}) done in {elapsed:.1f}s",
flush=True, flush=True,
) )
total = _time.time() - t0 total = time.time() - t0
print(f"[annotate] phase={name} complete: {processed}/{n} in {total:.1f}s", flush=True) print(f"[annotate] phase={name} complete: {processed}/{n} in {total:.1f}s", flush=True)
return PhaseResult(name=name, episodes_processed=processed, episodes_skipped=0) return PhaseResult(name=name, episodes_processed=processed, episodes_skipped=0)
def _run_plan_update_phase( # noqa: PLR0915 def _run_plan_update_phase( # noqa: PLR0915
self, records: list[EpisodeRecord], staging_dir: Path self, records: list[EpisodeRecord], staging_dir: Path
) -> PhaseResult: ) -> PhaseResult:
"""Re-emit ``plan`` rows at each interjection timestamp from Module 2. """Re-emit ``plan`` rows at each timestamp the ``interjections`` module produced.
Module 1 owns the prompt; Module 2 produced the timestamps. This phase The ``plan`` module owns the prompt; the ``interjections`` module
therefore calls back into Module 1 with the interjection timestamps so produced the timestamps. This phase therefore calls back into the
Module 1's existing prompt path is reused. ``plan`` module with the interjection timestamps so its existing
prompt path is reused.
""" """
if not self.module_1.enabled or not self.module_2.enabled: if not self.plan.enabled or not self.interjections.enabled:
return PhaseResult( return PhaseResult(
name="module_1_plan_update", episodes_processed=0, episodes_skipped=len(records) name="plan_update", episodes_processed=0, episodes_skipped=len(records)
) )
processed = 0 processed = 0
for record in records: for record in records:
staging = EpisodeStaging(staging_dir, record.episode_index) staging = EpisodeStaging(staging_dir, record.episode_index)
interjection_rows = [ interjection_rows = [
row for row in staging.read("module_2") if row.get("style") == "interjection" row for row in staging.read("interjections") if row.get("style") == "interjection"
] ]
interjection_times = [float(row["timestamp"]) for row in interjection_rows] interjection_times = [float(row["timestamp"]) for row in interjection_rows]
interjection_texts = [str(row.get("content") or "") for row in interjection_rows] interjection_texts = [str(row.get("content") or "") for row in interjection_rows]
if interjection_times: if interjection_times:
self.module_1.run_plan_updates(record, staging, interjection_times, interjection_texts) self.plan.run_plan_updates(record, staging, interjection_times, interjection_texts)
processed += 1 processed += 1
# Episodes without any interjections are skipped (no plan refresh # Episodes without any interjections are skipped (no plan refresh
# needed); count them so the summary's processed+skipped == total. # needed); count them so the summary's processed+skipped == total.
return PhaseResult( return PhaseResult(
name="module_1_plan_update", name="plan_update",
episodes_processed=processed, episodes_processed=processed,
episodes_skipped=len(records) - processed, episodes_skipped=len(records) - processed,
) )
src/lerobot/annotations/steerable_pipeline/frames.py
+16 -16
View File
@@ -49,7 +49,7 @@ class FrameProvider(Protocol):
Empty list if the camera is unavailable. ``camera_key=None`` falls back Empty list if the camera is unavailable. ``camera_key=None`` falls back
to the provider's default camera so existing single-camera callers to the provider's default camera so existing single-camera callers
(Module 1, Module 2) keep working unchanged. (the ``plan`` and ``interjections`` modules) keep working unchanged.
""" """
def video_for_episode( def video_for_episode(
@@ -100,10 +100,11 @@ def null_provider() -> FrameProvider:
class VideoFrameProvider: class VideoFrameProvider:
"""Decodes frames from the dataset's ``observation.images.*`` streams. """Decodes frames from the dataset's ``observation.images.*`` streams.
By default the *first* camera key is used for Module 1 (subtask By default the *first* camera key is used for the ``plan`` module
decomposition) and Module 2 (interjection scenarios) — those prompts care (subtask decomposition) and the ``interjections`` module (interjection
about *what is happening*, not which angle. Module 3 (VQA) instead scenarios) — those prompts care about *what is happening*, not which
iterates over every camera in :attr:`camera_keys` so each frame's 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 answer (bbox/keypoint/...) is tagged with the camera it was
grounded against. grounded against.
@@ -112,7 +113,7 @@ class VideoFrameProvider:
``video_for_episode`` to read a non-default stream. ``video_for_episode`` to read a non-default stream.
Caches up to ``cache_size`` decoded frames per process to keep Caches up to ``cache_size`` decoded frames per process to keep
co-timestamped Module 2 + Module 1 plan-update calls cheap. co-timestamped ``interjections`` + ``plan`` plan-update calls cheap.
""" """
root: Path root: Path
@@ -122,7 +123,7 @@ class VideoFrameProvider:
_meta: Any = field(default=None, init=False, repr=False) _meta: Any = field(default=None, init=False, repr=False)
_cache: dict = field(default_factory=dict, 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) _camera_keys: list[str] = field(default_factory=list, init=False, repr=False)
# Pipeline runs Module 1/2/3 phases under a ThreadPoolExecutor (see # Pipeline runs the three module phases under a ThreadPoolExecutor (see
# ``ExecutorConfig.episode_parallelism``); guard the dict cache and the # ``ExecutorConfig.episode_parallelism``); guard the dict cache and the
# one-shot warn flag against concurrent updates from worker threads. # one-shot warn flag against concurrent updates from worker threads.
_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
@@ -131,11 +132,10 @@ class VideoFrameProvider:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata # noqa: PLC0415 from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata # noqa: PLC0415
self._meta = LeRobotDatasetMetadata(repo_id="local", root=self.root) self._meta = LeRobotDatasetMetadata(repo_id="local", root=self.root)
# ``camera_keys`` covers both image- and video-stored cameras # ``camera_keys`` covers both image- and video-stored cameras and is
# (``video_keys`` is video-only). Some datasets declare cameras with # always defined on the metadata (``[]`` in the worst case), so it is
# ``dtype=image``, which would otherwise look empty here and silently # the single source we need here.
# disable Module 3 even though the videos are there. keys = list(self._meta.camera_keys)
keys = list(getattr(self._meta, "camera_keys", None) or self._meta.video_keys or [])
# Last-resort fallback: if metadata didn't surface anything but the # Last-resort fallback: if metadata didn't surface anything but the
# caller explicitly named a camera (``--vlm.camera_key=...``), trust # caller explicitly named a camera (``--vlm.camera_key=...``), trust
# them — the key is by definition known to exist on the dataset. # them — the key is by definition known to exist on the dataset.
@@ -275,10 +275,10 @@ class VideoFrameProvider:
try: try:
return _decode_pyav_direct(video_path, shifted, self.tolerance_s) return _decode_pyav_direct(video_path, shifted, self.tolerance_s)
except Exception as exc: except Exception as exc:
# Log loudly the first time decoding fails so silent # Log loudly the first time decoding fails so a silent
# Module-3-no-op (every prompt skipped because frames_at returned # vqa-module no-op (every prompt skipped because frames_at
# []) is debuggable from the job log instead of post-hoc parquet # returned []) is debuggable from the job log instead of
# inspection. Subsequent failures stay quiet. # post-hoc parquet inspection. Subsequent failures stay quiet.
with self._lock: with self._lock:
already_warned = getattr(self, "_warned_decode_fail", False) already_warned = getattr(self, "_warned_decode_fail", False)
if not already_warned: if not already_warned:
src/lerobot/annotations/steerable_pipeline/modules/general_vqa.py
+16 -26
View File
@@ -13,10 +13,12 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Module 3: general VQA at a timed cadence. """``vqa`` module: general VQA at a timed cadence.
Anchors ``K`` (question, answer) pairs to ``K`` consecutive frames per Every ``1/hz`` seconds an emission tick fires; each tick anchors ``K``
emission. For datasets with multiple cameras, every emission tick produces 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 one ``(vqa, user)`` + ``(vqa, assistant)`` pair *per camera*: each pair is
generated against that camera's frame and stamped with the matching generated against that camera's frame and stamped with the matching
``camera`` field on the emitted rows. The resolver disambiguates via ``camera`` field on the emitted rows. The resolver disambiguates via
@@ -26,7 +28,7 @@ per camera (see ``recipes/pi05_hirobot.yaml``).
Within a single (frame, camera) we still emit at most one ``(vqa, user)`` Within a single (frame, camera) we still emit at most one ``(vqa, user)``
and one ``(vqa, assistant)`` row, so the resolver contract stays scalar. and one ``(vqa, assistant)`` row, so the resolver contract stays scalar.
Question types covered (per the plan's Module 3 table): bbox, keypoint, Question types covered (per the plan's ``vqa`` table): bbox, keypoint,
count, attribute, spatial. The assistant's ``content`` is a JSON string count, attribute, spatial. The assistant's ``content`` is a JSON string
whose schema depends on the question type. Malformed JSON triggers one whose schema depends on the question type. Malformed JSON triggers one
retry inside :meth:`VlmClient.generate_json`. retry inside :meth:`VlmClient.generate_json`.
@@ -35,12 +37,13 @@ retry inside :meth:`VlmClient.generate_json`.
from __future__ import annotations from __future__ import annotations
import json import json
import logging
import random import random
from collections.abc import Sequence from collections.abc import Sequence
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any from typing import Any
from ..config import Module3Config from ..config import VqaConfig
from ..frames import FrameProvider, null_provider, to_image_blocks from ..frames import FrameProvider, null_provider, to_image_blocks
from ..prompts import load as load_prompt from ..prompts import load as load_prompt
from ..reader import EpisodeRecord from ..reader import EpisodeRecord
@@ -89,7 +92,7 @@ class GeneralVqaModule:
"""Emit grounded VQA pairs at a timed cadence.""" """Emit grounded VQA pairs at a timed cadence."""
vlm: VlmClient vlm: VlmClient
config: Module3Config config: VqaConfig
seed: int = 1729 seed: int = 1729
frame_provider: FrameProvider = field(default_factory=null_provider) frame_provider: FrameProvider = field(default_factory=null_provider)
@@ -99,7 +102,7 @@ class GeneralVqaModule:
def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None: def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
if not record.frame_timestamps: if not record.frame_timestamps:
staging.write("module_3", []) staging.write("vqa", [])
return return
rng = random.Random(f"{self.seed}:{record.episode_index}:vqa") rng = random.Random(f"{self.seed}:{record.episode_index}:vqa")
anchor_idx = _emission_anchor_indices( anchor_idx = _emission_anchor_indices(
@@ -111,17 +114,15 @@ class GeneralVqaModule:
# untagged rows that would fail validation. Surface a loud one- # untagged rows that would fail validation. Surface a loud one-
# time warning so this is never silently a no-op. # time warning so this is never silently a no-op.
if not getattr(self, "_warned_no_camera", False): if not getattr(self, "_warned_no_camera", False):
import logging # noqa: PLC0415
logging.getLogger(__name__).warning( logging.getLogger(__name__).warning(
"Module 3 (VQA) found no cameras on the frame provider — " "vqa module found no cameras on the frame provider — "
"every episode will emit zero VQA rows. Check that the " "every episode will emit zero VQA rows. Check that the "
"dataset declares observation.images.* features in " "dataset declares observation.images.* features in "
"meta/info.json; passing --vlm.camera_key=<key> at the " "meta/info.json; passing --vlm.camera_key=<key> at the "
"CLI now also seeds the cameras list as a fallback." "CLI now also seeds the cameras list as a fallback."
) )
self._warned_no_camera = True self._warned_no_camera = True
staging.write("module_3", []) staging.write("vqa", [])
return return
# Build all messages first (one per (frame, camera)), then issue them # Build all messages first (one per (frame, camera)), then issue them
@@ -140,13 +141,13 @@ class GeneralVqaModule:
per_call.append((ts, camera, qtype, messages)) per_call.append((ts, camera, qtype, messages))
if not per_call: if not per_call:
staging.write("module_3", []) staging.write("vqa", [])
return return
results = self.vlm.generate_json([m for _, _, _, m in per_call]) results = self.vlm.generate_json([m for _, _, _, m in per_call])
rows: list[dict[str, Any]] = [] rows: list[dict[str, Any]] = []
for (ts, camera, _qtype, _messages), result in zip(per_call, results): for (ts, camera, _qtype, _messages), result in zip(per_call, results, strict=True):
qa = self._postprocess(result) qa = self._postprocess(result)
if qa is None: if qa is None:
continue continue
@@ -171,10 +172,10 @@ class GeneralVqaModule:
"tool_calls": None, "tool_calls": None,
} }
) )
staging.write("module_3", rows) staging.write("vqa", rows)
def _target_cameras(self) -> list[str]: def _target_cameras(self) -> list[str]:
"""Return the cameras Module 3 should iterate per emission tick. """Return the cameras the ``vqa`` module should iterate per anchored frame.
Defaults to every camera the provider exposes. Datasets with no Defaults to every camera the provider exposes. Datasets with no
cameras (or test/null providers) yield an empty list, which makes cameras (or test/null providers) yield an empty list, which makes
@@ -214,17 +215,6 @@ class GeneralVqaModule:
return None return None
return question.strip(), answer return question.strip(), answer
def _generate_one(
self,
record: EpisodeRecord,
question_type: str,
frame_timestamp: float,
camera_key: str,
) -> tuple[str, dict[str, Any]] | None:
messages = self._build_messages(record, question_type, frame_timestamp, camera_key)
result = self.vlm.generate_json([messages])[0]
return self._postprocess(result)
def _has_image_block(messages: list[dict[str, Any]]) -> bool: def _has_image_block(messages: list[dict[str, Any]]) -> bool:
"""Return True if any user content block is a populated image block.""" """Return True if any user content block is a populated image block."""
src/lerobot/annotations/steerable_pipeline/modules/interjections_and_speech.py
+9 -9
View File
@@ -13,7 +13,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Module 2: interjections + paired speech (EVENT styles + speech atoms). """``interjections`` module: interjections + paired speech (EVENT styles + speech atoms).
Two sub-passes: Two sub-passes:
@@ -26,8 +26,8 @@ Two sub-passes:
speech atom (role:assistant, style:None, tool_calls=[say(...)]) speech atom (role:assistant, style:None, tool_calls=[say(...)])
Both rows go in ``language_events`` at the same timestamp. Both rows go in ``language_events`` at the same timestamp.
Module 1's :meth:`run_plan_updates` reuses Module 2's interjection The ``plan`` module's :meth:`run_plan_updates` reuses this module's
timestamps to refresh the ``plan`` row at the same instant. interjection timestamps to refresh the ``plan`` row at the same instant.
""" """
from __future__ import annotations from __future__ import annotations
@@ -37,7 +37,7 @@ from collections.abc import Sequence
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any from typing import Any
from ..config import Module2Config from ..config import InterjectionsConfig
from ..frames import FrameProvider, null_provider, to_image_blocks from ..frames import FrameProvider, null_provider, to_image_blocks
from ..prompts import load as load_prompt from ..prompts import load as load_prompt
from ..reader import EpisodeRecord, reconstruct_subtask_spans, snap_to_frame from ..reader import EpisodeRecord, reconstruct_subtask_spans, snap_to_frame
@@ -51,7 +51,7 @@ class InterjectionsAndSpeechModule:
"""Generate task-start speech and mid-episode interjection/speech pairs.""" """Generate task-start speech and mid-episode interjection/speech pairs."""
vlm: VlmClient vlm: VlmClient
config: Module2Config config: InterjectionsConfig
seed: int = 1729 seed: int = 1729
frame_provider: FrameProvider = field(default_factory=null_provider) frame_provider: FrameProvider = field(default_factory=null_provider)
@@ -66,13 +66,13 @@ class InterjectionsAndSpeechModule:
initial = self._initial_speech(record) initial = self._initial_speech(record)
if initial: if initial:
rows.append(speech_atom(t0, initial)) rows.append(speech_atom(t0, initial))
# Pull Module 1's subtask spans for this episode so the # Pull the ``plan`` module's subtask spans for this episode so the
# interjection prompt can ground itself in the actual current # interjection prompt can ground itself in the actual current
# subtask at each chosen timestamp. Module 1 ran first. # subtask at each chosen timestamp. The ``plan`` module ran first.
episode_end_t = float(record.frame_timestamps[-1]) if record.frame_timestamps else None episode_end_t = float(record.frame_timestamps[-1]) if record.frame_timestamps else None
subtask_spans = reconstruct_subtask_spans(staging.read("module_1"), episode_end_t=episode_end_t) subtask_spans = reconstruct_subtask_spans(staging.read("plan"), episode_end_t=episode_end_t)
rows.extend(self._mid_episode_interjections(record, subtask_spans)) rows.extend(self._mid_episode_interjections(record, subtask_spans))
staging.write("module_2", rows) staging.write("interjections", rows)
@staticmethod @staticmethod
def _subtask_at(spans: Sequence[dict[str, Any]], t: float) -> str | None: def _subtask_at(spans: Sequence[dict[str, Any]], t: float) -> str | None:
src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py
+14 -14
View File
@@ -13,7 +13,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""Module 1: subtask decomposition + plan + memory (PERSISTENT styles).""" """``plan`` module: subtask decomposition + plan + memory (PERSISTENT styles)."""
from __future__ import annotations from __future__ import annotations
@@ -22,7 +22,7 @@ from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
from ..config import Module1Config from ..config import PlanConfig
from ..frames import ( from ..frames import (
FrameProvider, FrameProvider,
VideoFrameProvider, VideoFrameProvider,
@@ -46,13 +46,13 @@ class PlanSubtasksMemoryModule:
(snapped to an exact frame). (snapped to an exact frame).
- ``plan`` rows: emitted at ``t=0``; refreshed at every interjection - ``plan`` rows: emitted at ``t=0``; refreshed at every interjection
timestamp via :meth:`run_plan_updates` (called by the executor after timestamp via :meth:`run_plan_updates` (called by the executor after
Module 2 completes). the ``interjections`` module completes).
- ``memory`` rows: emitted at each subtask boundary (= subtask start - ``memory`` rows: emitted at each subtask boundary (= subtask start
timestamp from the second subtask onward). timestamp from the second subtask onward).
""" """
vlm: VlmClient vlm: VlmClient
config: Module1Config config: PlanConfig
frame_provider: FrameProvider = field(default_factory=null_provider) frame_provider: FrameProvider = field(default_factory=null_provider)
@property @property
@@ -61,14 +61,14 @@ class PlanSubtasksMemoryModule:
def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None: def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
rows: list[dict[str, Any]] = [] rows: list[dict[str, Any]] = []
# Resolve the task that drives every other Module-1 prompt. May be # Resolve the task that drives every other ``plan``-module prompt.
# the canonical ``record.episode_task`` (default), or a fresh # May be the canonical ``record.episode_task`` (default), or a fresh
# description derived from the video when the canonical task is # description derived from the video when the canonical task is
# empty / placeholder / forced-off (see Module1Config.derive_task_*). # empty / placeholder / forced-off (see PlanConfig.derive_task_*).
effective_task = self._resolve_effective_task(record) effective_task = self._resolve_effective_task(record)
# ``task_aug`` rows at t=0 (role=user), one per rephrasing — the # ``task_aug`` rows at t=0 (role=user), one per rephrasing — the
# PR 1 renderer rotates ``${task}`` deterministically through them # message renderer rotates ``${task}`` deterministically through
# so the policy sees diverse phrasings during training. # them so the policy sees diverse phrasings during training.
t0 = float(record.frame_timestamps[0]) if record.frame_timestamps else 0.0 t0 = float(record.frame_timestamps[0]) if record.frame_timestamps else 0.0
if self.config.n_task_rephrasings > 0 and effective_task: if self.config.n_task_rephrasings > 0 and effective_task:
rephrasings = self._generate_task_rephrasings(effective_task, n=self.config.n_task_rephrasings) rephrasings = self._generate_task_rephrasings(effective_task, n=self.config.n_task_rephrasings)
@@ -134,7 +134,7 @@ class PlanSubtasksMemoryModule:
} }
) )
prior_memory = mem_text prior_memory = mem_text
staging.write("module_1", rows) staging.write("plan", rows)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# Task derivation + rephrasings # Task derivation + rephrasings
@@ -156,7 +156,7 @@ class PlanSubtasksMemoryModule:
) )
def _resolve_effective_task(self, record: EpisodeRecord) -> str: def _resolve_effective_task(self, record: EpisodeRecord) -> str:
"""Decide which task string drives Module 1 for this episode. """Decide which task string drives the ``plan`` module for this episode.
Returns the user-supplied ``record.episode_task`` unless Returns the user-supplied ``record.episode_task`` unless
``derive_task_from_video`` says otherwise (see config docstring). ``derive_task_from_video`` says otherwise (see config docstring).
@@ -182,7 +182,7 @@ class PlanSubtasksMemoryModule:
return task.lower() in self._PLACEHOLDER_TASKS return task.lower() in self._PLACEHOLDER_TASKS
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# VLM call helpers (factored out: every Module-1 prompt below follows # VLM call helpers (factored out: every ``plan``-module prompt below follows
# the same "build messages → single VLM call → pull a named field" # the same "build messages → single VLM call → pull a named field"
# shape, only differing in field name + post-processing). # shape, only differing in field name + post-processing).
# ------------------------------------------------------------------ # ------------------------------------------------------------------
@@ -258,7 +258,7 @@ class PlanSubtasksMemoryModule:
(the previous version told the model "an interjection happened" (the previous version told the model "an interjection happened"
without telling it what the user said). without telling it what the user said).
""" """
existing = staging.read("module_1") existing = staging.read("plan")
# Pass the episode's last frame timestamp so the final subtask # Pass the episode's last frame timestamp so the final subtask
# span is closed (otherwise its ``end`` equals its ``start``, # span is closed (otherwise its ``end`` equals its ``start``,
# zero duration, and the "current subtask at refresh_t" lookup # zero duration, and the "current subtask at refresh_t" lookup
@@ -289,7 +289,7 @@ class PlanSubtasksMemoryModule:
"tool_calls": None, "tool_calls": None,
} }
) )
staging.write("module_1", new_rows) staging.write("plan", new_rows)
def _generate_subtasks(self, record: EpisodeRecord, *, task: str | None = None) -> list[dict[str, Any]]: def _generate_subtasks(self, record: EpisodeRecord, *, task: str | None = None) -> list[dict[str, Any]]:
if record.row_count == 0 or not record.frame_timestamps: if record.row_count == 0 or not record.frame_timestamps:
src/lerobot/annotations/steerable_pipeline/reader.py
+17 -11
View File
@@ -38,6 +38,7 @@ from typing import Any
import pyarrow.parquet as pq import pyarrow.parquet as pq
from lerobot.datasets.io_utils import load_tasks
from lerobot.datasets.utils import DEFAULT_TASKS_PATH from lerobot.datasets.utils import DEFAULT_TASKS_PATH
@@ -83,8 +84,9 @@ def reconstruct_subtask_spans(
which is what downstream consumers (memory, interjection boundary which is what downstream consumers (memory, interjection boundary
selection) expect. selection) expect.
Used by Module 1 (plan-update pass) and Module 2 (interjection Used by the ``plan`` module (plan-update pass) and the
anchoring), which both need the same span shape. ``interjections`` module (interjection anchoring), which both need the
same span shape.
""" """
sorted_rows = sorted( sorted_rows = sorted(
(r for r in rows if r.get("style") == "subtask"), (r for r in rows if r.get("style") == "subtask"),
@@ -105,8 +107,9 @@ def snap_to_frame(t: float, frame_timestamps: Sequence[float]) -> float:
"""Snap an arbitrary float to the nearest exact source frame timestamp. """Snap an arbitrary float to the nearest exact source frame timestamp.
Modules use this when emitting event-style rows so the row's Modules use this when emitting event-style rows so the row's
timestamp matches a real parquet frame (event rows must land on an timestamp matches a real parquet frame: event rows must land on an
exact frame, see PR 1's "exact event matching" rule). exact frame, otherwise the per-frame event lookup the writer does
would never match them.
""" """
if not frame_timestamps: if not frame_timestamps:
return float(t) return float(t)
@@ -115,14 +118,17 @@ def snap_to_frame(t: float, frame_timestamps: Sequence[float]) -> float:
def _load_tasks_lookup(root: Path) -> dict[int, str]: def _load_tasks_lookup(root: Path) -> dict[int, str]:
tasks_path = root / DEFAULT_TASKS_PATH """Map ``task_index -> task`` from ``meta/tasks.parquet``.
if not tasks_path.exists():
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 {} return {}
table = pq.read_table(tasks_path) tasks = load_tasks(root)
cols = {name: table.column(name).to_pylist() for name in table.column_names} return {int(idx): str(task) for task, idx in zip(tasks.index, tasks["task_index"], strict=True)}
if "task_index" in cols and "task" in cols:
return dict(zip(cols["task_index"], cols["task"], strict=True))
raise ValueError(f"meta/tasks.parquet at {tasks_path} missing 'task_index' or 'task'")
def iter_episodes(root: Path, *, only_episodes: tuple[int, ...] | None = None) -> Iterator[EpisodeRecord]: def iter_episodes(root: Path, *, only_episodes: tuple[int, ...] | None = None) -> Iterator[EpisodeRecord]:
src/lerobot/annotations/steerable_pipeline/staging.py
+3 -3
View File
@@ -36,9 +36,9 @@ from typing import Any
ModuleName = str ModuleName = str
_MODULES: tuple[ModuleName, ...] = ( _MODULES: tuple[ModuleName, ...] = (
"module_1", "plan",
"module_2", "interjections",
"module_3", "vqa",
) )
src/lerobot/annotations/steerable_pipeline/validator.py
+4 -4
View File
@@ -15,7 +15,7 @@
# limitations under the License. # limitations under the License.
"""Pre-write validation against staged outputs. """Pre-write validation against staged outputs.
Runs after Modules 1–3 have all written their per-episode artifacts but Runs after all three modules have written their per-episode artifacts but
*before* the writer rewrites parquet shards. The validator never touches *before* the writer rewrites parquet shards. The validator never touches
parquet; it only inspects the staging tree and the source frame timestamps parquet; it only inspects the staging tree and the source frame timestamps
exposed by :class:`EpisodeRecord`. exposed by :class:`EpisodeRecord`.
@@ -218,11 +218,11 @@ class StagingValidator:
except ValueError: except ValueError:
report.add_error(f"ep={episode_index} module={module}: unknown style {style!r}") report.add_error(f"ep={episode_index} module={module}: unknown style {style!r}")
return return
if module == "module_1" and target_col != LANGUAGE_PERSISTENT: if module == "plan" and target_col != LANGUAGE_PERSISTENT:
report.add_error( report.add_error(
f"ep={episode_index} module=module_1 emitted style {style!r} that routes to {target_col} (must be persistent)" f"ep={episode_index} module=plan emitted style {style!r} that routes to {target_col} (must be persistent)"
) )
if module in {"module_2", "module_3"} and target_col != LANGUAGE_EVENTS: if module in {"interjections", "vqa"} and target_col != LANGUAGE_EVENTS:
report.add_error( report.add_error(
f"ep={episode_index} module={module} emitted style {style!r} that routes to {target_col} (must be events)" f"ep={episode_index} module={module} emitted style {style!r} that routes to {target_col} (must be events)"
) )
src/lerobot/annotations/steerable_pipeline/vlm_client.py
+15 -35
View File
@@ -32,10 +32,20 @@ The client speaks one method, :meth:`VlmClient.generate_json`, which:
from __future__ import annotations from __future__ import annotations
import atexit
import base64
import io
import json import json
import os import os
import shlex
import signal
import subprocess
import sys
import threading import threading
import time
import urllib.request
from collections.abc import Callable, Sequence from collections.abc import Callable, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Protocol from typing import Any, Protocol
@@ -212,10 +222,8 @@ def _make_vllm_client(config: VlmConfig) -> VlmClient:
# as CUDNN_STATUS_NOT_INITIALIZED in Qwen-VL vision-tower patch # as CUDNN_STATUS_NOT_INITIALIZED in Qwen-VL vision-tower patch
# embedders. Setting LEROBOT_DISABLE_CUDNN=1 forces native PyTorch # embedders. Setting LEROBOT_DISABLE_CUDNN=1 forces native PyTorch
# convolution kernels — slower but functional. # convolution kernels — slower but functional.
import os as _os # noqa: PLC0415 if os.environ.get("LEROBOT_DISABLE_CUDNN", "").lower() in {"1", "true", "yes"}:
import torch as _torch # noqa: PLC0415 - optional GPU dep, deferred
if _os.environ.get("LEROBOT_DISABLE_CUDNN", "").lower() in {"1", "true", "yes"}:
import torch as _torch # noqa: PLC0415
_torch.backends.cudnn.enabled = False _torch.backends.cudnn.enabled = False
llm_kwargs: dict[str, Any] = { llm_kwargs: dict[str, Any] = {
@@ -259,9 +267,7 @@ def _make_transformers_client(config: VlmConfig) -> VlmClient:
"for VL models." "for VL models."
) )
processor = AutoProcessor.from_pretrained(config.model_id, trust_remote_code=config.trust_remote_code) processor = AutoProcessor.from_pretrained(config.model_id, trust_remote_code=config.trust_remote_code)
import os as _os # noqa: PLC0415 use_accelerate = os.environ.get("LEROBOT_TRANSFORMERS_DEVICE_MAP", "manual") != "manual"
use_accelerate = _os.environ.get("LEROBOT_TRANSFORMERS_DEVICE_MAP", "manual") != "manual"
# ``device_map='auto'`` triggers a known std::bad_alloc on the Qwen3-VL # ``device_map='auto'`` triggers a known std::bad_alloc on the Qwen3-VL
# post-load dispatch path (the alloc fails in accelerate's hook setup # post-load dispatch path (the alloc fails in accelerate's hook setup
# even with TBs of host RAM). Default to manual: load on CPU with # even with TBs of host RAM). Default to manual: load on CPU with
@@ -276,7 +282,7 @@ def _make_transformers_client(config: VlmConfig) -> VlmClient:
trust_remote_code=config.trust_remote_code, trust_remote_code=config.trust_remote_code,
) )
else: else:
import torch as _torch # noqa: PLC0415 import torch as _torch # noqa: PLC0415 - optional GPU dep, deferred
model = auto_cls.from_pretrained( model = auto_cls.from_pretrained(
config.model_id, config.model_id,
@@ -390,8 +396,6 @@ def _make_openai_client(config: VlmConfig) -> VlmClient:
if len(batch) <= 1 or config.client_concurrency <= 1: if len(batch) <= 1 or config.client_concurrency <= 1:
return [_one_call(messages, max_tok, temp) for messages in batch] return [_one_call(messages, max_tok, temp) for messages in batch]
# Parallel fan-out — vllm batches these on the server side. # Parallel fan-out — vllm batches these on the server side.
from concurrent.futures import ThreadPoolExecutor # noqa: PLC0415
max_workers = min(config.client_concurrency, len(batch)) max_workers = min(config.client_concurrency, len(batch))
with ThreadPoolExecutor(max_workers=max_workers) as pool: with ThreadPoolExecutor(max_workers=max_workers) as pool:
futures = [pool.submit(_one_call, messages, max_tok, temp) for messages in batch] futures = [pool.submit(_one_call, messages, max_tok, temp) for messages in batch]
@@ -411,15 +415,6 @@ def _spawn_parallel_inference_servers(config: VlmConfig) -> list[str]:
Returns the list of ``api_base`` URLs the client should round-robin Returns the list of ``api_base`` URLs the client should round-robin
across. across.
""" """
import atexit # noqa: PLC0415
import os as _os # noqa: PLC0415
import shlex # noqa: PLC0415
import signal # noqa: PLC0415
import subprocess # noqa: PLC0415
import sys # noqa: PLC0415
import threading # noqa: PLC0415
import time # noqa: PLC0415
n = config.parallel_servers n = config.parallel_servers
api_bases: list[str] = [] api_bases: list[str] = []
procs: list[subprocess.Popen] = [] procs: list[subprocess.Popen] = []
@@ -449,7 +444,7 @@ def _spawn_parallel_inference_servers(config: VlmConfig) -> list[str]:
for i in range(n): for i in range(n):
port = config.serve_port + i port = config.serve_port + i
gpu = i % num_gpus gpu = i % num_gpus
env = _os.environ.copy() env = os.environ.copy()
env["CUDA_VISIBLE_DEVICES"] = str(gpu) env["CUDA_VISIBLE_DEVICES"] = str(gpu)
cmd = base_cmd.replace("{port}", str(port)) if "{port}" in base_cmd else f"{base_cmd} --port {port}" cmd = base_cmd.replace("{port}", str(port)) if "{port}" in base_cmd else f"{base_cmd} --port {port}"
api_base = f"http://localhost:{port}/v1" api_base = f"http://localhost:{port}/v1"
@@ -522,8 +517,6 @@ def _spawn_parallel_inference_servers(config: VlmConfig) -> list[str]:
def _server_is_up(api_base: str) -> bool: def _server_is_up(api_base: str) -> bool:
"""Return True if ``api_base/models`` answers 200 within 2 seconds.""" """Return True if ``api_base/models`` answers 200 within 2 seconds."""
import urllib.request # noqa: PLC0415
url = api_base.rstrip("/") + "/models" url = api_base.rstrip("/") + "/models"
# ``api_base`` is the user-configured local-server URL we just spawned # ``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 # or the user passed in via ``--vlm.api_base``; the bandit B310 warning
@@ -546,14 +539,6 @@ def _spawn_inference_server(config: VlmConfig) -> str:
Returns the full ``api_base`` URL the OpenAI client should use. Returns the full ``api_base`` URL the OpenAI client should use.
""" """
import atexit # noqa: PLC0415
import shlex # noqa: PLC0415
import signal # noqa: PLC0415
import subprocess # noqa: PLC0415
import sys # noqa: PLC0415
import threading # noqa: PLC0415
import time # noqa: PLC0415
cmd = config.serve_command cmd = config.serve_command
if not cmd: if not cmd:
cmd = ( cmd = (
@@ -695,8 +680,6 @@ def _to_openai_messages(
def _file_to_data_url(path: str) -> str: def _file_to_data_url(path: str) -> str:
"""Read a local video file and return a base64 ``data:video/mp4`` URL.""" """Read a local video file and return a base64 ``data:video/mp4`` URL."""
import base64 # noqa: PLC0415
with open(path, "rb") as f: with open(path, "rb") as f:
b64 = base64.b64encode(f.read()).decode("ascii") b64 = base64.b64encode(f.read()).decode("ascii")
return f"data:video/mp4;base64,{b64}" return f"data:video/mp4;base64,{b64}"
@@ -704,9 +687,6 @@ def _file_to_data_url(path: str) -> str:
def _pil_to_data_url(image: Any) -> str: def _pil_to_data_url(image: Any) -> str:
"""Encode a PIL.Image as a base64 data URL.""" """Encode a PIL.Image as a base64 data URL."""
import base64 # noqa: PLC0415
import io # noqa: PLC0415
buf = io.BytesIO() buf = io.BytesIO()
image.save(buf, format="PNG") image.save(buf, format="PNG")
b64 = base64.b64encode(buf.getvalue()).decode("ascii") b64 = base64.b64encode(buf.getvalue()).decode("ascii")
src/lerobot/annotations/steerable_pipeline/writer.py
+4 -4
View File
@@ -29,7 +29,7 @@ For every episode the writer:
The writer does NOT add a dataset-level ``tools`` column. Tool *calls* are 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 emitted per-row via the existing ``tool_calls`` field on the v3.1 row
struct (PR 1) for every speech atom. The tool *schema* (the description struct for every speech atom. The tool *schema* (the description
of the ``say`` function and its parameters) is a fixed code constant — of the ``say`` function and its parameters) is a fixed code constant —
``SAY_TOOL_SCHEMA`` below — and downstream chat-template consumers import ``SAY_TOOL_SCHEMA`` below — and downstream chat-template consumers import
it directly rather than reading a redundant per-row column. it directly rather than reading a redundant per-row column.
@@ -69,7 +69,7 @@ from .staging import EpisodeStaging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Tool schema constants moved to lerobot.datasets.language in PR 1 — single # Tool schema constants live in lerobot.datasets.language — single
# source of truth. Re-exported here so existing imports # source of truth. Re-exported here so existing imports
# (``from lerobot.annotations.steerable_pipeline.writer import SAY_TOOL_SCHEMA``) # (``from lerobot.annotations.steerable_pipeline.writer import SAY_TOOL_SCHEMA``)
# keep working. # keep working.
@@ -309,8 +309,8 @@ class LanguageColumnsWriter:
# uses `pa.json_()` for the `tool_calls` element type, which # uses `pa.json_()` for the `tool_calls` element type, which
# `pa.array(..., type=...)` cannot materialize from Python lists on # `pa.array(..., type=...)` cannot materialize from Python lists on
# current pyarrow versions. The inferred schema round-trips through # current pyarrow versions. The inferred schema round-trips through
# parquet and `LeRobotDataset` correctly — see PR 1's # parquet and `LeRobotDataset` correctly — `tests/datasets/test_language.py`
# `tests/datasets/test_language.py` which exercises the same flow. # exercises the same flow.
persistent_arr = pa.array(persistent) persistent_arr = pa.array(persistent)
events_arr = pa.array(events) events_arr = pa.array(events)
src/lerobot/scripts/lerobot_annotate.py
+18 -16
View File
@@ -24,7 +24,7 @@ Example:
--root=/path/to/dataset \\ --root=/path/to/dataset \\
--vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct --vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct
For distributed runs, see ``examples/annotation/run_hf_job.py``. For distributed runs, see ``examples/annotations/run_hf_job.py``.
""" """
import logging import logging
@@ -65,27 +65,27 @@ def annotate(cfg: AnnotationPipelineConfig) -> None:
vlm = make_vlm_client(cfg.vlm) vlm = make_vlm_client(cfg.vlm)
frame_provider = make_frame_provider(root, camera_key=cfg.vlm.camera_key) frame_provider = make_frame_provider(root, camera_key=cfg.vlm.camera_key)
# Surface the resolved cameras up front so silent Module-3-no-op # Surface the resolved cameras up front so a silent vqa-module no-op
# regressions are obvious in job output rather than discovered post-hoc # is obvious in job output rather than discovered post-hoc by counting
# by counting parquet rows. # parquet rows.
cam_keys = list(getattr(frame_provider, "camera_keys", []) or []) cam_keys = list(getattr(frame_provider, "camera_keys", []) or [])
logger.info( logger.info(
"annotate: frame_provider default camera=%r, all cameras=%s", "annotate: frame_provider default camera=%r, all cameras=%s",
getattr(frame_provider, "camera_key", None), getattr(frame_provider, "camera_key", None),
cam_keys, cam_keys,
) )
if cfg.module_3.enabled and not cam_keys: if cfg.vqa.enabled and not cam_keys:
logger.warning( logger.warning(
"annotate: Module 3 (VQA) is enabled but no cameras were " "annotate: the vqa module is enabled but no cameras were "
"resolved — Module 3 will produce zero VQA rows. Check " "resolved — it will produce zero VQA rows. Check "
"meta/info.json for observation.images.* features, or pass " "meta/info.json for observation.images.* features, or pass "
"--vlm.camera_key=<key> to seed the cameras list." "--vlm.camera_key=<key> to seed the cameras list."
) )
module_1 = PlanSubtasksMemoryModule(vlm=vlm, config=cfg.module_1, frame_provider=frame_provider) plan = PlanSubtasksMemoryModule(vlm=vlm, config=cfg.plan, frame_provider=frame_provider)
module_2 = InterjectionsAndSpeechModule( interjections = InterjectionsAndSpeechModule(
vlm=vlm, config=cfg.module_2, seed=cfg.seed, frame_provider=frame_provider vlm=vlm, config=cfg.interjections, seed=cfg.seed, frame_provider=frame_provider
) )
module_3 = GeneralVqaModule(vlm=vlm, config=cfg.module_3, seed=cfg.seed, frame_provider=frame_provider) vqa = GeneralVqaModule(vlm=vlm, config=cfg.vqa, seed=cfg.seed, frame_provider=frame_provider)
writer = LanguageColumnsWriter() writer = LanguageColumnsWriter()
validator = StagingValidator( validator = StagingValidator(
dataset_camera_keys=tuple(getattr(frame_provider, "camera_keys", []) or []) or None, dataset_camera_keys=tuple(getattr(frame_provider, "camera_keys", []) or []) or None,
@@ -93,9 +93,9 @@ def annotate(cfg: AnnotationPipelineConfig) -> None:
executor = Executor( executor = Executor(
config=cfg, config=cfg,
module_1=module_1, plan=plan,
module_2=module_2, interjections=interjections,
module_3=module_3, vqa=vqa,
writer=writer, writer=writer,
validator=validator, validator=validator,
) )
@@ -113,14 +113,16 @@ def annotate(cfg: AnnotationPipelineConfig) -> None:
logger.warning(w) logger.warning(w)
if cfg.push_to_hub: if cfg.push_to_hub:
if cfg.repo_id is None:
raise ValueError("--push_to_hub requires --repo_id (the dataset repo to push to).")
_push_to_hub(root, cfg) _push_to_hub(root, cfg)
def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None: def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
"""Upload the annotated dataset directory to the Hugging Face Hub.""" """Upload the annotated dataset directory back to ``cfg.repo_id`` on the Hub."""
from huggingface_hub import HfApi # noqa: PLC0415 from huggingface_hub import HfApi # noqa: PLC0415
repo_id = cfg.push_to_hub repo_id = cfg.repo_id
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)" commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
api = HfApi() api = HfApi()
print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True) print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
tests/annotations/conftest.py
+6 -67
View File
@@ -15,85 +15,24 @@
# limitations under the License. # limitations under the License.
"""Shared fixtures for annotation-pipeline tests. """Shared fixtures for annotation-pipeline tests.
Builds a minimal LeRobot-shaped dataset on disk so writer/validator tests The on-disk dataset builder lives with the other dataset factories in
can exercise real parquet reads and writes without needing a checked-in ``tests/fixtures/dataset_factories.py`` (:func:`build_annotation_dataset`);
LFS dataset. these fixtures only wire it into pytest.
""" """
from __future__ import annotations from __future__ import annotations
import json
from pathlib import Path from pathlib import Path
import pyarrow as pa
import pyarrow.parquet as pq
import pytest import pytest
from tests.fixtures.dataset_factories import build_annotation_dataset
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
@pytest.fixture @pytest.fixture
def fixture_dataset_root(tmp_path: Path) -> Path: def fixture_dataset_root(tmp_path: Path) -> Path:
"""A tiny dataset with two episodes, 12 frames each at 10 fps.""" """A tiny dataset with two episodes, 12 frames each at 10 fps."""
return _build_dataset( return build_annotation_dataset(
tmp_path / "ds", tmp_path / "ds",
episode_specs=[ episode_specs=[
(0, 12, "Could you tidy the kitchen please?"), (0, 12, "Could you tidy the kitchen please?"),
@@ -105,7 +44,7 @@ def fixture_dataset_root(tmp_path: Path) -> Path:
@pytest.fixture @pytest.fixture
def single_episode_root(tmp_path: Path) -> Path: def single_episode_root(tmp_path: Path) -> Path:
return _build_dataset( return build_annotation_dataset(
tmp_path / "ds_one", tmp_path / "ds_one",
episode_specs=[(0, 30, "Pour water from the bottle into the cup.")], episode_specs=[(0, 30, "Pour water from the bottle into the cup.")],
fps=10, fps=10,
tests/annotations/run_e2e_smoke.py
+14 -37
View File
@@ -15,22 +15,19 @@
# limitations under the License. # limitations under the License.
"""Opt-in E2E smoke run for ``make annotation-e2e``. """Opt-in E2E smoke run for ``make annotation-e2e``.
Builds the same fixture used by the pytest suite, runs the full Builds the shared annotation fixture (:func:`build_annotation_dataset`),
annotation pipeline against it with a stub VLM, and prints a short report. runs the full annotation pipeline against it with a stub VLM, and prints a
This is intentionally not a pytest test — it exercises the CLI plumbing short report. This is intentionally not a pytest test — it exercises the
without depending on conftest.py fixtures. 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 from __future__ import annotations
import json
import sys import sys
import tempfile import tempfile
from pathlib import Path 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.config import AnnotationPipelineConfig
from lerobot.annotations.steerable_pipeline.executor import Executor from lerobot.annotations.steerable_pipeline.executor import Executor
from lerobot.annotations.steerable_pipeline.modules import ( 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.validator import StagingValidator
from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
from tests.fixtures.dataset_factories import build_annotation_dataset
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
def _stub_responder(messages): def _stub_responder(messages):
@@ -102,14 +75,18 @@ def _stub_responder(messages):
def main() -> int: def main() -> int:
with tempfile.TemporaryDirectory() as tmp: 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) vlm = StubVlmClient(responder=_stub_responder)
cfg = AnnotationPipelineConfig() cfg = AnnotationPipelineConfig()
executor = Executor( executor = Executor(
config=cfg, config=cfg,
module_1=PlanSubtasksMemoryModule(vlm=vlm, config=cfg.module_1), plan=PlanSubtasksMemoryModule(vlm=vlm, config=cfg.plan),
module_2=InterjectionsAndSpeechModule(vlm=vlm, config=cfg.module_2, seed=cfg.seed), interjections=InterjectionsAndSpeechModule(vlm=vlm, config=cfg.interjections, seed=cfg.seed),
module_3=GeneralVqaModule(vlm=vlm, config=cfg.module_3, seed=cfg.seed), vqa=GeneralVqaModule(vlm=vlm, config=cfg.vqa, seed=cfg.seed),
writer=LanguageColumnsWriter(), writer=LanguageColumnsWriter(),
validator=StagingValidator(), validator=StagingValidator(),
) )
tests/annotations/test_modules.py
+16 -16
View File
@@ -23,9 +23,9 @@ from pathlib import Path
from typing import Any from typing import Any
from lerobot.annotations.steerable_pipeline.config import ( from lerobot.annotations.steerable_pipeline.config import (
Module1Config, InterjectionsConfig,
Module2Config, PlanConfig,
Module3Config, VqaConfig,
) )
from lerobot.annotations.steerable_pipeline.modules import ( from lerobot.annotations.steerable_pipeline.modules import (
GeneralVqaModule, GeneralVqaModule,
@@ -84,11 +84,11 @@ def test_module1_plan_memory_subtask_smoke(fixture_dataset_root: Path, tmp_path:
"Update the memory": {"memory": "wiped the counter once"}, "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)) record = next(iter_episodes(fixture_dataset_root))
staging = EpisodeStaging(tmp_path / "stage", record.episode_index) staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
module.run_episode(record, staging) module.run_episode(record, staging)
rows = staging.read("module_1") rows = staging.read("plan")
styles = {r["style"] for r in rows} styles = {r["style"] for r in rows}
assert {"subtask", "plan", "memory"}.issubset(styles) assert {"subtask", "plan", "memory"}.issubset(styles)
@@ -108,12 +108,12 @@ def test_module2_at_t0_emits_speech_only_no_interjection(fixture_dataset_root: P
) )
module = InterjectionsAndSpeechModule( module = InterjectionsAndSpeechModule(
vlm=vlm, vlm=vlm,
config=Module2Config(max_interjections_per_episode=0), config=InterjectionsConfig(max_interjections_per_episode=0),
) )
record = next(iter_episodes(fixture_dataset_root)) record = next(iter_episodes(fixture_dataset_root))
staging = EpisodeStaging(tmp_path / "stage", record.episode_index) staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
module.run_episode(record, staging) module.run_episode(record, staging)
rows = staging.read("module_2") rows = staging.read("interjections")
assert len(rows) == 1 assert len(rows) == 1
only = rows[0] only = rows[0]
assert only["role"] == "assistant" assert only["role"] == "assistant"
@@ -151,7 +151,7 @@ def test_module2_mid_episode_emits_paired_interjection_and_speech(
) )
module = InterjectionsAndSpeechModule( module = InterjectionsAndSpeechModule(
vlm=vlm, 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, seed=7,
) )
record = next(iter_episodes(fixture_dataset_root)) record = next(iter_episodes(fixture_dataset_root))
@@ -161,7 +161,7 @@ def test_module2_mid_episode_emits_paired_interjection_and_speech(
# production executor guarantees Module 1 ran first). # production executor guarantees Module 1 ran first).
boundary_ts = float(record.frame_timestamps[len(record.frame_timestamps) // 2]) boundary_ts = float(record.frame_timestamps[len(record.frame_timestamps) // 2])
staging.write( staging.write(
"module_1", "plan",
[ [
{ {
"role": "assistant", "role": "assistant",
@@ -180,7 +180,7 @@ def test_module2_mid_episode_emits_paired_interjection_and_speech(
], ],
) )
module.run_episode(record, staging) module.run_episode(record, staging)
rows = staging.read("module_2") rows = staging.read("interjections")
interjections = [r for r in rows if r["style"] == "interjection"] 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"] speeches = [r for r in rows if r["style"] is None and r["role"] == "assistant"]
@@ -198,14 +198,14 @@ def test_module3_vqa_unique_per_frame_and_camera(single_episode_root: Path, tmp_
vlm = make_canned_responder({"frame-grounded visual question": payload}) vlm = make_canned_responder({"frame-grounded visual question": payload})
module = GeneralVqaModule( module = GeneralVqaModule(
vlm=vlm, vlm=vlm,
config=Module3Config(vqa_emission_hz=1.0, K=3), config=VqaConfig(vqa_emission_hz=1.0, K=3),
seed=1, 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)) record = next(iter_episodes(single_episode_root))
staging = EpisodeStaging(tmp_path / "stage", record.episode_index) staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
module.run_episode(record, staging) 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 # every vqa row must carry a camera tag and one of the configured cameras
for r in rows: for r in rows:
assert r["style"] == "vqa" assert r["style"] == "vqa"
@@ -257,7 +257,7 @@ def test_module1_attaches_video_block_to_subtask_prompt(fixture_dataset_root: Pa
# call is the subtask one — keeps the assertions below focused on # call is the subtask one — keeps the assertions below focused on
# ``_generate_subtasks`` rather than fighting the order of unrelated # ``_generate_subtasks`` rather than fighting the order of unrelated
# text-only Module-1 sub-prompts. # text-only Module-1 sub-prompts.
config=Module1Config(max_video_frames=5, frames_per_second=10.0, n_task_rephrasings=0), config=PlanConfig(max_video_frames=5, frames_per_second=10.0, n_task_rephrasings=0),
frame_provider=provider, frame_provider=provider,
) )
record = next(iter_episodes(fixture_dataset_root)) record = next(iter_episodes(fixture_dataset_root))
@@ -304,7 +304,7 @@ def test_module3_attaches_frame_image_block_to_prompt(single_episode_root: Path,
provider = _StubFrameProvider() provider = _StubFrameProvider()
module = GeneralVqaModule( module = GeneralVqaModule(
vlm=_spy_responder(captured, payload), 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, seed=0,
frame_provider=provider, frame_provider=provider,
) )
@@ -336,14 +336,14 @@ def test_module3_assistant_content_is_valid_json(single_episode_root: Path, tmp_
vlm = make_canned_responder({"frame-grounded visual question": payload}) vlm = make_canned_responder({"frame-grounded visual question": payload})
module = GeneralVqaModule( module = GeneralVqaModule(
vlm=vlm, vlm=vlm,
config=Module3Config(vqa_emission_hz=1.0, K=2), config=VqaConfig(vqa_emission_hz=1.0, K=2),
seed=2, seed=2,
frame_provider=_StubFrameProvider(), frame_provider=_StubFrameProvider(),
) )
record = next(iter_episodes(single_episode_root)) record = next(iter_episodes(single_episode_root))
staging = EpisodeStaging(tmp_path / "stage", record.episode_index) staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
module.run_episode(record, staging) module.run_episode(record, staging)
rows = staging.read("module_3") rows = staging.read("vqa")
for row in rows: for row in rows:
if row["role"] == "assistant" and row["style"] == "vqa": if row["role"] == "assistant" and row["style"] == "vqa":
decoded = json.loads(row["content"]) decoded = json.loads(row["content"])
tests/annotations/test_pipeline_recipe_render.py
+9 -9
View File
@@ -23,9 +23,9 @@ import pyarrow.parquet as pq
from lerobot.annotations.steerable_pipeline.config import ( from lerobot.annotations.steerable_pipeline.config import (
AnnotationPipelineConfig, AnnotationPipelineConfig,
Module1Config, InterjectionsConfig,
Module2Config, PlanConfig,
Module3Config, VqaConfig,
) )
from lerobot.annotations.steerable_pipeline.executor import Executor from lerobot.annotations.steerable_pipeline.executor import Executor
from lerobot.annotations.steerable_pipeline.modules import ( from lerobot.annotations.steerable_pipeline.modules import (
@@ -115,15 +115,15 @@ def _build_executor() -> Executor:
}, },
) )
config = AnnotationPipelineConfig( config = AnnotationPipelineConfig(
module_1=Module1Config(), plan=PlanConfig(),
module_2=Module2Config(max_interjections_per_episode=1, interjection_min_t=0.5), interjections=InterjectionsConfig(max_interjections_per_episode=1, interjection_min_t=0.5),
module_3=Module3Config(vqa_emission_hz=1.0, K=2), vqa=VqaConfig(vqa_emission_hz=1.0, K=2),
) )
return Executor( return Executor(
config=config, config=config,
module_1=PlanSubtasksMemoryModule(vlm=vlm, config=config.module_1), plan=PlanSubtasksMemoryModule(vlm=vlm, config=config.plan),
module_2=InterjectionsAndSpeechModule(vlm=vlm, config=config.module_2, seed=config.seed), interjections=InterjectionsAndSpeechModule(vlm=vlm, config=config.interjections, seed=config.seed),
module_3=GeneralVqaModule(vlm=vlm, config=config.module_3, seed=config.seed), vqa=GeneralVqaModule(vlm=vlm, config=config.vqa, seed=config.seed),
writer=LanguageColumnsWriter(), writer=LanguageColumnsWriter(),
validator=StagingValidator(), validator=StagingValidator(),
) )
tests/annotations/test_validator.py
+6 -6
View File
@@ -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: def test_validator_catches_misaligned_timestamps(fixture_dataset_root: Path, tmp_path: Path) -> None:
staging_dir = tmp_path / "stage" staging_dir = tmp_path / "stage"
EpisodeStaging(staging_dir, 0).write( EpisodeStaging(staging_dir, 0).write(
"module_3", "vqa",
[ [
{ {
"role": "assistant", "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: def test_validator_catches_orphan_speech(fixture_dataset_root: Path, tmp_path: Path) -> None:
staging_dir = tmp_path / "stage" staging_dir = tmp_path / "stage"
EpisodeStaging(staging_dir, 0).write( EpisodeStaging(staging_dir, 0).write(
"module_2", "interjections",
[ [
speech_atom(0.0, "Got it."), speech_atom(0.0, "Got it."),
# interjection at 0.3s with NO paired speech # 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: def test_validator_catches_inconsistent_plan_memory(fixture_dataset_root: Path, tmp_path: Path) -> None:
staging_dir = tmp_path / "stage" staging_dir = tmp_path / "stage"
EpisodeStaging(staging_dir, 0).write( EpisodeStaging(staging_dir, 0).write(
"module_1", "plan",
[ [
{ {
"role": "assistant", "role": "assistant",
@@ -93,7 +93,7 @@ def test_validator_catches_inconsistent_plan_memory(fixture_dataset_root: Path,
], ],
) )
EpisodeStaging(staging_dir, 0).write( EpisodeStaging(staging_dir, 0).write(
"module_2", "interjections",
[ [
speech_atom(0.0, "Got it."), speech_atom(0.0, "Got it."),
speech_atom(0.4, "Replanning."), 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: def test_validator_catches_wrong_column(fixture_dataset_root: Path, tmp_path: Path) -> None:
staging_dir = tmp_path / "stage" staging_dir = tmp_path / "stage"
EpisodeStaging(staging_dir, 0).write( EpisodeStaging(staging_dir, 0).write(
"module_1", "plan",
[ [
{"role": "user", "content": "where?", "style": "vqa", "timestamp": 0.0, "tool_calls": None}, {"role": "user", "content": "where?", "style": "vqa", "timestamp": 0.0, "tool_calls": None},
], ],
) )
report = _validate(fixture_dataset_root, staging_dir) report = _validate(fixture_dataset_root, staging_dir)
assert not report.ok 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_writer.py
+17 -17
View File
@@ -35,17 +35,17 @@ def _stage_episode(
staging_dir: Path, staging_dir: Path,
episode_index: int, episode_index: int,
*, *,
module_1: list[dict] | None = None, plan: list[dict] | None = None,
module_2: list[dict] | None = None, interjections: list[dict] | None = None,
module_3: list[dict] | None = None, vqa: list[dict] | None = None,
) -> None: ) -> None:
staging = EpisodeStaging(staging_dir, episode_index) staging = EpisodeStaging(staging_dir, episode_index)
if module_1 is not None: if plan is not None:
staging.write("module_1", module_1) staging.write("plan", plan)
if module_2 is not None: if interjections is not None:
staging.write("module_2", module_2) staging.write("interjections", interjections)
if module_3 is not None: if vqa is not None:
staging.write("module_3", module_3) staging.write("vqa", vqa)
def test_writer_persistence_identity(fixture_dataset_root: Path, tmp_path: Path) -> None: 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( _stage_episode(
staging_dir, staging_dir,
0, 0,
module_1=[ plan=[
{ {
"role": "assistant", "role": "assistant",
"content": "grasp the sponge", "content": "grasp the sponge",
@@ -94,7 +94,7 @@ def test_writer_events_exact_timestamp(fixture_dataset_root: Path, tmp_path: Pat
_stage_episode( _stage_episode(
staging_dir, staging_dir,
0, 0,
module_2=[ interjections=[
speech_atom(0.0, "Got it."), speech_atom(0.0, "Got it."),
{ {
"role": "user", "role": "user",
@@ -127,7 +127,7 @@ def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> No
_stage_episode( _stage_episode(
staging_dir, staging_dir,
0, 0,
module_1=[ plan=[
{ {
"role": "assistant", "role": "assistant",
"content": "do X", "content": "do X",
@@ -150,7 +150,7 @@ def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> No
"tool_calls": None, "tool_calls": None,
}, },
], ],
module_2=[ interjections=[
speech_atom(0.0, "OK"), speech_atom(0.0, "OK"),
{ {
"role": "user", "role": "user",
@@ -161,7 +161,7 @@ def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> No
}, },
speech_atom(0.2, "Waiting"), speech_atom(0.2, "Waiting"),
], ],
module_3=[ vqa=[
{ {
"role": "user", "role": "user",
"content": "where is the cup?", "content": "where is the cup?",
@@ -201,7 +201,7 @@ def test_writer_drops_subtask_index_idempotent(fixture_dataset_root: Path, tmp_p
_stage_episode( _stage_episode(
staging_dir, staging_dir,
0, 0,
module_1=[ plan=[
{ {
"role": "assistant", "role": "assistant",
"content": "do X", "content": "do X",
@@ -277,7 +277,7 @@ def test_writer_does_not_add_tools_column(fixture_dataset_root: Path, tmp_path:
_stage_episode( _stage_episode(
staging_dir, staging_dir,
0, 0,
module_1=[ plan=[
{"role": "assistant", "content": "x", "style": "subtask", "timestamp": 0.0, "tool_calls": None} {"role": "assistant", "content": "x", "style": "subtask", "timestamp": 0.0, "tool_calls": None}
], ],
) )
@@ -316,7 +316,7 @@ def test_annotation_metadata_sync_allows_non_streaming_load(
_stage_episode( _stage_episode(
staging_dir, staging_dir,
0, 0,
module_1=[ plan=[
{"role": "assistant", "content": "do X", "style": "subtask", "timestamp": 0.0, "tool_calls": None} {"role": "assistant", "content": "do X", "style": "subtask", "timestamp": 0.0, "tool_calls": None}
], ],
) )
tests/fixtures/dataset_factories.py
Vendored
+61
View File
@@ -555,3 +555,64 @@ def lerobot_dataset_factory(
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
def empty_lerobot_dataset_factory() -> LeRobotDatasetFactory: def empty_lerobot_dataset_factory() -> LeRobotDatasetFactory:
return partial(LeRobotDataset.create, repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS) 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