admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-08-08 17:39:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(g05): use recipe message dropout

Browse Source
This commit is contained in:
Pepijn
2026-07-29 15:18:48 +02:00
parent e6ec909cc7
commit 6d2c3bb1e2
7 changed files with 166 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/g05.mdx
+4 -9
View File
@@ -156,15 +156,10 @@ 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 filters out unavailable formats before selecting one of sample timestamp. It applies independent `0.5` dropout to the optional BBox and
four author-compatible objectives: Subtask targets, producing Action-only, Subtask+Action, BBox+Action, and joint
BBox+Subtask+Action samples. An unavailable annotation is skipped before
| Assistant sequence | Weight | dropout through the recipe's `if_present` guard.
| -------------------------- | -----: |
| Action only | 1 |
| Subtask, then action | 2 |
| BBox, then action | 1 |
| BBox, subtask, then action | 1 |
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
docs/source/language_and_recipes.mdx
+14 -21
View File
@@ -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. 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.
Annotation-dependent blends can set `select_from_applicable: true` and declare Message turns can set `dropout` to a probability between `0` and `1`. Dropout is
`requires` on each component. The renderer first removes components whose deterministic for a given sample index and independent between turns. Combine it
required bindings resolve to `None`, then performs the deterministic weighted with `if_present` to sample optional supervision formats while gracefully
selection. This matches mixed-CoT policies where unavailable annotation formats skipping annotations that are unavailable:
must not consume probability:
```yaml ```yaml
select_from_applicable: true messages:
blend: - { role: user, content: "${task}", stream: low_level }
subtask: - {
weight: 2 role: assistant,
requires: [subtask] content: "${subtask}",
messages: stream: low_level,
- { target: true,
role: assistant, if_present: subtask,
content: "${subtask}", dropout: 0.5,
stream: low_level, }
target: true,
}
action:
weight: 1
messages:
- { role: user, content: "${task}", stream: low_level }
``` ```
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
src/lerobot/configs/recipe.py
+10 -3
View File
@@ -18,6 +18,7 @@ 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
@@ -49,8 +50,9 @@ 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, and ``if_present`` skips the turn when the named binding training target, ``if_present`` skips the turn when the named binding
resolves to ``None``. resolves to ``None``, and ``dropout`` independently omits it for a
deterministic fraction of samples.
""" """
role: MessageRole role: MessageRole
@@ -59,6 +61,7 @@ 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."""
@@ -86,6 +89,11 @@ 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:
@@ -106,7 +114,6 @@ class TrainingRecipe:
bindings: dict[str, str] | None = None bindings: dict[str, str] | None = None
requires: list[str] | None = None requires: list[str] | None = None
blend: dict[str, TrainingRecipe] | None = None blend: dict[str, TrainingRecipe] | None = None
select_from_applicable: bool = False
weight: float | None = None weight: float | None = None
def __post_init__(self) -> None: def __post_init__(self) -> None:
src/lerobot/configs/recipes/g05_bbox_subtask.yaml
+25 -57
View File
@@ -1,60 +1,28 @@
# G0.5's released MixedSamplesBuilder semantics, limited to the CoT signals # G0.5 CoT supervision from the signals LeRobot can currently provide:
# LeRobot can currently provide: active subtasks and grounded bounding boxes. # active subtasks and camera-grounded bounding boxes.
# #
# The renderer first removes branches whose required annotations are missing, # BBox and Subtask turns are dropped independently and deterministically. This
# then deterministically samples among the remaining weights. The G0.5 policy # produces Action-only, BBox+Action, Subtask+Action, and
# processor converts the selected target messages into the checkpoint's exact # BBox+Subtask+Action samples without a policy-specific recipe selector.
# BBox/Subtask/action template and supervises the complete assistant segment.
select_from_applicable: true bindings:
blend: bbox: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.exterior)"
no_cot_action:
weight: 1
messages:
- {role: user, content: "${task}", stream: low_level}
subtask_action: messages:
weight: 2 - {role: user, content: "${task}", stream: low_level}
requires: [subtask] - {
messages: role: assistant,
- {role: user, content: "${task}", stream: low_level} content: "BBoxJSON: ${bbox}",
- { stream: low_level,
role: assistant, target: true,
content: "Subtask: ${subtask}", if_present: bbox,
stream: low_level, dropout: 0.5,
target: true, }
} - {
role: assistant,
bbox_action: content: "Subtask: ${subtask}",
weight: 1 stream: low_level,
requires: [bbox] target: true,
bindings: if_present: subtask,
bbox: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.exterior)" dropout: 0.5,
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,
}
src/lerobot/datasets/language_render.py
+22 -62
View File
@@ -169,17 +169,6 @@ def render_sample(
persistent_rows = _normalize_rows(persistent or []) persistent_rows = _normalize_rows(persistent or [])
event_rows = _normalize_rows(events 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. # Route sparse VQA frames to a matching view-specific component before weighted selection.
# This avoids dropping annotated frames or selecting VQA without annotations. # This avoids dropping annotated frames or selecting VQA without annotations.
if recipe.blend is not None: if recipe.blend is not None:
@@ -205,55 +194,7 @@ def render_sample(
task=task, task=task,
dataset_ctx=dataset_ctx, dataset_ctx=dataset_ctx,
) )
return _render_message_recipe(selected_recipe, bindings) return _render_message_recipe(selected_recipe, bindings, sample_idx=sample_idx)
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]
def _render_vqa_if_present( def _render_vqa_if_present(
@@ -284,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) rendered = _render_message_recipe(component, bindings, sample_idx=sample_idx)
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))
@@ -441,6 +382,8 @@ 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
@@ -450,9 +393,11 @@ def _render_message_recipe(
streams: list[str | None] = [] streams: list[str | None] = []
target_indices: list[int] = [] 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: 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:
@@ -484,6 +429,21 @@ 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],
tests/configs/test_recipe.py
+24 -21
View File
@@ -47,6 +47,17 @@ 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(
@@ -152,33 +163,25 @@ 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_applicable_blend_round_trips_from_dict(): def test_message_dropout_round_trips_from_dict():
recipe = TrainingRecipe.from_dict( recipe = TrainingRecipe.from_dict(
{ {
"select_from_applicable": True, "messages": [
"blend": { {"role": "user", "content": "${task}", "stream": "low_level"},
"subtask": { {
"weight": 2, "role": "assistant",
"requires": ["subtask"], "content": "${subtask}",
"messages": [ "stream": "low_level",
{ "target": True,
"role": "assistant", "if_present": "subtask",
"content": "${subtask}", "dropout": 0.5,
"stream": "low_level",
"target": True,
}
],
}, },
"action": { ]
"weight": 1,
"messages": [{"role": "user", "content": "${task}", "stream": "low_level"}],
},
},
} }
) )
assert recipe.select_from_applicable assert recipe.messages[1].dropout == 0.5
assert recipe.blend["subtask"].requires == ["subtask"] assert recipe.messages[1].if_present == "subtask"
def test_from_yaml_round_trips_through_load_recipe(tmp_path: Path): def test_from_yaml_round_trips_through_load_recipe(tmp_path: Path):
tests/datasets/test_language_render.py
+67 -35
View File
@@ -176,27 +176,19 @@ def test_deterministic_blend_sampling():
assert first == second assert first == second
def test_applicable_blend_filters_missing_bindings_before_weighted_selection(): def test_message_dropout_skips_missing_optional_bindings():
recipe = TrainingRecipe( recipe = TrainingRecipe(
select_from_applicable=True, messages=[
blend={ MessageTurn(role="user", content="${task}", stream="low_level"),
"missing": TrainingRecipe( MessageTurn(
weight=1_000, role="assistant",
requires=["subtask"], content="${subtask}",
messages=[ stream="low_level",
MessageTurn( target=True,
role="assistant", if_present="subtask",
content="${subtask}", dropout=0.5,
stream="low_level",
target=True,
)
],
), ),
"action": TrainingRecipe( ]
weight=1,
messages=[MessageTurn(role="user", content="${task}", stream="low_level")],
),
},
) )
rendered = render_sample( rendered = render_sample(
@@ -212,7 +204,7 @@ def test_applicable_blend_filters_missing_bindings_before_weighted_selection():
assert rendered["target_message_indices"] == [] 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") 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 = [
@@ -224,24 +216,64 @@ def test_applicable_blend_can_select_joint_bbox_subtask_target():
} }
] ]
rendered = next( rendered_by_targets = {}
candidate for sample_idx in range(100):
for sample_idx in range(100) rendered = render_sample(
if ( recipe=recipe,
candidate := render_sample( persistent=persistent,
recipe=recipe, events=events,
persistent=persistent, t=0.0,
events=events, sample_idx=sample_idx,
t=0.0, task="pick the cup",
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:") first = render_sample(
assert rendered["messages"][2]["content"] == "Subtask: grasp the cup" 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(): def test_emitted_at_filters_vqa_by_camera():