admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-08-01 05:59:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(policies): early stopping + termination

Browse Source
This commit is contained in:
Steven Palma
2026-07-31 16:53:37 +02:00
parent 50192947fd
commit fdd9d92349
1 changed files with 19 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pi0_fast/modeling_pi0_fast.py
+19 -8
View File
@@ -604,6 +604,12 @@ class PI0FastPytorch(nn.Module): # see openpi `PI0Pytorch`
) -> torch.Tensor: ) -> torch.Tensor:
""" """
Optimized autoregressive decoding for FAST tokens using KV Caching. Optimized autoregressive decoding for FAST tokens using KV Caching.
Greedy decoding stops once every sequence emits the end-of-action marker. The
returned tensor keeps its fixed shape, with positions not generated after the
batch-wide stop left zero-filled. Stochastic decoding always runs to
``max_decoding_steps`` so early stopping does not change the RNG state used by
subsequent calls.
""" """
if max_decoding_steps is None: if max_decoding_steps is None:
max_decoding_steps = self.config.max_action_tokens max_decoding_steps = self.config.max_action_tokens
@@ -612,10 +618,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 # detokenize_actions() cuts at the first "|", so greedy decoding can stop once
# discarded anyway. # every sequence has emitted it. Keep stochastic decoding unchanged because
pipe_id = self._paligemma_tokenizer.convert_tokens_to_ids("|") # skipping multinomial calls would shift the RNG state for subsequent calls.
finished = torch.zeros(bsize, dtype=torch.bool, device=device) end_of_action_token_id = self._paligemma_tokenizer.convert_tokens_to_ids("|")
finished = torch.zeros(bsize, dtype=torch.bool, device=device) if temperature == 0 else None
# --- 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.
@@ -668,7 +675,10 @@ 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 if finished is not None:
finished |= next_token.squeeze(-1) == end_of_action_token_id
if bool(finished.all()):
return generated_action_tokens
# 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)
@@ -719,9 +729,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 finished is not None:
if bool(finished.all()): finished |= next_token.squeeze(-1) == end_of_action_token_id
break if bool(finished.all()):
break
return generated_action_tokens return generated_action_tokens