Compare commits

...

1 Commits

Author SHA1 Message Date
CarolinePascal 51ccac4ad4 fix(pi0fast): remove extra <bos> causing degenerate generation
pi0_fast decoding appended a second <bos> before the action block
(sample_actions_fast / sample_actions_fast_kv_cache), and
ActionTokenizerProcessorStep prepended the same <bos> to the training
target. The language prompt already begins with <bos> (standard PaliGemma
prefix "[image] <bos> prompt \n"), so this produced a ";\n <bos>" pattern
that does not exist in the openpi FAST format
("<bos> Task:...;\n Action: <fast> | <eos>", single leading <bos>,
"Action:" directly after ";\n").

On lerobot/pi0fast-base the extra <bos> pushed the model into a bos->bos
attractor, yielding degenerate generation (endless <bos> / random unicode)
instead of "Action: <fast>". Verified on stock PaliGemma with the same
weights: +bos -> <bos> collapse (logit ~51); -bos -> coherent, and on real
robot observations decoding now yields valid FAST action tokens.

Remove the extra <bos> in all three places so the prefix matches the
openpi format.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-10 15:21:44 +02:00
2 changed files with 16 additions and 17 deletions
@@ -613,15 +613,14 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
device = tokens.device
lm_head = self.paligemma_with_expert.paligemma.lm_head
# add bos token after tokens
bos_token = torch.full(
(bsize, 1), self._paligemma_tokenizer.bos_token_id, dtype=torch.long, device=device
)
tokens = torch.cat([tokens, bos_token], dim=1)
masks = torch.cat([masks, torch.ones((bsize, 1), dtype=torch.bool, device=device)], dim=1)
# NOTE (bug 2 fix): do NOT append a second <bos> here. The language tokens
# already begin with <bos> (standard PaliGemma prefix "[image] <bos> prompt \n").
# Appending another <bos> right before decoding pushes the checkpoint into a
# bos->bos attractor and yields degenerate generation. Generate directly after
# the prompt instead.
# 1. Initial Embedding (matches training prefix)
# prefix_embs will include [Images, Language Prompt, BOS]
# prefix_embs will include [Images, Language Prompt]
prefix_embs, prefix_pad_masks, prefix_att_masks, total_t_images, _ = self.embed_prefix_fast(
images, img_masks, tokens, masks, fast_action_tokens=None, fast_action_masks=None
)
@@ -709,14 +708,13 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
# --- 1. PREFILL PHASE ---
# Process Images + Text Prompt + BOS token once to populate the KV cache.
# Add BOS token to the prompt
bos_token = torch.full(
(bsize, 1), self._paligemma_tokenizer.bos_token_id, dtype=torch.long, device=device
)
tokens_in = torch.cat([tokens, bos_token], dim=1)
masks_in = torch.cat([masks, torch.ones((bsize, 1), dtype=torch.bool, device=device)], dim=1)
# NOTE (bug 2 fix): do NOT append a second <bos> here. The language tokens
# already begin with <bos> (standard PaliGemma prefix "[image] <bos> prompt \n").
# A second <bos> right before decoding causes degenerate bos->bos generation.
tokens_in = tokens
masks_in = masks
# Embed prefix [Images, Language, BOS]
# Embed prefix [Images, Language]
# fast_action_tokens=None means we are just embedding the condition (images+text)
prefix_embs, prefix_pad_masks, prefix_att_masks, total_t_images, _ = self.embed_prefix_fast(
images, img_masks, tokens_in, masks_in, fast_action_tokens=None, fast_action_masks=None
+4 -3
View File
@@ -476,11 +476,12 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
if tokens.dim() > 1:
tokens = tokens.flatten()
bos_id = self._paligemma_tokenizer.bos_token_id
# add bos
# NOTE (bug 2 fix): do NOT prepend a <bos> to the action target. The prompt
# already carries the leading <bos>; a second one before "Action:" mismatches
# the generation-time prefix (see sample_actions_fast*) and drives degenerate
# bos->bos decoding. Target is "Action: <fast tokens> |".
tokens = torch.cat(
[
torch.tensor([bos_id], device=action.device),
torch.tensor(
self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
device=action.device,