mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-01 05:59:44 +00:00
perf(pi0fast): stop FAST decoding at the end-of-action marker
The decode loop always runs all max_decoding_steps (256) tokens, but detokenize_actions() cuts the output at the first "|", so most of them get generated and then thrown away. On LIBERO the marker lands around token 25-35. Stopping there gives the same actions ~4x faster end to end. Checked on libero_10: 50/50 episodes came out identical to the old path, same success and same state trajectories.
This commit is contained in:
committed by
Steven Palma
parent
99443a936d
commit
50192947fd
@@ -612,6 +612,11 @@ 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
|
||||||
|
|
||||||
|
# detokenize_actions() cuts at the first "|", so anything generated past it is
|
||||||
|
# discarded anyway.
|
||||||
|
pipe_id = self._paligemma_tokenizer.convert_tokens_to_ids("|")
|
||||||
|
finished = torch.zeros(bsize, dtype=torch.bool, device=device)
|
||||||
|
|
||||||
# --- 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.
|
||||||
|
|
||||||
@@ -663,6 +668,7 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
|
|||||||
# Initialize storage for generated tokens
|
# Initialize storage for generated tokens
|
||||||
generated_action_tokens = torch.zeros((bsize, max_decoding_steps), dtype=torch.long, device=device)
|
generated_action_tokens = torch.zeros((bsize, max_decoding_steps), dtype=torch.long, device=device)
|
||||||
generated_action_tokens[:, 0] = next_token.squeeze(-1)
|
generated_action_tokens[:, 0] = next_token.squeeze(-1)
|
||||||
|
finished |= next_token.squeeze(-1) == pipe_id
|
||||||
|
|
||||||
# Track valid tokens mask (0 for pad, 1 for valid)
|
# Track valid tokens mask (0 for pad, 1 for valid)
|
||||||
# We need this to tell the new token what it can attend to (images + text + past actions)
|
# We need this to tell the new token what it can attend to (images + text + past actions)
|
||||||
@@ -713,6 +719,10 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
|
|||||||
|
|
||||||
generated_action_tokens[:, t] = next_token.squeeze(-1)
|
generated_action_tokens[:, t] = next_token.squeeze(-1)
|
||||||
|
|
||||||
|
finished |= next_token.squeeze(-1) == pipe_id
|
||||||
|
if bool(finished.all()):
|
||||||
|
break
|
||||||
|
|
||||||
return generated_action_tokens
|
return generated_action_tokens
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user