fix(pi0_fast): align FAST semantics with OpenPI

Use OpenPI-compatible normalization, token boundaries, balanced loss, and strict full-chunk decoding so training and inference share one sequence contract.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn
2026-07-17 11:50:37 +00:00
committed by Pepijn
parent a6f533a6dd
commit 7a05b31f83
7 changed files with 417 additions and 247 deletions
+10 -1
View File
@@ -43,6 +43,12 @@ This approach can transform **any existing VLM** into a VLA by training it to pr
pip install -e ".[pi]"
```
> [!WARNING]
> PI0-FAST now follows OpenPI's token and normalization semantics exactly. Checkpoints and
> saved preprocessors created with the previous LeRobot implementation are incompatible and
> must be retrained or regenerated. In particular, do not reuse `MEAN_STD` statistics or a
> processor that inserts a second BOS token before `Action:`.
## Training a Custom FAST Tokenizer
You have two options for the FAST tokenizer:
@@ -114,7 +120,7 @@ lerobot-train \
| `--policy.gradient_checkpointing=true` | Reduces memory usage significantly during training | `false` |
| `--policy.dtype=bfloat16` | Use mixed precision training for efficiency | `float32` |
| `--policy.chunk_size` | Number of action steps to predict (action horizon) | `50` |
| `--policy.n_action_steps` | Number of action steps to execute | `50` |
| `--policy.n_action_steps` | Number of decoded action steps to execute | `50` |
| `--policy.max_action_tokens` | Maximum number of FAST tokens per action chunk | `256` |
| `--policy.action_tokenizer_name` | FAST tokenizer to use | `lerobot/fast-action-tokenizer` |
| `--policy.auto_fit_fast_tokenizer=true` | Fit and cache a tokenizer for the training dataset | `false` |
@@ -157,6 +163,9 @@ actions = policy.predict_action_chunk(batch)
The model takes images, text instructions, and robot state as input, and outputs discrete FAST tokens that are decoded back to continuous actions.
PI0-FAST always decodes a complete `chunk_size` action chunk. `n_action_steps` controls only
how many actions from that chunk are executed before the policy predicts again.
## Configuration Options
| Parameter | Description | Default |