Pi0.5 模型架构

1. 整体架构概览

Pi0.5 (Physical Intelligence 0.5) 是 Physical Intelligence 推出的视觉-语言-动作 (VLA) 模型,在架构上与 Pi0 几乎完全相同,均采用 PaliGemma 视觉语言模型作为感知主干、Gemma 300M 作为动作专家 (Action Expert),通过双流 Transformer 实现视觉语言特征与动作特征的联合处理,并使用 Flow Matching 进行连续动作空间的生成。Pi0.5 相对于 Pi0 的核心改进在于三个方面:(1) 将状态和动作的归一化方式从 MEAN_STD 改为 QUANTILES,将所有值映射到 [-1, 1] 的有界空间,从根本上解决了零填充 (padding) 对 Flow Matching 稳定性的干扰问题;(2) 在动作专家中引入 AdaRMS (Adaptive Root Mean Square) 层归一化,以时间步信息对归一化进行条件调制;(3) 简化了 Time MLP 架构,将时间信息从与动作拼接融合改为通过 AdaRMS 注入。

graph TB subgraph Input["输入"] IMG["相机图像
(多视角支持)
[B, C, 224, 224]"] TXT["语言指令
(tokenized)
[B, max_length]"] end subgraph VLM["PaliGemma 视觉语言模型 (Prefix Stream)"] VE["SigLIP 视觉编码器
图像 → 视觉 token"] TE["Gemma 2B 文本嵌入
语言 token 嵌入"] LM["Gemma 2B 语言模型
18层 Transformer
hidden_size=2048"] end subgraph Expert["Action Expert (Suffix Stream)"] AIN["动作投影层
action_in_proj
Linear(32 → 1024)"] TMLP["Time MLP
正弦编码 → SiLU MLP
→ AdaRMS 条件"] GE["Gemma 300M 专家模型
18层 Transformer + AdaRMS
hidden_size=1024"] AOUT["动作输出投影
action_out_proj
Linear(1024 → 32)"] end subgraph DualStream["双流 Transformer (共享注意力)"] DS["逐层联合注意力
18层共享 Q/K/V 拼接
prefix + suffix 联合计算"] end subgraph FM["Flow Matching"] NOISE["高斯噪声 ~ N(0,1)"] INTERP["线性插值
x_t = t*noise + (1-t)*action"] EULER["欧拉积分去噪
(10步)"] end subgraph Output["输出"] ACT["预测动作轨迹
[B, 50, action_dim]"] end IMG --> VE --> LM TXT --> TE --> LM LM -->|"视觉语言特征"| DS NOISE --> INTERP INTERP --> AIN --> DS TMLP -->|"AdaRMS 条件"| GE GE --> DS DS --> AOUT EULER -.->|"推理时迭代"| DS AOUT --> ACT style Input fill:#e8f4fd,stroke:#2196F3 style VLM fill:#fff3e0,stroke:#FF9800 style Expert fill:#e8f5e9,stroke:#4CAF50 style DualStream fill:#f3e5f5,stroke:#9C27B0 style FM fill:#fff9c4,stroke:#FFC107 style Output fill:#fce4ec,stroke:#E91E63

2. 核心组件详解

2.1 PaliGemma 视觉语言模型

PaliGemma 由 Google 开发,是 Pi0/Pi0.5 的感知主干网络,负责将图像和语言指令编码为统一的视觉语言特征序列。它由 SigLIP 视觉编码器和 Gemma 2B 语言模型组成。

graph LR subgraph SigLIPBranch["SigLIP 视觉编码器"] PIX["pixel_values
[B, C, 224, 224]"] --> SIG["SigLIP
Vision Transformer
patch_size=14"] SIG --> PROJ_V["多模态投影层
intermediate=4304
projection_dim=2048"] end subgraph TextBranch["文本嵌入"] TOK["input_ids
[B, seq_len]"] --> EMB["Gemma 词元嵌入
vocab_size=257152"] EMB --> SCALE["* sqrt(hidden_size)
缩放因子"] end subgraph GemmaLM["Gemma 2B 语言模型"] MERGE["序列拼接
[img_tokens | lang_tokens]"] GEM["Gemma 2B
18层 Transformer
hidden_size=2048
num_heads=8, head_dim=256
MLP dim=16384
GQA: kv_heads=1"] end PROJ_V -->|"视觉 token
[B, N_patches, 2048]"| MERGE SCALE -->|"文本 token
[B, seq_len, 2048]"| MERGE MERGE --> GEM GEM -->|"视觉语言特征
[B, total_seq, 2048]"| OUT_VL["prefix_output"] style SigLIPBranch fill:#e3f2fd,stroke:#2196F3 style TextBranch fill:#fff3e0,stroke:#FF9800 style GemmaLM fill:#e8f5e9,stroke:#4CAF50

关键特性: - SigLIP 将 224x224 图像编码为一系列视觉 token,通过投影层映射到 Gemma 的 2048 维空间 - 多相机支持:每个相机图像独立编码,缺失的相机用全 -1 填充(SigLIP 的 padding 值),并通过 img_mask 标记为无效 - 图像预处理:图像从 [0,1] 归一化到 [-1,1],非方形图像通过 resize_with_pad 保持比例缩放后黑边填充 - Prefix-LM 注意力:图像 token 和语言 token 使用双向注意力(att_mask=0),即它们之间可以相互 attend

SigLIP vs CLIP/GLIP 对比

Pi0/Pi0.5 选择 SigLIP 而非更常见的 CLIP 作为视觉编码器,主要原因在于 SigLIP 的训练目标和效率优势:

特性 CLIP GLIP SigLIP
提出者 OpenAI (2021) Microsoft (2022) Google (2023)
训练目标 对比学习 (InfoNCE) 对比学习 + 短语级 Grounding Sigmoid 损失 (pairwise)
损失函数 softmax 归一化的交叉熵,需要全局 batch 内负样本 同 CLIP + 区域-短语对齐 每对样本独立的 sigmoid 二分类损失
batch 依赖 强:性能随 batch size 增大而提升,小 batch 性能下降明显 :每对样本独立计算,不需要大 batch 内的负样本对比
核心区别 对一个 batch 的 N 个图文对做 NxN 矩阵 softmax CLIP 基础上增加了目标检测和定位能力 对每个图文对独立做"是否匹配"的二分类
优势 开创性的图文对齐方法 支持开放词汇目标检测,细粒度区域-语言对齐 训练效率高、小模型性能好、与 Gemma 同属 Google 生态
典型应用 通用图文检索、零样本分类 开放词汇检测、视觉定位 (grounding) VLM 视觉编码器 (PaLI, PaliGemma)

为什么 Pi0.5 用 SigLIP 而不是 CLIP?

  1. Sigmoid vs Softmax 的关键区别:CLIP 使用 softmax,意味着每个图像的正确文本概率 = 1/N(N 为 batch size),所有负样本参与归一化。SigLIP 使用 sigmoid,每对图文独立判断匹配/不匹配,不依赖 batch 内其他样本。这使得 SigLIP 在小 batch 微调时更稳定——而机器人学习的 batch size 通常远小于互联网规模预训练
  2. Google 生态一致性:SigLIP 是 PaliGemma 的原生视觉编码器,与 Gemma 语言模型在预训练阶段已经对齐,直接使用避免了跨生态适配的额外工程
  3. 小模型效率:SigLIP 在 ViT-B/16 等较小模型上的性能显著优于同等规模的 CLIP,更适合实时机器人控制的延迟需求

简单理解:CLIP 像"从 N 个候选中选最匹配的"(排名问题),SigLIP 像"判断这一对是否匹配"(分类问题)。后者更简单、更稳定,尤其在 batch 不大的时候。


2.2 Action Expert (Gemma 300M)

动作专家是一个独立的 Gemma 300M 模型,专门处理动作序列。与 PaliGemma 的 Gemma 2B 共享相同的层数(18层),但每层的宽度更小。Pi0.5 的核心改进之一是在 Action Expert 中启用了 AdaRMS 层归一化

graph TB subgraph ActionInput["动作输入处理"] NA["含噪动作
[B, chunk_size, max_action_dim]"] AIP["action_in_proj
Linear(32 → 1024)"] NA --> AIP end subgraph TimeCond["时间步条件 (AdaRMS)"] T["时间步 t
[B]"] SIN["正弦位置编码
dim=1024
period=[4e-3, 4.0]"] MLP_IN["time_mlp_in
Linear(1024 → 1024)"] SILU["SiLU 激活"] MLP_OUT["time_mlp_out
Linear(1024 → 1024)"] SILU2["SiLU 激活"] T --> SIN --> MLP_IN --> SILU --> MLP_OUT --> SILU2 SILU2 -->|"adarms_cond
[B, 1024]"| ADARMSN["AdaRMS
层归一化条件"] end subgraph GemmaExpert["Gemma 300M 专家 (18层)"] direction TB L0["Layer 0: AdaRMS → Attn → AdaRMS → MLP"] L1["Layer 1: AdaRMS → Attn → AdaRMS → MLP"] DOTS["... (共18层)"] L17["Layer 17: AdaRMS → Attn → AdaRMS → MLP"] FN["Final AdaRMS Norm"] L0 --> L1 --> DOTS --> L17 --> FN end subgraph ActionOutput["动作输出"] AOP["action_out_proj
Linear(1024 → 32)"] PRED["预测速度 v_t
[B, chunk_size, action_dim]"] AOP --> PRED end AIP -->|"动作嵌入
[B, 50, 1024]"| GemmaExpert ADARMSN -->|"条件注入
每层 LayerNorm"| GemmaExpert FN --> AOP style ActionInput fill:#e3f2fd,stroke:#2196F3 style TimeCond fill:#fff9c4,stroke:#FFC107 style GemmaExpert fill:#e8f5e9,stroke:#4CAF50 style ActionOutput fill:#fce4ec,stroke:#E91E63

Gemma 300M 配置:

参数
hidden_size (width) 1024
num_hidden_layers (depth) 18
intermediate_size (mlp_dim) 4096
num_attention_heads 8
num_key_value_heads 1 (GQA)
head_dim 256
hidden_activation gelu_pytorch_tanh
use_adarms True (Pi0.5) / False (Pi0)

2.3 双流 Transformer

双流 Transformer 是 Pi0/Pi0.5 的核心计算机制。PaliGemma(prefix stream)和 Action Expert(suffix stream)各自拥有独立的权重,但在每一层的注意力计算中共享 Q/K/V 空间——将两个流的 query、key、value 沿序列维度拼接后进行联合注意力计算,再将结果拆分回各自的流中。

graph TB subgraph Layer["单层双流 Transformer (共18层)"] subgraph PrefixStream["Prefix Stream (PaliGemma 2B)"] P_LN1["input_layernorm
(标准 RMSNorm)"] P_QKV["Q/K/V 投影
(独立权重)"] P_RES1["+ 残差连接"] P_LN2["post_attention_layernorm
(标准 RMSNorm)"] P_MLP["MLP
(gate+up → down)"] P_RES2["+ 残差连接"] end subgraph SuffixStream["Suffix Stream (Action Expert 300M)"] S_LN1["input_layernorm
(AdaRMS, 条件=time_emb)"] S_QKV["Q/K/V 投影
(独立权重)"] S_RES1["+ 残差连接"] S_LN2["post_attention_layernorm
(AdaRMS, 条件=time_emb)"] S_MLP["MLP
(gate+up → down)"] S_RES2["+ 残差连接"] end subgraph SharedAttn["共享注意力计算"] CAT_Q["拼接 Q: [prefix_Q | suffix_Q]"] CAT_K["拼接 K: [prefix_K | suffix_K]"] CAT_V["拼接 V: [prefix_V | suffix_V]"] ROPE["RoPE 位置编码"] ATT["联合注意力
Attention(Q_cat, K_cat, V_cat)
+ 2D attention mask"] SPLIT["按序列位置拆分输出"] end end P_LN1 --> P_QKV --> CAT_Q P_QKV --> CAT_K P_QKV --> CAT_V S_LN1 --> S_QKV --> CAT_Q S_QKV --> CAT_K S_QKV --> CAT_V CAT_Q --> ROPE --> ATT CAT_K --> ROPE CAT_V --> ATT ATT --> SPLIT SPLIT -->|"prefix 部分"| P_RES1 SPLIT -->|"suffix 部分"| S_RES1 P_RES1 --> P_LN2 --> P_MLP --> P_RES2 S_RES1 --> S_LN2 --> S_MLP --> S_RES2 style PrefixStream fill:#fff3e0,stroke:#FF9800 style SuffixStream fill:#e8f5e9,stroke:#4CAF50 style SharedAttn fill:#f3e5f5,stroke:#9C27B0

注意力掩码机制:

双流注意力使用精心设计的 2D 注意力掩码来控制信息流向:

att_mask 构造:
  图像 tokens:   [0, 0, 0, ..., 0]     ← 双向注意力
  语言 tokens:   [0, 0, 0, ..., 0]     ← 双向注意力
  动作 token[0]: [1]                    ← 因果边界:之前的 prefix 不能 attend 到此
  动作 token[1:]: [0, 0, ..., 0]       ← 动作 tokens 之间双向注意���

这意味着: - 图像 token、语言 token 之间可以相互 attend(prefix-LM 风格) - 动作 token 可以 attend 到所有图像和语言 token(单向) - 图像和语言 token 不能 attend 到动作 token(保证前缀可以被缓存复用)


2.4 Flow Matching

Flow Matching 是一种基于常微分方程 (ODE) 的生成模型方法。与扩散模型不同,Flow Matching 定义了从噪声分布到数据分布的直线路径,模型学习在每个时间点的速度场。

graph TB subgraph Training["训练阶段"] direction TB A_GT["真实动作 a
[B, 50, action_dim]
(QUANTILES 归一化后 ∈ [-1,1])"] N["高斯噪声 noise
~ N(0, 1)"] TBETA["t ~ Beta(1.5, 1.0)
* 0.999 + 0.001"] INTERP["线性插值构造含噪样本
x_t = t * noise + (1-t) * action"] TARGET["速度目标
u_t = noise - action"] A_GT --> INTERP N --> INTERP TBETA --> INTERP A_GT --> TARGET N --> TARGET INTERP -->|"x_t 输入模型"| MODEL["Pi0.5 模型"] MODEL -->|"预测速度 v_t"| LOSS["MSE 损失
L = ||u_t - v_t||^2"] TARGET --> LOSS end subgraph Inference["推理阶段"] direction TB NOISE2["初始噪声
x_1 ~ N(0, 1)"] STEP1["步骤 1: t=1.0
v_1 = model(x_1, t=1.0)
x_0.9 = x_1 + (-0.1)*v_1"] STEP2["步骤 2: t=0.9
v_2 = model(x_0.9, t=0.9)
x_0.8 = x_0.9 + (-0.1)*v_2"] DOTS["... (共10步)"] STEP10["步骤 10: t=0.1
v_10 = model(x_0.1, t=0.1)
x_0 = x_0.1 + (-0.1)*v_10"] RESULT["最终预测
x_0 ≈ action"] NOISE2 --> STEP1 --> STEP2 --> DOTS --> STEP10 --> RESULT end style Training fill:#e8f5e9,stroke:#4CAF50 style Inference fill:#e3f2fd,stroke:#2196F3

时间采样策略: - 训练时,时间步 t 从 Beta(1.5, 1.0) 分布中采样,再缩放到 [0.001, 1.0] - Beta(1.5, 1.0) 偏向较大的 t 值(接近纯噪声),这意味着模型在训练时更多地关注"从噪声开始去噪"的早期阶段 - 推理时,从 t=1(纯噪声)等间距积分到 t=0(纯动作),默认 10 步


3. 与 Pi0 的关键区别

Pi0.5 在架构上与 Pi0 几乎完全相同(同样的 PaliGemma + Gemma 300M 双流 Transformer + Flow Matching),但有三个关键改进。这些改进看似微小,实则对模型的训练稳定性和泛化能力有深远影响。

3.1 QUANTILES 归一化 vs MEAN_STD 归一化

graph LR subgraph Pi0Norm["Pi0: MEAN_STD 归一化"] RAW1["原始值
e.g. 关节角度
[0.5, 1.2, -3.7, ...]"] MS["x_norm = (x - mean) / std"] UNBOUNDED["归一化后: 无界
理论范围 (-inf, +inf)
e.g. [-0.3, 1.5, -2.1, ...]"] PAD1["零填充 padding
[..., 0, 0, 0]
对应原始空间均值位置"] RAW1 --> MS --> UNBOUNDED UNBOUNDED --> PAD1 end subgraph Pi05Norm["Pi0.5: QUANTILES 归一化"] RAW2["原始值
e.g. 关节角度
[0.5, 1.2, -3.7, ...]"] QT["分位数映射
映射到 [-1, 1]"] BOUNDED["归一化后: 有界
严格范围 [-1, 1]
e.g. [-0.1, 0.6, -0.9, ...]"] PAD2["零填充 padding
[..., 0, 0, 0]
正好是范围中点!"] RAW2 --> QT --> BOUNDED BOUNDED --> PAD2 end style Pi0Norm fill:#ffcdd2,stroke:#E91E63 style Pi05Norm fill:#c8e6c9,stroke:#4CAF50

为什么 QUANTILES 归一化至关重要?

这是 Pi0.5 最重要的改进,直接影响 Flow Matching 的训练稳定性。核心问题在于 padding 零值的语义

问题 MEAN_STD (Pi0) QUANTILES (Pi0.5)
归一化后值域 无界 (-inf, +inf) 有界 [-1, 1]
padding 零值含义 对应数据均值,有具体物理含义 对应范围中点 0,是自然中立值
Flow Matching 噪声 N(0,1) 噪声与无界数据混合 N(0,1) 噪声与有界数据混合
padding 维度的梯度 非零梯度(零值有特定含义) 接近零的梯度(零值是中立的)
异常值影响 均值/标准差易被极端值偏移 分位数对异常值天然鲁棒

Flow Matching 中 padding 的关键影响——为什么零填充会从根本上干扰 Flow Matching?

要理解这个问题,需要先理解 Flow Matching 的工作原理。Flow Matching 学习一个"速度场",指导数据从噪声分布流向真实数据分布:

x_t = t * noise + (1-t) * action     ← 时间 t 处的插值样本
u_t = noise - action                  ← 模型需要学习预测的"速度"(真实动作到噪声的方向)

现在考虑 padding 维度。假设一个机器人只有 7 个关节(7 维动作),但 max_action_dim=32,那么第 8~32 维全部填零。这些零值本身没有物理意义——它们只是占位符

问题的根源:在 MEAN_STD 归一化下,padding 的零值变成了一个"有意义的目标"

举一个具体例子来说明:

假设某个 padding 维度的训练集统计量为: mean=0.5, std=0.2

MEAN_STD 归一化:
  - 原始 padding 值 = 0
  - 归一化后 = (0 - 0.5) / 0.2 = -2.5  ← 零值变成了 -2.5!

  Flow Matching 的速度目标:
    u_t = noise - (-2.5) = noise + 2.5
  模型必须学习"把噪声去噪到 -2.5"这个具体值
  这个 -2.5 还落在了分布的尾部(距均值 2.5 个标准差),是一个极端值

QUANTILES 归一化:
  - 原始 padding 值 = 0
  - 归一化后 = 0(因为 [-1,1] 的中点就是 0)  ← 零值保持为零!

  Flow Matching 的速度目标:
    u_t = noise - 0 = noise
  模型只需要学习"把噪声去噪到 0"——而 0 恰好是高斯噪声的均值
  这意味着 padding 维度几乎不产生学习信号,模型可以"忽略"它们

这为什么会造成训练不稳定?

  1. 梯度冲突:在 MEAN_STD 下,模型在 padding 维度上收到的梯度信号和真实动作维度的梯度信号混在一起。padding 维度要求模型去噪到一个极端值(如 -2.5),但这些维度没有实际意义,这些"虚假梯度"会干扰模型对真实动作的学习

  2. 异常值放大:MEAN_STD 的均值和标准差容易被极端值偏移。如果训练数据中某些动作维度存在少数异常值,标准差会被拉大,归一化后的值域变得不可预测。而 QUANTILES(分位数映射)天然对异常值鲁棒——无论离群点多极端,映射后的值始终在 [-1, 1] 内

  3. 噪声-数据尺度匹配:Flow Matching 的噪声来自 N(0,1),值域大致在 [-3, 3]。QUANTILES 归一化后数据在 [-1, 1],两者尺度匹配良好。MEAN_STD 归一化后数据的实际范围不可控,可能出现噪声比数据小得多(模型难以学习去噪方向)或噪声比数据大得多(插值路径不平滑)的情况

一句话总结:QUANTILES 归一化让 padding 的零值变成真正的"无信息"中性值,让 Flow Matching 可以优雅地忽略这些维度;而 MEAN_STD 下的零值被映射到一个有具体含义的非零位置,迫使模型浪费容量去"学习"这些没有意义的 padding 维度。

配置项 Pi0 Pi0.5
STATE 归一化 NormalizationMode.MEAN_STD NormalizationMode.QUANTILES
ACTION 归一化 NormalizationMode.MEAN_STD NormalizationMode.QUANTILES
VISUAL 归一化 NormalizationMode.IDENTITY NormalizationMode.IDENTITY

3.2 AdaRMS 层归一化

Pi0.5 在 Action Expert (Gemma 300M) 中使用 AdaRMS (Adaptive Root Mean Square) 层归一化替代标准 RMSNorm。这使得时间步信息能够直接调制每一层的归一化行为。

前置知识:什么是 RMSNorm?

在理解 AdaRMS 之前,先回顾标准 RMSNorm (Root Mean Square Normalization)。它是 LayerNorm 的简化版本,被 Gemma/LLaMA 等现代 LLM 广泛采用:

标准 LayerNorm:
  y = (x - mean(x)) / sqrt(var(x) + eps) * gamma + beta
  → 需要计算均值和方差,有可学习的 scale (gamma) 和 bias (beta)

RMSNorm (简化版):
  y = x / sqrt(mean(x^2) + eps) * weight
  → 只除以均方根 (RMS),不减均值,没有 bias
  → 计算量更小,效果相当

RMSNorm 的作用是稳定每层特征的数值范围,防止深层网络中特征值越来越大或越来越小。weight 是一个可学习参数,为每个特征维度提供独立的缩放因子。

AdaRMS:让归一化"感知"时间步

AdaRMS (Adaptive RMS) 在 RMSNorm 的基础上增加了条件调制 (conditioning)——根据外部条件信号(这里是 Flow Matching 的时间步 t)动态调整归一化后的特征幅度。具体来说:

# 标准 RMSNorm(Pi0 的 Action Expert 使用):
def rmsnorm(x, weight):
    rms = sqrt(mean(x^2) + eps)
    return (x / rms) * weight          # weight 是固定的可学习参数

# AdaRMS(Pi0.5 的 Action Expert 使用):
def adarms(x, weight, time_embedding):
    rms = sqrt(mean(x^2) + eps)
    x_norm = (x / rms) * weight        # 先做标准 RMSNorm

    # 从时间嵌入中生成 scale 和 gate 两个调制信号
    scale, gate = Dense(time_embedding)  # Dense: Linear(1024 → 2048), 然后拆成两半

    # 用 scale 调制归一化后的特征幅度
    x_modulated = x_norm * (1 + scale)   # scale ≈ 0 时退化为标准 RMSNorm
                                          # scale > 0 时放大特征
                                          # scale < 0 时缩小特征

    return x_modulated
    # gate 用于后续的残差连接门控(控制"这一层的输出有多大比例通过")

AdaRMS 的直观理解:

想象 Transformer 的 18 层就像一条流水线,每层对特征做归一化 → 注意力 → 归一化 → MLP 四步处理。在标准 RMSNorm 下,归一化的 weight 参数是固定的,不管当前处于 Flow Matching 的哪个时间步,归一化行为完全相同。

但在去噪过程中,不同时间步需要完全不同的处理策略: - t=1.0(纯噪声):输入几乎是随机噪声,模型需要大幅度调整,识别出"这大概是个什么动作" - t=0.5(半噪声半信号):模型需要在已有的粗略轮廓上精修 - t=0.1(接近真实动作):模型只需做微小的精细修正

AdaRMS 让每一层的归一化行为能够根据时间步动态变化。通过 scale 参数,模型可以在 t=1.0 时放大特征(大幅修正),在 t=0.1 时保持特征不变(微调)。这种"自适应"是通过 time_embedding → Dense → scale 这条路径学习出来的,而非人工设计。

graph TB subgraph Pi0LN["Pi0: 标准 RMSNorm (Action Expert)"] direction LR H1["hidden_states"] --> RMS1["RMSNorm(x)
x / sqrt(mean(x^2) + eps)
* weight"] RMS1 --> OUT1["归一化输出"] TIME1["时间信息"] -.->|"不参与归一化
仅通过 action 嵌入传入"| H1 end subgraph Pi05LN["Pi0.5: AdaRMS (Action Expert)"] direction LR H2["hidden_states"] --> RMS2["RMSNorm(x)
x / sqrt(mean(x^2) + eps)"] COND["adarms_cond
(time_emb)
[B, 1024]"] --> DENSE["Dense 层
→ (scale, gate)"] DENSE -->|"scale"| MODULATE["x * (1 + scale)"] DENSE -->|"gate"| GATE["门控残差连接"] RMS2 --> MODULATE MODULATE --> OUT2["归一化输出"] end style Pi0LN fill:#ffcdd2,stroke:#E91E63 style Pi05LN fill:#c8e6c9,stroke:#4CAF50

AdaRMS 的设计意图:

在 Flow Matching 中,模型需要在不同时间步 t 上表现出不同的行为: - t 接近 1(纯噪声):模型需要做大幅度的去噪,预测粗略方向 - t 接近 0(接近真实动作):模型需要做精细的修正

AdaRMS 允许时间步信息在每一层的归一化环节直接调制特征的幅度和门控,使 Action Expert 能够根据当前去噪阶段动态调整其内部表示。这比 Pi0 中仅通过初始嵌入传递时间信息更加直接和有效。

代码层面的差异:

# Pi0: use_adarms=[False, False]  — 两个流都不使用 AdaRMS
self.paligemma_with_expert = PaliGemmaWithExpertModel(
    ..., use_adarms=[False, False], ...)

# Pi0.5: use_adarms=[False, True]  — 仅 Action Expert 使用 AdaRMS
self.paligemma_with_expert = PaliGemmaWithExpertModel(
    ..., use_adarms=[False, True], ...)

3.3 简化的 Time MLP 架构

Pi0.5 对时间步嵌入的处理流程进行了简化,移除了状态投影层 (state_proj) 和时间-动作拼接操作。

前置概念说明

graph TB subgraph Pi0Time["Pi0: Time + Action + State 融合"] direction TB T_P0["时间步 t"] --> SIN_P0["正弦编码
[B, 1024]"] SIN_P0 --> EXP_P0["expand → [B, 50, 1024]"] A_P0["含噪动作"] --> AIN_P0["action_in_proj
Linear(32→1024)
[B, 50, 1024]"] EXP_P0 --> CAT_P0["cat([action_emb, time_emb])
[B, 50, 2048]"] AIN_P0 --> CAT_P0 CAT_P0 --> MLP_P0["action_time_mlp_in
Linear(2048→1024) → SiLU"] MLP_P0 --> MLP2_P0["action_time_mlp_out
Linear(1024→1024)"] MLP2_P0 --> EMBS_P0["action_time_emb
[B, 50, 1024]"] S_P0["机器人状态"] --> SP_P0["state_proj
Linear(32→1024)
[B, 1, 1024]"] SP_P0 --> SCAT["[state_emb | action_time_emb]
[B, 51, 1024]"] EMBS_P0 --> SCAT NOTE_P0["adarms_cond = None"] end subgraph Pi05Time["Pi0.5: 简化的 Time MLP + AdaRMS"] direction TB T_P5["时间步 t"] --> SIN_P5["正弦编码
[B, 1024]"] SIN_P5 --> MLP_P5["time_mlp_in
Linear(1024→1024) → SiLU"] MLP_P5 --> MLP2_P5["time_mlp_out
Linear(1024→1024) → SiLU"] MLP2_P5 --> COND_P5["adarms_cond
[B, 1024]"] A_P5["含噪动作"] --> AIN_P5["action_in_proj
Linear(32→1024)"] AIN_P5 --> EMBS_P5["action_emb (直接使用)
[B, 50, 1024]"] NOTE_P5["无 state_proj
无 cat 操作
时间信息通过 AdaRMS 注入"] end style Pi0Time fill:#ffcdd2,stroke:#E91E63 style Pi05Time fill:#c8e6c9,stroke:#4CAF50

关键区别总结:

特性 Pi0 Pi0.5
状态投影层 state_proj: Linear(32 → 1024) (已移除)
状态 token 作为独立 token 加入 suffix 序列 (Pi0.5 不处理显式状态)
Time MLP 输入维度 2048 (action_emb + time_emb 拼接) 1024 (仅 time_emb)
Time MLP 输出 融合后的 action_time_emb adarms_cond (用于 AdaRMS 条件)
时间信息注入方式 与动作嵌入拼接后通过 MLP 通过 AdaRMS 在每层归一化中注入
suffix 序列长度 chunk_size + 1 (state + actions) chunk_size (仅 actions)

设计动机:

Pi0 中时间信息通过拼接方式"一次性"注入动作嵌入,之后在 18 层 Transformer 中逐渐稀释。Pi0.5 的 AdaRMS 设计让时间信息在每一层都直接参与特征调制,实现了更持久、更直接的时间条件注入。同时移除 state_proj 简化了模型结构——Pi0.5 依赖语言指令和视觉观测来隐式传递状态信息。

类比理解:Pi0 的方式就像在信的开头写上"这是紧急信件",然后希望读者(18 层 Transformer)一直记得这个信息。Pi0.5 的方式则是在每一页的页眉上都印着"紧急"——每一层处理特征时都能直接看到当前的时间条件,不会遗忘。


4. 训练流水线

graph TB subgraph DataPrep["数据准备"] RAW["LeRobot 数据集
(图像, 语言, 动作)"] NORM["QUANTILES 归一化
动作 → [-1, 1]"] PAD["零填充
action_dim → max_action_dim=32"] TOK["语言 Tokenize
max_length=200"] RAW --> NORM --> PAD RAW --> TOK end subgraph ImgPrep["图像预处理"] IMG_RAW["原始图像
[B, C, H, W]"] RESIZE["resize_with_pad
→ [B, C, 224, 224]"] SCALE["[0,1] → [-1,1]"] IMG_RAW --> RESIZE --> SCALE end subgraph FlowMatch["Flow Matching 采样"] ACT_NORM["归一化后动作
[B, 50, 32]"] NOISE["noise ~ N(0,1)
[B, 50, 32]"] TIME["t ~ Beta(1.5, 1.0)
* 0.999 + 0.001"] XT["x_t = t*noise + (1-t)*action"] UT["u_t = noise - action"] ACT_NORM --> XT NOISE --> XT TIME --> XT ACT_NORM --> UT NOISE --> UT end subgraph Forward["模型前向传播"] EMB_PRE["embed_prefix
图像 → SigLIP → tokens
语言 → 嵌入 → tokens"] EMB_SUF["embed_suffix
x_t → action_in_proj
t → Time MLP → adarms_cond"] DUAL["双流 Transformer
18层联合注意力"] PROJ_OUT["action_out_proj
→ 预测速度 v_t"] end subgraph Loss["损失计算"] MSE["MSE 损失
L = mean(||u_t - v_t||^2)
仅计算真实维度,忽略 padding"] BACK["反向传播
AdamW + Cosine Decay"] end PAD --> FlowMatch SCALE --> Forward TOK --> Forward XT --> EMB_SUF TIME --> EMB_SUF EMB_PRE --> DUAL EMB_SUF --> DUAL DUAL --> PROJ_OUT PROJ_OUT --> MSE UT --> MSE MSE --> BACK style DataPrep fill:#e8f4fd,stroke:#2196F3 style ImgPrep fill:#fff3e0,stroke:#FF9800 style FlowMatch fill:#fff9c4,stroke:#FFC107 style Forward fill:#e8f5e9,stroke:#4CAF50 style Loss fill:#fce4ec,stroke:#E91E63

5. 推理流水线

推理时,Pi0.5 使用 KV Cache 优化:prefix(图像+语言)只需编码一次并缓存 KV,后续 10 步欧拉去噪只需运行 Action Expert 的 suffix 部分。

graph TB subgraph Encode["Step 0: Prefix 编码 (只执行一次)"] IMG_IN["相机图像"] --> SIG_INF["SigLIP 编码"] TXT_IN["语言指令"] --> TOK_INF["Token 嵌入"] SIG_INF --> PRE_FWD["PaliGemma 前向
(仅 prefix stream)
use_cache=True"] TOK_INF --> PRE_FWD PRE_FWD --> KV["past_key_values
(KV Cache 缓存)"] end subgraph Denoise["迭代去噪 (10步欧拉积分)"] X1["x_1 ~ N(0,1)
[B, 50, 32]"] S1["Step 1: t=1.0"] S2["Step 2: t=0.9"] S3["Step 3: t=0.8"] DOTS["..."] S10["Step 10: t=0.1"] X1 --> S1 --> S2 --> S3 --> DOTS --> S10 end subgraph SingleStep["单步去噪详情"] XT_IN["x_t"] --> EMB_S["embed_suffix
action_in_proj + Time MLP"] EMB_S --> ATT_S["Action Expert 前向
attention_mask 包含 prefix
past_key_values = KV Cache"] KV -->|"复用 prefix KV"| ATT_S ATT_S --> PROJ_S["action_out_proj → v_t"] PROJ_S --> EULER_S["x_{t-dt} = x_t + dt * v_t
dt = -1/10 = -0.1"] end subgraph PostProcess["后处理"] X0["x_0 ≈ 预测动作
[B, 50, 32]"] UNPAD["去除 padding
[B, 50, 32] → [B, 50, action_dim]"] DENORM["QUANTILES 反归一化
[-1,1] → 原始动作空间"] X0 --> UNPAD --> DENORM --> FINAL["最终动作轨迹"] end S10 --> X0 style Encode fill:#e3f2fd,stroke:#2196F3 style Denoise fill:#e8f5e9,stroke:#4CAF50 style SingleStep fill:#fff9c4,stroke:#FFC107 style PostProcess fill:#fce4ec,stroke:#E91E63

推理优化要点: - KV Cache 复用:prefix 的 KV 在第一步计算后缓存,后续 9 步去噪直接复用,大幅减少计算量 - Prefix-only 编码:第一步只运行 inputs_embeds=[prefix_embs, None],不启动 Action Expert - Suffix-only 去噪:后续每步只运行 inputs_embeds=[None, suffix_embs],使用缓存的 prefix KV - Euler 积分方向:从 t=1.0 到 t=0.0,步长 dt = -1/num_steps


6. 关键超参数表

参数 说明
模型结构
paligemma_variant gemma_2b PaliGemma 使用 Gemma 2B
action_expert_variant gemma_300m Action Expert 使用 Gemma 300M
Transformer 层数 18 两个流共享层数
PaliGemma hidden_size 2048 Gemma 2B 隐藏维度
Expert hidden_size 1024 Gemma 300M 隐藏维度
head_dim 256 注意力头维度
num_heads 8 注意力头数
num_kv_heads 1 GQA 键值头数
vocab_size 257152 词汇表大小
动作空间
chunk_size 50 预测动作步数 (action horizon)
n_action_steps 50 执行动作步数
max_action_dim 32 最大动作维度 (零填充)
max_state_dim 32 最大状态维度
Flow Matching
num_inference_steps 10 推理去噪步数
time_sampling_beta_alpha 1.5 Beta 分布 alpha
time_sampling_beta_beta 1.0 Beta 分布 beta
time_sampling_scale 0.999 时间缩放
time_sampling_offset 0.001 时间偏移
min_period 4e-3 正弦编码最小周期
max_period 4.0 正弦编码最大周期
归一化
VISUAL IDENTITY 图像不归一化
STATE QUANTILES 分位数归一化到 [-1,1]
ACTION QUANTILES 分位数归一化到 [-1,1]
训练
optimizer_lr 2.5e-5 峰值学习率
optimizer_betas (0.9, 0.95) AdamW beta
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 最终学习率
image_resolution (224, 224) 输入图像分辨率
tokenizer_max_length 200 语言 token 最大长度
dtype float32 默认精度

7. 关键源文件表

组件 类名 文件路径
主策略类 PI05Policy lerobot/policies/pi05/modeling_pi05.py
核心模型 PI05Pytorch lerobot/policies/pi05/modeling_pi05.py
双流 Transformer PaliGemmaWithExpertModel lerobot/policies/pi05/modeling_pi05.py
逐层联合注意力 compute_layer_complete() lerobot/policies/pi05/modeling_pi05.py
模型配置 PI05Config lerobot/policies/pi05/configuration_pi05.py
Pi0 主策略类 (对比) PI0Policy lerobot/policies/pi0/modeling_pi0.py
Pi0 核心模型 (对比) PI0Pytorch lerobot/policies/pi0/modeling_pi0.py
Pi0 配置 (对比) PI0Config lerobot/policies/pi0/configuration_pi0.py
Gemma 变体配置 GemmaConfig / get_gemma_config() lerobot/policies/pi05/modeling_pi05.py
注意力掩码工具 make_att_2d_masks() lerobot/policies/pi05/modeling_pi05.py
正弦位置编码 create_sinusoidal_pos_embedding() lerobot/policies/pi05/modeling_pi05.py
图像预处理 resize_with_pad_torch() lerobot/policies/pi05/modeling_pi05.py
RTC 处理器 RTCProcessor lerobot/policies/rtc/modeling_rtc.py
RTC 配置 RTCConfig lerobot/policies/rtc/configuration_rtc.py