feat(train): debug dump runs inference too, with parity check

Extends the periodic LM-head dump (LEROBOT_DEBUG_PREDS_EVERY) to ALSO
run select_message autoregressively on the same prompt prefix and show:

  prompt                          : '<bos>User: ... Assistant: '
  target  (ground truth)          : ' close the gripper ...'
  training argmax (teacher-fed)   : ' close the gri lift ...'  acc=12/15=80%
  inference (autoregressive)      : ' close the gripper around ...'
  first-token parity              : train=3387 (' close') vs infer=3387 (' close')  ✓ MATCH

The first-token parity check is decisive: training-side argmax at the
prompt-end position and inference's first generated token both compute
``argmax(lm_head(h_last_prompt))`` on identical context, so they MUST
match. Any divergence signals a training↔inference bug (mask, dtype,
KI routing, embedding scale, etc.). Subsequent tokens can diverge
because training uses teacher forcing while inference free-runs.

debug_text_predictions now also returns an ``inference`` list keyed
by sample, each entry carrying ``first_sup_pos`` and ``decoded``.
Limited to 24 new tokens per sample to keep the dump fast.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

src/lerobot/configs/recipes/subtask_mem.yaml

@@ -0,0 +1,57 @@
+# subtask_mem_vqa_speech — Hi-Robot blend + memory + spoken responses.
+#
+# Superset of subtasks_vqa.yaml. Keeps the core subtask + action + VQA
+# training, and adds two text-supervised tasks:
+#
+#   high_level_subtask         — predict the subtask from the task.
+#   low_level_execution        — flow loss with [images, subtask, state].
+#   memory_update              — compress progress into a memory note.
+#   user_interjection_response — reply to a user interjection with a
+#                                spoken `say` tool call (no plan, no
+#                                subtask text — just the spoken reply).
+#   ask_vqa_{top,wrist}        — camera-grounded VQA.
+#
+# Plan is intentionally left out — memory is the only persistent
+# high-level state here, keeping the prompt short.
+#
+# Requires the dataset to carry `memory`, `interjection` and `say`-tool
+# annotations (the annotation pipeline's memory + interjection modules)
+# in addition to `subtask` and `vqa`. Sub-recipes whose `if_present`
+# bindings are missing simply don't render for that sample, so a
+# dataset without interjections still trains the rest of the blend.
+#
+# SmolVLA2 note: the `say` tool call on the interjection-response turn
+# is flattened to a `<say>...</say>` text marker by the chat tokenizer
+# (`_flatten_say_tool_calls`) before `apply_chat_template`, so the LM
+# head learns to emit exactly the marker the runtime parses back
+# (`_split_plan_and_say`).
+
+blend:
+
+  high_level_subtask:
+    weight: 0.30
+    messages:
+      - {role: user, content: "${task}", stream: high_level}
+      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
+
+  low_level_execution:
+    weight: 0.40
+    messages:
+      # The action expert is conditioned on the SUBTASK — at inference
+      # `HighLevelSubtaskFwd` generates it via the LM head and feeds it
+      # here. `stream: low_level` flips `predict_actions=True` so the
+      # flow loss fires; no text-CE target (subtask prediction is owned
+      # by `high_level_subtask`).
+      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
+
+  memory_update:
+    weight: 0.30
+    bindings:
+      prior_memory: "nth_prev(style=memory, offset=1)"
+      current_memory: "emitted_at(t, style=memory)"
+      completed_subtask: "nth_prev(style=subtask, offset=1)"
+    messages:
+      - {role: user, content: "${task}", stream: high_level}
+      - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
+      - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
+      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}

src/lerobot/policies/pi05/modeling_pi05.py

@@ -617,10 +617,13 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             )
         return func(*args, **kwargs)
 
-    def _prepare_attention_masks_4d(self, att_2d_masks):
+    def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
         """Helper method to prepare 4D attention masks for transformer."""
         att_2d_masks_4d = att_2d_masks[:, None, :, :]
-        return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
+        result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
+        if dtype is not None:
+            result = result.to(dtype=dtype)
+        return result
 
     def sample_noise(self, shape, device):
         return torch.normal(
@@ -756,7 +759,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
         position_ids = torch.cumsum(pad_masks, dim=1) - 1
 
-        att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks)
+        att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks, dtype=prefix_embs.dtype)
 
         def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond):
             (_, suffix_out), _ = self.paligemma_with_expert.forward(
@@ -814,7 +817,9 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
         prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
 
-        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks)
+        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            prefix_att_2d_masks, dtype=prefix_embs.dtype
+        )
         self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"  # noqa: SLF001
 
         _, past_key_values = self.paligemma_with_expert.forward(
@@ -884,7 +889,9 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
         position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
 
-        full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            full_att_2d_masks, dtype=suffix_embs.dtype
+        )
         self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
 
         past_key_values = copy.deepcopy(past_key_values)

src/lerobot/policies/pi052/modeling_pi052.py

@@ -630,7 +630,9 @@ class PI052Policy(PI05Policy):
             att_2d_masks[:, fast_end:, fast_start:fast_end] = False
 
         position_ids = torch.cumsum(pad_masks, dim=1) - 1
-        att_2d_masks_4d = self.model._prepare_attention_masks_4d(att_2d_masks)
+        att_2d_masks_4d = self.model._prepare_attention_masks_4d(
+            att_2d_masks, dtype=prefix_embs.dtype
+        )
 
         # ---- forward (capture BOTH expert outputs) ------------------
         (prefix_out, suffix_out), _ = self.model.paligemma_with_expert.forward(
@@ -740,7 +742,7 @@ class PI052Policy(PI05Policy):
 
         att_2d = make_att_2d_masks(full_pad, full_att)
         position_ids = torch.cumsum(full_pad, dim=1) - 1
-        att_2d_4d = self.model._prepare_attention_masks_4d(att_2d)
+        att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=full_embs.dtype)
 
         (vlm_out, _), _ = self.model.paligemma_with_expert.forward(
             attention_mask=att_2d_4d,
@@ -780,6 +782,133 @@ class PI052Policy(PI05Policy):
 
         return text_loss, fast_loss
 
+    # ------------------------------------------------------------------
+    # Diagnostic: forward + argmax for supervised text positions
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def debug_text_predictions(
+        self, batch: dict[str, Tensor], max_samples: int = 5
+    ) -> dict[str, Tensor]:
+        """Run the text-loss forward but return argmax predictions instead of CE.
+
+        Lets a periodic training-loop hook compare what the LM head emits
+        right now against what it *should* emit at every supervised
+        position — the cheapest "is text training actually working"
+        diagnostic. Returns CPU tensors keyed by ``input_ids``,
+        ``attention_mask``, ``labels``, ``predictions``; predictions are
+        aligned with input positions (``predictions[t]`` is the head's
+        argmax after seeing ``input_ids[:t+1]``, so it should match
+        ``input_ids[t+1]`` for next-token prediction). Returns ``{}``
+        when the batch has no supervised text positions.
+        """
+        from ..pi05.modeling_pi05 import make_att_2d_masks  # noqa: PLC0415
+
+        text_labels = batch.get("text_labels")
+        if text_labels is None or not bool((text_labels != -100).any().item()):
+            return {}
+
+        was_training = self.training
+        self.eval()
+        try:
+            n = min(max_samples, int(text_labels.shape[0]))
+            sub: dict[str, Any] = {
+                OBS_LANGUAGE_TOKENS: batch[OBS_LANGUAGE_TOKENS][:n],
+                OBS_LANGUAGE_ATTENTION_MASK: batch[OBS_LANGUAGE_ATTENTION_MASK][:n],
+            }
+            for k, v in batch.items():
+                if isinstance(k, str) and k.startswith("observation.images.") and torch.is_tensor(v):
+                    sub[k] = v[:n]
+
+            sub_labels = text_labels[:n]
+            images, img_masks = self._preprocess_images(sub)
+            lang_tokens = sub[OBS_LANGUAGE_TOKENS]
+            lang_masks = sub[OBS_LANGUAGE_ATTENTION_MASK]
+
+            prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix(
+                images, img_masks, lang_tokens, lang_masks
+            )
+            lang_start = prefix_embs.shape[1] - sub_labels.shape[1]
+            if lang_start >= 0:
+                prefix_att = _mark_target_span_causal(
+                    prefix_att, sub_labels, lang_start, prefix_embs.shape[1]
+                )
+
+            att_2d = make_att_2d_masks(prefix_pad, prefix_att)
+            position_ids = torch.cumsum(prefix_pad, dim=1) - 1
+            att_2d_4d = self.model._prepare_attention_masks_4d(att_2d)
+            backbone = self.model.paligemma_with_expert
+            backbone_dtype = (
+                backbone.paligemma.model.language_model.layers[0]
+                .self_attn.q_proj.weight.dtype
+            )
+            if att_2d_4d.dtype != backbone_dtype:
+                att_2d_4d = att_2d_4d.to(dtype=backbone_dtype)
+
+            (vlm_out, _), _ = backbone.forward(
+                attention_mask=att_2d_4d,
+                position_ids=position_ids,
+                past_key_values=None,
+                inputs_embeds=[prefix_embs, None],
+                use_cache=False,
+            )
+            text_hidden = vlm_out[:, -sub_labels.shape[1]:, :]
+            lm_head = backbone.paligemma.lm_head
+            text_logits = lm_head(text_hidden.to(lm_head.weight.dtype))
+            preds = text_logits.argmax(dim=-1)
+
+            # Train/inference parity check — run select_message on the
+            # *same* prompt prefix (the language up to but not including
+            # the supervised span) and capture the auto-regressive
+            # generation. The first generated token MUST match the
+            # training-side argmax at the prompt-end position (both are
+            # ``argmax lm_head(h_last_prompt)`` over identical context);
+            # any divergence is a parity bug (mask, dtype, KI routing
+            # difference). Later tokens can diverge because training
+            # uses teacher forcing while inference free-runs.
+            inference_outputs: list[dict[str, Any]] = []
+            for s in range(n):
+                row_labels = sub_labels[s]
+                sup_pos = (row_labels != -100).nonzero(as_tuple=True)[0]
+                if sup_pos.numel() == 0:
+                    inference_outputs.append({"first_token": None, "decoded": ""})
+                    continue
+                first_sup = int(sup_pos[0].item())
+                # Build a single-sample batch with attention zeroed past
+                # the supervised span — that gives ``embed_prefix`` only
+                # the user-prompt portion to attend over.
+                prompt_mask = sub[OBS_LANGUAGE_ATTENTION_MASK][s : s + 1].clone()
+                prompt_mask[:, first_sup:] = 0
+                inf_batch: dict[str, Any] = {
+                    OBS_LANGUAGE_TOKENS: sub[OBS_LANGUAGE_TOKENS][s : s + 1],
+                    OBS_LANGUAGE_ATTENTION_MASK: prompt_mask,
+                }
+                for k, v in sub.items():
+                    if isinstance(k, str) and k.startswith("observation.images."):
+                        inf_batch[k] = v[s : s + 1]
+                if "observation.state" in batch and torch.is_tensor(batch["observation.state"]):
+                    inf_batch["observation.state"] = batch["observation.state"][s : s + 1]
+                try:
+                    # Tight budget — we just want to see the model's
+                    # opening continuation, not the full sequence.
+                    decoded = self.select_message(
+                        inf_batch, max_new_tokens=24, temperature=0.0, top_p=1.0
+                    )
+                except Exception as exc:  # noqa: BLE001
+                    decoded = f"<inference failed: {type(exc).__name__}: {exc}>"
+                inference_outputs.append({"first_sup_pos": first_sup, "decoded": decoded})
+
+            return {
+                "input_ids": lang_tokens.detach().cpu(),
+                "attention_mask": lang_masks.detach().cpu(),
+                "labels": sub_labels.detach().cpu(),
+                "predictions": preds.detach().cpu(),
+                "inference": inference_outputs,
+            }
+        finally:
+            if was_training:
+                self.train()
+
     # ------------------------------------------------------------------
     # select_message — AR text generation at inference
     # ------------------------------------------------------------------
@@ -864,9 +993,7 @@ class PI052Policy(PI05Policy):
         for _ in range(max_new_tokens):
             att_2d = make_att_2d_masks(current_pad, current_att)
             position_ids = torch.cumsum(current_pad, dim=1) - 1
-            att_2d_4d = self.model._prepare_attention_masks_4d(att_2d)
-            if att_2d_4d.dtype != backbone_dtype:
-                att_2d_4d = att_2d_4d.to(dtype=backbone_dtype)
+            att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=backbone_dtype)
             (vlm_out, _), _ = backbone.forward(
                 attention_mask=att_2d_4d,
                 position_ids=position_ids,

src/lerobot/policies/pi_gemma.py

@@ -272,6 +272,8 @@ class PiGemmaModel(GemmaModel):  # type: ignore[misc]
         # Convert to bfloat16 if the first layer uses bfloat16
         if len(self.layers) > 0 and self.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16:
             hidden_states = hidden_states.to(torch.bfloat16)
+        if causal_mask is not None and torch.is_floating_point(causal_mask):
+            causal_mask = causal_mask.to(dtype=hidden_states.dtype)
 
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)

src/lerobot/scripts/lerobot_train.py

@@ -20,6 +20,7 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
+import os
 import time
 from contextlib import nullcontext
 from pprint import pformat
@@ -156,6 +157,122 @@ def update_policy(
     return train_metrics, output_dict
 
 
+def _print_debug_text_predictions(
+    policy: Any, batch: dict[str, Any], step: int, n_samples: int = 5
+) -> None:
+    """Forward the current batch and print head-argmax vs label per supervised position.
+
+    Opt-in via ``LEROBOT_DEBUG_PREDS_EVERY=<step_interval>``. Only the
+    policy types that expose ``debug_text_predictions`` participate
+    (currently PI052); others are silently skipped. Pretty-prints up to
+    ``n_samples`` samples from the current batch, showing the prompt,
+    every supervised position's (label, prediction, ✓/✗), and a
+    per-sample token-accuracy summary — the cheapest "is text training
+    actually learning anything" signal.
+    """
+    if not hasattr(policy, "debug_text_predictions"):
+        return
+    try:
+        debug = policy.debug_text_predictions(batch, max_samples=n_samples)
+    except Exception as exc:  # noqa: BLE001
+        logging.warning("debug_text_predictions failed: %s", exc)
+        return
+    if not debug:
+        return
+
+    # Build a tokenizer for decoding — match training side exactly.
+    try:
+        from transformers import AutoTokenizer  # noqa: PLC0415
+
+        from lerobot.policies.pi052.text_processor_pi052 import (  # noqa: PLC0415
+            register_paligemma_loc_tokens,
+        )
+
+        tok_name = (
+            getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
+        )
+        tokenizer = register_paligemma_loc_tokens(AutoTokenizer.from_pretrained(tok_name))
+    except Exception as exc:  # noqa: BLE001
+        logging.warning("debug preds: tokenizer load failed: %s", exc)
+        return
+
+    ids = debug["input_ids"]
+    labels = debug["labels"]
+    preds = debug["predictions"]
+    attn = debug["attention_mask"]
+    inference = debug.get("inference") or []
+
+    n = ids.shape[0]
+    print(
+        f"\n========== STEP {step} DEBUG PREDICTIONS ({n} samples) ==========",
+        flush=True,
+    )
+    for s in range(n):
+        a = attn[s].tolist()
+        real = sum(a)
+        sid = ids[s].tolist()
+        sl = labels[s].tolist()
+        sp = preds[s].tolist()
+        prompt = tokenizer.decode(sid[:real], skip_special_tokens=False)
+        print(f"\n  --- sample {s + 1}/{n} ---", flush=True)
+        print(f"  prompt: {prompt!r}", flush=True)
+
+        # Ground-truth target (the contiguous supervised label span).
+        sup_ids = [int(sid[i]) for i in range(real) if sl[i] != -100]
+        if sup_ids:
+            print(
+                f"  target  (ground truth)        : {tokenizer.decode(sup_ids, skip_special_tokens=False)!r}",
+                flush=True,
+            )
+
+        # Training-side teacher-forced argmax on the same prompt+target.
+        n_sup = n_ok = 0
+        first_sup_pred: int | None = None
+        teacher_chars: list[int] = []
+        for i in range(1, real):
+            label = sl[i]
+            if label == -100:
+                continue
+            n_sup += 1
+            pred = int(sp[i - 1])
+            if first_sup_pred is None:
+                first_sup_pred = pred
+            teacher_chars.append(pred)
+            if label == pred:
+                n_ok += 1
+        teacher_text = (
+            tokenizer.decode(teacher_chars, skip_special_tokens=False) if teacher_chars else ""
+        )
+        acc = n_ok / max(n_sup, 1)
+        print(
+            f"  training argmax (teacher-fed) : {teacher_text!r}   acc={n_ok}/{n_sup}={acc:.1%}",
+            flush=True,
+        )
+
+        # Inference-side autoregressive output from the same prompt prefix.
+        inf_entry = inference[s] if s < len(inference) else None
+        if inf_entry:
+            inf_decoded = inf_entry.get("decoded", "")
+            print(f"  inference (autoregressive)    : {inf_decoded!r}", flush=True)
+            # First-token parity: training-side argmax at the prompt-end
+            # position MUST equal inference's first generated token —
+            # both compute argmax(lm_head(h_last_prompt)) on identical
+            # context. Any divergence signals a training↔inference bug.
+            if first_sup_pred is not None and inf_decoded and not inf_decoded.startswith("<inference"):
+                inf_ids = tokenizer(inf_decoded, add_special_tokens=False)["input_ids"]
+                if inf_ids:
+                    inf_first = int(inf_ids[0])
+                    match = inf_first == first_sup_pred
+                    print(
+                        f"  first-token parity            : "
+                        f"train={first_sup_pred} ({tokenizer.decode([first_sup_pred])!r}) "
+                        f"vs infer={inf_first} ({tokenizer.decode([inf_first])!r})  "
+                        f"{'✓ MATCH' if match else '✗ DIVERGED — training/inference mismatch'}",
+                        flush=True,
+                    )
+    print("=" * 60 + "\n", flush=True)
+
+
 def _build_vqa_oversample_weights(dataset: Any, target_fraction: float) -> "torch.Tensor | None":
     """Build per-frame sampling weights that oversample VQA-annotated frames.
 
@@ -542,6 +659,27 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
         is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
 
+        # Optional periodic head-prediction dump for the LM head:
+        # ``LEROBOT_DEBUG_PREDS_EVERY=1000`` prints 5 samples + per-token
+        # (label, argmax, ✓/✗) every 1000 steps. Cheap diagnostic to see
+        # whether the text head is actually learning what we expect, vs
+        # collapsing to a fixed token. Refilling the recipe-sample dump
+        # budget at the same cadence also redumps the raw input shapes.
+        _debug_preds_every = int(os.environ.get("LEROBOT_DEBUG_PREDS_EVERY", "0"))
+        if (
+            _debug_preds_every > 0
+            and step % _debug_preds_every == 0
+            and is_main_process
+        ):
+            try:
+                from lerobot.policies.pi052 import text_processor_pi052 as _tp  # noqa: PLC0415
+
+                _tp._DUMPED_SO_FAR = 0
+                _tp._DUMP_BUDGET = max(_tp._DUMP_BUDGET, 5)
+            except Exception:  # noqa: BLE001
+                pass
+            _print_debug_text_predictions(policy, batch, step, n_samples=5)
+
         if is_log_step:
             logging.info(train_tracker)
             if wandb_logger: