From 50192947fdef06563a00c420060354b62c4c88d2 Mon Sep 17 00:00:00 2001 From: Thomas Landeg Date: Thu, 16 Jul 2026 13:19:45 -0700 Subject: [PATCH] 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. --- src/lerobot/policies/pi0_fast/modeling_pi0_fast.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py b/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py index 494d25644..475563d7c 100644 --- a/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py +++ b/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py @@ -612,6 +612,11 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch` device = tokens.device 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 --- # 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 generated_action_tokens = torch.zeros((bsize, max_decoding_steps), dtype=torch.long, device=device) 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) # 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) + finished |= next_token.squeeze(-1) == pipe_id + if bool(finished.all()): + break + return generated_action_tokens