From 279c6c7af36183508f5e5c3b2d4bf9c7bb4cc3e6 Mon Sep 17 00:00:00 2001 From: Pepijn <138571049+pkooij@users.noreply.github.com> Date: Wed, 15 Jul 2026 11:38:49 +0200 Subject: [PATCH] feat(annotate): improve VLM subtask annotation (legible contact sheets, seeded relabeling, self-hosted vLLM recipe) (#3896) * feat(annotate): WGO-tuned subtask prompt (atomic completed-events + duration prior) Rework the plan-module subtask segmentation prompt toward the WGO-Bench atomic annotation protocol: segment by completed world-state changes (grasp/place/open/close/pour/insert), fold approach+retreat into their event, keep separate events separate, and add a 2-10s duration prior. Drops the pi0.7 "fewer larger composites preferred" bias that drove under-segmentation on the benchmark. Output JSON shape unchanged. Co-authored-by: Cursor * feat(annotate): seeded-relabeling second pass for subtasks Add an opt-in relabel pass (plan.subtask_seeded_relabel) that, after segmentation, re-labels each span using previous/current/next segment contact sheets and the seed label as a strong prior, minimally correcting it. Mirrors macrodata's best end-to-end labeling step. Boundaries are untouched; one extra VLM call per span. Off by default. Co-authored-by: Cursor * feat(annotate): robust OpenAI-compat client for hosted VLMs Guard against a choice with no message (safety filter or a thinking model that spends its whole budget before emitting content) so one empty reply no longer crashes the whole annotation run; treat it as an empty response and let the existing JSON-retry path handle it. Add an optional `reasoning_effort` knob on VlmConfig, forwarded to the server when set, to cap a thinking model's reasoning (needed for Gemini via its OpenAI-compatible endpoint). Co-authored-by: Cursor * feat(annotate): legible tile-scaled timestamp on contact sheets The burned-in timestamp used the ~10px bitmap default font, which blurs once the model downsamples a full contact sheet into 768px tiles, so the VLM can no longer read the exact source time a boundary depends on. Scale the timestamp to the tile height (with a graceful fallback on older Pillow) so the visual time cue stays readable at sheet resolution. Co-authored-by: Cursor * feat(annotate): lean GEPA-aligned subtask segmentation prompt Replace the verbose, label-heavy segmentation prompt with a lean adaptation of the blog's GEPA-found completed_events_duration_prior recipe: focus on completed manipulation events, explicit no-split / no-merge rules, a 2-10s duration prior, and an instruction to prioritize temporally correct boundaries over label wording. The previous prompt over-weighted label guidance, which traded away boundary precision. Co-authored-by: Cursor * revert: restore original subtask segmentation prompt The lean GEPA-aligned paraphrase (dd4b0110d) regressed Gemini on the 30-ep subset: Seg F1 0.259 -> 0.189 and E2E 0.184 -> 0.135, driven by worse under-segmentation (224 -> 188 preds). The blog's 0.306 came from the actual GEPA-search artifact, which a hand paraphrase does not reproduce. Restore the original prompt, which remains our best config. Co-authored-by: Cursor * feat(annotate): env-var override for prompt templates Allow LEROBOT_PROMPT_OVERRIDE_ to supersede the packaged prompt file at load time. Enables prompt search (GEPA) to inject candidate segmentation prompts into a remote annotate job via an env secret, without committing a branch per candidate. Co-authored-by: Cursor * docs(annotate): genericize hosted-VLM comments (no model name) Co-authored-by: Cursor * docs(annotate): document seeded-relabel and reasoning_effort flags Co-authored-by: Cursor * test(annotate): update subtask-prompt marker to match WGO-tuned prompt The three plan-module tests keyed the canned VLM responder on the literal 'atomic subtasks', which the WGO-tuned segmentation prompt no longer contains (it now segments 'COMPLETED manipulation events'). Point the fixture markers at the current wording so the subtask call is matched again. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- docs/source/annotation_pipeline.mdx | 51 +++--- .../annotations/steerable_pipeline/config.py | 13 ++ .../annotations/steerable_pipeline/frames.py | 11 +- .../modules/plan_subtasks_memory.py | 47 ++++++ .../steerable_pipeline/prompts/__init__.py | 13 +- .../prompts/plan_subtask_relabel.txt | 35 ++++ .../prompts/plan_subtasks.txt | 154 +++++++----------- .../steerable_pipeline/vlm_client.py | 10 +- tests/annotations/test_modules.py | 6 +- 9 files changed, 214 insertions(+), 126 deletions(-) create mode 100644 src/lerobot/annotations/steerable_pipeline/prompts/plan_subtask_relabel.txt diff --git a/docs/source/annotation_pipeline.mdx b/docs/source/annotation_pipeline.mdx index 02658ec9a..3c389663a 100644 --- a/docs/source/annotation_pipeline.mdx +++ b/docs/source/annotation_pipeline.mdx @@ -81,6 +81,12 @@ merged. Both prompts also carry a causal **event-boundary** definition (a new event starts when an object becomes held / is released / reaches a new location / a lid changes state / contents move) to sharpen where cuts land. +Optionally, a third **seeded-relabel** pass (`--plan.subtask_seeded_relabel`) +revisits each span with its previous/current/next segment contact sheets and +minimally corrects the label, using the first label as a prior — it keeps the +boundaries fixed and only sharpens wording, at the cost of one extra call per +subtask. + The resulting spans are then stitched into a gap-free, full-episode cover, so **every frame has exactly one active subtask**. See [`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py) @@ -157,30 +163,33 @@ Every module is on by default and can be toggled independently (set to ### The VLM (`--vlm.*`) -| Flag | Default | What it does | -| -------------------------- | ------------------ | ----------------------------------------------------------------------------------- | -| `--vlm.model_id` | `Qwen/Qwen3.6-27B` | The model to serve and prompt. | -| `--vlm.camera_key` | first `images.*` | Which camera every prompt is grounded on. | -| `--vlm.serve_command` | auto | The exact `vllm serve …` command (set TP size, GPU memory, `--max-model-len` here). | -| `--vlm.parallel_servers` | `1` | Independent servers for round-robin routing (one per GPU). | -| `--vlm.num_gpus` | `0` | GPUs per server (`0` = one each). | -| `--vlm.client_concurrency` | `16` | In-flight requests across all servers. | -| `--vlm.max_new_tokens` | `512` | Generation cap per call. | -| `--vlm.temperature` | `0.2` | Sampling temperature. | +| Flag | Default | What it does | +| -------------------------- | ------------------ | ------------------------------------------------------------------------------------ | +| `--vlm.model_id` | `Qwen/Qwen3.6-27B` | The model to serve and prompt. | +| `--vlm.camera_key` | first `images.*` | Which camera every prompt is grounded on. | +| `--vlm.serve_command` | auto | The exact `vllm serve …` command (set TP size, GPU memory, `--max-model-len` here). | +| `--vlm.parallel_servers` | `1` | Independent servers for round-robin routing (one per GPU). | +| `--vlm.num_gpus` | `0` | GPUs per server (`0` = one each). | +| `--vlm.client_concurrency` | `16` | In-flight requests across all servers. | +| `--vlm.max_new_tokens` | `512` | Generation cap per call. | +| `--vlm.temperature` | `0.2` | Sampling temperature. | +| `--vlm.reasoning_effort` | `null` | Thinking-budget hint (`low`/`medium`/`high`) forwarded to OpenAI-compatible servers. | ### Subtasks / plan / memory (`--plan.*`) -| Flag | Default | What it does | -| ------------------------------- | ---------- | ------------------------------------------------------------------------------------------------------------------------- | -| `--plan.frames_per_second` | `2.0` | Frame sampling rate for the contact sheets (`2.0` = one frame every 0.5s). | -| `--plan.max_frames_per_prompt` | `60` | Frame budget per VLM call. Episodes whose sampling exceeds this are auto-windowed at the same density, then stitched. | -| `--plan.contact_sheet_columns` | `5` | Columns per contact-sheet grid (`contact_sheet_frames_per_sheet` tiles, time row-major). | -| `--plan.plan_max_steps` | `8` | Upper bound on subtasks per episode. | -| `--plan.subtask_describe_first` | `true` | Run the describe→segment grounding pass (best subtask quality; +1 call/episode). | -| `--plan.emit_plan` | `true` | Emit the numbered `plan` rows (`false` = subtasks + memory only). | -| `--plan.emit_memory` | `true` | Emit the `memory` rows (`false` = subtasks + plan only); symmetric to `emit_plan`. | -| `--plan.n_task_rephrasings` | `10` | How many `task_aug` rephrasings to emit (`0` disables). | -| `--plan.derive_task_from_video` | `if_short` | Use the dataset task as-is (`off`), only when it's missing/short (`if_short`), or always re-derive from video (`always`). | +| Flag | Default | What it does | +| ------------------------------- | ---------- | ---------------------------------------------------------------------------------------------------------------------------- | +| `--plan.frames_per_second` | `2.0` | Frame sampling rate for the contact sheets (`2.0` = one frame every 0.5s). | +| `--plan.max_frames_per_prompt` | `60` | Frame budget per VLM call. Episodes whose sampling exceeds this are auto-windowed at the same density, then stitched. | +| `--plan.contact_sheet_columns` | `5` | Columns per contact-sheet grid (`contact_sheet_frames_per_sheet` tiles, time row-major). | +| `--plan.plan_max_steps` | `8` | Upper bound on subtasks per episode. | +| `--plan.subtask_describe_first` | `true` | Run the describe→segment grounding pass (best subtask quality; +1 call/episode). | +| `--plan.subtask_seeded_relabel` | `false` | Second pass: re-label each subtask from its prev/current/next contact sheets, seeded with the first label (+1 call/subtask). | +| `--plan.subtask_relabel_frames` | `5` | Frames sampled uniformly per segment sheet in the relabel pass (only used when `subtask_seeded_relabel=true`). | +| `--plan.emit_plan` | `true` | Emit the numbered `plan` rows (`false` = subtasks + memory only). | +| `--plan.emit_memory` | `true` | Emit the `memory` rows (`false` = subtasks + plan only); symmetric to `emit_plan`. | +| `--plan.n_task_rephrasings` | `10` | How many `task_aug` rephrasings to emit (`0` disables). | +| `--plan.derive_task_from_video` | `if_short` | Use the dataset task as-is (`off`), only when it's missing/short (`if_short`), or always re-derive from video (`always`). | ### Interjections + VQA diff --git a/src/lerobot/annotations/steerable_pipeline/config.py b/src/lerobot/annotations/steerable_pipeline/config.py index 86d6cadd9..73f41f4f7 100644 --- a/src/lerobot/annotations/steerable_pipeline/config.py +++ b/src/lerobot/annotations/steerable_pipeline/config.py @@ -65,6 +65,14 @@ class PlanConfig: # invented from the task text (+1 VLM call/episode). subtask_describe_first: bool = True + # Seeded relabeling: after segmentation, re-label each span with a focused + # pass that sees the previous / current / next segment contact sheets and + # minimally corrects the seed label (macrodata's best end-to-end labeling + # step). Costs +1 VLM call per subtask; off by default. + subtask_seeded_relabel: bool = False + # Frames sampled uniformly per segment sheet in the relabel pass. + subtask_relabel_frames: int = 5 + # Emit ``style="plan"`` rows at each boundary; False = subtasks + memory only. emit_plan: bool = True @@ -160,6 +168,11 @@ class VlmConfig: # Forwarded as extra_body.chat_template_kwargs (e.g. {"enable_thinking": false}). chat_template_kwargs: dict[str, Any] | None = None + # OpenAI-style thinking budget hint ("low"/"medium"/"high"); forwarded to + # the server when set. Used to cap a thinking model's reasoning so it + # leaves tokens for the actual JSON answer on OpenAI-compatible endpoints. + reasoning_effort: str | None = None + @dataclass class ExecutorConfig: diff --git a/src/lerobot/annotations/steerable_pipeline/frames.py b/src/lerobot/annotations/steerable_pipeline/frames.py index 5a6a5879c..403581132 100644 --- a/src/lerobot/annotations/steerable_pipeline/frames.py +++ b/src/lerobot/annotations/steerable_pipeline/frames.py @@ -413,7 +413,16 @@ def _draw_timestamp_badge(image: PIL.Image.Image, timestamp: float) -> PIL.Image result = image.copy() draw = ImageDraw.Draw(result) - font = ImageFont.load_default() + # Scale the timestamp to the tile so it stays legible after the model + # downsamples the full sheet into 768px tiles — a tiny bitmap font blurs + # at contact-sheet resolution and the VLM can no longer read the exact + # source time, which is what the boundary score depends on. ``size=`` is + # supported by Pillow's bitmap default since 10.1; fall back otherwise. + badge_px = max(14, round(image.height * 0.12)) + try: + font = ImageFont.load_default(size=badge_px) + except TypeError: + font = ImageFont.load_default() label = f"{timestamp:06.2f}s" left, top, right, bottom = draw.textbbox((0, 0), label, font=font) text_w, text_h = right - left, bottom - top diff --git a/src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py b/src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py index b6df6551c..fa39b5f89 100644 --- a/src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py +++ b/src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py @@ -116,6 +116,8 @@ class PlanSubtasksMemoryModule: rows.extend(self._task_aug_rows([effective_task, *variants], t0)) subtask_spans = self._generate_subtasks(record, task=effective_task) + if self.config.subtask_seeded_relabel and subtask_spans: + subtask_spans = self._seeded_relabel(record, subtask_spans, effective_task) # subtask rows for span in subtask_spans: @@ -509,6 +511,51 @@ class PlanSubtasksMemoryModule: return cleaned + def _seeded_relabel( + self, record: EpisodeRecord, spans: list[dict[str, Any]], task: str + ) -> list[dict[str, Any]]: + """Re-label each span using prev/current/next segment contact sheets. + + Boundaries are kept fixed; only ``text`` is refined. The original + ("seed") label is passed as a strong prior so the model verifies and + minimally corrects it rather than re-describing from scratch — the + macrodata seeded-relabeling step. One VLM call per span. + """ + n = len(spans) + out: list[dict[str, Any]] = [] + for i, span in enumerate(spans): + content: list[dict[str, Any]] = [] + if i > 0: + content += self._segment_sheet(record, spans[i - 1]) + content += self._segment_sheet(record, span) + if i < n - 1: + content += self._segment_sheet(record, spans[i + 1]) + prompt = load_prompt("plan_subtask_relabel").format( + episode_task=task, + seed_label=span["text"], + segment_index=i + 1, + segment_count=n, + start=float(span["start"]), + end=float(span["end"]), + ) + content.append({"type": "text", "text": prompt}) + label = self._vlm_field([{"role": "user", "content": content}], "label") + text = label.strip() if isinstance(label, str) and label.strip() else span["text"] + out.append({**span, "text": text}) + return out + + def _segment_sheet(self, record: EpisodeRecord, span: dict[str, Any]) -> list[dict[str, Any]]: + """Contact-sheet block(s) for one span: up to N frames sampled uniformly.""" + s, e = float(span["start"]), float(span["end"]) + n = max(1, int(self.config.subtask_relabel_frames)) + if e <= s or n == 1: + timestamps = [s] + else: + step = (e - s) / (n - 1) + timestamps = [s + i * step for i in range(n)] + frames = self.frame_provider.frames_at(record, timestamps) + return self._contact_sheet_blocks(frames, timestamps[: len(frames)]) + def _generate_subtasks_windowed( self, record: EpisodeRecord, task: str, window_s: float ) -> list[dict[str, Any]]: diff --git a/src/lerobot/annotations/steerable_pipeline/prompts/__init__.py b/src/lerobot/annotations/steerable_pipeline/prompts/__init__.py index 5ce8e163b..de884f11d 100644 --- a/src/lerobot/annotations/steerable_pipeline/prompts/__init__.py +++ b/src/lerobot/annotations/steerable_pipeline/prompts/__init__.py @@ -22,12 +22,23 @@ plain editors and roundtrip cleanly through ``ruff format``. from __future__ import annotations +import os from pathlib import Path _DIR = Path(__file__).parent def load(name: str) -> str: - """Read prompt template ``name.txt`` from the ``prompts/`` directory.""" + """Read prompt template ``name.txt`` from the ``prompts/`` directory. + + A ``LEROBOT_PROMPT_OVERRIDE_`` environment variable, when set to a + non-empty value, takes precedence over the packaged file. This lets prompt + search (e.g. GEPA) inject candidate templates into a remote job without + rebuilding the package; the override must keep the same ``{placeholder}`` + fields the call site formats in. + """ + override = os.environ.get(f"LEROBOT_PROMPT_OVERRIDE_{name}") + if override and override.strip(): + return override path = _DIR / f"{name}.txt" return path.read_text(encoding="utf-8") diff --git a/src/lerobot/annotations/steerable_pipeline/prompts/plan_subtask_relabel.txt b/src/lerobot/annotations/steerable_pipeline/prompts/plan_subtask_relabel.txt new file mode 100644 index 000000000..336abfe05 --- /dev/null +++ b/src/lerobot/annotations/steerable_pipeline/prompts/plan_subtask_relabel.txt @@ -0,0 +1,35 @@ +Annotate one fixed segment from a longer robot demonstration. + +Return only JSON: + {{"label": ""}} + +You are shown up to three timestamped contact sheets, in order: +- The FIRST sheet is the PREVIOUS segment (context only); it may be absent. +- The SECOND sheet is the CURRENT target segment. +- The THIRD sheet is the NEXT segment (context only); it may be absent. +Each tile has its timestamp (seconds, absolute video time) burned into its +top-left corner. + +Episode instruction: "{episode_task}" +Target segment: {segment_index} of {segment_count} +Target time: {start:.2f}s to {end:.2f}s +Original predicted label for this exact segment: "{seed_label}" + +Rules: +- Label ONLY the current target segment (the second sheet). Use the + previous/next sheets only to disambiguate what changed. +- Treat the original predicted label as a STRONG PRIOR, not ground truth: + verify it against the current segment and correct it minimally. +- If it already names the right action and main object, keep it; only fix + grammar or add a clearly visible essential detail. +- If it is vague but directionally correct, make it more specific. +- If it describes the previous/next segment, the wrong action, wrong + object, wrong destination, or a wrong state change, replace it. +- Do not describe the previous or next segment, and do not split, merge, + or move the fixed segment. +- Do not introduce an action that is not clearly visible in the current + target segment. +- Use one concise imperative phrase. Name the manipulated object and the + action / state change. Include source, destination, side, direction, + final placement, or opened/closed state when visible and central. +- Do not mention timestamps, frame numbers, uncertainty, or intent. diff --git a/src/lerobot/annotations/steerable_pipeline/prompts/plan_subtasks.txt b/src/lerobot/annotations/steerable_pipeline/prompts/plan_subtasks.txt index e6a5260a7..8227245a4 100644 --- a/src/lerobot/annotations/steerable_pipeline/prompts/plan_subtasks.txt +++ b/src/lerobot/annotations/steerable_pipeline/prompts/plan_subtasks.txt @@ -1,112 +1,68 @@ -You are labeling a teleoperated robot demonstration. +You are annotating a teleoperated robot demonstration shown as +timestamped contact sheets (each tile has its time in seconds burned +into the top-left corner). The operator's goal was: "{episode_task}" -The user originally asked: "{episode_task}" +{observation_block}Reconstruct the sequence of COMPLETED manipulation events the robot +performs, in chronological order. Output one segment per event with a +[start, end] time in seconds and a short action label. -You are shown the entire demonstration as a single video. Watch the -whole clip, then segment it into a list of consecutive atomic subtasks -the robot performs. +GROUNDING — read first, it overrides everything below: +- Label ONLY events you can SEE in the frames. The instruction is the + goal; the VIDEO is the ground truth for what actually happened. +- Do NOT invent, anticipate, or pad steps that are not shown. -{observation_block}GROUNDING — read this first, it overrides everything below: -- Label ONLY what the robot actually does in the video. Every subtask - you emit must correspond to motion you can SEE in specific frames. -- Do NOT invent, anticipate, or pad. If the robot only does one thing - (e.g. it just navigates to a location and the clip ends), emit - EXACTLY ONE subtask. Many demonstrations are a single atomic skill. -- ``max_steps`` below is a hard CEILING, not a target. Emitting fewer - subtasks than the ceiling is not just allowed, it is expected for - short / atomic demonstrations. One correct subtask is far better - than several invented ones. -- If the video does not clearly show the action implied by the task, - describe what you actually see — do NOT fabricate the task's steps - from the instruction text. The instruction tells you the goal; the - VIDEO is the ground truth for what happened. +Granularity — segment by completed events, not by motion: +- Start a NEW segment whenever the world state changes: an object is + grasped, lifted, transported, placed, or released; a held object + changes; a drawer/door/lid/container opens or closes; contents move + between containers (poured); a tool starts or stops acting on a + surface. Watch the gripper open/close transitions — they usually mark + boundaries. +- Do NOT split approach, reach, grasp adjustment, small repositioning, + hesitation, or retreat into their own segments. Fold each into the + event it belongs to (the approach is part of the pick; the retreat is + part of the place). +- Do NOT merge separate completed events. Each distinct pick, place, + open, close, pour, push, wipe, or insert is its own segment, even when + they repeat on different objects or locations. +- Most segments last 2-10 seconds. Shorter segments are okay ONLY for + fast pick / place / open / close / release events. Never emit a + segment shorter than {min_subtask_seconds} seconds; merge a too-short + candidate into its neighbour instead. +- Skip idle time, pure camera motion, and tiny hand jitter. -Authoring rules — Hi Robot atom granularity, pi0.7-style short prompts: +Labels — short imperative phrases: +- One concise command naming the action and the manipulated object, e.g. + "pick up the red cup", "put the cup on the shelf", "open the top + drawer", "pour water into the glass", "insert the plug into the + socket". +- Include source, destination, side, direction, or the final + open/closed state when it is visible and central to the event. +- Prefer these verbs (extend only when none fits): pick up, put, place, + push, pull, turn, press, open, close, pour, insert, wipe, stack. + Disambiguate by what you SEE: + * STACK vs PUT: object placed ON TOP OF another object -> "stack". + * INSERT vs PUT: object pushed INTO a fitted slot/hole/socket -> "insert". + * PICK UP vs PUT (direction): gripper CLOSES and object moves WITH + the hand -> "pick up"; gripper OPENS and object stays -> "put". + * POUR vs PUT: source is tilted and contents flow -> "pour". +- Use the exact object nouns implied by the task; stay consistent across + the episode (don't switch "cube" to "block"). +- Write imperative commands, never third person ("the robot ..."), and + drop articles/adverbs. -- Each subtask = one COMPOSITE atomic skill the low-level policy can - execute end-to-end. A "skill" bundles its own approach motion with - its terminal action — do NOT split the approach off as its own - subtask. The whole-arm policy already learns to reach as part of - every manipulation primitive. -- Write each subtask as an IMPERATIVE COMMAND, starting with one of - these verbs (extend only when none fits): - pick up — approach + grasp + lift in one subtask - put on/in — transport + release in one subtask - place on/in — synonym of "put"; pick one and stay consistent - push — contact + linear shove - pull — contact + linear retract - turn — rotary actuation - press