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:
Thomas Landeg
2026-07-16 13:19:45 -07:00
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