refactor(g05): keep recipe runtime unchanged

This commit is contained in:
Pepijn
2026-07-29 15:59:38 +02:00
parent 3dee8e640a
commit fe4497ba27
7 changed files with 19 additions and 157 deletions
+3 -4
View File
@@ -149,10 +149,9 @@ lerobot-train \
The bundled `g05_bbox_subtask.yaml` recipe resolves the active The bundled `g05_bbox_subtask.yaml` recipe resolves the active
`language_persistent` `subtask` and camera-scoped grounded `vqa` event at each `language_persistent` `subtask` and camera-scoped grounded `vqa` event at each
sample timestamp. It applies independent `0.5` dropout to the optional BBox and sample timestamp. It emits each optional BBox or Subtask target when that
Subtask targets, producing Action-only, Subtask+Action, BBox+Action, and joint annotation is present; unavailable annotations are skipped through the recipe's
BBox+Subtask+Action samples. An unavailable annotation is skipped before existing `if_present` guard.
dropout through the recipe's `if_present` guard.
Grounded VQA boxes are converted from pixel-space `xyxy` JSON using the source Grounded VQA boxes are converted from pixel-space `xyxy` JSON using the source
camera dimensions captured before image resizing, then serialized as G0.5 camera dimensions captured before image resizing, then serialized as G0.5
-18
View File
@@ -146,24 +146,6 @@ The renderer does not apply a tokenizer chat template. Policy processors decide
Blend recipes select one weighted sub-recipe deterministically from the sample index. 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. `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.
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
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 A message recipe with a supervised assistant turn on the `low_level` stream trains
the π0.5 paper's joint sequence instead of a blend: the target span gets text CE the π0.5 paper's joint sequence instead of a blend: the target span gets text CE
while also conditioning the action losses in the same forward. while also conditioning the action losses in the same forward.
+2 -10
View File
@@ -18,7 +18,6 @@ from __future__ import annotations
import re import re
from dataclasses import dataclass from dataclasses import dataclass
from math import isfinite
from pathlib import Path from pathlib import Path
from typing import Any, Literal, get_args from typing import Any, Literal, get_args
@@ -50,9 +49,8 @@ class MessageTurn:
``content`` may be a plain string, a list of HF-style multimodal blocks, or ``content`` may be a plain string, a list of HF-style multimodal blocks, or
``None`` when ``tool_calls_from`` supplies tool-call payloads instead. ``None`` when ``tool_calls_from`` supplies tool-call payloads instead.
``stream`` tags the turn for downstream filtering, ``target`` flags it as a ``stream`` tags the turn for downstream filtering, ``target`` flags it as a
training target, ``if_present`` skips the turn when the named binding training target, and ``if_present`` skips the turn when the named binding
resolves to ``None``, and ``dropout`` independently omits it for a resolves to ``None``.
deterministic fraction of samples.
""" """
role: MessageRole role: MessageRole
@@ -61,7 +59,6 @@ class MessageTurn:
target: bool = False target: bool = False
if_present: str | None = None if_present: str | None = None
tool_calls_from: str | None = None tool_calls_from: str | None = None
dropout: float = 0.0
def __post_init__(self) -> None: def __post_init__(self) -> None:
"""Validate role, stream, and content after dataclass construction.""" """Validate role, stream, and content after dataclass construction."""
@@ -89,11 +86,6 @@ class MessageTurn:
raise ValueError( raise ValueError(
"Multimodal content blocks must be HF-style dictionaries with a type key." "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 @classmethod
def from_dict(cls, data: dict[str, Any]) -> MessageTurn: def from_dict(cls, data: dict[str, Any]) -> MessageTurn:
@@ -1,9 +1,6 @@
# G0.5 CoT supervision from the signals LeRobot can currently provide: # G0.5 CoT supervision from the signals LeRobot can currently provide:
# active subtasks and camera-grounded bounding boxes. # active subtasks and camera-grounded bounding boxes.
# # Each target is emitted when its corresponding annotation is present.
# 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.
bindings: bindings:
bbox: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.exterior)" bbox: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.exterior)"
@@ -16,7 +13,6 @@ messages:
stream: low_level, stream: low_level,
target: true, target: true,
if_present: bbox, if_present: bbox,
dropout: 0.5,
} }
- { - {
role: assistant, role: assistant,
@@ -24,5 +20,4 @@ messages:
stream: low_level, stream: low_level,
target: true, target: true,
if_present: subtask, if_present: subtask,
dropout: 0.5,
} }
+3 -22
View File
@@ -194,7 +194,7 @@ def render_sample(
task=task, task=task,
dataset_ctx=dataset_ctx, dataset_ctx=dataset_ctx,
) )
return _render_message_recipe(selected_recipe, bindings, sample_idx=sample_idx) return _render_message_recipe(selected_recipe, bindings)
def _render_vqa_if_present( def _render_vqa_if_present(
@@ -225,7 +225,7 @@ def _render_vqa_if_present(
task=task, task=task,
dataset_ctx=dataset_ctx, dataset_ctx=dataset_ctx,
) )
rendered = _render_message_recipe(component, bindings, sample_idx=sample_idx) rendered = _render_message_recipe(component, bindings)
if rendered is not None: if rendered is not None:
renderable.append((float(component.weight or 0.0), rendered)) renderable.append((float(component.weight or 0.0), rendered))
@@ -382,8 +382,6 @@ def _parse_resolver_args(args: str) -> dict[str, Any]:
def _render_message_recipe( def _render_message_recipe(
recipe: TrainingRecipe, recipe: TrainingRecipe,
bindings: dict[str, LanguageRow | str | None], bindings: dict[str, LanguageRow | str | None],
*,
sample_idx: int,
) -> RenderedMessages | None: ) -> RenderedMessages | None:
"""Expand ``recipe.messages`` into rendered chat messages using ``bindings``.""" """Expand ``recipe.messages`` into rendered chat messages using ``bindings``."""
assert recipe.messages is not None assert recipe.messages is not None
@@ -393,11 +391,9 @@ def _render_message_recipe(
streams: list[str | None] = [] streams: list[str | None] = []
target_indices: list[int] = [] target_indices: list[int] = []
for turn_idx, turn in enumerate(recipe.messages): for turn in recipe.messages:
if turn.if_present is not None and bindings.get(turn.if_present) is None: if turn.if_present is not None and bindings.get(turn.if_present) is None:
continue continue
if _drop_message_turn(turn.dropout, sample_idx=sample_idx, turn_idx=turn_idx):
continue
message = {"role": turn.role} message = {"role": turn.role}
if turn.content is not None: if turn.content is not None:
@@ -429,21 +425,6 @@ def _render_message_recipe(
return rendered 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( def _render_content(
content: str | list[dict[str, Any]], content: str | list[dict[str, Any]],
bindings: dict[str, LanguageRow | str | None], bindings: dict[str, LanguageRow | str | None],
+1 -14
View File
@@ -47,17 +47,6 @@ def test_message_turn_requires_a_stream():
MessageTurn(role="user", content="${task}") 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(): def test_message_recipe_requires_at_least_one_target():
with pytest.raises(ValueError, match="target"): with pytest.raises(ValueError, match="target"):
TrainingRecipe( TrainingRecipe(
@@ -163,7 +152,7 @@ def test_from_dict_with_nested_blend():
assert isinstance(recipe.blend["a"].messages[0], MessageTurn) assert isinstance(recipe.blend["a"].messages[0], MessageTurn)
def test_message_dropout_round_trips_from_dict(): def test_message_if_present_round_trips_from_dict():
recipe = TrainingRecipe.from_dict( recipe = TrainingRecipe.from_dict(
{ {
"messages": [ "messages": [
@@ -174,13 +163,11 @@ def test_message_dropout_round_trips_from_dict():
"stream": "low_level", "stream": "low_level",
"target": True, "target": True,
"if_present": "subtask", "if_present": "subtask",
"dropout": 0.5,
}, },
] ]
} }
) )
assert recipe.messages[1].dropout == 0.5
assert recipe.messages[1].if_present == "subtask" assert recipe.messages[1].if_present == "subtask"
+9 -83
View File
@@ -176,35 +176,7 @@ def test_deterministic_blend_sampling():
assert first == second assert first == second
def test_message_dropout_skips_missing_optional_bindings(): def test_g05_recipe_emits_available_bbox_and_subtask():
recipe = TrainingRecipe(
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,
),
]
)
rendered = render_sample(
recipe=recipe,
persistent=[],
events=[],
t=0.0,
sample_idx=0,
task="pick the cup",
)
assert rendered["messages"] == [{"role": "user", "content": "pick the cup"}]
assert rendered["target_message_indices"] == []
def test_g05_dropout_produces_all_bbox_subtask_combinations():
recipe = TrainingRecipe.from_yaml("src/lerobot/configs/recipes/g05_bbox_subtask.yaml") recipe = TrainingRecipe.from_yaml("src/lerobot/configs/recipes/g05_bbox_subtask.yaml")
persistent = [persistent_row("assistant", "grasp the cup", "subtask", 0.0)] persistent = [persistent_row("assistant", "grasp the cup", "subtask", 0.0)]
events = [ events = [
@@ -216,64 +188,18 @@ def test_g05_dropout_produces_all_bbox_subtask_combinations():
} }
] ]
rendered_by_targets = {} rendered = render_sample(
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",
)
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,
),
]
)
first = render_sample(
recipe=recipe, recipe=recipe,
persistent=PERSISTENT, persistent=persistent,
events=[], events=events,
t=0.0, t=0.0,
sample_idx=42, sample_idx=0,
task="pick the cup", task="pick the cup",
) )
second = render_sample(
recipe=recipe, assert rendered["target_message_indices"] == [1, 2]
persistent=PERSISTENT, assert rendered["messages"][1]["content"].startswith("BBoxJSON:")
events=[], assert rendered["messages"][2]["content"] == "Subtask: grasp the cup"
t=0.0,
sample_idx=42,
task="pick the cup",
)
assert first == second
def test_emitted_at_filters_vqa_by_camera(): def test_emitted_at_filters_vqa_by_camera():