From 6d2c3bb1e2cb81ba001e15c648be7efc94c3e9f3 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 29 Jul 2026 15:18:48 +0200 Subject: [PATCH] refactor(g05): use recipe message dropout --- docs/source/g05.mdx | 13 +-- docs/source/language_and_recipes.mdx | 35 +++--- src/lerobot/configs/recipe.py | 13 ++- .../configs/recipes/g05_bbox_subtask.yaml | 82 +++++--------- src/lerobot/datasets/language_render.py | 84 ++++----------- tests/configs/test_recipe.py | 45 ++++---- tests/datasets/test_language_render.py | 102 ++++++++++++------ 7 files changed, 166 insertions(+), 208 deletions(-) diff --git a/docs/source/g05.mdx b/docs/source/g05.mdx index 012e020f7..ff120ca74 100644 --- a/docs/source/g05.mdx +++ b/docs/source/g05.mdx @@ -156,15 +156,10 @@ lerobot-train \ The bundled `g05_bbox_subtask.yaml` recipe resolves the active `language_persistent` `subtask` and camera-scoped grounded `vqa` event at each -sample timestamp. It filters out unavailable formats before selecting one of -four author-compatible objectives: - -| Assistant sequence | Weight | -| -------------------------- | -----: | -| Action only | 1 | -| Subtask, then action | 2 | -| BBox, then action | 1 | -| BBox, subtask, then action | 1 | +sample timestamp. It applies independent `0.5` dropout to the optional BBox and +Subtask targets, producing Action-only, Subtask+Action, BBox+Action, and joint +BBox+Subtask+Action samples. An unavailable annotation is skipped before +dropout through the recipe's `if_present` guard. Grounded VQA boxes are converted from pixel-space `xyxy` JSON using the source camera dimensions captured before image resizing, then serialized as G0.5 diff --git a/docs/source/language_and_recipes.mdx b/docs/source/language_and_recipes.mdx index b6ebaae75..470338f97 100644 --- a/docs/source/language_and_recipes.mdx +++ b/docs/source/language_and_recipes.mdx @@ -146,29 +146,22 @@ The renderer does not apply a tokenizer chat template. Policy processors decide Blend recipes select one weighted sub-recipe deterministically from the sample index. `recipes/subtask_mem.yaml` trains the compact core blend — high-level subtask prediction, low-level execution, and memory. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds VQA and spoken interjection responses. -Annotation-dependent blends can set `select_from_applicable: true` and declare -`requires` on each component. The renderer first removes components whose -required bindings resolve to `None`, then performs the deterministic weighted -selection. This matches mixed-CoT policies where unavailable annotation formats -must not consume probability: +Message turns can set `dropout` to a probability between `0` and `1`. Dropout is +deterministic for a given sample index and independent between turns. Combine it +with `if_present` to sample optional supervision formats while gracefully +skipping annotations that are unavailable: ```yaml -select_from_applicable: true -blend: - subtask: - weight: 2 - requires: [subtask] - messages: - - { - role: assistant, - content: "${subtask}", - stream: low_level, - target: true, - } - action: - weight: 1 - messages: - - { role: user, content: "${task}", stream: low_level } +messages: + - { role: user, content: "${task}", stream: low_level } + - { + role: assistant, + content: "${subtask}", + stream: low_level, + target: true, + if_present: subtask, + dropout: 0.5, + } ``` A message recipe with a supervised assistant turn on the `low_level` stream trains diff --git a/src/lerobot/configs/recipe.py b/src/lerobot/configs/recipe.py index 3b6ab2557..fe5c4c183 100644 --- a/src/lerobot/configs/recipe.py +++ b/src/lerobot/configs/recipe.py @@ -18,6 +18,7 @@ from __future__ import annotations import re from dataclasses import dataclass +from math import isfinite from pathlib import Path from typing import Any, Literal, get_args @@ -49,8 +50,9 @@ class MessageTurn: ``content`` may be a plain string, a list of HF-style multimodal blocks, or ``None`` when ``tool_calls_from`` supplies tool-call payloads instead. ``stream`` tags the turn for downstream filtering, ``target`` flags it as a - training target, and ``if_present`` skips the turn when the named binding - resolves to ``None``. + training target, ``if_present`` skips the turn when the named binding + resolves to ``None``, and ``dropout`` independently omits it for a + deterministic fraction of samples. """ role: MessageRole @@ -59,6 +61,7 @@ class MessageTurn: target: bool = False if_present: str | None = None tool_calls_from: str | None = None + dropout: float = 0.0 def __post_init__(self) -> None: """Validate role, stream, and content after dataclass construction.""" @@ -86,6 +89,11 @@ class MessageTurn: raise ValueError( "Multimodal content blocks must be HF-style dictionaries with a type key." ) + if isinstance(self.dropout, bool) or not isinstance(self.dropout, int | float): + raise TypeError("MessageTurn.dropout must be a probability.") + self.dropout = float(self.dropout) + if not isfinite(self.dropout) or not 0.0 <= self.dropout <= 1.0: + raise ValueError("MessageTurn.dropout must be between 0 and 1.") @classmethod def from_dict(cls, data: dict[str, Any]) -> MessageTurn: @@ -106,7 +114,6 @@ class TrainingRecipe: bindings: dict[str, str] | None = None requires: list[str] | None = None blend: dict[str, TrainingRecipe] | None = None - select_from_applicable: bool = False weight: float | None = None def __post_init__(self) -> None: diff --git a/src/lerobot/configs/recipes/g05_bbox_subtask.yaml b/src/lerobot/configs/recipes/g05_bbox_subtask.yaml index c3bcf2528..4738232b9 100644 --- a/src/lerobot/configs/recipes/g05_bbox_subtask.yaml +++ b/src/lerobot/configs/recipes/g05_bbox_subtask.yaml @@ -1,60 +1,28 @@ -# G0.5's released MixedSamplesBuilder semantics, limited to the CoT signals -# LeRobot can currently provide: active subtasks and grounded bounding boxes. +# G0.5 CoT supervision from the signals LeRobot can currently provide: +# active subtasks and camera-grounded bounding boxes. # -# The renderer first removes branches whose required annotations are missing, -# then deterministically samples among the remaining weights. The G0.5 policy -# processor converts the selected target messages into the checkpoint's exact -# BBox/Subtask/action template and supervises the complete assistant segment. +# BBox and Subtask turns are dropped independently and deterministically. This +# produces Action-only, BBox+Action, Subtask+Action, and +# BBox+Subtask+Action samples without a policy-specific recipe selector. -select_from_applicable: true -blend: - no_cot_action: - weight: 1 - messages: - - {role: user, content: "${task}", stream: low_level} +bindings: + bbox: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.exterior)" - subtask_action: - weight: 2 - requires: [subtask] - messages: - - {role: user, content: "${task}", stream: low_level} - - { - role: assistant, - content: "Subtask: ${subtask}", - stream: low_level, - target: true, - } - - bbox_action: - weight: 1 - requires: [bbox] - bindings: - bbox: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.exterior)" - messages: - - {role: user, content: "${task}", stream: low_level} - - { - role: assistant, - content: "BBoxJSON: ${bbox}", - stream: low_level, - target: true, - } - - bbox_subtask_action: - weight: 1 - requires: [bbox, subtask] - bindings: - bbox: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.exterior)" - messages: - - {role: user, content: "${task}", stream: low_level} - - { - role: assistant, - content: "BBoxJSON: ${bbox}", - stream: low_level, - target: true, - } - - { - role: assistant, - content: "Subtask: ${subtask}", - stream: low_level, - target: true, - } +messages: + - {role: user, content: "${task}", stream: low_level} + - { + role: assistant, + content: "BBoxJSON: ${bbox}", + stream: low_level, + target: true, + if_present: bbox, + dropout: 0.5, + } + - { + role: assistant, + content: "Subtask: ${subtask}", + stream: low_level, + target: true, + if_present: subtask, + dropout: 0.5, + } diff --git a/src/lerobot/datasets/language_render.py b/src/lerobot/datasets/language_render.py index c4420b1cd..c67c0f6a2 100644 --- a/src/lerobot/datasets/language_render.py +++ b/src/lerobot/datasets/language_render.py @@ -169,17 +169,6 @@ def render_sample( persistent_rows = _normalize_rows(persistent or []) event_rows = _normalize_rows(events or []) - if recipe.blend is not None and recipe.select_from_applicable: - return _render_applicable_blend( - recipe, - persistent=persistent_rows, - events=event_rows, - t=t, - sample_idx=sample_idx, - task=task, - dataset_ctx=dataset_ctx, - ) - # Route sparse VQA frames to a matching view-specific component before weighted selection. # This avoids dropping annotated frames or selecting VQA without annotations. if recipe.blend is not None: @@ -205,55 +194,7 @@ def render_sample( task=task, dataset_ctx=dataset_ctx, ) - return _render_message_recipe(selected_recipe, bindings) - - -def _render_applicable_blend( - recipe: TrainingRecipe, - *, - persistent: Sequence[LanguageRow], - events: Sequence[LanguageRow], - t: float, - sample_idx: int, - task: str | None, - dataset_ctx: Any | None, -) -> RenderedMessages | None: - """Select deterministically among components whose required bindings resolve. - - Ordinary blends preserve their historical select-then-resolve behavior. - Annotation-dependent policies such as G0.5 instead need the author - ``MixedSamplesBuilder`` contract: discard unavailable CoT formats first, - then draw according to the remaining relative weights. - """ - - assert recipe.blend is not None - renderable: list[tuple[float, RenderedMessages]] = [] - for component in recipe.blend.values(): - bindings = _resolve_bindings( - component, - persistent=persistent, - events=events, - t=t, - sample_idx=sample_idx, - task=task, - dataset_ctx=dataset_ctx, - ) - rendered = _render_message_recipe(component, bindings) - if rendered is not None: - renderable.append((float(component.weight or 0.0), rendered)) - - if not renderable: - return None - - total_weight = sum(weight for weight, _ in renderable) - digest = hashlib.blake2b(f"applicable:{sample_idx}".encode(), digest_size=8).digest() - draw = int.from_bytes(digest, "big") / 2**64 * total_weight - cumulative = 0.0 - for weight, rendered in renderable: - cumulative += weight - if draw < cumulative: - return rendered - return renderable[-1][1] + return _render_message_recipe(selected_recipe, bindings, sample_idx=sample_idx) def _render_vqa_if_present( @@ -284,7 +225,7 @@ def _render_vqa_if_present( task=task, dataset_ctx=dataset_ctx, ) - rendered = _render_message_recipe(component, bindings) + rendered = _render_message_recipe(component, bindings, sample_idx=sample_idx) if rendered is not None: renderable.append((float(component.weight or 0.0), rendered)) @@ -441,6 +382,8 @@ def _parse_resolver_args(args: str) -> dict[str, Any]: def _render_message_recipe( recipe: TrainingRecipe, bindings: dict[str, LanguageRow | str | None], + *, + sample_idx: int, ) -> RenderedMessages | None: """Expand ``recipe.messages`` into rendered chat messages using ``bindings``.""" assert recipe.messages is not None @@ -450,9 +393,11 @@ def _render_message_recipe( streams: list[str | None] = [] target_indices: list[int] = [] - for turn in recipe.messages: + for turn_idx, turn in enumerate(recipe.messages): if turn.if_present is not None and bindings.get(turn.if_present) is None: continue + if _drop_message_turn(turn.dropout, sample_idx=sample_idx, turn_idx=turn_idx): + continue message = {"role": turn.role} if turn.content is not None: @@ -484,6 +429,21 @@ def _render_message_recipe( return rendered +def _drop_message_turn(dropout: float, *, sample_idx: int, turn_idx: int) -> bool: + """Apply reproducible, independent dropout to one recipe message.""" + + if dropout <= 0.0: + return False + if dropout >= 1.0: + return True + digest = hashlib.blake2b( + f"message_dropout:{sample_idx}:{turn_idx}".encode(), + digest_size=8, + ).digest() + draw = int.from_bytes(digest, "big") / 2**64 + return draw < dropout + + def _render_content( content: str | list[dict[str, Any]], bindings: dict[str, LanguageRow | str | None], diff --git a/tests/configs/test_recipe.py b/tests/configs/test_recipe.py index a39212c21..88b72525f 100644 --- a/tests/configs/test_recipe.py +++ b/tests/configs/test_recipe.py @@ -47,6 +47,17 @@ def test_message_turn_requires_a_stream(): MessageTurn(role="user", content="${task}") +@pytest.mark.parametrize("dropout", [-0.1, 1.1, float("inf"), float("nan")]) +def test_message_turn_rejects_invalid_dropout(dropout): + with pytest.raises(ValueError, match="between 0 and 1"): + MessageTurn(role="user", content="${task}", stream="high_level", dropout=dropout) + + +def test_message_turn_rejects_non_numeric_dropout(): + with pytest.raises(TypeError, match="probability"): + MessageTurn(role="user", content="${task}", stream="high_level", dropout="half") + + def test_message_recipe_requires_at_least_one_target(): with pytest.raises(ValueError, match="target"): TrainingRecipe( @@ -152,33 +163,25 @@ def test_from_dict_with_nested_blend(): assert isinstance(recipe.blend["a"].messages[0], MessageTurn) -def test_applicable_blend_round_trips_from_dict(): +def test_message_dropout_round_trips_from_dict(): recipe = TrainingRecipe.from_dict( { - "select_from_applicable": True, - "blend": { - "subtask": { - "weight": 2, - "requires": ["subtask"], - "messages": [ - { - "role": "assistant", - "content": "${subtask}", - "stream": "low_level", - "target": True, - } - ], + "messages": [ + {"role": "user", "content": "${task}", "stream": "low_level"}, + { + "role": "assistant", + "content": "${subtask}", + "stream": "low_level", + "target": True, + "if_present": "subtask", + "dropout": 0.5, }, - "action": { - "weight": 1, - "messages": [{"role": "user", "content": "${task}", "stream": "low_level"}], - }, - }, + ] } ) - assert recipe.select_from_applicable - assert recipe.blend["subtask"].requires == ["subtask"] + assert recipe.messages[1].dropout == 0.5 + assert recipe.messages[1].if_present == "subtask" def test_from_yaml_round_trips_through_load_recipe(tmp_path: Path): diff --git a/tests/datasets/test_language_render.py b/tests/datasets/test_language_render.py index 4cc8bf947..84376f8c8 100644 --- a/tests/datasets/test_language_render.py +++ b/tests/datasets/test_language_render.py @@ -176,27 +176,19 @@ def test_deterministic_blend_sampling(): assert first == second -def test_applicable_blend_filters_missing_bindings_before_weighted_selection(): +def test_message_dropout_skips_missing_optional_bindings(): recipe = TrainingRecipe( - select_from_applicable=True, - blend={ - "missing": TrainingRecipe( - weight=1_000, - requires=["subtask"], - messages=[ - MessageTurn( - role="assistant", - content="${subtask}", - stream="low_level", - target=True, - ) - ], + messages=[ + MessageTurn(role="user", content="${task}", stream="low_level"), + MessageTurn( + role="assistant", + content="${subtask}", + stream="low_level", + target=True, + if_present="subtask", + dropout=0.5, ), - "action": TrainingRecipe( - weight=1, - messages=[MessageTurn(role="user", content="${task}", stream="low_level")], - ), - }, + ] ) rendered = render_sample( @@ -212,7 +204,7 @@ def test_applicable_blend_filters_missing_bindings_before_weighted_selection(): assert rendered["target_message_indices"] == [] -def test_applicable_blend_can_select_joint_bbox_subtask_target(): +def test_g05_dropout_produces_all_bbox_subtask_combinations(): recipe = TrainingRecipe.from_yaml("src/lerobot/configs/recipes/g05_bbox_subtask.yaml") persistent = [persistent_row("assistant", "grasp the cup", "subtask", 0.0)] events = [ @@ -224,24 +216,64 @@ def test_applicable_blend_can_select_joint_bbox_subtask_target(): } ] - rendered = next( - candidate - for sample_idx in range(100) - if ( - candidate := render_sample( - recipe=recipe, - persistent=persistent, - events=events, - t=0.0, - sample_idx=sample_idx, - task="pick the cup", - ) + rendered_by_targets = {} + for sample_idx in range(100): + rendered = render_sample( + recipe=recipe, + persistent=persistent, + events=events, + t=0.0, + sample_idx=sample_idx, + task="pick the cup", ) - and candidate["target_message_indices"] == [1, 2] + targets = tuple( + rendered["messages"][idx]["content"].split(":", 1)[0] + for idx in rendered["target_message_indices"] + ) + rendered_by_targets.setdefault(targets, rendered) + + assert set(rendered_by_targets) == { + (), + ("BBoxJSON",), + ("Subtask",), + ("BBoxJSON", "Subtask"), + } + joint = rendered_by_targets[("BBoxJSON", "Subtask")] + assert joint["messages"][1]["content"].startswith("BBoxJSON:") + assert joint["messages"][2]["content"] == "Subtask: grasp the cup" + + +def test_message_dropout_is_deterministic_for_sample_index(): + recipe = TrainingRecipe( + messages=[ + MessageTurn(role="user", content="${task}", stream="low_level"), + MessageTurn( + role="assistant", + content="${subtask}", + stream="low_level", + target=True, + dropout=0.5, + ), + ] ) - assert rendered["messages"][1]["content"].startswith("BBoxJSON:") - assert rendered["messages"][2]["content"] == "Subtask: grasp the cup" + first = render_sample( + recipe=recipe, + persistent=PERSISTENT, + events=[], + t=0.0, + sample_idx=42, + task="pick the cup", + ) + second = render_sample( + recipe=recipe, + persistent=PERSISTENT, + events=[], + t=0.0, + sample_idx=42, + task="pick the cup", + ) + assert first == second def test_emitted_at_filters_vqa_by_camera():