From 2749cf7767c49308ef837cd3e6325ecbf753d3fc Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 15 Jul 2026 15:35:08 +0200 Subject: [PATCH] refactor(pi052): remove debug prediction dumps --- src/lerobot/policies/pi052/debug_utils.py | 107 ------------------- src/lerobot/policies/pi052/modeling_pi052.py | 75 ------------- src/lerobot/scripts/lerobot_train.py | 8 -- 3 files changed, 190 deletions(-) delete mode 100644 src/lerobot/policies/pi052/debug_utils.py diff --git a/src/lerobot/policies/pi052/debug_utils.py b/src/lerobot/policies/pi052/debug_utils.py deleted file mode 100644 index 569dfbed0..000000000 --- a/src/lerobot/policies/pi052/debug_utils.py +++ /dev/null @@ -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) diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index 4114b88e4..24aa3834f 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/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], diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index d9172b562..a74a545a5 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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=``): 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()