From 17c08004617690d1efcaea289df4943b9440bf97 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 13 May 2026 12:05:37 +0200 Subject: [PATCH] fix(pi052): FAST loss masking + predict_actions gating + smolvla2 review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FAST loss changes ----------------- 1. Gate by ``predict_actions`` (same routing as flow loss). The ActionTokenizerProcessorStep tokenises actions for *every* sample regardless of which sub-recipe rendered it; for text-only recipes (high_level_subtask, memory_update, ...) the action tokens are still in the batch but mustn't be supervised. Skip the FAST forward+CE entirely when no sample in the batch has ``predict_actions=True``. 2. Switch from "multiply-by-mask" masking to ``ignore_index=-100``. The old pattern computed per-token CE for all positions, then zeroed out invalid ones. Two issues: (a) any out-of-vocab target id at a padded position would have crashed cross_entropy before the mask got a chance to zero it out, and (b) the pattern is needlessly clever. Now ``shift_targets.masked_fill(~mask, -100)`` followed by ``ignore_index=-100`` cleanly drops invalid positions. Matches the smolvla2 text-loss convention. 3. Clean up unused ``bsize`` variable in _compute_fast_action_loss and expand the attention-mask docstring with the ``make_att_2d_masks`` mask_ar convention spec (causal vs bidirectional blocks). smolvla2 audit (reference review, no code change) ------------------------------------------------- Compared smolvla2/modeling_smolvla2.py against pi052/modeling_pi052.py to catch parallel bugs. Findings: * No ``paligemma.language_model`` vs ``paligemma.model.language_model`` issue — smolvla2 uses SmolVLM (different class, different attribute layout) so the bug doesn't apply. * ``fill_kv_cache=True`` is correctly passed to smolvla's ``vlm_with_expert.forward`` — that class *does* accept the kwarg (unlike pi05's PaliGemmaWithExpertModel.forward, which is why pi052 must omit it). * Text-loss alignment is correct: ``_compute_text_loss`` computes ``lang_start`` / ``lang_end`` from the known prefix layout (``[image_blocks..., lang, state]``) and slices ``prefix_out`` to just the language positions before applying ``lm_head``. The parallel bug I fixed in pi052 (lm_head over the full prefix, shape-mismatched against text_labels) was *not* present in smolvla2. * Per-sample flow routing via ``predict_actions``: correctly masks per-sample by calling the parent ``forward(..., reduction='none')`` and applying the predict_actions mask before the mean. pi052 only has the batch-level any() gate — a parallel improvement for pi052 would require modifying PI05Pytorch.forward to support per-sample reduction, deferred. * ``reduction="none"`` returns ``total.expand(bsize)``: identical scalar-broadcast limitation in both policies. Acknowledged but low priority (only RA-BC weighting uses the per-sample path and it's documented as a known approximation in smolvla2). * Chat tokenizer correctly handles batched/unbatched messages, pads with -100 for label positions, builds attention masks. No bugs found. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/lerobot/policies/pi052/modeling_pi052.py | 56 ++++++++++++++------ 1 file changed, 39 insertions(+), 17 deletions(-) diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index 52349c614..657df36f7 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -339,7 +339,18 @@ class PI052Policy(PI05Policy): # ACTION_TOKENS / ACTION_TOKEN_MASK into the batch — we # forward them through the PaliGemma backbone alongside the # language prefix and compute CE on the action positions. - if getattr(self.config, "enable_fast_action_loss", False): + # + # Gated on ``predict_actions`` (same routing the flow loss + # uses): for text-only recipes the action_tokens are still + # present in the batch but shouldn't be supervised. Skip the + # entire FAST forward when no sample in the batch wants action + # supervision. + run_fast = ( + getattr(self.config, "enable_fast_action_loss", False) + and self.config.fast_action_loss_weight > 0 + and (predict_actions_t is None or bool(predict_actions_t.any().item())) + ) + if run_fast: from lerobot.utils.constants import ACTION_TOKEN_MASK, ACTION_TOKENS # noqa: PLC0415 action_tokens = batch.get(ACTION_TOKENS) @@ -399,14 +410,19 @@ class PI052Policy(PI05Policy): fast_emb = self.model.paligemma_with_expert.embed_language_tokens(action_tokens) fast_emb = fast_emb * math.sqrt(emb_dim) - # Concat onto the prefix. Pad masks: language uses - # ``lang_masks``; FAST uses ``action_mask`` (True at valid - # token positions). Attention masks add ``True`` (causal) - # for FAST so they can attend to the bidirectional prefix - # but only causally among themselves. - bsize, fast_len = action_tokens.shape + # Concat onto the prefix. + # pad masks: language uses ``lang_masks``; FAST uses + # ``action_mask`` (True at valid token positions). + # att masks: prefix is 0 (bidirectional block); FAST is 1 + # (each token starts its own causal block). Per + # ``make_att_2d_masks``'s mask_ar convention this + # yields prefix-LM attention: FAST tokens attend + # bidirectionally to images+language and causally + # among themselves, while prefix tokens *cannot* + # see FAST tokens. Matches pi05_full §III.C. + fast_len = action_tokens.shape[1] device = prefix_embs.device - ones_att = torch.ones((bsize, fast_len), dtype=torch.bool, device=device) + ones_att = torch.ones((action_tokens.shape[0], fast_len), dtype=torch.bool, device=device) full_embs = torch.cat([prefix_embs, fast_emb], dim=1) full_pad = torch.cat([prefix_pad, action_mask.to(prefix_pad.dtype)], dim=1) full_att = torch.cat([prefix_att, ones_att], dim=1) @@ -430,17 +446,23 @@ class PI052Policy(PI05Policy): lm_head = self.model.paligemma_with_expert.paligemma.lm_head fast_logits = lm_head(fast_hidden.to(lm_head.weight.dtype)) - # Shift for next-token prediction. + # Shift for next-token prediction. Replace targets at padded + # positions with -100 so ``ignore_index`` in cross_entropy + # cleanly drops them rather than relying on a post-hoc + # multiply-by-mask (which still computes the CE numerator at + # invalid positions and could crash if a padded target id + # falls outside the vocab). shift_logits = fast_logits[:, :-1, :].contiguous() - shift_targets = action_tokens[:, 1:].contiguous() - shift_mask = action_mask[:, 1:].contiguous().float() + shift_targets = action_tokens[:, 1:].contiguous().long() + shift_valid = action_mask[:, 1:].contiguous().bool() + shift_targets = shift_targets.masked_fill(~shift_valid, -100) - per_tok = F.cross_entropy( - shift_logits.view(-1, shift_logits.size(-1)), - shift_targets.view(-1), - reduction="none", - ).view(shift_targets.shape) - loss = (per_tok * shift_mask).sum() / shift_mask.sum().clamp_min(1.0) + # Mean over valid positions via ``ignore_index``. + loss = F.cross_entropy( + shift_logits.reshape(-1, shift_logits.shape[-1]), + shift_targets.reshape(-1), + ignore_index=-100, + ) return loss def _compute_text_loss(self, batch: dict[str, Tensor], text_labels: Tensor) -> Tensor: