pi_0

Abstract

A Vision-Language-Action Flow Model for General Robot Control

1. 背景与动机

三个瓶颈：

1. 现有的具身智能模型依赖大量的对特定 Task 的针对性训练，才有可能在这个 Task 上表现良好。因此，其泛化能力是取决于各 Domain Task 数据量的。
2. 其次，能否泛化依赖于模型架构能否有效利用多种数据（网络文本，视频演示，robot 状态等），从而高效学习。
3. 最后，能否提出一种标准的、类似于 LLM 的训练流程，方便规模化和精细调节。

我们能不能设计一种模型架构，使得模型在很少的数据集上能快速微调？我们把这种模型叫做 “机器人基础模型” / “通用机器人策略”。

这篇文章给出了一种解答: ${\pi }_0$

2. 模型架构与设计

2.1 大致思路

• VLM（PaliGemma 3B） + Action Expert (300M) = VLA: VLM 能利用互联网的海量数据学习表征；Action Expert 能使模型学会操作。
• Flow-matching: 与 Auto-regressive 不同，Flow matching 作为 diffusion model 的一种变体，学习的是连续向量的变化方向。

Note

有点像 MOE，但只有两个 Expert。 二者有独立的 MLP，只在 Self-Attention 层有交互。

2.2 深入代码（pi0.py）

2.2.1 基础模块

先从 big_vision 抄了一个注意力掩码模块和 positional Embedding 模块，这里暂且按下不表。

# 这个有点帅，可以根据不同的 mask_ar 输出不同的掩码
def make_attn_mask(input_mask, mask_ar);
 
def posemb_sincos(
    pos: at.Real[at.Array, " b"], embedding_dim: int, min_period: float, max_period: float
) -> at.Float[at.Array, "b {embedding_dim}"];

2.2.2 主体类

都在一个 Pi0 class 下：

class Pi0(_model.BaseModel):
    def __init__(self, config: pi0_config.Pi0Config, rngs: nnx.Rngs):
    # init 就略了，比较繁杂，源代码还有很多考虑是不是 Pi0.5 的 if-else，就没必要细看了，只需要知道这个初始化了 Pi0Config 中规定的那些参数

2.2.3 Prefix / Suffix 切分

数据分四种：文字 + 图像 + 状态 + 动作，前两者称作 prefix，后两者称作 suffix。

这样分方便：

• 喂数据
• KV Cache（只更新后缀，前缀可以 cached）
• 不同 attention mask 处理

因为 suffix 是模型要学的分布，所以 prefix 不能 attend suffix，只能 suffix 单向看 prefix；但注意，二者内部可以互相看。

训练时：
prefix: 观测 -> suffix: 真实动作
    ↓                ↓
   已知           学习目标
 
推理时：
prefix: 观测 -> suffix: 噪声->动作
    ↓                ↓
    已知     迭代生成

# 先把文字和图像 embed 一下。注意掩码细节
def embed_prefix(
        self, obs: _model.Observation # Observation 是一个三元数组，[视觉，状态，语言], 但 prefix 这里只处理 images and texts，故 state 并非 prefix!
    ) -> tuple[at.Float[at.Array, "b s emb"], at.Bool[at.Array, "b s"], at.Bool[at.Array, " s"]]:
    # 返回 `(embeddings, mask_ar, mask_loss)`, 其中 mask_loss 标记哪些位置需要计算loss（一般是只对 action 算 loss）
        input_mask = []
        ar_mask = []
        tokens = []
        # embed images
        for name in obs.images:
            image_tokens, _ = self.PaliGemma.img(obs.images[name], train=False) ## 调用 PaliGemma 的视觉编码器，提取图像 token （ViT就在这里）
 
            tokens.append(image_tokens)
            # 把每个 batch 元素的 mask 复制 s 份（s = image token 数量），让每张图的所有 token 共享同一个 mask 值。
            input_mask.append(
                einops.repeat(
                    obs.image_masks[name],
                    "b -> b s",
                    s=image_tokens.shape[1],
                )
            )
            # image tokens attend to each other
            ar_mask += [False] * image_tokens.shape[1]
 
        # add language (aka tokenized inputs)
        if obs.tokenized_prompt is not None:
            tokenized_inputs = self.PaliGemma.llm(obs.tokenized_prompt, method="embed")
            tokens.append(tokenized_inputs)
            input_mask.append(obs.tokenized_prompt_mask)
            # full attention between image and language inputs
            ar_mask += [False] * tokenized_inputs.shape[1]
        tokens = jnp.concatenate(tokens, axis=1)
        input_mask = jnp.concatenate(input_mask, axis=1)
        ar_mask = jnp.array(ar_mask)
        return tokens, input_mask, ar_mask

# 这里 suffix 处理 State + Action + Timestep:
def embed_suffix(
        # 输入：Observation(只用state)，加噪后的动作 [b, horizon, action_dim], 扩散时间步[b] \in [0,1]
        self, obs: _model.Observation, noisy_actions: _model.Actions, timestep: at.Float[at.Array, " b"]
        # 输出：tokens, input_mask（全是有效token，故为 all true）, ar_mask (因果边界标记)
    ) -> tuple[
        at.Float[at.Array, "b s emb"],
        at.Bool[at.Array, "b s"],
        at.Bool[at.Array, " s"],
        at.Float[at.Array, "b emb"] | None,
    ]:
        input_mask = []
        ar_mask = []
        tokens = []
 
 
        if not self.pi05:
            # 1. add a single state token （把机器人当前状态变成一个token）
            state_token = self.state_proj(obs.state)[:, None, :]
            tokens.append(state_token)
            input_mask.append(jnp.ones((obs.state.shape[0], 1), dtype=jnp.bool_))
            # image/language inputs do not attend to state or actions
            ar_mask += [True]
 
 
            # 2. 把 noisy_action 投影到 embedding 空间
        action_tokens = self.action_in_proj(noisy_actions)
        # 3. embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1]
        time_emb = posemb_sincos(timestep, self.action_in_proj.out_features, min_period=4e-3, max_period=4.0)
        if self.pi05:
            # time MLP (for adaRMS)
            time_emb = self.time_mlp_in(time_emb)
            time_emb = nnx.swish(time_emb)
            time_emb = self.time_mlp_out(time_emb)
            time_emb = nnx.swish(time_emb)
            action_expert_tokens = action_tokens
            adarms_cond = time_emb
        else:
            # 4. mix timestep + action information using an MLP (no adaRMS)
            time_tokens = einops.repeat(time_emb, "b emb -> b s emb", s=self.action_horizon)
            action_time_tokens = jnp.concatenate([action_tokens, time_tokens], axis=-1)
            action_time_tokens = self.action_time_mlp_in(action_time_tokens)
            action_time_tokens = nnx.swish(action_time_tokens)
            action_time_tokens = self.action_time_mlp_out(action_time_tokens)
            action_expert_tokens = action_time_tokens
            adarms_cond = None
        tokens.append(action_expert_tokens)
        input_mask.append(jnp.ones(action_expert_tokens.shape[:2], dtype=jnp.bool_))
        # image/language/state inputs do not attend to action tokens
        ar_mask += [True] + ([False] * (self.action_horizon - 1))
        tokens = jnp.concatenate(tokens, axis=1)
        input_mask = jnp.concatenate(input_mask, axis=1)
        ar_mask = jnp.array(ar_mask)
        return tokens, input_mask, ar_mask, adarms_cond
 
 
        # 最后返回带噪动作的 token 序列 + 因果掩码，告诉模型"当前的带噪动作长什么样，需要去噪生成真实动作"

Prefix + Suffix 构成了完整 embedding 序列：

完整序列: [图像 | 语言 | state | action_0 | action_1 | action_2 | action_3]
cumsum:   [0    | 0    | 1     | 2        | 2        | 2        | 2       ]
           ↑              ↑
         prefix        suffix
        (cumsum=0)    (cumsum≥1)
注意力规则：只能看到 cumsum ≤ 自己的位置
- prefix 内部互相看 ✓
- suffix 内部互相看 ✓
- prefix → suffix 可以看 ✓
- suffix → prefix 不能看 ✗

2.2.4 训练过程怎么算 Loss？ → Flow Matching

整体流程：

干净 action → 加噪得到 x_t → 模型预测去噪方向 v_t → 和真实去噪方向 u_t 算 MSE

def compute_loss(
        self, rng: at.KeyArrayLike, observation: _model.Observation, actions: _model.Actions, *, train: bool = False
    ) -> at.Float[at.Array, "*b ah"]:
        preprocess_rng, noise_rng, time_rng = jax.random.split(rng, 3)
        observation = _model.preprocess_observation(preprocess_rng, observation, train=train)
 
        batch_shape = actions.shape[:-2]
 
        # 1. 生成噪声：
        noise = jax.random.normal(noise_rng, actions.shape) # 正态分布的噪声
        time = jax.random.beta(time_rng, 1.5, 1, batch_shape) * 0.999 + 0.001 # beta分布采样
        time_expanded = time[..., None, None]
        # 得到噪声 x_t
        x_t = time_expanded * noise + (1 - time_expanded) * actions
        # 真实去噪方向u_t, 注意是直线轨迹，所以直接向量相减
        u_t = noise - actions
 
 
        # one big forward pass of prefix + suffix at once
        # 2. 先 Embed
        prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
        suffix_tokens, suffix_mask, suffix_ar_mask, adarms_cond = self.embed_suffix(observation, x_t, time)
        input_mask = jnp.concatenate([prefix_mask, suffix_mask], axis=1)
        ar_mask = jnp.concatenate([prefix_ar_mask, suffix_ar_mask], axis=0)
        attn_mask = make_attn_mask(input_mask, ar_mask)
        positions = jnp.cumsum(input_mask, axis=1) - 1
        (prefix_out, suffix_out), _ = self.PaliGemma.llm(
            [prefix_tokens, suffix_tokens], mask=attn_mask, positions=positions, adarms_cond=[None, adarms_cond]
        )
        # 3. 模型预测 v_t
        v_t = self.action_out_proj(suffix_out[:, -self.action_horizon :])
 
        # 4. Loss = MSE(v_t, u_t)
        return jnp.mean(jnp.square(v_t - u_t), axis=-1)

2.2.5 推理过程（又名“动作采样”）

欧拉积分法（其实就是走一小步）

def sample_actions(
       self,
       rng: at.KeyArrayLike,
       observation: _model.Observation,
       *,
       num_steps: int | at.Int[at.Array, ""] = 10,
       noise: at.Float[at.Array, "b ah ad"] | None = None,
   ) -> _model.Actions:
       observation = _model.preprocess_observation(None, observation, train=False)
       # note that we use the convention more common in diffusion literature, where t=1 is noise and t=0 is the target
       # distribution. yes, this is the opposite of the pi0 paper, and I'm sorry. 给我整笑了，还挺真诚
       dt = -1.0 / num_steps # 定义一小步的步长
       batch_size = observation.state.shape[0]
       if noise is None:
           noise = jax.random.normal(rng, (batch_size, self.action_horizon, self.action_dim))
 
       # first fill KV cache with a forward pass of the prefix
       prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
       prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask)
       positions = jnp.cumsum(prefix_mask, axis=1) - 1
       _, kv_cache = self.PaliGemma.llm([prefix_tokens, None], mask=prefix_attn_mask, positions=positions)
 
       def step(carry):
           x_t, time = carry
           suffix_tokens, suffix_mask, suffix_ar_mask, adarms_cond = self.embed_suffix(
               observation, x_t, jnp.broadcast_to(time, batch_size)
           )
           # `suffix_attn_mask` is shape (b, suffix_len, suffix_len) indicating how the suffix tokens can attend to each
           # other
           suffix_attn_mask = make_attn_mask(suffix_mask, suffix_ar_mask)
           # `prefix_attn_mask` is shape (b, suffix_len, prefix_len) indicating how the suffix tokens can attend to the
           # prefix tokens
           prefix_attn_mask = einops.repeat(prefix_mask, "b p -> b s p", s=suffix_tokens.shape[1])
           # `combined_mask` is shape (b, suffix_len, prefix_len + suffix_len) indicating how the suffix tokens (which
           # generate the queries) can attend to the full prefix + suffix sequence (which generates the keys and values)
           full_attn_mask = jnp.concatenate([prefix_attn_mask, suffix_attn_mask], axis=-1)
           assert full_attn_mask.shape == (
               batch_size,
               suffix_tokens.shape[1],
               prefix_tokens.shape[1] + suffix_tokens.shape[1],
           )
           # `positions` is shape (b, suffix_len) indicating the positions of the suffix tokens
           positions = jnp.sum(prefix_mask, axis=-1)[:, None] + jnp.cumsum(suffix_mask, axis=-1) - 1
 
           (prefix_out, suffix_out), _ = self.PaliGemma.llm(
               [None, suffix_tokens],
               mask=full_attn_mask,
               positions=positions,
               kv_cache=kv_cache,
               adarms_cond=[None, adarms_cond],
           )
           assert prefix_out is None
           v_t = self.action_out_proj(suffix_out[:, -self.action_horizon :])
           # 高斯积分法：
           return x_t + dt * v_t, time + dt
 
       def cond(carry):
           # 定义结束条件
           x_t, time = carry
           # robust to floating-point error
           return time >= -dt / 2
 
       x_0, _ = jax.lax.while_loop(cond, step, (noise, 1.0))
       return x_0

3. 数据管道与训练

To be continued…

4. 微调

To be continued…