diff --git a/docs/source/language_and_recipes.mdx b/docs/source/language_and_recipes.mdx index 7df272973..8d2aed59d 100644 --- a/docs/source/language_and_recipes.mdx +++ b/docs/source/language_and_recipes.mdx @@ -146,6 +146,12 @@ 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. +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. +`recipes/subtask_joint.yaml` is the provided example; pair it with +`--policy.joint_subtask_conditioning=true` at inference. + ## Graceful absence If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op. diff --git a/docs/source/pi052.mdx b/docs/source/pi052.mdx index 9308124f8..7184cafaf 100644 --- a/docs/source/pi052.mdx +++ b/docs/source/pi052.mdx @@ -55,9 +55,20 @@ The provided recipes are: | Recipe | Required annotations | Trains | | ------------------------------------- | ----------------------------------------------------------------------- | -------------------------------------------------- | | `recipes/subtask.yaml` | `subtask` | Subtask prediction and subtask-conditioned actions | +| `recipes/subtask_joint.yaml` | `subtask` | Paper-style joint sequence: subtask text and actions in one sample | | `recipes/subtask_mem.yaml` | `subtask`, `memory` | Subtasks, actions, and memory updates | | `recipes/subtask_mem_vqa_speech.yaml` | `subtask`, `memory`, `vqa`; interjection/speech rows for those branches | Subtasks, actions, memory, VQA, and spoken replies | +The blend recipes factorize training into separate high-level (task → subtask) +and low-level (subtask → actions) samples, matching how inference decomposes +π(a|o, subtask)·π(subtask|o, task). `recipes/subtask_joint.yaml` instead uses +the π0.5 paper's single-sequence layout — the supervised subtask span is +attended causally and conditions the FAST and flow losses in the same forward. +Checkpoints trained with the joint recipe must set +`--policy.joint_subtask_conditioning=true` at inference so the flow prefix +rebuilds the same layout (task turn with state, then the generated subtask as a +causal assistant turn); leave it `false` for the blend recipes. + Use `lerobot-annotate` to generate these columns. The repository includes a Hugging Face Jobs launcher that you can edit for your source and destination datasets. For a local annotation run, first install @@ -117,6 +128,14 @@ the expected prompt, text target, and action endpoints before scaling up. | `policy.knowledge_insulation` | `true` | Blocks action-loss gradients through the VLM K/V path | | `policy.flow_num_repeats` | `5` | Reuses one VLM prefix for independent denoising targets | | `policy.lm_head_lr_scale` | `1.0` | Scales language-head learning rate; `1.0` uses the base rate | +| `policy.fast_skip_tokens` | `1152` | FAST id offset; skips ``+`` so VQA and FAST never collide | +| `policy.joint_subtask_conditioning` | `false` | Rebuilds the joint-sequence prefix at inference (see recipes) | + +`fast_skip_tokens=1152` places FAST codes below PaliGemma's `` range. +openpi's pi0-FAST convention is `128` (FAST occupies the `` ids); use that +value only to stay weight-compatible with checkpoints trained that way, and +avoid combining it with the VQA recipe, whose `` targets would share +embedding rows with FAST codes. The loss weights are starting points, not dataset-independent constants. Track flow loss and text/FAST losses separately, and inspect generated subtasks rather diff --git a/src/lerobot/configs/recipes/subtask_joint.yaml b/src/lerobot/configs/recipes/subtask_joint.yaml new file mode 100644 index 000000000..5258ccf23 --- /dev/null +++ b/src/lerobot/configs/recipes/subtask_joint.yaml @@ -0,0 +1,13 @@ +# Paper-style joint sequence (pi0.5 §IV-B): one sample supervises the subtask +# text with CE and, because the assistant turn is part of the prefix, conditions +# the FAST and flow action losses on the same annotated subtask in one forward. +# The supervised span is attended causally; the action losses see task + subtask. +# +# Pair with `--policy.joint_subtask_conditioning=true` at inference so the flow +# prefix reproduces this layout (task turn with state + causal generated subtask). +# Samples without a `subtask` annotation fall back to a plain task-prompt +# low-level sample via `if_present`. + +messages: + - {role: user, content: "${task}", stream: low_level} + - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask} diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index a0a88e560..c7fe2d260 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -63,9 +63,14 @@ class PI052Config(PI05Config): max_action_tokens: int = 256 """Maximum number of FAST tokens per action chunk.""" - fast_skip_tokens: int = 128 - """Number of low-vocab tokens the FAST tokenizer skips to avoid - collisions with PaliGemma's text vocabulary.""" + fast_skip_tokens: int = 1152 + """Number of top-of-vocab tokens the FAST id mapping skips. + + 1152 skips PaliGemma's 128 ```` and 1024 ```` special tokens so + FAST codes land in plain-text ids below 256000 and never collide with the + ```` targets used for VQA. openpi's pi0-FAST convention is 128 (FAST + occupies the ```` range); use 128 only to stay weight-compatible with + checkpoints trained that way.""" fast_action_loss_weight: float = 1.0 """Weight on FAST action-token CE relative to continuous-flow supervision.""" @@ -76,6 +81,16 @@ class PI052Config(PI05Config): Non-positive values regenerate each action chunk while still refreshing the action prompt every chunk. """ + joint_subtask_conditioning: bool = False + """Condition low-level action inference on the task plus the generated subtask. + + Matches paper-style joint-sequence recipes (``recipes/subtask_joint.yaml``) + where one sample supervises the subtask text and conditions the action + losses on it: the inference prefix becomes + ``User: {task}, State: ...;\\nAssistant: {subtask}`` with the subtask + span attended causally, exactly as trained. Leave ``False`` for the blend + recipes, whose low-level samples use ``User: {subtask}, State: ...;``.""" + auto_fit_fast_tokenizer: bool = False """Fit and cache a dataset-specific FAST tokenizer before training. diff --git a/src/lerobot/policies/pi052/inference/pi052_adapter.py b/src/lerobot/policies/pi052/inference/pi052_adapter.py index 72737c4e8..a570ef1af 100644 --- a/src/lerobot/policies/pi052/inference/pi052_adapter.py +++ b/src/lerobot/policies/pi052/inference/pi052_adapter.py @@ -41,22 +41,53 @@ class PI052PolicyAdapter(BaseLanguageAdapter): subtask = state.language_context.get("subtask") or state.task or "" # Match the training prompt by conditioning on both subtask and discretized state. - content = subtask + state_str = None obs_state = observation.get(OBS_STATE) if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0: from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str # noqa: PLC0415 state_row = obs_state[0] if obs_state.ndim > 1 else obs_state - content = f"{subtask}, State: {discretize_state_str(state_row)};" + state_str = discretize_state_str(state_row) - text_batch = _build_text_batch( - self.policy, - [{"role": "user", "content": content}], - add_generation_prompt=False, - ) batch = dict(observation) - batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"] - batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"] + if getattr(self.policy.config, "joint_subtask_conditioning", False): + # Joint sequences keep the task turn (with state) and render the + # subtask as a causal assistant turn, exactly as trained. + from transformers import AutoTokenizer # noqa: PLC0415 + + from lerobot.policies.pi052.text_processor_pi052 import ( # noqa: PLC0415 + encode_prompt_with_targets, + register_paligemma_loc_tokens, + ) + from lerobot.utils.constants import OBS_LANGUAGE_CAUSAL_MARKS # noqa: PLC0415 + + task = state.task or "" + task_content = task if state_str is None else f"{task}, State: {state_str};" + tok_name = getattr(self.policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224" + tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens) + ids, attn, marks = encode_prompt_with_targets( + tokenizer, + [ + {"role": "user", "content": task_content}, + {"role": "assistant", "content": subtask}, + ], + target_indices=[1], + ) + device = getattr(self.policy.config, "device", None) + if device is not None: + ids, attn, marks = ids.to(device), attn.to(device), marks.to(device) + batch[OBS_LANGUAGE_TOKENS] = ids + batch[OBS_LANGUAGE_ATTENTION_MASK] = attn + batch[OBS_LANGUAGE_CAUSAL_MARKS] = marks + else: + content = subtask if state_str is None else f"{subtask}, State: {state_str};" + text_batch = _build_text_batch( + self.policy, + [{"role": "user", "content": content}], + add_generation_prompt=False, + ) + batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"] + batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"] return self.policy.predict_action_chunk(batch) def generate_text( @@ -67,6 +98,21 @@ class PI052PolicyAdapter(BaseLanguageAdapter): user_text: str | None = None, ) -> str: messages = self.build_messages(kind, state, user_text=user_text) + if kind == "subtask" and getattr(self.policy.config, "joint_subtask_conditioning", False): + # Joint samples carry state on the task turn, so the subtask must be + # generated from the same state-bearing prompt. + import torch # noqa: PLC0415 + + from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str # noqa: PLC0415 + from lerobot.utils.constants import OBS_STATE # noqa: PLC0415 + + obs_state = (observation or {}).get(OBS_STATE) + if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0: + state_row = obs_state[0] if obs_state.ndim > 1 else obs_state + for m in reversed(messages): + if m.get("role") == "user": + m["content"] = f"{m.get('content', '')}, State: {discretize_state_str(state_row)};" + break return _generate_with_policy( self.policy, messages, @@ -141,7 +187,10 @@ def _build_text_batch( messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in prompt_messages] prompt, _spans = _format_messages(messages) if add_generation_prompt: - prompt = prompt + "Assistant: " + # No trailing space: SentencePiece folds it into the first target token + # ("▁move"), so a space-suffixed prefill ends in a lone "▁" the model + # never saw at this position during training. + prompt = prompt + "Assistant:" encoded = tokenizer(prompt, return_tensors="pt") ids = encoded["input_ids"] diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index 9b838a79f..b22423f46 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -28,6 +28,7 @@ from torch.nn import functional from lerobot.utils.constants import ( ACTION, OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_CAUSAL_MARKS, OBS_LANGUAGE_TOKENS, OBS_STATE, ) @@ -143,9 +144,42 @@ class PI05Pytorch(PI05PytorchBase): # see openpi `PI0Pytorch` suffix_out = suffix_out.to(dtype=torch.float32) return self.action_out_proj(suffix_out) + def sample_actions( + self, + images, + img_masks, + tokens, + masks, + noise=None, + num_steps=None, + lang_causal_marks=None, + **kwargs, + ) -> Tensor: + """Sample actions, optionally marking trailing language positions causal. -# FAST tokens occupy the high vocabulary range and must be masked during text generation. -_FAST_ACTION_VOCAB_SIZE = 2048 + ``lang_causal_marks`` (B, L_lang bool) flags generated-subtask tokens so + joint-sequence checkpoints see the same causal prefix layout at + inference as during training (``_mark_target_span_causal``). + """ + self._lang_causal_marks = lang_causal_marks + try: + return super().sample_actions( + images, img_masks, tokens, masks, noise=noise, num_steps=num_steps, **kwargs + ) + finally: + self._lang_causal_marks = None + + def embed_prefix(self, images, img_masks, tokens, masks): + prefix_embs, prefix_pad, prefix_att = super().embed_prefix(images, img_masks, tokens, masks) + marks = getattr(self, "_lang_causal_marks", None) + if marks is not None: + prefix_att = _apply_causal_language_marks(prefix_att, marks.to(prefix_att.device)) + return prefix_embs, prefix_pad, prefix_att + + +# The universal `physical-intelligence/fast` tokenizer (and dataset refits of it) +# uses 1024 BPE codes; text generation must mask any that map below the range. +_FAST_ACTION_VOCAB_SIZE = 1024 _HF_KERNELS_ENABLED = False @@ -331,6 +365,17 @@ def _mark_target_span_causal( return att +def _apply_causal_language_marks(prefix_att_masks: Tensor, marks: Tensor) -> Tensor: + """OR per-token causal marks into the trailing language segment of a prefix.""" + att = prefix_att_masks.clone() + n = min(marks.shape[1], att.shape[1]) + if n <= 0: + return att + seg = att[:, -n:].bool() + att[:, -n:] = (seg | marks[:, -n:].bool()).to(att.dtype) + return att + + def _fast_lin_ce( hidden: Tensor, lm_head_weight: Tensor, @@ -1622,8 +1667,12 @@ class PI052Policy(PI05Policy): return decoded def _prepare_action_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: - from .inference.pi052_adapter import _build_text_batch # noqa: PLC0415 - from .text_processor_pi052 import discretize_state_str # noqa: PLC0415 + from .inference.pi052_adapter import _build_text_batch, _get_loc_tokenizer # noqa: PLC0415 + from .text_processor_pi052 import ( # noqa: PLC0415 + discretize_state_str, + encode_prompt_with_targets, + register_paligemma_loc_tokens, + ) n = self._batch_size_from_observation(batch) self._ensure_subtask_state(n) @@ -1631,6 +1680,14 @@ class PI052Policy(PI05Policy): # Mirror training by appending the already normalized state to low-level prompts. state_all = batch.get(OBS_STATE) + joint = bool(getattr(self.config, "joint_subtask_conditioning", False)) + joint_tokenizer = None + if joint: + from transformers import AutoTokenizer # noqa: PLC0415 + + tok_name = getattr(self.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224" + joint_tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens) + # Hold subtasks for the configured interval to match training and avoid rapid replanning. replan = int(getattr(self.config, "subtask_replan_steps", 0) or 0) hold_chunks = max(1, round(replan / self.config.n_action_steps)) if replan > 0 else 1 @@ -1638,7 +1695,7 @@ class PI052Policy(PI05Policy): self._subtask_chunk_counter += 1 # Generate and batch one independently conditioned subtask per environment. - rows: list[tuple[Tensor, Tensor | None]] = [] + rows: list[tuple[Tensor, Tensor | None, Tensor | None]] = [] tokenizer = None for i in range(n): if regenerate or not self.last_subtasks[i]: @@ -1648,32 +1705,62 @@ class PI052Policy(PI05Policy): # Hold the previously generated subtask; only the state in the # prompt below is refreshed to the current observation. subtask = self.last_subtasks[i] - content = subtask - if torch.is_tensor(state_all): - content = f"{subtask}, State: {discretize_state_str(state_all[i])};" - text_batch = _build_text_batch( - self, - [{"role": "user", "content": content}], - add_generation_prompt=False, - ) - rows.append((text_batch["lang_tokens"], text_batch["lang_masks"])) - tokenizer = text_batch["tokenizer"] + state_str = discretize_state_str(state_all[i]) if torch.is_tensor(state_all) else None + if joint: + # Joint sequences keep the task turn (with state) and render the + # subtask as a causal assistant turn, exactly as trained. + task_content = tasks[i] + if state_str is not None: + task_content = f"{task_content}, State: {state_str};" + ids, attn, marks = encode_prompt_with_targets( + joint_tokenizer, + [ + {"role": "user", "content": task_content}, + {"role": "assistant", "content": subtask}, + ], + target_indices=[1], + ) + device = getattr(self.config, "device", None) + if device is not None: + ids, attn, marks = ids.to(device), attn.to(device), marks.to(device) + rows.append((ids, attn, marks)) + tokenizer = joint_tokenizer + else: + content = subtask if state_str is None else f"{subtask}, State: {state_str};" + text_batch = _build_text_batch( + self, + [{"role": "user", "content": content}], + add_generation_prompt=False, + ) + rows.append((text_batch["lang_tokens"], text_batch["lang_masks"], None)) + tokenizer = text_batch["tokenizer"] - tokens, masks = self._stack_token_rows(rows, tokenizer) + tokens, masks, marks = self._stack_token_rows(rows, tokenizer) out = dict(batch) out[OBS_LANGUAGE_TOKENS] = tokens out[OBS_LANGUAGE_ATTENTION_MASK] = masks + if marks is not None: + out[OBS_LANGUAGE_CAUSAL_MARKS] = marks return out def _generate_low_level_subtask(self, obs_i: dict[str, Tensor], task: str, i: int) -> str: from .inference.pi052_adapter import _generate_with_policy # noqa: PLC0415 + from .text_processor_pi052 import discretize_state_str # noqa: PLC0415 msg = "" if task: + content = task + if getattr(self.config, "joint_subtask_conditioning", False): + # Joint samples carry state on the task turn, so the subtask + # must be generated from the same state-bearing prompt. + state = obs_i.get(OBS_STATE) + if torch.is_tensor(state) and state.numel() > 0: + state_row = state[0] if state.ndim > 1 else state + content = f"{task}, State: {discretize_state_str(state_row)};" msg = _generate_with_policy( self, - [{"role": "user", "content": task}], + [{"role": "user", "content": content}], observation=obs_i, label=f"eval subtask gen[{i}]", suppress_loc_tokens=True, @@ -1740,21 +1827,27 @@ class PI052Policy(PI05Policy): return out @staticmethod - def _stack_token_rows(rows: list[tuple[Tensor, Tensor | None]], tokenizer: Any) -> tuple[Tensor, Tensor]: - """Right-pad per-env ``(1, L_i)`` token/mask rows and stack to ``(n, L)``. + def _stack_token_rows( + rows: list[tuple[Tensor, Tensor | None, Tensor | None]], tokenizer: Any + ) -> tuple[Tensor, Tensor, Tensor | None]: + """Right-pad per-env ``(1, L_i)`` token/mask/marks rows and stack to ``(n, L)``. Right-padding with a False attention mask matches the training-time tokenizer (``padding_side="right"``), so the action expert treats pad - positions as masked. + positions as masked. Causal marks (third element, optional) pad False. """ - max_len = max(t.shape[1] for t, _ in rows) + max_len = max(t.shape[1] for t, _, _ in rows) pad_id = getattr(tokenizer, "pad_token_id", None) or 0 + has_marks = any(m is not None for _, _, m in rows) tok_rows: list[Tensor] = [] mask_rows: list[Tensor] = [] - for tokens, masks in rows: + marks_rows: list[Tensor] = [] + for tokens, masks, marks in rows: length = tokens.shape[1] if masks is None: masks = torch.ones((1, length), dtype=torch.bool, device=tokens.device) + if has_marks and marks is None: + marks = torch.zeros((1, length), dtype=torch.bool, device=tokens.device) if length < max_len: pad = max_len - length tokens = torch.cat( @@ -1765,9 +1858,17 @@ class PI052Policy(PI05Policy): [masks, torch.zeros((1, pad), dtype=masks.dtype, device=masks.device)], dim=1, ) + if has_marks: + marks = torch.cat( + [marks, torch.zeros((1, pad), dtype=marks.dtype, device=marks.device)], + dim=1, + ) tok_rows.append(tokens) mask_rows.append(masks) - return torch.cat(tok_rows, dim=0), torch.cat(mask_rows, dim=0) + if has_marks: + marks_rows.append(marks) + stacked_marks = torch.cat(marks_rows, dim=0) if has_marks else None + return torch.cat(tok_rows, dim=0), torch.cat(mask_rows, dim=0), stacked_marks @staticmethod def _fallback_subtask_from_task(task: str) -> str: @@ -1915,4 +2016,22 @@ class PI052Policy(PI05Policy): if self.config.use_flashrt_fp8_mlp and not getattr(self, "_fp8_applied", False): self._fp8_applied = True self.apply_flashrt_fp8_mlp(batch) - return super().predict_action_chunk(batch, **kwargs) + marks = batch.get(OBS_LANGUAGE_CAUSAL_MARKS) + if marks is None: + return super().predict_action_chunk(batch, **kwargs) + return self._predict_action_chunk_with_marks(batch, marks, **kwargs) + + @torch.no_grad() + def _predict_action_chunk_with_marks( + self, batch: dict[str, Tensor], marks: Tensor, **kwargs: Unpack[ActionSelectKwargs] + ) -> Tensor: + """Base ``predict_action_chunk`` plus causal marks on the generated-subtask span.""" + self.eval() + images, img_masks = self._preprocess_images(batch) + tokens = batch[OBS_LANGUAGE_TOKENS] + masks = batch[OBS_LANGUAGE_ATTENTION_MASK] + actions = self.model.sample_actions( + images, img_masks, tokens, masks, lang_causal_marks=marks, **kwargs + ) + original_action_dim = self.config.output_features[ACTION].shape[0] + return actions[:, :, :original_action_dim] diff --git a/src/lerobot/policies/pi052/text_processor_pi052.py b/src/lerobot/policies/pi052/text_processor_pi052.py index 45288a8e2..9b6dd8bb2 100644 --- a/src/lerobot/policies/pi052/text_processor_pi052.py +++ b/src/lerobot/policies/pi052/text_processor_pi052.py @@ -251,6 +251,44 @@ def _format_messages( return "".join(parts), spans +def encode_prompt_with_targets( + tokenizer: Any, messages: list[dict[str, Any]], target_indices: list[int] +) -> tuple[Tensor, Tensor, Tensor]: + """Tokenize a flat prompt and mark the token positions of target spans. + + Inference-side twin of ``PI052TextTokenizerStep._encode_messages``: same + serialization (role headers, target EOS) and the same offset-overlap span + arithmetic, but unpadded and returning a boolean target mask instead of + labels. Used to rebuild joint-sequence prompts whose target spans must be + attended causally, matching ``_mark_target_span_causal`` at train time. + + Returns ``(input_ids, attention_mask, target_marks)``, each ``(1, L)``. + """ + prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None)) + encoded = tokenizer(prompt, return_tensors="pt", return_offsets_mapping=True) + input_ids = encoded["input_ids"][0] + attention_mask = encoded.get("attention_mask") + if attention_mask is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + else: + attention_mask = attention_mask[0].bool() + offsets = encoded["offset_mapping"][0] + + marks = torch.zeros_like(input_ids, dtype=torch.bool) + for idx in target_indices: + if idx >= len(spans): + continue + char_start, char_end = spans[idx] + for token_pos in range(input_ids.shape[0]): + if not attention_mask[token_pos]: + continue + tok_start, tok_end = int(offsets[token_pos, 0]), int(offsets[token_pos, 1]) + if tok_end <= char_start or tok_start >= char_end: + continue + marks[token_pos] = True + return input_ids.unsqueeze(0), attention_mask.unsqueeze(0), marks.unsqueeze(0) + + @dataclass @ProcessorStepRegistry.register(name="pi052_text_tokenizer") class PI052TextTokenizerStep(ProcessorStep): diff --git a/src/lerobot/utils/constants.py b/src/lerobot/utils/constants.py index 6a44a4c8e..fae865b16 100644 --- a/src/lerobot/utils/constants.py +++ b/src/lerobot/utils/constants.py @@ -26,6 +26,7 @@ OBS_IMAGES = OBS_IMAGE + "s" OBS_LANGUAGE = OBS_STR + ".language" OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens" OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask" +OBS_LANGUAGE_CAUSAL_MARKS = OBS_LANGUAGE + ".causal_marks" OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask" OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens" OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask" diff --git a/tests/policies/pi052/test_pi052_joint_sequence.py b/tests/policies/pi052/test_pi052_joint_sequence.py new file mode 100644 index 000000000..a0fa22892 --- /dev/null +++ b/tests/policies/pi052/test_pi052_joint_sequence.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for PI052 joint-sequence (paper-style) subtask conditioning. + +Joint recipes train the subtask text and the action losses in one sequence, +with the supervised subtask span attended causally. At inference the same +layout is rebuilt around the *generated* subtask, so these tests pin: + +- the inference-side encoder produces the same token ids and target positions + as the training-time tokenizer step for the same messages; +- OR-ing causal marks into a prefix reproduces the training-time attention + pattern (prompt cannot see the subtask; subtask is causal over itself); +- the joint recipe file stays a valid message recipe; +- the FAST id mapping with the default ``fast_skip_tokens`` stays clear of + PaliGemma's ```` range so VQA and FAST supervision never collide. +""" + +from pathlib import Path + +import torch + +from lerobot.configs.recipe import TrainingRecipe +from lerobot.policies.pi052.text_processor_pi052 import ( + PI052TextTokenizerStep, + encode_prompt_with_targets, +) + + +class _CharTokenizer: + """Char-level stub: 1 char = 1 token, so offsets are trivially aligned.""" + + pad_token_id = 0 + eos_token = "\x1f" # unit separator — a 1-char "EOS" for testing + + def __call__(self, text, max_length=None, padding=None, return_tensors=None, **kwargs): + limit = max_length if max_length is not None else len(text) + ids = [ord(c) % 251 + 1 for c in text[:limit]] + offsets = [(i, i + 1) for i in range(len(ids))] + attention = [1] * len(ids) + if padding == "max_length" and max_length is not None and len(ids) < max_length: + pad = max_length - len(ids) + ids += [self.pad_token_id] * pad + offsets += [(0, 0)] * pad + attention += [0] * pad + return { + "input_ids": torch.tensor([ids], dtype=torch.long), + "attention_mask": torch.tensor([attention], dtype=torch.long), + "offset_mapping": torch.tensor([offsets], dtype=torch.long), + } + + +_MESSAGES = [ + {"role": "user", "content": "fold the towel"}, + {"role": "assistant", "content": "grab the near corner"}, +] + + +def test_encode_prompt_with_targets_matches_training_labels(): + tokenizer = _CharTokenizer() + + step = PI052TextTokenizerStep(max_length=120) + step._tokenizer = tokenizer + train_ids, train_attn, labels, predict_actions, _prompt = step._encode_messages( + tokenizer, + [dict(m) for m in _MESSAGES], + message_streams=["low_level", "low_level"], + target_indices=[1], + complementary={}, + ) + assert bool(predict_actions) + + ids, attn, marks = encode_prompt_with_targets(tokenizer, [dict(m) for m in _MESSAGES], [1]) + + n = int(attn.sum()) + assert n == int(train_attn.sum()) + assert torch.equal(ids[0, :n], train_ids[:n]) + # Causal marks at inference must cover exactly the supervised label span. + assert torch.equal(marks[0, :n], labels[:n] != -100) + assert marks.any(), "the assistant target span must be marked" + # The user turn must stay unmarked (bidirectional prompt). + user_len = len("User: fold the towel\n") + assert not marks[0, :user_len].any() + + +def test_apply_causal_language_marks_reproduces_training_mask(): + from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks + from lerobot.policies.pi052.modeling_pi052 import _apply_causal_language_marks + + n_img, n_lang = 4, 8 + prefix_len = n_img + n_lang + pad = torch.ones((1, prefix_len), dtype=torch.bool) + att = torch.zeros((1, prefix_len), dtype=torch.bool) + # Subtask span = language positions 5..7 (prefix positions 9..11). + marks = torch.zeros((1, n_lang), dtype=torch.bool) + marks[0, 5:8] = True + + att_marked = _apply_causal_language_marks(att, marks) + att_2d = make_att_2d_masks(pad, att_marked)[0] + + subtask = [n_img + 5, n_img + 6, n_img + 7] + # Prompt and images never see the subtask. + for q in range(n_img + 5): + for k in subtask: + assert not att_2d[q, k], f"prompt position {q} must not attend subtask position {k}" + # Subtask tokens see the full prompt and earlier subtask tokens only. + for qi, q in enumerate(subtask): + for k in range(n_img + 5): + assert att_2d[q, k] + for ki, k in enumerate(subtask): + assert bool(att_2d[q, k]) == (ki <= qi) + + +def test_joint_recipe_is_a_valid_message_recipe(): + recipe_path = Path(__file__).parents[3] / "src" / "lerobot" / "configs" / "recipes" / "subtask_joint.yaml" + recipe = TrainingRecipe.from_yaml(recipe_path) + assert recipe.messages is not None and len(recipe.messages) == 2 + assert all(turn.stream == "low_level" for turn in recipe.messages) + assert not recipe.messages[0].target + assert recipe.messages[1].target + assert recipe.messages[1].if_present == "subtask" + + +def test_default_fast_mapping_clears_loc_and_seg_ranges(): + from lerobot.policies.pi052.configuration_pi052 import PI052Config + from lerobot.policies.pi052.modeling_pi052 import _FAST_ACTION_VOCAB_SIZE + + skip = PI052Config.__dataclass_fields__["fast_skip_tokens"].default + assert skip == 1152 + + paligemma_vocab = 257152 + fast_ids = paligemma_vocab - 1 - skip - torch.arange(_FAST_ACTION_VOCAB_SIZE) + # Below the range [256000, 257024) and the range [257024, 257152). + assert int(fast_ids.max()) < 256000 + assert int(fast_ids.min()) >= 0