pi0.5

1
  π₀.₅ (Pi0.5) 模型架构
2

3
  基于代码库的分析，π₀.₅ 是一个基于流匹配的视觉-语言-动作模型，是 π₀ 的升级版本，具有更好的开放世界泛化能力。
4

5
  核心架构图
6

7
  ┌─────────────────────────────────────────────────────────────────────────┐
8
  │                        π₀.₅ 模型架构 (Pi0.5)                            │
9
  └─────────────────────────────────────────────────────────────────────────┘
10

11
  输入层
12
  ├─────────────────────────────────────────────────────────────────────┤
13
  │ 前缀 - 上下文编码                                                     │
14
  │ ├─ 多视角图像                                                         │
15
  │ │  ├─ base_0_rgb (224×224×3) - 基础相机视角                        │
144 collapsed lines
16
  │ │  ├─ left_wrist_0_rgb (224×224×3) - 左腕相机视角                  │
17
  │ │  └─ right_wrist_0_rgb (224×224×3) - 右腕相机视角                 │
18
  │ ├─ 语言指令                                                          │
19
  │ │  └─ tokenized_prompt (最大 200 tokens) - 离散语言标记             │
20
  │ └─ 状态表示                                                          │
21
  │    └─ state (32-dim) - **作为离散语言标记处理** (π₀.₅ 特点)         │
22
  │                                                                       │
23
  │ 后缀 - 动作生成                                                       │
24
  │ ├─ 带噪声动作 (noisy_actions) - 流匹配输入                          │
25
  │ ├─ 时间步嵌入                                                         │
26
  │ │  └─ timestep ∈ [0.001, 0.999] - 正弦-余弦位置编码                 │
27
  │ └─ **adaRMSNorm 条件** (π₀.₅ 特点)                                  │
28
  └─────────────────────────────────────────────────────────────────────┘
29
                      ↓
30
  视觉编码器
31
  ├─────────────────────────────────────────────────────────────────────┤
32
  │ SigLIP Vision Tower (So400m/14)                                      │
33
  │ ├─ 2D sincos 位置编码                                                │
34
  │ ├─ Vision Transformer 处理                                          │
35
  │ └─ 输出维度: PaliGemma width (取决于 gemma_2b 变体)                 │
36
  └─────────────────────────────────────────────────────────────────────┘
37
                      ↓
38
  双专家语言模型
39
  ├─────────────────────────────────────────────────────────────────────┤
40
  │ PaliGemma Dual Experts (基于 Gemma 架构)                           │
41
  │                                                                       │
42
  │ 专家 1: PaliGemma 专家 (gemma_2b)                                   │
43
  │ ├─ 主要语言理解和视觉-语言融合                                       │
44
  │ ├─ 处理图像和语言输入                                                │
45
  │ └─ **不使用 adaRMSNorm**                                            │
46
  │                                                                       │
47
  │ 专家 2: 动作专家 (gemma_300m)                                       │
48
  │ ├─ 专门用于动作预测                                                  │
49
  │ ├─ 使用 **adaRMSNorm** 注入时间步条件 (π₀.₅ 特点)                   │
50
  │ └─ 接收动作+时间混合表示                                            │
51
  └─────────────────────────────────────────────────────────────────────┘
52
                      ↓
53
  动作输出投影
54
  ├─────────────────────────────────────────────────────────────────────┤
55
  │ Action Projection Layer                                             │
56
  │ ├─ Linear: action_expert.width → action_dim (32)                    │
57
  │ └─ 输出: 清洁动作 v_t (去噪后的速度场)                              │
58
  └─────────────────────────────────────────────────────────────────────┘
59

60
  π₀.₅ 与 π₀ 的关键区别
61

62
  | 特性          | π₀                            | π₀.₅                     |
63
  |-------------|-------------------------------|--------------------------|
64
  | 状态输入方式      | 连续状态 token (单独的 suffix token) | 离散语言标记 (作为 prefix 的一部分)  |
65
  | 时间条件注入      | MLP 混合 (action_time_mlp)      | adaRMSNorm (自适应 RMS 归一化) |
66
  | 最大 token 长度 | 48                            | 200                      |
67
  | 动作专家配置      | 不使用 adaRMS                    | 使用 adaRMS 注入流匹配时间步       |
68
  | 开放世界泛化      | 基线性能                          | 更好 (通过知识隔离训练)            |
69

70
  数据流详解
71

72
  1. 前缀编码 - embed_prefix()
73

74
  # src/openpi/models/pi0.py:106-137
75
  输入: Observation
76
  ├─ Images: 3 个相机视角 (224×224×3)
77
  ├─ Language: tokenized_prompt (最多 200 tokens)
78
  └─ State: (32-dim) - 在 π₀.₅ 中作为语言标记
79

80
  处理流程:
81
  1. SigLIP 编码每个图像 → image_tokens
82
  2. PaliGemma LLM 嵌入语言标记 → tokenized_inputs
83
  3. 图像和语言标记完全注意 (full attention between image and language)
84
  4. 拼接所有 tokens → prefix_tokens
85

86
  2. 后缀编码 - embed_suffix()
87

88
  # src/openpi/models/pi0.py:140-186
89
  输入: noisy_actions, timestep
90

91
  处理流程:
92
  1. Action 投影: Linear(action_dim, width) → action_tokens
93
  2. 时间嵌入: posemb_sincos(timestep) → time_emb
94
  3. **时间 MLP** (π₀.₅ 特点):
95
     - time_mlp_in: Linear(width, width)
96
     - Swish 激活
97
     - time_mlp_out: Linear(width, width)
98
     - Swish 激活
99
  4. **adaRMS 条件**: adarms_cond = time_emb (用于动作专家)
100
  5. 动作标记因果自回归 (causal attention)
101

102
  3. 流匹配训练
103

104
  # src/openpi/models/pi0.py:189-214
105
  损失计算:
106
  1. 采样时间步: time ~ Beta(1.5, 1) * 0.999 + 0.001
107
  2. 构造噪声状态: x_t = time * noise + (1-time) * actions
108
  3. 计算目标速度场: u_t = noise - actions
109
  4. 前向传播: v_t = model(x_t, time, observation)
110
  5. MSE 损失: mean((v_t - u_t)²)
111

112
  4. 推理采样
113

114
  # src/openpi/models/pi0.py:217-279
115
  Flow Matching 采样 (ODE Solver):
116
  1. 初始化: x_T ~ N(0, I), time = 1.0
117
  2. 迭代去噪 (num_steps 步):
118
     - 计算速度场: v_t = model(x_t, time)
119
     - 更新状态: x_{t-1} = x_t - dt * v_t (dt = 1/num_steps)
120
     - 更新时间: time ← time - dt
121
  3. 输出: x_0 (预测的动作序列)
122

123
  配置示例
124

125
  基础配置 (pi0_config.py:18-53):
126
  Pi0Config(
127
      dtype="bfloat16",
128
      paligemma_variant="gemma_2b",        # PaliGemma 专家
129
      action_expert_variant="gemma_300m",  # 动作专家
130
      action_dim=32,
131
      action_horizon=50,
132
      max_token_len=200,  # π₀.₅: 200, π₀: 48
133
      pi05=True,  # 启用 π₀.₅ 模式
134
      discrete_state_input=True,  # 状态作为离散标记
135
  )
136

137
  预训练检查点
138

139
  | 模型          | 用途        | 检查点路径                                      |
140
  |-------------|-----------|--------------------------------------------|
141
  | π₀.₅-Base   | 微调基础模型    | gs://openpi-assets/checkpoints/pi05_base   |
142
  | π₀.₅-LIBERO | LIBERO 基准 | gs://openpi-assets/checkpoints/pi05_libero |
143
  | π₀.₅-DROID  | DROID 数据集 | gs://openpi-assets/checkpoints/pi05_droid  |
144

145
  核心创新点
146

147
  1. 知识隔离: 通过特殊训练策略提升开放世界泛化
148
  2. adaRMSNorm: 自适应归一化，将流匹配时间步注入动作专家层
149
  3. 离散状态表示: 将状态作为语言标记处理，而非连续后缀 token
150
  4. 双专家架构: 分离语言理解和动作预测，提升专门化性能
151

152
  文件位置
153

154
  - 配置: src/openpi/models/pi0_config.py:18-110
155
  - 模型实现: src/openpi/models/pi0.py:66-280
156
  - PyTorch 版本: src/openpi/models_pytorch/pi0_pytorch.py
157
  - 训练配置: src/openpi/training/config.py (搜索 pi05_)
158

159
  这个架构设计使得 π₀.₅ 在保持 π₀ 的流匹配框架优势的同时，通过更精细的状态表示和条件注入机制实现了更好的泛化性能。

1
  | 维度      | π₀               | π₀.₅             |
2
  |---------|------------------|------------------|
3
  | 状态表示能力  | 连续向量，1个token     | 离散tokens，多个token |
4
  | 状态-视觉交互 | ❌ 无交互（单向注意）      | ✅ 完全交互           |
5
  | 时间条件注入  | MLP混合（token级）    | adaRMS（层级条件）     |
6
  | 开放世界泛化  | 基线               | 更好               |
7
  | 训练稳定性   | 需要careful tuning | 更稳定（知识隔离）        |
8
  | 推理速度    | 相同               | 相同               |
9
  | 显存占用    | 略少               | 略多（更长序列）         |

1
  1. 预训练 vs 后训练 - 会冻结什么模块？
2

3
  根据代码分析，取决于训练模式（全量微调 vs LoRA）：
4

5
  预训练
6

7
  - 全参数训练：所有模块都训练
8
    - SigLIP Vision Encoder
9
    - PaliGemma 专家
10
    - Action Expert (gemma_300m)
11
    - 所有投影层
12

13
  后训练/微调
14

15
  有两种策略：
109 collapsed lines
16

17
  a) 全量微调 (默认)
18

19
  # 示例: pi0_libero, pi05_libero
20
  model=Pi0Config(pi05=True)  # 无 lora 后缀
21
  weight_loader=CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params")
22
  # freeze_filter=None → 全部训练
23
  - 不冻结任何模块
24
  - 从 base checkpoint 加载全部权重
25
  - 所有参数继续更新
26

27
  b) LoRA 微调 (低内存)
28

29
  # 示例: pi0_libero_low_mem_finetune
30
  model=Pi0Config(
31
      paligemma_variant="gemma_2b_lora",      # ← LoRA
32
      action_expert_variant="gemma_300m_lora"  # ← LoRA
33
  )
34
  freeze_filter=Pi0Config(...).get_freeze_filter()
35

36
  冻结逻辑 (pi0_config.py:80-109):
37
  def get_freeze_filter(self):
38
      # 当使用 lora 变体时，冻结主权重，只训练 LoRA 参数
39

40
      if "lora" in paligemma_variant:
41
          冻结: 所有 llm 参数 (.*llm.*)
42
          排除: lora 参数本身 (.*lora.*)
43

44
      if "lora" in action_expert_variant:
45
          冻结: 动作专家参数 (.*llm.*_1.*)
46
          排除: lora 参数 (.*lora.*)
47

48
  冻结模块总结:
49

50
  | 模块                   | 预训练  | 全量微调 | LoRA 微调           |
51
  |----------------------|------|------|-------------------|
52
  | SigLIP Vision        | ✅ 训练 | ✅ 训练 | ❌ 冻结              |
53
  | PaliGemma (gemma_2b) | ✅ 训练 | ✅ 训练 | ❌ 冻结主权重✅ 只训练 LoRA |
54
  | Action Expert        | ✅ 训练 | ✅ 训练 | ❌ 冻结主权重✅ 只训练 LoRA |
55
  | 投影层                  | ✅ 训练 | ✅ 训练 | ❌ 冻结              |
56

57
  ---
58
  2. PaliGemma 的来源 - 天生 VLM 还是纯 LLM？
59

60
  PaliGemma 本来就是 VLM（Vision-Language Model），并非 π₀ 发现的"意外可用"。
61

62
  证据 1: PaliGemma 架构设计
63

64
  # pi0.py:73-91
65
  # PaliGemma = SigLIP + Gemma (已经是 VLM 结构)
66
  llm = nnx_bridge.ToNNX(_gemma.Module(
67
      configs=[paligemma_config, action_expert_config],  # ← 双专家
68
      ...
69
  ))
70

71
  img = nnx_bridge.ToNNX(_siglip.Module(
72
      num_classes=paligemma_config.width,  # ← 输出给 LLM 的维度
73
      ...
74
  ))
75
  self.PaliGemma = nnx.Dict(llm=llm, img=img)
76

77
  PaliGamma 的原始设计:
78
  - SigLIP (Vision Encoder) → 图像嵌入
79
  - Gemma (Language Model) → 语言处理
80
  - 直接拼接 vision tokens + language tokens
81
  - 跨模态注意力融合
82

83
  证据 2: 官方描述
84

85
  从 README:
86
  π₀ 是基于 PaliGemma 的 flow-based VLA
87

88
  PaliGemma 是 Google 发布的开源 VLM 模型，专门用于视觉-语言任务。
89

90
  π₀ 的创新点在于:
91

92
  1. 双专家架构:
93
    - PaliGemma 专家 (原 VLM 能力)
94
    - Action Expert (新增，专门预测动作)
95
  2. 流匹配头:
96
    - 将 VLM 改造成 VLA
97
    - 用 flow matching 替代传统的语言建模目标
98
  3. 并非"发现 vision token 可用":
99
    - PaliGemma 本来就处理 vision + language
100
    - π₀ 只是复用了 PaliGemma 的视觉编码和语言理解
101
    - 然后加了第二个专家专门输出动作
102

103
  架构对比:
104

105
  原始 PaliGemma (VLM):
106
  ┌─────────────┐    ┌──────────────┐
107
  │ SigLIP Img  │───→│ Gemma LLM    │───→ 文本输出
108
  └─────────────┘    └──────────────┘
109
         ↑
110
      Image Input
111
         +
112
      Language Tokens
113

114
  π₀/π₀.₅ (VLA):
115
  ┌─────────────┐    ┌──────────────────┬─────────────┐
116
  │ SigLIP Img  │───→│ PaliGemma Expert │ Action Expert│
117
  └─────────────┘    └──────────────────┴─────────────┘
118
                            ↓                    ↓
119
                       (理解)               (动作输出)
120

121
  总结:
122
  - PaliGemma = 天生 VLM，不是 π₀ 发现的
123
  - π₀ = VLM → VLA 的改造
124
  - 核心创新 = 双专家 + 流匹配，而非"视觉 tokens 可用"