feat(annotate): add plan toggle, drop subtask verify pass, 4xH200 job

- PlanConfig.emit_plan (default True): keep subtasks + memory but skip
  the per-boundary "plan" rows and their VLM call when False.
- Remove the subtask_verify pass entirely: pruning dropped legitimate
  subtasks and the stitch step already guarantees full-episode coverage.
  Deletes _verify_subtasks, both call sites, and the now-unused
  module_1_subtask_verify prompt.
- run_hf_job example: 4xH200 (4 vllm servers), emit_plan=false, vqa off.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Pepijn
2026-06-02 18:00:41 +02:00
parent 1417fd69b2
commit 4c86332fe3
4 changed files with 56 additions and 163 deletions
+18 -21
View File
@@ -1,10 +1,10 @@
#!/usr/bin/env python
"""Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6-27B VLM).
Spawns one ``h200x2`` job that:
Spawns one ``h200x4`` job that:
1. installs this branch of ``lerobot`` plus the annotation extras,
2. boots two vllm servers (one per GPU) with Qwen3.6-27B (dense VLM),
2. boots four vllm servers (one per GPU) with Qwen3.6-27B (dense VLM),
3. runs the plan / interjections / vqa modules across the dataset
in free-form mode (each episode generates its own subtasks +
memory),
@@ -36,13 +36,13 @@ CMD = (
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
"export VLLM_VIDEO_BACKEND=pyav && "
"lerobot-annotate "
"--repo_id=pepijn223/robocasa_smoke_2atomic_v3 "
"--dest_repo_id=pepijn223/robocasa_smoke_2atomic_v3_ann "
"--repo_id=pepijn223/robocasa_pretrain_human300_v4 "
"--dest_repo_id=pepijn223/robocasa_pretrain_human300_v4_annotated5 "
"--push_to_hub=true "
"--vlm.backend=openai "
"--vlm.model_id=Qwen/Qwen3.6-27B "
"--vlm.parallel_servers=2 "
"--vlm.num_gpus=2 "
"--vlm.parallel_servers=4 "
"--vlm.num_gpus=4 "
'--vlm.serve_command="vllm serve Qwen/Qwen3.6-27B '
"--tensor-parallel-size 1 --max-model-len 32768 "
'--gpu-memory-utilization 0.8 --uvicorn-log-level warning --port {port}" '
@@ -63,7 +63,7 @@ CMD = (
# CONTEXT BUDGET: with embedded frames, each frame is ~250-320 vision
# tokens. The model's context is 32768 (see --max-model-len). 32
# frames sampled uniformly across the episode (~8-10k tokens) fits
# comfortably alongside the prompt and the describe/verify passes.
# comfortably alongside the prompt and the describe pass.
# Do NOT raise max_video_frames toward 128 with embedded frames — that
# is ~33-39k tokens and overflows the context (BadRequestError 400,
# "Input length exceeds maximum context length").
@@ -73,7 +73,7 @@ CMD = (
# Constant 1 fps density via windowing: episodes longer than 32s are
# split into 32-second windows (each 32 frames @ 1 fps, fits context),
# so long episodes get MORE subtasks instead of a sparser whole-episode
# view. describe->segment->verify runs per window; spans are merged +
# view. describe->segment runs per window; spans are merged +
# stitched to a contiguous whole-episode cover. 0 disables.
"--plan.subtask_window_seconds=32 "
# IMPORTANT for RoboCasa: the dataset's task string ("Navigate to the
@@ -95,26 +95,23 @@ CMD = (
# the subtask text — useful only for long composite manipulation
# tasks. Leave off for RoboCasa atomic / navigation.
# Keep subtask decomposition tight for atomic tasks:
"--plan.plan_max_steps=6 "
# NOTE: the multi-call subtask quality chain (describe -> segment ->
# verify, 3 VLM calls/episode) is ON BY DEFAULT now. Pass
# --plan.subtask_describe_first=false / --plan.subtask_verify=false to
# disable on datasets you've verified are easy and want fewer calls.
"--plan.plan_max_steps=10 "
# Only annotate subtasks + memory — skip the numbered "plan" rows
# (and their per-boundary VLM call). Flip to true to re-enable plan.
"--plan.emit_plan=false "
# NOTE: the grounding pass (describe -> segment, +1 VLM call/episode)
# is ON BY DEFAULT. Pass --plan.subtask_describe_first=false to disable
# on datasets you've verified are easy and want fewer calls.
# Phase 2 — interjections + speech.
"--interjections.max_interjections_per_episode=6 "
# Phase 4 — general VQA.
# Ground VQA on the SAME single camera as plan/interjections
# (--vlm.camera_key) instead of iterating every camera. The whole
# pipeline then focuses on one view, e.g. observation.images.base.
"--vqa.restrict_to_default_camera=true "
"--vqa.K=1 "
"--vqa.vqa_emission_hz=1.0"
# Phase 4 — general VQA: DISABLED for this run.
"--vqa.enabled=false"
)
job = run_job(
image="vllm/vllm-openai:latest",
command=["bash", "-c", CMD],
flavor="h200x2",
flavor="h200x4",
secrets={"HF_TOKEN": token},
timeout="2h",
)
@@ -74,27 +74,19 @@ class PlanConfig:
min_subtask_seconds: float = 1.5
plan_max_steps: int = 8
# Multi-call subtask quality chain. ON by default — the single-call
# 'watch video -> emit subtask JSON' pattern makes the VLM commit to
# structured output before reasoning about the video, so it
# pattern-matches the task text and hallucinates steps. The chain
# costs 2 extra VLM calls/episode (3 total for subtasks) but is the
# difference between trustworthy and fabricated labels. Set either to
# False to trade quality for fewer calls on datasets you've verified
# are easy.
#
# ``subtask_describe_first``: run a grounding pass that narrates ONLY
# what is visible in the video (no subtask JSON yet), then inject that
# description into the segmentation prompt. Forces the model to
# observe before committing to structured output — the strongest
# lever against subtasks invented from the task text. +1 VLM call/ep.
# description into the segmentation prompt. Forces the model to observe
# before committing to structured output — the strongest lever against
# subtasks invented from the task text. ON by default; +1 VLM call/ep.
# Set False to trade quality for fewer calls on easy datasets.
subtask_describe_first: bool = True
# ``subtask_verify``: after segmentation, re-watch the video and drop
# any proposed subtask that can't be verified as visible. Prunes
# hallucinations; can only remove subtasks, never add/rewrite them.
# Fail-open (keeps un-verified spans if the verify call returns
# nothing). +1 VLM call/ep.
subtask_verify: bool = True
# Emit ``style="plan"`` rows (the numbered still-todo list re-emitted at
# every subtask boundary). Set False to keep only subtasks + memory and
# skip the plan rows entirely — saves one ``_generate_plan`` VLM call per
# subtask boundary. Subtask and memory generation are unaffected.
emit_plan: bool = True
# NOTE: subtask spans are ALWAYS stitched into a contiguous
# full-episode cover (first subtask pulled back to t0, gaps closed,
@@ -170,19 +170,22 @@ class PlanSubtasksMemoryModule:
# contains exactly the subtasks that started at or after the
# current span. Saves the runtime from having to derive
# "what's still left" at inference time.
for span in subtask_spans:
boundary_t = snap_to_frame(span["start"], record.frame_timestamps)
plan_text = self._generate_plan(record, subtask_spans, refresh_t=boundary_t, task=effective_task)
if plan_text is not None:
rows.append(
{
"role": "assistant",
"content": plan_text,
"style": "plan",
"timestamp": float(boundary_t),
"tool_calls": None,
}
if self.config.emit_plan:
for span in subtask_spans:
boundary_t = snap_to_frame(span["start"], record.frame_timestamps)
plan_text = self._generate_plan(
record, subtask_spans, refresh_t=boundary_t, task=effective_task
)
if plan_text is not None:
rows.append(
{
"role": "assistant",
"content": plan_text,
"style": "plan",
"timestamp": float(boundary_t),
"tool_calls": None,
}
)
# memory rows at every subtask boundary except the very first start
prior_memory = ""
for i, span in enumerate(subtask_spans[1:], start=1):
@@ -450,7 +453,7 @@ class PlanSubtasksMemoryModule:
def _episode_video_block(
self, record: EpisodeRecord, window: tuple[float, float] | None = None
) -> list[dict[str, Any]]:
"""Video block for the segmentation / describe / verify prompts.
"""Video block for the segmentation / describe prompts.
Always returns a block that actually carries the video. When
``use_video_url`` is set we try the server-side ``video_url``
@@ -514,6 +517,8 @@ class PlanSubtasksMemoryModule:
(the previous version told the model "an interjection happened"
without telling it what the user said).
"""
if not self.config.emit_plan:
return
existing = staging.read("plan")
# Pass the episode's last frame timestamp so the final subtask
# span is closed (otherwise its ``end`` equals its ``start``,
@@ -559,9 +564,6 @@ class PlanSubtasksMemoryModule:
the model segments its own grounded observations instead of
pattern-matching the task text.
2. segmentation emit subtask JSON (as before).
3. ``subtask_verify`` an adversarial pass that re-watches the
video and drops any proposed subtask it cannot actually see,
pruning hallucinations.
"""
if record.row_count == 0 or not record.frame_timestamps:
return []
@@ -573,8 +575,8 @@ class PlanSubtasksMemoryModule:
# than one window, process the episode in fixed-length windows so
# the VLM always sees ``frames_per_second`` density (instead of a
# sparse 32-frame whole-episode view). Each window runs the full
# describe -> segment -> verify chain on its own frames; results
# are merged + stitched into a contiguous whole-episode cover.
# describe -> segment chain on its own frames; results are merged +
# stitched into a contiguous whole-episode cover.
window_s = float(getattr(self.config, "subtask_window_seconds", 0.0) or 0.0)
if window_s > 0.0 and episode_duration > window_s:
return self._generate_subtasks_windowed(record, effective_task, window_s)
@@ -606,18 +608,11 @@ class PlanSubtasksMemoryModule:
if not cleaned:
return []
# ---- Pass 3 (optional): verification / pruning ---------------
if getattr(self.config, "subtask_verify", False):
cleaned = self._verify_subtasks(record, effective_task, cleaned)
if not cleaned:
return []
# ---- Full-episode coverage stitch ----------------------------
# The VLM (especially after the verify pass prunes spans) can
# leave the first subtask starting after t0 or leave gaps between
# spans, so the subtask timeline no longer tiles the whole
# episode and frames fall through with no active subtask. Always
# stitch the surviving spans into a contiguous cover of
# The VLM can leave the first subtask starting after t0 or leave
# gaps between spans, so the subtask timeline no longer tiles the
# whole episode and frames fall through with no active subtask.
# Always stitch the surviving spans into a contiguous cover of
# [t0, t_last] — there is no scenario where a sparse, gap-ridden
# subtask timeline is desirable for conditioning.
cleaned = self._stitch_full_coverage(cleaned, record)
@@ -630,8 +625,8 @@ class PlanSubtasksMemoryModule:
"""Subtask generation in fixed-length windows at constant fps.
Splits ``[t0, t_last]`` into consecutive windows of ``window_s``
seconds, runs the describe -> segment -> verify chain on each
window's own frames (sampled at ``frames_per_second``), offsets
seconds, runs the describe -> segment chain on each window's own
frames (sampled at ``frames_per_second``), offsets
each window's spans back to absolute episode time, then merges +
stitches into a contiguous whole-episode cover.
"""
@@ -662,7 +657,7 @@ class PlanSubtasksMemoryModule:
def _subtasks_for_window(
self, record: EpisodeRecord, task: str, w0: float, w1: float
) -> list[dict[str, Any]]:
"""Run describe -> segment -> verify on one ``[w0, w1]`` window.
"""Run describe -> segment on one ``[w0, w1]`` window.
The model works in window-RELATIVE time ``[0, L]`` (it perceives
the window as a clip starting at 0); spans are offset back to
@@ -698,11 +693,6 @@ class PlanSubtasksMemoryModule:
if not cleaned:
return []
if getattr(self.config, "subtask_verify", False):
cleaned = self._verify_subtasks(record, task, cleaned, window=window)
if not cleaned:
return []
# Offset window-relative spans back to absolute episode time.
for s in cleaned:
s["start"] = w0 + float(s["start"])
@@ -722,7 +712,7 @@ class PlanSubtasksMemoryModule:
subtask's ``end`` extends to the last frame ``t_last``.
Starts are otherwise left as the (already frame-snapped, distinct)
values the VLM + verify produced only the FIRST start is pulled
values the VLM produced only the FIRST start is pulled
back to ``t0``, which can't collide with a later span because it
was already the earliest. Purely deterministic; runs after the
VLM passes.
@@ -792,59 +782,6 @@ class PlanSubtasksMemoryModule:
text = self._vlm_field(self._video_message(record, prompt, window=window), "description")
return text.strip() if isinstance(text, str) and text.strip() else ""
def _verify_subtasks(
self,
record: EpisodeRecord,
task: str,
spans: list[dict[str, Any]],
window: tuple[float, float] | None = None,
) -> list[dict[str, Any]]:
"""Adversarial pass: drop proposed subtasks not visible in the video.
Keeps the original span on a verified ``text`` match (the verify
prompt is told not to rewrite text), so verification can only
PRUNE never invent or mutate. If the verify call fails or
returns nothing parseable, the un-verified spans are kept (fail
open: better to keep a possibly-good label than silently drop
everything on a transient VLM hiccup).
"""
import json # noqa: PLC0415
subtasks_json = json.dumps(
{
"subtasks": [
{"text": s["text"], "start": round(s["start"], 3), "end": round(s["end"], 3)}
for s in spans
]
},
indent=2,
)
prompt = load_prompt("module_1_subtask_verify").format(episode_task=task, subtasks_json=subtasks_json)
kept_raw = self._vlm_field(self._video_message(record, prompt, window=window), "subtasks")
# Windowed verify: the video is sampled from the absolute window
# ``[w0, w1]`` but the model perceives it as a clip starting at 0,
# so proposed + returned times are window-RELATIVE in ``[0, L]``.
# Clamp to that relative range and skip the absolute frame-snap
# dedupe (done once later on the merged absolute-time set).
clamp = (0.0, float(window[1] - window[0])) if window is not None else None
kept = self._clean_spans(kept_raw, record, bounds=clamp, dedupe=window is None)
if not kept:
logger.info(
"episode %d: verify pass returned nothing — keeping the %d "
"un-verified subtask(s) (fail-open)",
record.episode_index,
len(spans),
)
return spans
if len(kept) < len(spans):
logger.info(
"episode %d: verify pass pruned %d -> %d subtask(s)",
record.episode_index,
len(spans),
len(kept),
)
return kept
@staticmethod
def _dedupe_starts_to_distinct_frames(
spans: list[dict[str, Any]], record: EpisodeRecord
@@ -1,33 +0,0 @@
You previously segmented a teleoperated robot demonstration into these
candidate subtasks (JSON):
{subtasks_json}
The user's task was: "{episode_task}"
This is a VERIFICATION pass. Re-watch the video. For EACH candidate
subtask, decide whether the robot can ACTUALLY be seen performing that
action within its [start, end] time window.
Rules:
- KEEP a subtask only if its action is clearly visible in the video in
roughly that time window.
- DROP any subtask whose action you cannot see, that describes
something not actually present in the video, that was inferred from
the task instruction rather than observed, or that duplicates another
kept subtask.
- Do NOT add new subtasks. Do NOT rewrite the text of kept subtasks.
Do NOT change the start/end timestamps of kept subtasks.
- It is correct and expected to return FEWER subtasks than you were
given — even just one — if that is all the video supports. Returning
zero is allowed if none can be verified.
Output strictly valid JSON of the SAME shape, containing only the kept
subtasks in chronological order:
{{
"subtasks": [
{{"text": "<kept verbatim>", "start": <float>, "end": <float>}},
...
]
}}