mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-16 22:41:49 +00:00
refactor(pi052): remove debug prediction dumps
This commit is contained in:
Pepijn
@@ -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)
|
||||
@@ -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],
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user