diff --git a/docs/source/g05.mdx b/docs/source/g05.mdx index 0ee62c62a..6d11464c9 100644 --- a/docs/source/g05.mdx +++ b/docs/source/g05.mdx @@ -123,9 +123,11 @@ For example, fine-tune the private SO-101 checkpoint on a LeRobot dataset: export HF_USER=your_hf_username 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.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.private=true \ --output_dir=outputs/train/g05_so101 \ @@ -135,6 +137,31 @@ lerobot-train \ --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 +`` 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 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 diff --git a/docs/source/language_and_recipes.mdx b/docs/source/language_and_recipes.mdx index 8d2aed59d..b6ebaae75 100644 --- a/docs/source/language_and_recipes.mdx +++ b/docs/source/language_and_recipes.mdx @@ -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. `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 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. diff --git a/src/lerobot/configs/recipe.py b/src/lerobot/configs/recipe.py index 43d90c1ce..3b6ab2557 100644 --- a/src/lerobot/configs/recipe.py +++ b/src/lerobot/configs/recipe.py @@ -104,7 +104,9 @@ class TrainingRecipe: messages: list[MessageTurn] | None = None bindings: dict[str, str] | None = None + requires: list[str] | None = None blend: dict[str, TrainingRecipe] | None = None + select_from_applicable: bool = False weight: float | None = None def __post_init__(self) -> None: @@ -151,6 +153,9 @@ class TrainingRecipe: assert self.messages is not None 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: missing = self._referenced_bindings(turn) - known_bindings if missing: @@ -168,6 +173,8 @@ class TrainingRecipe: def _validate_blend_recipe(self) -> None: """Ensure each blend component is a non-empty, weighted message recipe.""" 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: raise ValueError("Blend recipes must contain at least one component.") diff --git a/src/lerobot/configs/recipes/g05_bbox_subtask.yaml b/src/lerobot/configs/recipes/g05_bbox_subtask.yaml new file mode 100644 index 000000000..c3bcf2528 --- /dev/null +++ b/src/lerobot/configs/recipes/g05_bbox_subtask.yaml @@ -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, + } diff --git a/src/lerobot/datasets/language_render.py b/src/lerobot/datasets/language_render.py index 624f19e9b..c4420b1cd 100644 --- a/src/lerobot/datasets/language_render.py +++ b/src/lerobot/datasets/language_render.py @@ -169,6 +169,17 @@ def render_sample( persistent_rows = _normalize_rows(persistent or []) event_rows = _normalize_rows(events or []) + if recipe.blend is not None and recipe.select_from_applicable: + return _render_applicable_blend( + recipe, + persistent=persistent_rows, + events=event_rows, + t=t, + sample_idx=sample_idx, + task=task, + dataset_ctx=dataset_ctx, + ) + # Route sparse VQA frames to a matching view-specific component before weighted selection. # This avoids dropping annotated frames or selecting VQA without annotations. if recipe.blend is not None: @@ -197,6 +208,54 @@ def render_sample( 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( recipe: TrainingRecipe, *, @@ -385,6 +444,8 @@ def _render_message_recipe( ) -> RenderedMessages | None: """Expand ``recipe.messages`` into rendered chat messages using ``bindings``.""" 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]] = [] streams: list[str | None] = [] target_indices: list[int] = [] diff --git a/src/lerobot/policies/g05/configuration_g05.py b/src/lerobot/policies/g05/configuration_g05.py index 385c2b2a3..e21441a93 100644 --- a/src/lerobot/policies/g05/configuration_g05.py +++ b/src/lerobot/policies/g05/configuration_g05.py @@ -91,6 +91,39 @@ def make_g05_prompt_template(num_images: int, *, predict_cot: bool, flow_only: b return f"{prefix}Action: |" +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": "|", + "subtask": "|", + } + images = "".join(f"" for index in range(num_images)) + prefix = ( + f"{images}" + "Embodiment: ; Task: " + "State: ;" + "" + "\n" + ) + action = "Action: " if flow_only else "Action: |" + 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: # 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 @@ -198,6 +231,8 @@ class G05Config(PreTrainedConfig): processor_metadata: dict[str, Any] = field(default_factory=dict) action_codec_metadata: dict[str, Any] = field(default_factory=dict) prompt_template: str = "" + recipe_path: str | None = None + cot_bbox_camera: str | None = None normalization_mapping: dict[str, NormalizationMode] = field( default_factory=lambda: { @@ -254,6 +289,8 @@ class G05Config(PreTrainedConfig): raise ValueError("runtime_system must be 'system1' or 'system2'.") if self.runtime_system == "system2" and not self.predict_cot: 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: raise ValueError("n_action_steps must be between 1 and chunk_size.") 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.") if not set(self.optional_camera_keys) <= set(self.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: raise ValueError( "num_input_images must equal len(camera_order) * n_obs_steps for the selected checkpoint." diff --git a/src/lerobot/policies/g05/modeling_g05.py b/src/lerobot/policies/g05/modeling_g05.py index 9131a5033..47697f5db 100644 --- a/src/lerobot/policies/g05/modeling_g05.py +++ b/src/lerobot/policies/g05/modeling_g05.py @@ -11,6 +11,7 @@ from __future__ import annotations import importlib +import json import shutil from collections import deque from collections.abc import Mapping @@ -26,7 +27,12 @@ from lerobot.optim.optimizers import OptimizerParams from lerobot.policies.pretrained import PreTrainedPolicy 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: @@ -229,6 +235,158 @@ class G05Policy(PreTrainedPolicy): return value[index] 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"" + + 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 = " dict[str, Any]: prepare = getattr(self.backend, "prepare_lerobot_batch", None) if callable(prepare): @@ -284,13 +442,30 @@ class G05Policy(PreTrainedPolicy): if frequency is not None: sample["frequency"] = frequency if self.config.predict_cot: - 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}" - ) + 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=" 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( features: dict[PipelineFeatureType, dict[str, PolicyFeature]], ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: 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 @ProcessorStepRegistry.register(name="g05_image_transform") class G05ImageTransformStep(ProcessorStep): @@ -690,17 +733,26 @@ def make_g05_pre_post_processors( action_names=list(config.action_feature_names) or None, num_obs_steps=config.n_obs_steps, ) - steps: list[ProcessorStep] = [ - RenameObservationsProcessorStep(rename_map={}), - AddBatchDimensionProcessorStep(), + steps: list[ProcessorStep] = [RenameObservationsProcessorStep(rename_map={})] + if config.recipe_path: + 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( camera_order=config.camera_order, camera_sizes=config.camera_sizes, mean=config.image_mean, std=config.image_std, optional_camera_keys=config.optional_camera_keys, - ), - ] + ) + ) action_filter = config.processor_metadata.get("action_filter") or {} if str(action_filter.get("_target_", "")).endswith("R1LiteJointActionFilter"): action_parts = tuple( diff --git a/tests/configs/test_recipe.py b/tests/configs/test_recipe.py index 53520bfa3..a39212c21 100644 --- a/tests/configs/test_recipe.py +++ b/tests/configs/test_recipe.py @@ -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(): with pytest.raises(ValueError, match="only one"): TrainingRecipe( @@ -144,6 +152,35 @@ def test_from_dict_with_nested_blend(): 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): yaml_text = dedent( """ diff --git a/tests/datasets/test_language_render.py b/tests/datasets/test_language_render.py index 4ab4a2f2c..4cc8bf947 100644 --- a/tests/datasets/test_language_render.py +++ b/tests/datasets/test_language_render.py @@ -176,6 +176,74 @@ def test_deterministic_blend_sampling(): 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(): top = emitted_at( 3.0, diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py index e69e549bf..693a5dbf6 100644 --- a/tests/policies/g05/test_g05.py +++ b/tests/policies/g05/test_g05.py @@ -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" +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 "|Action: |" 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 " + assert sample["atomic_task"] == "Subtask: grasp the cup" + assert "||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 "Action: |" 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 " + assert sample["atomic_task"] == "Subtask: grasp the cup" + assert "||Action:" in sample["template"] + + def test_author_inference_payload_synthesizes_required_dummy_action(): policy = G05Policy(_config(), backend=TinyG05Backend()) batch = _policy_batch()