refactor(pi052): remove debug prediction dumps

src/lerobot/policies/pi052/debug_utils.py

@@ -1,107 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 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.
-"""Training-time debug helpers for PI052's language head."""
-
-import logging
-from typing import Any
-
-
-def print_debug_text_predictions(policy: Any, batch: dict[str, Any], step: int, n_samples: int = 5) -> None:
-    """Print supervised text predictions and token accuracy for up to ``n_samples`` rows."""
-    # Unwrap distributed wrappers that do not proxy custom policy methods.
-    inner = policy
-    while hasattr(inner, "module") and not hasattr(inner, "debug_text_predictions"):
-        inner = inner.module
-    if not hasattr(inner, "debug_text_predictions"):
-        logging.warning(
-            "LEROBOT_DEBUG_PREDS_EVERY set but policy %s has no "
-            "debug_text_predictions method — skipping dump.",
-            type(inner).__name__,
-        )
-        return
-    try:
-        debug = inner.debug_text_predictions(batch, max_samples=n_samples)
-    except Exception as exc:  # noqa: BLE001
-        logging.warning("debug_text_predictions failed: %s", exc, exc_info=True)
-        return
-    if not debug:
-        logging.warning(
-            "debug_text_predictions returned no supervised samples — current batch has no text labels."
-        )
-        return
-    policy = inner  # used below for select_message-style decoding parity
-
-    # 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"]
-
-    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
-        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])
-            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,
-        )
-    print("=" * 60 + "\n", flush=True)

src/lerobot/policies/pi052/modeling_pi052.py

@@ -1548,81 +1548,6 @@ class PI052Policy(PreTrainedPolicy):
 
         return text_loss, fast_loss
 
-    @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.
-        """
-
-        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)
-
-            return {
-                "input_ids": lang_tokens.detach().cpu(),
-                "attention_mask": lang_masks.detach().cpu(),
-                "labels": sub_labels.detach().cpu(),
-                "predictions": preds.detach().cpu(),
-            }
-        finally:
-            if was_training:
-                self.train()
-
     def select_message(
         self,
         batch: dict[str, Tensor],

src/lerobot/scripts/lerobot_train.py

@@ -678,14 +678,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         is_env_eval_step = cfg.env_eval_freq > 0 and step % cfg.env_eval_freq == 0
         is_eval_step = cfg.eval_steps > 0 and eval_dataloader is not None and step % cfg.eval_steps == 0
 
-        # Optional LM-head diagnostic (``LEROBOT_DEBUG_PREDS_EVERY=<steps>``): prints
-        # per-token (label, argmax) for a few samples to check the text head is learning.
-        _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:
-            from lerobot.policies.pi052.debug_utils import print_debug_text_predictions  # noqa: PLC0415
-
-            print_debug_text_predictions(policy, batch, step, n_samples=5)
-
         if is_log_step:
             # Collective reduce must run on every rank, before the main-process gate below.
             train_tracker.reduce_across_ranks()