Pi0-FAST 模型架构#
1. 整体架构概览#
Pi0-FAST (Fast Action Sequence Tokenizer) 是 Physical Intelligence 提出的一种高效视觉-语言-动作 (VLA) 模型。与 Pi0 使用 Flow Matching 扩散生成连续动作不同，Pi0-FAST 将连续动作离散化为 token，然后通过 PaliGemma 视觉语言模型进行 自回归 next-token 预测。这种设计避免了多步去噪推理，实现了更快的动作生成。
核心思路：连续动作 → DCT 变换 → BPE 分词 → 离散 token → 自回归预测
graph TB
subgraph Input["输入"]
IMG["相机图像<br/>(pixel_values)"]
TXT["语言指令<br/>(input_ids)"]
ACT_TRAIN["动作序列<br/>[B, chunk_size, action_dim]<br/>(仅训练)"]
end
subgraph Tokenizer["FAST 动作分词器"]
DCT["DCT 变换<br/>连续动作 → 频域系数"]
BPE["BPE 分词<br/>频域系数 → 离散 token ID"]
end
subgraph VLM["PaliGemma 视觉语言模型"]
VE["SigLIP 视觉编码器"]
EMB["词元嵌入层<br/>(共享: 语言 + 动作 token)"]
GEMMA["Gemma 2B<br/>18 层 Transformer"]
LM_HEAD["LM Head<br/>(Linear → vocab logits)"]
end
subgraph Output["输出"]
LOGITS["下一个 token 预测<br/>(训练: 交叉熵损失)"]
ACTIONS["预测动作序列<br/>(推理: 自回归解码 → IDCT)"]
end
IMG --> VE --> GEMMA
TXT --> EMB --> GEMMA
ACT_TRAIN --> DCT --> BPE --> EMB
GEMMA --> LM_HEAD
LM_HEAD --> LOGITS
LM_HEAD --> ACTIONS
style Input fill:#e8f4fd,stroke:#2196F3
style Tokenizer fill:#fff3e0,stroke:#FF9800
style VLM fill:#e8f5e9,stroke:#4CAF50
style Output fill:#fce4ec,stroke:#E91E63
2. 核心组件详解#
2.1 PaliGemma 视觉语言模型#
Pi0-FAST 使用 PaliGemma（SigLIP + Gemma）作为统一的视觉语言主干。与 Pi0 不同，Pi0-FAST 没有独立的 Action Expert——动作 token 直接复用语言模型的词元嵌入和 LM Head 进行预测。
graph LR
subgraph Vision["SigLIP 视觉编码器"]
PIX["pixel_values<br/>[B, C, H, W]"] --> SIG["SigLIP<br/>视觉 Transformer"]
SIG --> IMG_EMB["图像嵌入<br/>[B, N_patches, 2048]"]
end
subgraph Language["词元嵌入（共享）"]
LANG_TOK["语言 token<br/>[B, seq_len]"] --> EMBED["embed_tokens()"]
ACT_TOK["动作 token<br/>[B, num_tokens]"] --> EMBED
EMBED --> LANG_EMB["嵌入 × √dim<br/>[B, *, 2048]"]
end
subgraph LM["Gemma 2B (18 层)"]
direction TB
L0["Layer 0"] --> L1["Layer 1~16"] --> L17["Layer 17"]
end
IMG_EMB --> LM
LANG_EMB --> LM
LM --> HEAD["LM Head<br/>2048 → 257152 (vocab)"]
style Vision fill:#e3f2fd,stroke:#2196F3
style Language fill:#f3e5f5,stroke:#9C27B0
style LM fill:#e8f5e9,stroke:#4CAF50
关键设计：
- 动作 token 和语言 token 共享同一个嵌入层 (embed_tokens)
- 动作 token 和语言 token 共享同一个 LM Head 进行预测
- 嵌入后乘以 √hidden_dim 进行缩放（与 Gemma 标准做法一致）
- 无需额外的状态投影或动作投影层
2.2 FAST 动作分词器 (Action Tokenizer)#
FAST 分词器将连续动作序列转换为离散 token，是 Pi0-FAST 的核心创新。
graph LR
subgraph Encode["编码（训练时）"]
A["连续动作<br/>[B, T, action_dim]"] --> DCT["DCT 变换<br/>(离散余弦变换)"]
DCT --> COEFF["频域系数<br/>(压缩表示)"]
COEFF --> BPE["BPE 分词<br/>(字节对编码)"]
BPE --> TOKENS["离散 token ID<br/>[B, num_tokens]"]
end
subgraph Decode["解码（推理时）"]
PRED_TOK["预测 token ID"] --> DE_BPE["BPE 解码"]
DE_BPE --> IDCT["IDCT 逆变换<br/>(scipy.fftpack.idct)"]
IDCT --> PRED_A["连续动作<br/>[B, T, action_dim]"]
end
style Encode fill:#e8f5e9,stroke:#4CAF50
style Decode fill:#fff3e0,stroke:#FF9800
DCT 变换的优势： - 将时间域动作序列转换为频域系数，实现 有损压缩 - 低频成分（平滑运动）用少量系数即可表示 - 天然规避了连续值域的归一化和 padding 不稳定问题 - BPE 分词进一��将系数编码为可预测的离散 token
Token 空间：
- 总词表大小: 257,152（Gemma 原始词表）
- fast_skip_tokens: 128（跳过前 128 个 token）
- 最大动作 token 数: 256（max_action_tokens）
2.3 自定义注意力掩码#
Pi0-FAST 使用精心设计的注意力掩码，确保图像/语言和动作 token 之间正确的信息流动。
graph TB
subgraph Mask["注意力掩码结构"]
direction LR
subgraph QueryAxis["Query (行)"]
QI["图像 token"]
QL["语言 token"]
QA["动作 token"]
end
subgraph Rules["注意力规则"]
R1["图像 ↔ 图像: ✅ 双向"]
R2["图像 ↔ 语言: ✅ 双向"]
R3["语言 ↔ 语言: ✅ 双向"]
R4["动作 → 图像: ✅ 允许"]
R5["动作 → 语言: ✅ 允许"]
R6["图像 → 动作: ❌ 禁止"]
R7["语言 → 动作: ❌ 禁止"]
R8["动作 ↔ 动作: ⬇️ 因果"]
end
end
style Mask fill:#e0f2f1,stroke:#009688
graph LR
subgraph AttentionMatrix["注意力矩阵可视化"]
M[" 图像 语言 动作<br/>图像 [✅✅✅ ✅✅ ❌❌❌]<br/>语言 [✅✅✅ ✅✅ ❌❌❌]<br/>动作₁ [✅✅✅ ✅✅ ✅❌❌]<br/>动作₂ [✅✅✅ ✅✅ ✅✅❌]<br/>动作₃ [✅✅✅ ✅✅ ✅✅✅]"]
end
style AttentionMatrix fill:#f5f5f5,stroke:#9E9E9E
实现细节 (_create_custom_attention_mask_fast):
1. 初始化全零布尔掩码 [B, total_len, total_len]
2. 图像/语言段之间：设为 True（双向注意力）
3. 动作段 → 图像/语言段：设为 True（动作可看到观测）
4. 动作段内部：设下三角 True（因果注意力，自回归生成）
5. 最后与 padding mask 相交：att_masks &= pad_2d_masks
3. 训练流水线#
graph TB
subgraph DataPrep["数据准备"]
RAW_A["原始连续动作<br/>[B, chunk_size, action_dim]"]
RAW_A --> NORM["归一化<br/>(MEAN_STD)"]
NORM --> PAD["零填充至 max_action_dim=32"]
PAD --> FAST_ENC["FAST Tokenizer 编码<br/>DCT → BPE"]
FAST_ENC --> ACT_TOKENS["动作 token ID<br/>[B, max_action_tokens]"]
RAW_IMG["原始图像<br/>[B, n_cameras, C, H, W]"]
RAW_IMG --> RESIZE["resize_with_pad<br/>→ 224×224"]
RAW_LANG["语言指令<br/>(文本)"]
RAW_LANG --> TOK["PaliGemma Tokenizer<br/>max_length=200"]
TOK --> LANG_TOKENS["语言 token ID<br/>[B, seq_len]"]
end
subgraph Embedding["嵌入构建"]
RESIZE --> VIS_EMB["SigLIP 编码<br/>→ 图像嵌入"]
LANG_TOKENS --> LANG_EMB["embed_tokens × √dim<br/>→ 语言嵌入"]
ACT_TOKENS --> ACT_EMB["embed_tokens × √dim<br/>→ 动作嵌入"]
VIS_EMB --> CONCAT["拼接<br/>[图像, 语言, 动作]"]
LANG_EMB --> CONCAT
ACT_EMB --> CONCAT
end
subgraph Forward["前向传播"]
CONCAT --> ATT_MASK["构建注意力掩码<br/>(双向 + 因果)"]
ATT_MASK --> GEMMA["Gemma 2B<br/>18 层 Transformer"]
GEMMA --> HIDDEN["隐藏状态<br/>[B, total_len, 2048]"]
end
subgraph Loss["损失计算"]
HIDDEN --> EXTRACT["提取动作位置<br/>hidden[:, -num_fast:]"]
EXTRACT --> LM_HEAD["LM Head<br/>→ vocab logits"]
LM_HEAD --> SHIFT["左移 logits / 右移 targets<br/>(next-token prediction)"]
SHIFT --> CE["CrossEntropyLoss<br/>(带 padding mask)"]
CE --> LOSS["fast_loss = masked_sum / mask_count"]
end
style DataPrep fill:#e3f2fd,stroke:#2196F3
style Embedding fill:#f3e5f5,stroke:#9C27B0
style Forward fill:#e8f5e9,stroke:#4CAF50
style Loss fill:#fce4ec,stroke:#E91E63
训练目标详解：
输入 token: [img_1, ..., img_N, lang_1, ..., lang_M, act_1, act_2, ..., act_{T-1}]
目标 token: [ act_2, act_3, ..., act_T ]
- logits 左移一位，targets 右移一位，实现 next-token prediction
- 只在动作 token 位置计算损失（图像和语言 token 不参与损失）
- 使用
fast_action_masks屏蔽 padding token - 损失函数：
CrossEntropyLoss(reduction="none")，然后 mask 加权平均
4. 推理流水线#
Pi0-FAST 支持两种推理模式：无缓存 和 KV 缓存 加速。
4.1 KV 缓存推理（默认，高效）#
graph TB
subgraph Prefill["阶段一: Prefill（一次性）"]
IMG["图像"] --> VIS["SigLIP 编码"]
LANG["语言 + BOS"] --> EMB["词元嵌入"]
VIS --> CAT["拼接 [图像, 语言, BOS]"]
EMB --> CAT
CAT --> FWD1["Gemma 前向<br/>use_cache=True"]
FWD1 --> KV["KV Cache<br/>(18 层缓存)"]
FWD1 --> FIRST["LM Head → 采样<br/>→ 第一个动作 token"]
end
subgraph Decode["阶段二: 自回归解码（逐 token）"]
FIRST --> LOOP["循环 t = 1 → max_decoding_steps"]
LOOP --> EMB_NEW["嵌入上一个 token<br/>[B, 1, 2048]"]
EMB_NEW --> FWD2["Gemma 前向<br/>+ KV Cache<br/>只处理 1 个 token"]
KV --> FWD2
FWD2 --> SAMPLE["LM Head → 采样"]
SAMPLE -->|"temperature=0: argmax<br/>temperature>0: multinomial"| NEXT["下一个 token"]
NEXT --> LOOP
end
subgraph PostProcess["后处理"]
LOOP -->|"生成完毕"| ALL_TOKENS["所有动作 token<br/>[B, max_decoding_steps]"]
ALL_TOKENS --> DEBPE["BPE 解码"]
DEBPE --> IDCT["IDCT 逆变换"]
IDCT --> FINAL["连续动作<br/>[B, chunk_size, action_dim]"]
end
style Prefill fill:#e3f2fd,stroke:#2196F3
style Decode fill:#e8f5e9,stroke:#4CAF50
style PostProcess fill:#fff3e0,stroke:#FF9800
KV 缓存优化要点： - Prefill 阶段：处理完整前缀序列，生成 KV Cache - 解码阶段：每步只处理 1 个新 token，重用缓存的 K/V - 注意力掩码：新 token 可注意到所有历���有效 token - 复杂度从 O(T²) 降低到 O(T)（T = 总序列长度）
4.2 无缓存推理（简单但慢）#
graph LR
subgraph NaiveLoop["逐步完整前向传播"]
T0["t=0: [img, lang] → Gemma → token₁"]
T1["t=1: [img, lang, token₁] → Gemma → token₂"]
T2["t=2: [img, lang, token₁, token₂] → Gemma → token₃"]
DOTS["..."]
TN["t=N: [img, lang, token₁...N] → Gemma → token_{N+1}"]
T0 --> T1 --> T2 --> DOTS --> TN
end
style NaiveLoop fill:#ffebee,stroke:#f44336
- 每步重新计算完整序列（无缓存复用）
- 序列逐步增长，注意力掩码动态扩展
- 仅作为参考实现，实际使用 KV 缓存版本
5. 与 Pi0 的关键区别#
| 特性 | Pi0 | Pi0-FAST |
|---|---|---|
| 动作表示 | 连续值 (Flow Matching) | 离散 token (FAST Tokenizer) |
| 动作生成方式 | 迭代去噪 (10 步) | 自回归 next-token 预测 |
| Action Expert | 独立 Gemma 300M | 无（复用 PaliGemma） |
| 模型结构 | 双流 (PaliGemma + Expert) | 单流 (仅 PaliGemma) |
| 训练损失 | MSE (flow matching) | CrossEntropy (next-token) |
| 推理速度 | 10 次完整前向传播 | ~256 次轻量解码步 (KV cache) |
| 状态输入 | 投影到隐藏维 → suffix | 不直接使用（通过语言描述） |
| 归一化敏感性 | 高（MEAN_STD + padding 不稳定） | 低（token 化天然规避） |
| Padding 问题 | 严重（30 维零值影响 flow matching） | 无（token 化后无连续 padding） |
graph LR
subgraph Pi0Path["Pi0: Flow Matching 路径"]
P_IN["连续动作"] --> P_NOISE["加噪 x_t"] --> P_DIT["双流 Transformer<br/>× 10 步去噪"] --> P_OUT["连续动作"]
end
subgraph FastPath["Pi0-FAST: Token 路径"]
F_IN["连续动作"] --> F_DCT["DCT"] --> F_BPE["BPE"] --> F_AR["单流 Transformer<br/>自回归解码"] --> F_IDCT["IDCT"] --> F_OUT["连续动作"]
end
style Pi0Path fill:#e3f2fd,stroke:#2196F3
style FastPath fill:#e8f5e9,stroke:#4CAF50
6. 关键超参数#
模型结构#
| 参数 | 值 | 说明 |
|---|---|---|
paligemma_variant |
gemma_2b |
PaliGemma 使用 Gemma 2B |
| Gemma 2B width | 2048 | 隐藏层维度 |
| Gemma 2B depth | 18 | Transformer 层数 |
| Gemma 2B mlp_dim | 16,384 | FFN 中间层维度 |
| Gemma 2B num_heads | 8 | 注意力头数 |
| Gemma 2B head_dim | 256 | 每头维度 |
| Gemma 2B num_kv_heads | 1 | GQA KV 头数 |
| 词表大小 | 257,152 | Gemma 原始词表 |
动作空间#
| 参数 | 值 | 说明 |
|---|---|---|
chunk_size |
50 | 预测动作步数 |
n_action_steps |
50 | 执行动作步数 |
max_state_dim |
32 | 状态向量填充维度 |
max_action_dim |
32 | 动作向量填充维度 |
max_action_tokens |
256 | 最大离散动作 token 数 |
fast_skip_tokens |
128 | 跳过词表前 128 个 token |
推理参数#
| 参数 | 值 | 说明 |
|---|---|---|
temperature |
0.0 | 采样温度 (0=argmax) |
max_decoding_steps |
256 | 最大自回归步数 |
use_kv_cache |
True | 启用 KV 缓存加速 |
validate_action_token_prefix |
True | 验证 "Action: " 前缀 |
训练参数#
| 参数 | 值 | 说明 |
|---|---|---|
optimizer_lr |
2.5e-5 | AdamW 学习率 |
optimizer_betas |
(0.9, 0.95) | Adam β 参数 |
optimizer_weight_decay |
0.01 | 权重衰减 |
optimizer_grad_clip_norm |
1.0 | 梯度裁剪阈值 |
scheduler_warmup_steps |
1,000 | 预热步数 |
scheduler_decay_steps |
30,000 | 余弦衰减总步数 |
scheduler_decay_lr |
2.5e-6 | 最终学习率 |
tokenizer_max_length |
200 | 语言 token 最大长度 |
image_resolution |
(224, 224) | 输入图像分辨率 |
7. 关键源文件#
| 组件 | 类名 | 文件 |
|---|---|---|
| 策略封装 | PI0FastPolicy |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:804 |
| 核心模型 | PI0FastPytorch |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:278 |
| PaliGemma 封装 | PI0FastPaliGemma |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:186 |
| Gemma 配置 | GemmaConfig / get_gemma_config |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:150 |
| 配置类 | PI0FastConfig |
lerobot/policies/pi0_fast/configuration_pi0_fast.py:31 |
| 嵌入构建 | embed_prefix_fast() |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:355 |
| 注意力掩码 | _create_custom_attention_mask_fast() |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:443 |
| 训练前向 | forward() |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:484 |
| 推理 (无缓存) | sample_actions_fast() |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:583 |
| 推理 (KV 缓存) | sample_actions_fast_kv_cache() |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:680 |
| 图像预处理 | resize_with_pad_torch() |
lerobot/policies/pi0_fast/modeling_pi0_fast.py:76 |
| FAST 分词器 | physical-intelligence/fast |
HuggingFace Hub (外部) |
| 文本分词器 | google/paligemma-3b-pt-224 |
HuggingFace Hub (外部) |