feat(g05): train with LeRobot language recipes

This commit is contained in:
Pepijn
2026-07-29 11:08:03 +02:00
parent 274ee585b4
commit 0675920df9
11 changed files with 678 additions and 14 deletions
+28 -1
View File
@@ -123,9 +123,11 @@ For example, fine-tune the private SO-101 checkpoint on a LeRobot dataset:
export HF_USER=your_hf_username export HF_USER=your_hf_username
lerobot-train \ lerobot-train \
--dataset.repo_id=${HF_USER}/my_so101_dataset \ --dataset.repo_id=${HF_USER}/my_so101_dataset_annotated \
--policy.path=lerobot/g05_so101 \ --policy.path=lerobot/g05_so101 \
--policy.device=cuda \ --policy.device=cuda \
--policy.recipe_path=recipes/g05_bbox_subtask.yaml \
--policy.cot_bbox_camera=observation.images.exterior \
--policy.repo_id=${HF_USER}/g05_so101_finetuned \ --policy.repo_id=${HF_USER}/g05_so101_finetuned \
--policy.private=true \ --policy.private=true \
--output_dir=outputs/train/g05_so101 \ --output_dir=outputs/train/g05_so101 \
@@ -135,6 +137,31 @@ lerobot-train \
--save_freq=1000 --save_freq=1000
``` ```
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 |
Grounded VQA boxes are converted from pixel-space `xyxy` JSON using the source
camera dimensions captured before image resizing, then serialized as G0.5
`<locXXXX>` tokens. Joint samples preserve the released checkpoint's
`BBox → Subtask → Action` order. The user/task conditioning tokens remain
masked; the author backend applies its language/action objective to the
assistant sequence.
Generate the required `subtask` and grounded `vqa` language columns with
`lerobot-annotate` as described in the
[annotation pipeline](./annotation_pipeline). The bundled recipe targets
`observation.images.exterior`; copy the YAML and change its camera-filtered
bindings when training an embodiment with a different grounded camera.
The SO-101 recipe uses AdamW at `8e-5` with 1,000 warmup steps. The packaged The SO-101 recipe uses AdamW at `8e-5` with 1,000 warmup steps. The packaged
LIBERO and RoboTwin configurations use their released `1e-5` recipe, with LIBERO and RoboTwin configurations use their released `1e-5` recipe, with
1,000 and 500 warmup steps respectively. All profiles preserve G0.5's six 1,000 and 500 warmup steps respectively. All profiles preserve G0.5's six
+25
View File
@@ -146,6 +146,31 @@ 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
`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:
```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 }
```
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.
+7
View File
@@ -104,7 +104,9 @@ class TrainingRecipe:
messages: list[MessageTurn] | None = None messages: list[MessageTurn] | None = None
bindings: dict[str, str] | None = None bindings: dict[str, 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:
@@ -151,6 +153,9 @@ class TrainingRecipe:
assert self.messages is not None assert self.messages is not None
known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"} known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"}
missing_requirements = set(self.requires or ()) - known_bindings
if missing_requirements:
raise ValueError(f"TrainingRecipe requires unknown binding(s): {sorted(missing_requirements)}")
for turn in self.messages: for turn in self.messages:
missing = self._referenced_bindings(turn) - known_bindings missing = self._referenced_bindings(turn) - known_bindings
if missing: if missing:
@@ -168,6 +173,8 @@ class TrainingRecipe:
def _validate_blend_recipe(self) -> None: def _validate_blend_recipe(self) -> None:
"""Ensure each blend component is a non-empty, weighted message recipe.""" """Ensure each blend component is a non-empty, weighted message recipe."""
assert self.blend is not None assert self.blend is not None
if self.requires:
raise ValueError("A blend recipe cannot declare requires; set it on its components.")
if not self.blend: if not self.blend:
raise ValueError("Blend recipes must contain at least one component.") raise ValueError("Blend recipes must contain at least one component.")
@@ -0,0 +1,60 @@
# G0.5's released MixedSamplesBuilder semantics, limited to the CoT signals
# LeRobot can currently provide: active subtasks and 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.
select_from_applicable: true
blend:
no_cot_action:
weight: 1
messages:
- {role: user, content: "${task}", stream: low_level}
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,
}
+61
View File
@@ -169,6 +169,17 @@ 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:
@@ -197,6 +208,54 @@ def render_sample(
return _render_message_recipe(selected_recipe, bindings) 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]
def _render_vqa_if_present( def _render_vqa_if_present(
recipe: TrainingRecipe, recipe: TrainingRecipe,
*, *,
@@ -385,6 +444,8 @@ def _render_message_recipe(
) -> 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
if any(bindings.get(name) is None for name in recipe.requires or ()):
return None
messages: list[dict[str, Any]] = [] messages: list[dict[str, Any]] = []
streams: list[str | None] = [] streams: list[str | None] = []
target_indices: list[int] = [] target_indices: list[int] = []
@@ -91,6 +91,39 @@ def make_g05_prompt_template(num_images: int, *, predict_cot: bool, flow_only: b
return f"{prefix}Action: <EOV><EOC><action_action>|<eos>" return f"{prefix}Action: <EOV><EOC><action_action>|<eos>"
def make_g05_cot_prompt_template(
num_images: int,
*,
fields: tuple[str, ...],
flow_only: bool,
) -> str:
"""Build the exact author template for a selected Subtask/BBox CoT format."""
supported_fields = {"bbox", "subtask"}
if not fields or not set(fields) <= supported_fields:
raise ValueError(
f"G0.5 CoT fields must be a non-empty subset of {sorted(supported_fields)}, got {fields}."
)
# The released BBoxSubtaskCoTBuilder emits BBox before Subtask. Preserve
# checkpoint serialization even though the paper's schematic orders the
# independently composable labels differently.
ordered_fields = tuple(field for field in ("bbox", "subtask") if field in fields)
placeholders = {
"bbox": "<bbox_text>|",
"subtask": "<atomic_task_text>|",
}
images = "".join(f"<image{index}_image_!>" for index in range(num_images))
prefix = (
f"<chat_user_prefix>{images}<bos>"
"Embodiment: <embodiment_text_!>; Task: <command_text_!_200> "
"State: <proprio_proprio_!>;"
"<chat_user_suffix><chat_assistant_prefix>"
"<prompt_text_!>\n<EOC>"
)
action = "Action: <EOV><eos>" if flow_only else "Action: <EOV><action_action>|<eos>"
return prefix + "".join(placeholders[field] for field in ordered_fields) + action
# Raw dimensions are inserted in these exact policy slots. The G0.5 shared layout is: # Raw dimensions are inserted in these exact policy slots. The G0.5 shared layout is:
# left_control[9] | left_gripper[1] | right_control[9] | right_gripper[1] | lower_body[7]. # left_control[9] | left_gripper[1] | right_control[9] | right_gripper[1] | lower_body[7].
# LIBERO uses only the right EEF delta and right gripper. atomic_4 is a single-arm mobile # LIBERO uses only the right EEF delta and right gripper. atomic_4 is a single-arm mobile
@@ -198,6 +231,8 @@ class G05Config(PreTrainedConfig):
processor_metadata: dict[str, Any] = field(default_factory=dict) processor_metadata: dict[str, Any] = field(default_factory=dict)
action_codec_metadata: dict[str, Any] = field(default_factory=dict) action_codec_metadata: dict[str, Any] = field(default_factory=dict)
prompt_template: str = "" prompt_template: str = ""
recipe_path: str | None = None
cot_bbox_camera: str | None = None
normalization_mapping: dict[str, NormalizationMode] = field( normalization_mapping: dict[str, NormalizationMode] = field(
default_factory=lambda: { default_factory=lambda: {
@@ -254,6 +289,8 @@ class G05Config(PreTrainedConfig):
raise ValueError("runtime_system must be 'system1' or 'system2'.") raise ValueError("runtime_system must be 'system1' or 'system2'.")
if self.runtime_system == "system2" and not self.predict_cot: if self.runtime_system == "system2" and not self.predict_cot:
raise ValueError("G0.5 System 2 requires predict_cot=True in the packaged checkpoint.") raise ValueError("G0.5 System 2 requires predict_cot=True in the packaged checkpoint.")
if self.recipe_path is not None and not self.predict_cot:
raise ValueError("G0.5 recipe-driven CoT training requires predict_cot=True.")
if not 1 <= self.n_action_steps <= self.chunk_size: if not 1 <= self.n_action_steps <= self.chunk_size:
raise ValueError("n_action_steps must be between 1 and chunk_size.") raise ValueError("n_action_steps must be between 1 and chunk_size.")
if self.action_head == "actioncodec" and not self.discrete_action: if self.action_head == "actioncodec" and not self.discrete_action:
@@ -314,6 +351,8 @@ class G05Config(PreTrainedConfig):
raise ValueError("camera_sizes must contain exactly the ordered checkpoint camera keys.") raise ValueError("camera_sizes must contain exactly the ordered checkpoint camera keys.")
if not set(self.optional_camera_keys) <= set(self.camera_order): if not set(self.optional_camera_keys) <= set(self.camera_order):
raise ValueError("optional_camera_keys must be a subset of camera_order.") raise ValueError("optional_camera_keys must be a subset of camera_order.")
if self.cot_bbox_camera is not None and self.cot_bbox_camera not in self.camera_order:
raise ValueError("cot_bbox_camera must be one of camera_order.")
if self.num_input_images != len(self.camera_order) * self.n_obs_steps: if self.num_input_images != len(self.camera_order) * self.n_obs_steps:
raise ValueError( raise ValueError(
"num_input_images must equal len(camera_order) * n_obs_steps for the selected checkpoint." "num_input_images must equal len(camera_order) * n_obs_steps for the selected checkpoint."
+177 -2
View File
@@ -11,6 +11,7 @@
from __future__ import annotations from __future__ import annotations
import importlib import importlib
import json
import shutil import shutil
from collections import deque from collections import deque
from collections.abc import Mapping from collections.abc import Mapping
@@ -26,7 +27,12 @@ from lerobot.optim.optimizers import OptimizerParams
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.constants import ACTION, OBS_STATE from lerobot.utils.constants import ACTION, OBS_STATE
from .configuration_g05 import G05_POLICY_PARTS, G05Config from .configuration_g05 import (
G05_POLICY_PARTS,
G05Config,
make_g05_cot_prompt_template,
make_g05_prompt_template,
)
def _author_backend(config: G05Config) -> nn.Module: def _author_backend(config: G05Config) -> nn.Module:
@@ -229,6 +235,158 @@ class G05Policy(PreTrainedPolicy):
return value[index] return value[index]
return value return value
def _recipe_cot_targets(
self,
batch: Mapping[str, Any],
index: int,
batch_size: int,
) -> tuple[str | None, str | None]:
"""Read the selected recipe's supervised Subtask/BBox messages."""
messages = batch.get("messages")
target_indices = batch.get("target_message_indices")
if messages is None or target_indices is None:
return None, None
sample_messages = messages
if (
isinstance(messages, list | tuple)
and len(messages) == batch_size
and (not messages or isinstance(messages[0], list | tuple))
):
sample_messages = messages[index]
sample_target_indices = target_indices
has_batched_target_indices = (isinstance(target_indices, Tensor) and target_indices.ndim > 1) or (
isinstance(target_indices, list | tuple)
and len(target_indices) == batch_size
and (not target_indices or isinstance(target_indices[0], list | tuple | Tensor))
)
if has_batched_target_indices:
sample_target_indices = target_indices[index]
if isinstance(sample_messages, Mapping):
sample_messages = [sample_messages]
if isinstance(sample_target_indices, Tensor):
sample_target_indices = sample_target_indices.detach().cpu().tolist()
if not isinstance(sample_messages, list | tuple) or not isinstance(
sample_target_indices, list | tuple
):
return None, None
subtask: str | None = None
bbox_json: str | None = None
for target_index in sample_target_indices:
message = sample_messages[int(target_index)]
content = message.get("content") if isinstance(message, Mapping) else None
if not isinstance(content, str):
continue
if content.startswith("Subtask:"):
value = content.removeprefix("Subtask:").strip()
if value:
subtask = value
elif content.startswith("BBoxJSON:"):
value = content.removeprefix("BBoxJSON:").strip()
if value:
bbox_json = value
return subtask, bbox_json
@staticmethod
def _format_bbox_target(bbox_json: str | None, image_size: tuple[int, int]) -> str | None:
"""Convert LeRobot grounded-VQA JSON into G0.5's location-token format."""
if not bbox_json:
return None
try:
payload = json.loads(bbox_json)
if isinstance(payload, str):
payload = json.loads(payload)
except (json.JSONDecodeError, TypeError):
return None
if not isinstance(payload, Mapping):
return None
if isinstance(payload.get("answer"), Mapping):
payload = payload["answer"]
height, width = image_size
boxes: list[tuple[str, list[float]]] = []
detections = payload.get("detections")
if isinstance(detections, list):
for detection in detections:
if not isinstance(detection, Mapping) or detection.get("bbox_format", "xyxy") != "xyxy":
continue
coords = detection.get("bbox")
if not isinstance(coords, list | tuple) or len(coords) != 4:
continue
label = str(detection.get("label") or "object")
boxes.append((label, [float(value) for value in coords]))
else:
for label, coords in payload.items():
if isinstance(coords, list | tuple) and len(coords) == 4:
boxes.append((str(label), [float(value) for value in coords]))
if not boxes:
return None
def normalize(coords: list[float]) -> list[float]:
if max(abs(value) for value in coords) <= 1.0:
return coords
x1, y1, x2, y2 = coords
return [x1 / width, y1 / height, x2 / width, y2 / height]
def location_token(value: float) -> str:
location = max(0, min(1023, round(value * 1024)))
return f"<loc{location:04d}>"
formatted = []
for label, raw_coords in boxes:
x1, y1, x2, y2 = normalize(raw_coords)
locations = "".join(location_token(value) for value in (y1, x1, y2, x2))
formatted.append(f"{label} {locations}")
return "BBox: " + "; ".join(formatted)
def _apply_recipe_cot(
self,
sample: dict[str, Any],
batch: Mapping[str, Any],
index: int,
batch_size: int,
) -> bool:
"""Populate one author sample from recipe-rendered CoT targets."""
subtask, bbox_json = self._recipe_cot_targets(batch, index, batch_size)
image_size = batch.get("g05_bbox_image_size")
if (
isinstance(image_size, list | tuple)
and len(image_size) == batch_size
and image_size
and isinstance(image_size[0], list | tuple | Tensor)
):
image_size = image_size[index]
if isinstance(image_size, Tensor):
image_size = image_size.detach().cpu().tolist()
if not isinstance(image_size, list | tuple) or len(image_size) != 2:
camera = self.config.cot_bbox_camera or self.config.camera_order[0]
image_size = self.config.camera_sizes[camera]
bbox = self._format_bbox_target(bbox_json, (int(image_size[0]), int(image_size[1])))
fields = tuple(field for field, value in (("bbox", bbox), ("subtask", subtask)) if value)
if not fields:
return False
flow_only = "<action_action" not in self.config.prompt_template
sample["template"] = make_g05_cot_prompt_template(
self.config.num_prompt_images,
fields=fields,
flow_only=flow_only,
)
if bbox is not None:
sample["bbox"] = bbox
if subtask is not None:
sample["atomic_task"] = f"Subtask: {subtask}"
sample["prompt"] = {
("bbox",): "predict bbox",
("subtask",): "predict subtask",
("bbox", "subtask"): "predict bbox, subtask and action",
}[fields]
return True
def _prepare_author_batch(self, batch: Mapping[str, Any], task: str | None = None) -> dict[str, Any]: def _prepare_author_batch(self, batch: Mapping[str, Any], task: str | None = None) -> dict[str, Any]:
prepare = getattr(self.backend, "prepare_lerobot_batch", None) prepare = getattr(self.backend, "prepare_lerobot_batch", None)
if callable(prepare): if callable(prepare):
@@ -284,12 +442,29 @@ class G05Policy(PreTrainedPolicy):
if frequency is not None: if frequency is not None:
sample["frequency"] = frequency sample["frequency"] = frequency
if self.config.predict_cot: if self.config.predict_cot:
rendered_recipe = "messages" in batch
applied_recipe_cot = rendered_recipe and self._apply_recipe_cot(
sample, batch, index, batch_size
)
if not applied_recipe_cot:
# During mixed-recipe training an applicable no-CoT branch is a
# genuine target format. At inference, where actions are absent,
# retain the checkpoint's configured System 2 prompt.
if rendered_recipe and isinstance(batch.get(ACTION), Tensor):
sample["template"] = make_g05_prompt_template(
self.config.num_prompt_images,
predict_cot=False,
flow_only="<action_action" not in self.config.prompt_template,
)
else:
sample["prompt"] = "predict subtask" sample["prompt"] = "predict subtask"
atomic_task = batch.get("atomic_task") atomic_task = batch.get("atomic_task")
if atomic_task is not None: if atomic_task is not None:
atomic_task = str(self._batch_item(atomic_task, index, batch_size)) atomic_task = str(self._batch_item(atomic_task, index, batch_size))
sample["atomic_task"] = ( sample["atomic_task"] = (
atomic_task if atomic_task.startswith("Subtask:") else f"Subtask: {atomic_task}" atomic_task
if atomic_task.startswith("Subtask:")
else f"Subtask: {atomic_task}"
) )
for image_index in range(self.config.num_prompt_images): for image_index in range(self.config.num_prompt_images):
camera = self.config.camera_order[image_index % len(self.config.camera_order)] camera = self.config.camera_order[image_index % len(self.config.camera_order)]
+57 -5
View File
@@ -11,6 +11,7 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path
from typing import Any from typing import Any
import torch import torch
@@ -48,12 +49,54 @@ from lerobot.utils.constants import (
from .configuration_g05 import G05_EMBODIMENT_MAPPINGS, G05_POLICY_PARTS, G05Config from .configuration_g05 import G05_EMBODIMENT_MAPPINGS, G05_POLICY_PARTS, G05Config
def _load_recipe(path_str: str) -> Any:
"""Load an absolute recipe path or one relative to ``lerobot/configs``."""
from lerobot.configs.recipe import TrainingRecipe
path = Path(path_str)
if not path.is_absolute() and not path.exists():
from lerobot.configs import recipe as recipe_module
candidate = Path(recipe_module.__file__).resolve().parent / path
if candidate.exists():
path = candidate
return TrainingRecipe.from_yaml(path)
def _copy_feature_tree( def _copy_feature_tree(
features: dict[PipelineFeatureType, dict[str, PolicyFeature]], features: dict[PipelineFeatureType, dict[str, PolicyFeature]],
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return {kind: values.copy() for kind, values in features.items()} return {kind: values.copy() for kind, values in features.items()}
@dataclass
@ProcessorStepRegistry.register(name="g05_bbox_image_size")
class G05BBoxImageSizeStep(ProcessorStep):
"""Preserve the annotated camera's source size before checkpoint resizing."""
camera_key: str
def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition.get(TransitionKey.OBSERVATION) or {}
image = observation.get(self.camera_key)
if image is None:
return transition
image = torch.as_tensor(image)
if image.ndim < 3:
raise ValueError(f"G0.5 bbox camera {self.camera_key!r} has invalid shape {image.shape}.")
transition = transition.copy()
complementary = dict(transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
complementary["g05_bbox_image_size"] = (int(image.shape[-2]), int(image.shape[-1]))
transition[TransitionKey.COMPLEMENTARY_DATA] = complementary
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
@dataclass @dataclass
@ProcessorStepRegistry.register(name="g05_image_transform") @ProcessorStepRegistry.register(name="g05_image_transform")
class G05ImageTransformStep(ProcessorStep): class G05ImageTransformStep(ProcessorStep):
@@ -690,17 +733,26 @@ def make_g05_pre_post_processors(
action_names=list(config.action_feature_names) or None, action_names=list(config.action_feature_names) or None,
num_obs_steps=config.n_obs_steps, num_obs_steps=config.n_obs_steps,
) )
steps: list[ProcessorStep] = [ steps: list[ProcessorStep] = [RenameObservationsProcessorStep(rename_map={})]
RenameObservationsProcessorStep(rename_map={}), if config.recipe_path:
AddBatchDimensionProcessorStep(), from lerobot.processor.render_messages_processor import RenderMessagesStep
steps.extend(
[
G05BBoxImageSizeStep(camera_key=config.cot_bbox_camera or config.camera_order[0]),
RenderMessagesStep(recipe=_load_recipe(config.recipe_path)),
]
)
steps.append(AddBatchDimensionProcessorStep())
steps.append(
G05ImageTransformStep( G05ImageTransformStep(
camera_order=config.camera_order, camera_order=config.camera_order,
camera_sizes=config.camera_sizes, camera_sizes=config.camera_sizes,
mean=config.image_mean, mean=config.image_mean,
std=config.image_std, std=config.image_std,
optional_camera_keys=config.optional_camera_keys, optional_camera_keys=config.optional_camera_keys,
), )
] )
action_filter = config.processor_metadata.get("action_filter") or {} action_filter = config.processor_metadata.get("action_filter") or {}
if str(action_filter.get("_target_", "")).endswith("R1LiteJointActionFilter"): if str(action_filter.get("_target_", "")).endswith("R1LiteJointActionFilter"):
action_parts = tuple( action_parts = tuple(
+37
View File
@@ -57,6 +57,14 @@ def test_message_recipe_requires_at_least_one_target():
) )
def test_message_recipe_requires_known_bindings():
with pytest.raises(ValueError, match="requires unknown binding"):
TrainingRecipe(
messages=[_minimal_target_turn()],
requires=["not_a_binding"],
)
def test_recipe_rejects_both_messages_and_blend(): def test_recipe_rejects_both_messages_and_blend():
with pytest.raises(ValueError, match="only one"): with pytest.raises(ValueError, match="only one"):
TrainingRecipe( TrainingRecipe(
@@ -144,6 +152,35 @@ 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():
recipe = TrainingRecipe.from_dict(
{
"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"}],
},
},
}
)
assert recipe.select_from_applicable
assert recipe.blend["subtask"].requires == ["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):
yaml_text = dedent( yaml_text = dedent(
""" """
+68
View File
@@ -176,6 +176,74 @@ def test_deterministic_blend_sampling():
assert first == second assert first == second
def test_applicable_blend_filters_missing_bindings_before_weighted_selection():
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,
)
],
),
"action": TrainingRecipe(
weight=1,
messages=[MessageTurn(role="user", content="${task}", stream="low_level")],
),
},
)
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_applicable_blend_can_select_joint_bbox_subtask_target():
recipe = TrainingRecipe.from_yaml("src/lerobot/configs/recipes/g05_bbox_subtask.yaml")
persistent = [persistent_row("assistant", "grasp the cup", "subtask", 0.0)]
events = [
{
"role": "assistant",
"content": '{"detections": [{"label": "cup", "bbox": [1, 2, 3, 4]}]}',
"style": "vqa",
"camera": "observation.images.exterior",
}
]
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",
)
)
and candidate["target_message_indices"] == [1, 2]
)
assert rendered["messages"][1]["content"].startswith("BBoxJSON:")
assert rendered["messages"][2]["content"] == "Subtask: grasp the cup"
def test_emitted_at_filters_vqa_by_camera(): def test_emitted_at_filters_vqa_by_camera():
top = emitted_at( top = emitted_at(
3.0, 3.0,
+113
View File
@@ -492,6 +492,119 @@ def test_system2_training_target_is_forwarded_without_replacing_operator_task():
assert prepared["samples"][0]["atomic_task"] == "Subtask: grasp the cup" assert prepared["samples"][0]["atomic_task"] == "Subtask: grasp the cup"
def test_system2_recipe_subtask_target_selects_author_template():
policy = G05Policy(_config(predict_cot=True, runtime_system="system2"), backend=TinyG05Backend())
batch = _policy_batch("operator task")
batch["messages"] = [
[
{"role": "user", "content": "operator task"},
{"role": "assistant", "content": "Subtask: grasp the cup"},
]
]
batch["target_message_indices"] = [[1]]
sample = policy._prepare_author_batch(batch)["samples"][0]
assert sample["command"] == "operator task"
assert sample["prompt"] == "predict subtask"
assert sample["atomic_task"] == "Subtask: grasp the cup"
assert "<EOC><atomic_task_text>|Action: <EOV><action_action>|<eos>" in sample["template"]
def test_system2_recipe_bbox_and_subtask_use_checkpoint_field_order():
policy = G05Policy(_config(predict_cot=True, runtime_system="system2"), backend=TinyG05Backend())
batch = _policy_batch("operator task")
batch["messages"] = [
[
{"role": "user", "content": "operator task"},
{
"role": "assistant",
"content": (
'BBoxJSON: {"detections": [{"label": "cup", "bbox_format": "xyxy", '
'"bbox": [20, 10, 100, 50]}]}'
),
},
{"role": "assistant", "content": "Subtask: grasp the cup"},
]
]
batch["target_message_indices"] = [[1, 2]]
batch["g05_bbox_image_size"] = (100, 200)
sample = policy._prepare_author_batch(batch)["samples"][0]
assert sample["prompt"] == "predict bbox, subtask and action"
assert sample["bbox"] == "BBox: cup <loc0102><loc0102><loc0512><loc0512>"
assert sample["atomic_task"] == "Subtask: grasp the cup"
assert "<EOC><bbox_text>|<atomic_task_text>|Action:" in sample["template"]
def test_system2_recipe_no_cot_branch_uses_action_only_training_template():
policy = G05Policy(_config(predict_cot=True, runtime_system="system2"), backend=TinyG05Backend())
batch = _policy_batch("operator task")
batch["messages"] = [[{"role": "user", "content": "operator task"}]]
batch["target_message_indices"] = [[]]
sample = policy._prepare_author_batch(batch)["samples"][0]
assert "prompt" not in sample
assert "atomic_task" not in sample
assert "<chat_assistant_prefix>Action: <EOV><EOC><action_action>|<eos>" in sample["template"]
def test_recipe_preprocessor_resolves_lerobot_subtask_and_bbox_annotations():
pytest.importorskip("datasets", reason="recipe rendering requires lerobot[dataset]")
config = _config(
predict_cot=True,
runtime_system="system2",
recipe_path="recipes/g05_bbox_subtask.yaml",
)
preprocessor, _ = make_pre_post_processors(config)
policy = G05Policy(config, backend=TinyG05Backend())
raw = {
OBS_STATE: torch.zeros(7),
ACTION: torch.zeros(4, 7),
"observation.images.image": torch.zeros(3, 100, 200, dtype=torch.uint8),
"observation.images.wrist_image": torch.zeros(3, 100, 200, dtype=torch.uint8),
"task": "operator task",
"timestamp": torch.tensor(0.0),
"language_persistent": [
{
"role": "assistant",
"content": "grasp the cup",
"style": "subtask",
"timestamp": 0.0,
"camera": None,
"tool_calls": None,
}
],
"language_events": [
{
"role": "assistant",
"content": (
'{"detections": [{"label": "cup", "bbox_format": "xyxy", "bbox": [20, 10, 100, 50]}]}'
),
"style": "vqa",
"camera": "observation.images.exterior",
"tool_calls": None,
}
],
}
processed = next(
candidate
for sample_index in range(100)
if (candidate := preprocessor({**raw, "index": torch.tensor(sample_index)}))["target_message_indices"]
== [[1, 2]]
)
sample = policy._prepare_author_batch(processed)["samples"][0]
assert "language_persistent" not in processed
assert "language_events" not in processed
assert sample["bbox"] == "BBox: cup <loc0102><loc0102><loc0512><loc0512>"
assert sample["atomic_task"] == "Subtask: grasp the cup"
assert "<EOC><bbox_text>|<atomic_task_text>|Action:" in sample["template"]
def test_author_inference_payload_synthesizes_required_dummy_action(): def test_author_inference_payload_synthesizes_required_dummy_action():
policy = G05Policy(_config(), backend=TinyG05Backend()) policy = G05Policy(_config(), backend=TinyG05Backend())
batch = _policy_batch() batch = _policy_batch()