diff --git a/examples/annotations/run_hf_job.py b/examples/annotations/run_hf_job.py index 86575f72f..6af40a268 100644 --- a/examples/annotations/run_hf_job.py +++ b/examples/annotations/run_hf_job.py @@ -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", ) diff --git a/src/lerobot/annotations/steerable_pipeline/config.py b/src/lerobot/annotations/steerable_pipeline/config.py index 63fbe83b0..cdcf38072 100644 --- a/src/lerobot/annotations/steerable_pipeline/config.py +++ b/src/lerobot/annotations/steerable_pipeline/config.py @@ -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, 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 0d9c1a7dd..fecd42d3a 100644 --- a/src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py +++ b/src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py @@ -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 diff --git a/src/lerobot/annotations/steerable_pipeline/prompts/module_1_subtask_verify.txt b/src/lerobot/annotations/steerable_pipeline/prompts/module_1_subtask_verify.txt deleted file mode 100644 index e52dc0aeb..000000000 --- a/src/lerobot/annotations/steerable_pipeline/prompts/module_1_subtask_verify.txt +++ /dev/null @@ -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": "", "start": , "end": }}, - ... - ] - }}