mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
feat(g05): train with LeRobot language recipes
This commit is contained in:
+28
-1
@@ -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
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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,
|
||||||
|
}
|
||||||
@@ -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."
|
||||||
|
|||||||
@@ -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,13 +442,30 @@ 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:
|
||||||
sample["prompt"] = "predict subtask"
|
rendered_recipe = "messages" in batch
|
||||||
atomic_task = batch.get("atomic_task")
|
applied_recipe_cot = rendered_recipe and self._apply_recipe_cot(
|
||||||
if atomic_task is not None:
|
sample, batch, index, batch_size
|
||||||
atomic_task = str(self._batch_item(atomic_task, index, batch_size))
|
)
|
||||||
sample["atomic_task"] = (
|
if not applied_recipe_cot:
|
||||||
atomic_task if atomic_task.startswith("Subtask:") else f"Subtask: {atomic_task}"
|
# 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"
|
||||||
|
atomic_task = batch.get("atomic_task")
|
||||||
|
if atomic_task is not None:
|
||||||
|
atomic_task = str(self._batch_item(atomic_task, index, batch_size))
|
||||||
|
sample["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)]
|
||||||
sample[f"image{image_index}"] = self.config.camera_sizes[camera]
|
sample[f"image{image_index}"] = self.config.camera_sizes[camera]
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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(
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user