mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 10:16:09 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 51ccac4ad4 |
@@ -613,15 +613,14 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
|
|||||||
device = tokens.device
|
device = tokens.device
|
||||||
lm_head = self.paligemma_with_expert.paligemma.lm_head
|
lm_head = self.paligemma_with_expert.paligemma.lm_head
|
||||||
|
|
||||||
# add bos token after tokens
|
# NOTE (bug 2 fix): do NOT append a second <bos> here. The language tokens
|
||||||
bos_token = torch.full(
|
# already begin with <bos> (standard PaliGemma prefix "[image] <bos> prompt \n").
|
||||||
(bsize, 1), self._paligemma_tokenizer.bos_token_id, dtype=torch.long, device=device
|
# Appending another <bos> right before decoding pushes the checkpoint into a
|
||||||
)
|
# bos->bos attractor and yields degenerate generation. Generate directly after
|
||||||
tokens = torch.cat([tokens, bos_token], dim=1)
|
# the prompt instead.
|
||||||
masks = torch.cat([masks, torch.ones((bsize, 1), dtype=torch.bool, device=device)], dim=1)
|
|
||||||
|
|
||||||
# 1. Initial Embedding (matches training prefix)
|
# 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(
|
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
|
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 ---
|
# --- 1. PREFILL PHASE ---
|
||||||
# Process Images + Text Prompt + BOS token once to populate the KV cache.
|
# Process Images + Text Prompt + BOS token once to populate the KV cache.
|
||||||
|
|
||||||
# Add BOS token to the prompt
|
# NOTE (bug 2 fix): do NOT append a second <bos> here. The language tokens
|
||||||
bos_token = torch.full(
|
# already begin with <bos> (standard PaliGemma prefix "[image] <bos> prompt \n").
|
||||||
(bsize, 1), self._paligemma_tokenizer.bos_token_id, dtype=torch.long, device=device
|
# A second <bos> right before decoding causes degenerate bos->bos generation.
|
||||||
)
|
tokens_in = tokens
|
||||||
tokens_in = torch.cat([tokens, bos_token], dim=1)
|
masks_in = masks
|
||||||
masks_in = torch.cat([masks, torch.ones((bsize, 1), dtype=torch.bool, device=device)], dim=1)
|
|
||||||
|
|
||||||
# Embed prefix [Images, Language, BOS]
|
# Embed prefix [Images, Language]
|
||||||
# fast_action_tokens=None means we are just embedding the condition (images+text)
|
# 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(
|
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
|
images, img_masks, tokens_in, masks_in, fast_action_tokens=None, fast_action_masks=None
|
||||||
|
|||||||
@@ -476,11 +476,12 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
|||||||
if tokens.dim() > 1:
|
if tokens.dim() > 1:
|
||||||
tokens = tokens.flatten()
|
tokens = tokens.flatten()
|
||||||
|
|
||||||
bos_id = self._paligemma_tokenizer.bos_token_id
|
# NOTE (bug 2 fix): do NOT prepend a <bos> to the action target. The prompt
|
||||||
# add bos
|
# 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(
|
tokens = torch.cat(
|
||||||
[
|
[
|
||||||
torch.tensor([bos_id], device=action.device),
|
|
||||||
torch.tensor(
|
torch.tensor(
|
||||||
self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
|
self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
|
||||||
device=action.device,
|
device=action.device,
|
||||||
|
|||||||
Reference in New Issue
Block a user