From 51ccac4ad407463e062f980fd5f8d24ea2d070b7 Mon Sep 17 00:00:00 2001
From: CarolinePascal <caroline8.pascal@gmail.com>
Date: Fri, 10 Jul 2026 15:21:44 +0200
Subject: [PATCH] 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>
---
 .../policies/pi0_fast/modeling_pi0_fast.py    | 26 +++++++++----------
 src/lerobot/processor/tokenizer_processor.py  |  7 ++---
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py b/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py
index d9342eb24..7b922a712 100644
--- a/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py
+++ b/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py
@@ -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
diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py
index a808e6127..ed3f8d07d 100644
--- a/src/lerobot/processor/tokenizer_processor.py
+++ b/src/lerobot/processor/tokenizer_processor.py
@@ -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,