chore(dependencies): update uv.lock

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

@@ -613,14 +613,15 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         device = tokens.device
         lm_head = self.paligemma_with_expert.paligemma.lm_head
 
-        # 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.
+        # 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)
 
         # 1. Initial Embedding (matches training prefix)
-        # prefix_embs will include [Images, Language Prompt]
+        # prefix_embs will include [Images, Language Prompt, BOS]
         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
         )
@@ -708,13 +709,14 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         # --- 1. PREFILL PHASE ---
         # Process Images + Text Prompt + BOS token once to populate the KV cache.
 
-        # 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
+        # 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)
 
-        # Embed prefix [Images, Language]
+        # Embed prefix [Images, Language, BOS]
         # 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

src/lerobot/processor/tokenizer_processor.py

@@ -476,12 +476,11 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             if tokens.dim() > 1:
                 tokens = tokens.flatten()
 
-            # 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> |".
+            bos_id = self._paligemma_tokenizer.bos_token_id
+            # add bos
             tokens = torch.cat(
                 [
+                    torch.tensor([bos_id], device=action.device),
                     torch.tensor(
                         self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
                         device=action.device,