This Is Where We Start🗽pi0 : Our First Generalist Policy (Part. 1)
type
status
date
slug
summary
tags
category
icon
password
本质上 pi0 提出的是一个架构，是组合现有的模型实现 robotics 位姿的输出。需要做的，就是理解这个架构，看 VLA 这个模型是怎么构建起来的。
Part 1: Paper
Introduction
这张图描述了整个 policy 的基本流程。核心的 VLA 模型用一个 pre-trained VLM 模型（我的理解是提供强大的推理能力）和一个 action expert 拼接而成。用 cross-embodiment 的数据做训练，用 high-quality 的数据做 fine-tuned。
以 VLM 为基础的原因
By basing our model on a VLM, we inherit the general knowledge, semantic reasoning, and problem-solving abilities of language and vision-language models.
论文提出了现在 generalist robot policies 的三个 bottlenecks ，分别是
- very large scale
- a right model architectures make use of diverse data source, represent intricate and subtle behaviors
- a right training recipe
论文分别给出了如下解决方案
- Problem 1：以已经预训练好的 VLM 为 base
- Problem 2：
- diverse data: VLM, cross-embodiment training
- intricate and subtle behaviors: an action chunking architecture with flow matching
- Problem 3：pre-training/post-training separation
pre-training/post-training separation
Intuitively, training only on high-quality data does not teach the model how to recover from mistakes, since mistakes are rarely seen in such data. Training on only lower-quality pretraining data does not teach the model to act efficiently and robustly. Combining both provides the desired behavior: the model attempts insofar as possible to act in a manner similar to the high-quality data, but still has a repertoire of recoveries and corrections that it can deploy in the case of a mistake.
Related Work
之前的模型会采用 autoregressive discretization 来预测动作，这里使用的是一种 diffusion model 的变体，称为 flow matching（主要卖点）。flow matching 可以提供高精度和多模态的能力。
diffusion model for action generation
论文里列举了两篇
训练模式与之前不同
… we train our model via a diffusion-style (flow matching) loss applied on individual sequence elements, in lieu of the standard cross-entropy loss for decoder-only transformers.
… we use a separate set of weights for the tokens corresponding to diffusion.
Overview
Overview of pi0’s framework
We start with a pre-training mixture, which consists of both our own dexterous manipulation datasets and open-source data. We use this mixture to train our flow matching VLA model, which consists of a larger VLM backbone and a smaller action expert for processing robot states and actions. The VLM backbone weights are initialized from PaliGemma, providing representations learned from large-scale Internet pre-training. The resulting π 0 model can be used to control multiple robot embodiments with differing action spaces to accomplish a wide variety of tasks.
训练总体分两步：pre-training 希望得到一个有多种能力、良好泛化性的模型，post-training 希望让动作更精细
The model
整个模型的骨架是 VLM。VLM 的底层是 language model transformer。在机器人上 image encoder 做的是将机器人观察到的图像 embed 到 language tokens 的语义空间中。
late fusion VLM recipe
VLM 是需要搞懂的，论文中列举了 3 篇
在这个 backbone 上扩展，对于机器人我们需要加入特定的输入 proprioceptive state （本体状态）和特定的输出 robot action。
使用 conditional flow matching 来为连续的动作分布函数建模，基本框架遵循 Transfusion，用 flow matching loss 来监督 “tokens corresponding to continuous outputs”，用 cross-entropy loss 来监督 “tokens corresponding to discrete outputs”。
Transfusion
trains a single transformer using multiple objective
对于机器人的 state 和 action 单独使用一组不同的参数（应该就是指把 action 的预测分离出来）能获得更好的性能，类似 a mixture of experts（使用两套 mixture elements，第一套针对 image 和 text 的输入，第二套针对 robotics-specific inputs and outputs）。论文中将这一套参数称为 action expert。
Mathematical representation
训练部分
问题的本质是想建模出一个 action 的概率密度函数 ，对这个函数做随机抽样来预测下一个 action。
其中，action chunk 表示未来的一系列动作；
observation 。 是 RGB 图像输入， 是 language tokens 序列， 是 a vector of joint angle。
用各自的 encoder + a linear projection layer 将 image 和 state 映射到和 language tokens 一样的 embedding space 中。
接下来的数学推导有点复杂。
对于 中的每个动作 ，有一个对应的输入给 action expert 的 action token。使用 conditional flow matching 的 loss（这玩意儿就离谱）来监督
与其说生成过程是一个 conditional flow matching ，我觉得就是扩散模型。
我们从标准高斯分布中采样一个随机噪声 ，与原始动作线性组合得到加噪后的动作 。根据 flow matching ，我们训练一个网络，该网络可以生成 ，这个 velocity 要和目标向量场 尽可能匹配，其中 （不太理解）。
action expert 采用 full bidirectional attention mask ，让所有的 action token 能互相关注（类似于文本生成中 token 上下文的信息都能注意到）。
flow matching 时间步 是从 分布中采样得到，更注重嘈杂的时间步（在相对嘈杂的部分采样更多），有助于提升模型处理噪声的能力。
推理部分
预测一个动作是从噪声变换到目标动作，从 flow matching 时间步的角度来看是 。起点 ，中间由模型预测的 velocity 给出变换方向，变化量是 velocity 的积分。
从数值计算的角度，使用 forward Euler integration rule ，论文里取 ，即将积分分为 小步去做，每步
Model Architecture Details
在论文的 Appendix B 中，作者给出了对模型各部分更详细的描述，包括 的创新点，也是对论文对总结。
Additional inputs and outputs
在 PaliGemma 原始的 输入基础上，为适配机器人增加了描述本体状态的 。除此之外还有 noisy action chunk 。
- 输入
图像序列，语言提示，本体状态 表征当前状态
noisy action chunk ，未来一段时间（ ， the action horizon ）的动作规划，但此时是有噪声的，有待 flow matching 做去噪。
- 输出
将输入一股脑塞进一个 transformer 中， transformer 的任务是从当前状态的信息中提取对未来动作有用的信息，并会特别注意 noisy action chunk 。
Incorporating the flow matching timestep
将 action chunk 输入 transformer 中，需要将它映射到一个统一的 embedding space 。除此之外， action chunk 有一个自己的 flow matching 时间步 ，用于标记 flow matching 的过程。
映射通过 MLP （Multilayer Perceptron ，多层感知机）实现，可以用如下公式表示
其中， 对 做 positional encoding ，将 转化为一个维度为 （ 是 embedding dimension ）的向量。 是一个线性变换矩阵，将 noisy action 映射到一个维度为 的空间。 concat 操作将两个向量拼接起来，形成维度为 的向量，再用 将该向量映射回 维。 swish 是一个激活函数，能够提供更平滑的非线性变换。最后一个 又映射了一次（意义是？？？）。
Attention mask
使用一种 blockwise casual attention mask 。将输入序列分为 3 块： ， ， 。
- The first block
这个块被限制不能注意到另外两个块，目的是 “minimize distribution shift from said pre-training”，大概是模型在预训练时学习到的输入模态的分布应该尽可能保持不变。
- The second block
机器人本体状态不随着 flow matching integrate step 的进行而改变。阻止它注意第三个块（动作规划）可以在采样过程中缓冲对应的键和值，提高运算效率。
- The third block
这个块可以注意到前面的两个块，目的是基于所有的信息生成动作规划。
在每个块内部使用 full bidirectional attention ，即所有的 token 可以相互注意，但当前的块都不能注意到未来的块。
Action expert
使用一个 transformer 但有两套权重，即两个“专家”，每个 token 会被路由到其中一个“专家”进行处理。
处理 的权重直接从 PaliGemma 初始化。 和 在 VLM 预训练期间没有见过，将被路由到 action expert 。
两个专家仅通过 transformer’s self-attention layers 交互，因此它们的宽度
width 和 MLP 维度 mlp_dim 不需要完全匹配，这样可以降低 action expert 的参数量。Sampling the flow matching timestep
flow matching 原论文对时间步 的采样满足均匀分布 ，即假设每个时间步长对训练的作用都是等同的。
与图像生成任务中对采样分布函数的优化不同，论文认为模型在低时间步长（高噪声水平）时需要更精细地学习如何从噪声中恢复动作分布（在论文中有解释和对比，这里就不赘述了）。所以如前文提及的， 中对 的采样服从 ，实验中阈值 取为 0.999 。
上一篇
Object- Oriented Programming Notebook
下一篇
pi0 : Our First Generalist Policy (Part. 2)
Loading...