mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-10 03:21:54 +00:00
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:
@@ -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>}},
|
||||
...
|
||||
]
|
||||
}}
|
||||
Reference in New Issue
Block a user