From 29a85509cb38f2918dfab427e9d422b83130d96a Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 30 Jun 2026 11:56:32 +0200 Subject: [PATCH] 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 --- .../annotations/steerable_pipeline/config.py | 8 ++++ .../modules/plan_subtasks_memory.py | 47 +++++++++++++++++++ .../prompts/plan_subtask_relabel.txt | 35 ++++++++++++++ 3 files changed, 90 insertions(+) create mode 100644 src/lerobot/annotations/steerable_pipeline/prompts/plan_subtask_relabel.txt diff --git a/src/lerobot/annotations/steerable_pipeline/config.py b/src/lerobot/annotations/steerable_pipeline/config.py index 86d6cadd9..52188183e 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 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/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.