openpi论文及代码解析(A Vision-Language-Action Flow Model for General Robot Control) (一)

目前一说到具身算法不得不提两种经典模型一个是rdt模型一个是pi0模型, rdt模型在之前的博客介绍过了RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation 论文及代码总结(三), 今天记录下pi0模型,当然现在已经出了pi0.5, 但是介于目前代码没有公布, 因此我们先以pi0作为主要模型介绍, 后面也会详细说明下pi-fast. 这里让我很苦恼的是, 代码是基于flax框架而不是传统我们用到的pytorch框架, 所以在学习起来相当费力, 这里也很感谢这篇博客π0源码(openpi)剖析——从π0模型架构的实现:如何基于PaLI-Gemma和扩散策略去噪生成动作,到基于C/S架构下的模型训练与部署 该作者分析很全面, 让我对模型有进一步细致了解. 因为这里涉及到的背景还是不少的, 所以我们还是先巩固下背景知识.

1️⃣ 论文地址: π0: A Vision-Language-Action Flow Model for General Robot Control
2️⃣ 源码地址: openpi

一、 背景知识


1. 流匹配(FlowMatching)

说到流匹配, 网上说的很多, 尤其是一对数学公式和推导, 我是看的云里雾里的,不是很理解, 当我看到b站有个大佬解释的非常简单易于理解(【复现】扩散模型Flow Matching), 同时给出了对应的代码, 看完以后我大概就知道是怎么回事了, 这里我简单说下流匹配是怎么回事(同时可以参考这篇文章直观理解Flow Matching算法(带源码)).

首先来说流匹配做的事情和扩散模型做的事情是一样的都是由噪音变成图像, 但是扩散模型由高斯噪音不断进行迭代去噪音变成清晰的图片, 但是流匹配从一张高斯噪音纯图片中的像素向一个目标图像分布去流动, 流动可以理解为对图像的像素进行有方向的加加减减经过多轮的迭代, 最终像素流到它需要去的地方形成一个需要的图像. 所以说都是由噪音进行迭代到图片, 但是理念不一样, 需要求的东西不一样, 所以可以简单理解

1️⃣ 扩散模型: 求出噪音是什么
2️⃣ 流匹配模型: 求出像素流动的向量, 其目的是让模型学会一个流场(vector field), 这个流场定义了数据从起点(比如噪声分布)到终点(目标分布如何逐步变化的方式)
具体也可以参考这篇文章说的非常好: 通俗易懂理解Flow Matching

通俗理解FlowMatching算法.png

通俗理解数学原理(一)
通俗理解数学原理(二)

为了让大家更容易理解, 这里我引用了代码, 方便大家理解

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np

# 超参数
dim = 2         # 数据维度(2D点)
num_samples = 1000
num_steps = 50  # ODE求解步数
lr = 1e-3
epochs = 5000

# 目标分布:正弦曲线上的点(x1坐标)
x1_samples = torch.rand(num_samples, 1) * 4 * torch.pi  # 0到4π
y1_samples = torch.sin(x1_samples)                      # y=sin(x)
target_data = torch.cat([x1_samples, y1_samples], dim=1)

# 噪声分布:高斯噪声(x0坐标)
noise_data = torch.randn(num_samples, dim) * 2

class VectorField(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim + 1, 64),  # 输入维度: x (2) + t (1) = 3
            nn.ReLU(),
            nn.Linear(64, dim)
        )
  
    def forward(self, x, t):
        # 直接拼接x和t(t的形状需为(batch_size, 1))
        return self.net(torch.cat([x, t], dim=1))
        
model = VectorField()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

for epoch in range(epochs):
    # 随机采样噪声点和目标点
    idx = torch.randperm(num_samples)
    x0 = noise_data[idx]  # 起点:噪声
    x1 = target_data[idx] # 终点:正弦曲线

    # 时间t的形状为 (batch_size, 1)
    t = torch.rand(x0.size(0), 1)  # 例如:shape (1000, 1)
  
    # 线性插值生成中间点
    xt = (1 - t) * x0 + t * x1
  
    # 模型预测向量场(直接传入t,无需squeeze)
    vt_pred = model(xt, t)  # t的维度保持不变
  
    # 目标向量场:x1 - x0
    vt_target = x1 - x0
  
    # 损失函数
    loss = torch.mean((vt_pred - vt_target)**2)
  
    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

x = noise_data[0:1]  # 初始噪声点
trajectory = [x.detach().numpy()]

tag = torch.from_numpy(np.array([1]))
# 数值求解ODE(欧拉法)
t = 0
delta_t = 1 / num_steps
with torch.no_grad():
    for i in range(num_steps):
        vt = model(x, torch.tensor([[t]], dtype=torch.float32))
        t += delta_t
        x = x + vt * delta_t  # x(t+Δt) = x(t) + v(t)Δt
        trajectory.append(x.detach().numpy())

trajectory = torch.tensor(trajectory).squeeze()

print(trajectory[-1] / (torch.pi / 10 * 4))

# 绘制向量场和生成轨迹
plt.figure(figsize=(10, 5))
plt.scatter(target_data[:,0], target_data[:,1], c='blue', label='Target (sin(x))')
plt.scatter(noise_data[:,0], noise_data[:,1], c='red', alpha=0.3, label='Noise')
plt.plot(trajectory[:,0], trajectory[:,1], 'g-', linewidth=2, label='Generated Path')
plt.legend()
plt.title("Flow Matching: From Noise to Target Distribution")
plt.show()
image.png

在一个带条件的数字生成模型代码


import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from torchdiffeq import odeint
import matplotlib.pyplot as plt

# ================== 配置参数 ==================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
image_size = 28
channels = 1
batch_size = 256
lr = 1e-4
epochs = 100
num_classes = 10

# ================== 数据加载 ==================
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: 2 * x - 1)  # [-1, 1] 归一化
])

train_dataset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)

# ================== 模型架构 ==================
class ConditionedDoubleConv(nn.Module):
    """带条件注入的双卷积模块"""
    def __init__(self, in_channels, out_channels, cond_dim):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_channels)
        self.conv2 = nn.Conv2d(out_channels + cond_dim, out_channels, kernel_size=3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_channels)
      
    def forward(self, x, cond):
        x = F.silu(self.norm1(self.conv1(x)))
        cond = cond.expand(-1, -1, x.size(2), x.size(3))  # 动态广播条件
        x = torch.cat([x, cond], dim=1)
        return F.silu(self.norm2(self.conv2(x)))

class Down(nn.Module):
    """下采样模块"""
    def __init__(self, in_channels, out_channels, cond_dim):
        super().__init__()
        self.maxpool = nn.MaxPool2d(2)
        self.conv = ConditionedDoubleConv(in_channels, out_channels, cond_dim)
      
    def forward(self, x, cond):
        x = self.maxpool(x)
        return self.conv(x, cond)

class Up(nn.Module):
    """上采样模块"""
    def __init__(self, in_channels, out_channels, cond_dim):
        super().__init__()
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv = ConditionedDoubleConv(in_channels, out_channels, cond_dim)
      
    def forward(self, x1, x2, cond):
        x1 = self.up(x1)
        # 尺寸对齐
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]
        x1 = F.pad(x1, [diffX//2, diffX - diffX//2,
                        diffY//2, diffY - diffY//2])
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x, cond)

class ConditionalUNet(nn.Module):
    """维度安全的条件生成UNet"""
    def __init__(self):
        super().__init__()
        # 统一条件编码维度
        self.t_dim = 16
        self.label_dim = 16
        self.cond_dim = self.t_dim + self.label_dim  # 32
      
        # 时间嵌入
        self.time_embed = nn.Sequential(
            nn.Linear(1, 32),
            nn.SiLU(),
            nn.Linear(32, self.t_dim)
        )
        # 标签嵌入
        self.label_embed = nn.Embedding(num_classes, self.label_dim)
      
        # 编码路径
        self.inc = ConditionedDoubleConv(1, 64, self.cond_dim)
        self.down1 = Down(64, 128, self.cond_dim)
        self.down2 = Down(128, 256, self.cond_dim)
      
        # 解码路径
        self.up1 = Up(256 + 128, 128, self.cond_dim)  # 输入通道修正
        self.up2 = Up(128 + 64, 64, self.cond_dim)
        self.outc = nn.Conv2d(64, 1, kernel_size=1)
      
    def forward(self, x, t, labels):
        # 条件编码 (统一维度)
        t_emb = self.time_embed(t.view(-1, 1))  # [B, 16]
        lbl_emb = self.label_embed(labels)      # [B, 16]
        cond = torch.cat([t_emb, lbl_emb], dim=1)  # [B, 32]
        cond = cond.unsqueeze(-1).unsqueeze(-1)    # [B, 32, 1, 1]
      
        # 编码器
        x1 = self.inc(x, cond)
        x2 = self.down1(x1, cond)
        x3 = self.down2(x2, cond)
      
        # 解码器
        x = self.up1(x3, x2, cond)
        x = self.up2(x, x1, cond)
        return self.outc(x)

# ================== 训练与生成 ==================
model = ConditionalUNet().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

@torch.no_grad()
def generate_with_label(label, num_samples=16, device="cuda"):
    """生成指定标签的样本(修复条件维度问题)"""
    model.eval()
  
    # 初始噪声和标签
    x0 = torch.randn(num_samples, 1, 28, 28, device=device)
    labels = torch.full((num_samples,), label, device=device, dtype=torch.long)
  
    # 定义ODE函数
    def ode_func(t: torch.Tensor, x: torch.Tensor):
        t_expanded = t.expand(x.size(0))  # [1] -> [num_samples]
        vt = model(x, t_expanded, labels)
        return vt
  
    # 时间点(从0到1)
    t_eval = torch.tensor([0.0, 1.0], device=device)
  
    # 解ODE(自适应步长)
    generated = odeint(
        ode_func,
        x0,
        t_eval,
        rtol=1e-5,
        atol=1e-5,
        method='dopri5'
    )
  
    # 后处理
    images = (generated[-1].clamp(-1, 1) + 1) / 2  # [0,1]
    return images.cpu().squeeze(1)  # 移除通道维度

def visualize_samples(samples, title="Generated Samples"):
    """可视化生成结果"""
    plt.figure(figsize=(10, 10))
    for i in range(16):
        plt.subplot(4, 4, i+1)
        plt.imshow(samples[i].squeeze().cpu().numpy(), cmap='gray', vmin=0, vmax=1)
        plt.axis('off')
    plt.suptitle(title)
    plt.show()

def plot_100_digits(image_size=28, device="cuda"):
    """
    生成0-9各10张数字并绘制在10x10网格中
    Args:
        model: 训练好的生成模型
        image_size: 图像尺寸(默认MNIST为28)
        device: 计算设备
    """
    plt.figure(figsize=(8, 8))
  
    # 为每个数字0-9生成10张图
    for label in range(10):
        # 生成当前数字的10个样本
        generated = generate_with_label(
            label=label,
            num_samples=10
        ).numpy()  # 形状 (10, 28, 28)
      
        # 在当前行绘制
        for i in range(10):
            ax = plt.subplot(10, 10, i * 10 + 1 + label)
            plt.imshow(generated[i], cmap='gray')
            ax.axis('off')
            # 在每列第一行添加标签
            if i == 0:
                ax.text(14, -10, str(label), fontsize=20, ha='center')
  
    plt.tight_layout()
    plt.show()

def train():
    """训练循环"""
    for epoch in range(epochs):
        model.train()
        total_loss = 0
      
        for images, labels in train_loader:
            images = images.to(device)
            labels = labels.to(device)
          
            # 动态噪声生成
            noise = torch.randn_like(images)
            t = torch.rand(images.size(0), device=device)
            xt = (1 - t.view(-1,1,1,1)) * noise + t.view(-1,1,1,1) * images
          
            # 前向计算
            vt_pred = model(xt, t, labels)
            loss = F.mse_loss(vt_pred, images - noise)
          
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
          
            total_loss += loss.item()
      
        # 每10个epoch生成示例
        if epoch % 10 == 0:
            plot_100_digits()
        
        print(f"Epoch {epoch} Loss: {total_loss/len(train_loader):.4f}")

if __name__ == "__main__":
    train()
image.png

可以看下面代码中框起来的关键部分(t是在0~1之间的数)


模型训练

模型推理

x = (x+1)/2 ➡️ -1到1的值转为0-1的值

2. 修正流(Rectified Flow)


具体可以参考这篇文章ICLR 2023 | 扩散生成模型新方法:极度简化,一步生成 就是在流匹配的基础上希望路径更加直
image.png

下面这篇文章解释了流匹配和修正流之间的主要区别


上述两者的推理代码逻辑一样, 修正流代码可以参考图像生成的Rectified Flow方法之一:原理

这篇文章也说的特别好可以参考: SD3的采样下篇——Rectified Flow








3. SigLIP以及Gemma模型理解

(1) SigLIP模型
SigLIP【全称:Sigmoid Loss for Language Image Pre-Training】,SigLIP是在batch内,利用sigmod对文-图对做二分类;CLIP是在batch内,利用softmax对文-图对做多分类。

SigLIP不需要对两两相似进行全局归一化,这样的做法允许扩大batch的大小,同时在较小的batch下也能表现的好。

如下图所示:


SigLIP利用sigmod对文-图对做二分类,是在指导模型朝着文字Tokens和图像Tokens的两个序列的对角线上值越来越大,非对角线上的值越来越小的方向前进。即,希望配对的文-图对越来越匹配,非配对的文-图对越来越不匹配, 只需要计算文-图矩阵,对角线的得分就可以了,相比softmax计算更加简单,快捷。


(2) Gemma模型
Gemma 3 是 Google 最新的开放权重大型语言模型。它有四种尺寸,分别是 10 亿、40 亿、120 亿 和 270 亿 参数,包含基础(预训练)和指令调优版本。Gemma 3 支持 多模态! 4B亿、12B和 27B参数的模型可以处理 图像 和 文本,而1B参数的模型仅限于文本. 详细可以参考这篇文章Gemma 3 27B版本超越DeepSeek V3:技术要点分析!

(3) PaliGemma(可参考多模态PaliGemma 2(含1代):Google推出的基于SigLIP和Gemma 2的视觉语言模型(附SigLIP详解))

PaliGemma

PaliGemma将 400M SigLIP 和 2B Gemma 模型结合成一个小于 3 B 的 VLM
PaliGemma

由一下三个组件组成:

  • 一个图像编码器,使用SigLIP,其shape optimized ViT-So400m图像编码器,该模型通过sigmoid损失在大规模上进行了对比预训练,且其在小尺寸上也表现出色
  • 一个仅解码器的语言模型,使用gemma-2B v1.0,该模型可以匹配或超越使用相对更大些的语言模型的VLMs的性能,包括之前的PaLIs
  • 一个线性层,将SigLIP的输出投影到与gemma-2B的词汇token相同的维度,以便它们可以被连接「其实PaliGemma的作者也不是没考虑别的选择,而是在他们的早期实验中,他们发现更复杂的替代方案——MLPs,并没有提供明显的优势,因此决定使用最简单的选项」


4. Tranfusion(Transfusion: Predict the Next Token andDiffuse Images with One Multi-Modal Model)

tranfusion

本质是将LLM的transformer和图像中的diffusion结合了起来,使用同一个transformer来同时处理文本和图像信息.之前的DiT架构都是使用一个预训练的TextEncoder来提取文本信息,,并通过Concat、AdaLN、 CrossAttention、MMDit等方式将文本信息融入模型,而本文的方式直接同时训练文本和图像信息,并且是使用同一个模型来进行处理.

(1) 模型训练
  • 训练数据由文本和图像数据组成,通常以 50% 文本和 50% 图像的比例混合,每个图像序列前后都会添加特殊的 BOI 和 EOI token。
  • 将 LLM 训练的损失和扩散模型的损失相加,通过一个平衡系数来调整两者的权重。
    a. 文本:同正常的 next-token prediction 模型
    b. 图像:使用 Diffusion 的 loss (使用 [DDPM](https://zhida.zhihu.com/search?content_id=253033481&content_type=Article&match_order=1&q=DDPM&zhida_source=entity) 的加噪去噪方式进行训练)
(2)模型推理

在两种模式之间切换

  • LM 模式:按照标准做法,从预测分布中逐个 token 地采样,对于文本使用因果注意力
  • 扩散模式:对于图像则使用双向注意力(需要感知全部 patch 的信息)
    a. 当采样到一个 <BOI> token 时,解码算法切换到扩散模式 以n个图像 patch 的形式将纯噪声$x_T$附加到输入序列, 在 T 个 step 进行去噪,每个 step t,模型预测 noise,并根据 noise 更新 $x_{t-1}$
    b. 在 <EOI> 后,切换回 LM 模式。
  • 左图经过一个VAE来得到tokens,并插入到文本token中,文本也会在经过一个tokenizer之后通过一个轻量级的模块进行处理,然后再通过一个transformer来处理文本和图像的信息.
  • 右图文本的attention方式和图像不一致,文本因为要采用causal的方式,而图像则需要采用bidirectional的方式.

transfusion: Predict the next token and diffuse images with one multi-modal model,一个既可以预测下一个token又可以生成图像的多模态模型(相当于训练单个模型来同时预测离散文本token和扩散连续图像),总之,其通过在50%文本和50%图像数据上预训练一个Transformer模型来展示Transfusion, 其通过在单个序列元素上应用扩散风格(流匹配)损失训练他们的模型,而不是仅解码器的transformers的标准交叉熵损失

5. Flax框架简单理解

参考这篇文章从0开始】使用Flax NNX API 构建简单神经网络并训练
拟合公式: y = 2x^2+1, 代码如下

import jax.numpy as jnp
import jax.random as jrm
import optax as ox
from jax import Array
from flax import nnx
from typing import Generator


class Network(nnx.Module):
    """def a simple MLP"""

    def __init__(self, in_dim: int, out_dim: int, rng: nnx.Rngs, hidden_dim: int):
        super().__init__()
        self.linear1 = nnx.Linear(in_dim, hidden_dim, rngs=rng)
        self.linear2 = nnx.Linear(hidden_dim, hidden_dim, rngs=rng)
        self.linear3 = nnx.Linear(hidden_dim, out_dim, rngs=rng)

    def __call__(self, x) -> Array:
        x = self.linear1(x)
        x = nnx.relu(x)
        x = self.linear2(x)
        x = nnx.relu(x)
        x = self.linear3(x)
        return x


def make_dataset(
    X: Array, Y: Array, batch: int, seed: int = 0
) -> Generator[tuple[jnp.ndarray, jnp.ndarray], None, None]:
    "dataset sample function"
    combined = jnp.stack((X, Y), axis=1)[..., None]
    key = jrm.key(seed)
    while True:
        selected = jrm.choice(key, combined, shape=(batch,))
        yield selected[:, 0], selected[:, 1]


def loss_fn(model: Network, batch):
    x, y = batch
    predicted = model(x)
    return ox.l2_loss(predicted, y).mean()


# hyper parameter
seed = 0
batch = 16

# make dataset
X = jnp.arange(0, 10, 0.005)
Y = 2 * X**2 + 1.0

# build model & optimizer
model = Network(1, 1, hidden_dim=20, rng=nnx.Rngs(seed))
optimizer = nnx.Optimizer(model, ox.adamw(0.001, 0.90))

# train
for i, (x, y) in enumerate(make_dataset(X, Y, batch)):
    loss, grads = nnx.value_and_grad(loss_fn)(model, (x, y))
    optimizer.update(grads)
    print(i, loss)
    if i >= 6000:
        break
6. LeRobot

参考博客如下:
[1] Hugging Face 开发的 LeRobot 框架介绍
[2] Hugging Face 中 LeRobot 使用的入门指南
[3] Hugging Face 开发的 LeRobot 框架安装说明
[4] Hugging Face 中 LeRobot 学习的策略
[5] LeRobot 框架的核心架构概念和组件(上)
[6] LeRobot 框架的核心架构概念和组件(中)
[7] LeRobot 框架的核心架构概念和组件(下)
[8] LeRobot的数据集系统(上)
[9] LeRobot的数据集系统(下) (⭐️)
[10] LeRobot的机器人控制系统(上)
[11] LeRobot的机器人控制系统(下)
[12] LeRobot 框架的开发指南 (上)
[13] LeRobot 框架的开发指南 (下)
[14] LeRobot 实现的 PI 0 策略
[15] LeRobot 实现的 PI 0 FAST 策略

二、 论文阅读


该论文提出了一种基于预训练视觉 - 语言模型(VLM)的新型流匹配架构,以继承互联网规模的语义知识. 整体模型框架如下所示:

在文章中说明了几处我觉得有些意义的说明:

[1] 在AI大型语言和视觉 - 语言模型中:这些系统在来自网络的大量多样化图像和文本语料库上进行预训练,然后使用更精心策划的数据集进行微调(“对齐”),以诱导所需的行为和响应模式, 尽管已经证明了这种模型可以表现出广泛的指导和解决问题的能力,但它们并不像人们那样真正位于身体世界中,并且他们对身体互动的理解完全基于抽象的描述。

[2] 在自然语言和计算机视觉领域,在多样化多任务数据上预训练的通用基础模型往往优于狭隘定制的专门解决方案。例如,如果目标是识别照片中的鸟类,在许多不同的图像 - 语言关联上预训练,然后微调或提示进行鸟类识别任务,可能比仅在鸟类识别数据上训练更好. 类似地,我们可能会发现,对于有效的专门机器人系统,先在高度多样化的机器人数据上预训练,然后微调或提示所需任务更为有效。这可以解决数据稀缺问题,因为通用模型可以获取更多数据来源(包括来自其他任务、其他机器人甚至非机器人来源的数据) 并且可以解决鲁棒性和泛化挑战,因为多样化数据覆盖了更广泛的观察和动作,提供了各种场景、校正和恢复行为,而这些可能在专业化的专门数据中不存在

[3] 为了使利用各种多样化机器人数据源变得可行,我们采用跨实体训练,将来自许多机器人类型的数据组合到同一模型中。这些不同的机器人类型具有不同的配置空间和动作表示,包括单臂和双臂系统,以及移动操作器。此外,为了能够执行高度灵巧和复杂的物理任务,我们使用带有流匹配(扩散的一种变体)的动作分块架构来表示复杂的连续动作分布。这使我们的模型能够以高达 50Hz 的频率控制机器人执行灵巧任务,如衣物折叠.

[4] 与语言模型一样,我们的方法类似于亿级规模语言和图像 - 语言模型中常见的预训练 / 后训练分离,即模型首先在非常大且多样化的语料库上预训练,然后在更狭窄且更精心策划的数据上微调,以诱导所需的行为模式 —— 在我们的案例中是灵巧性、效率和鲁棒性。直观地说,仅在高质量数据上训练不会教会模型如何从错误中恢复,因为此类数据中很少出现错误。仅在较低质量的预训练数据上训练不会教会模型高效且鲁棒地行动。两者结合提供了所需的行为:模型尽可能尝试以类似于高质量数据的方式行动,但在出现错误时仍有一系列恢复和纠正措施可供部署。

[5] 我们通过在超过 10,000 小时的机器人数据上预训练, RDT-1B:在超过 100 万次的多机器人剧集上进行预训练,其预训练数据集包含 46 个数据集,总大小达 21TB。π0:使用的数据集包含来自 7 种不同机器人和 68 个任务的灵巧操作数据,包括 106M 步的单臂数据和 797M 的双臂数据,此外还结合了来自开放源代码数据集(如 OXE、DROID 和 Bridge)的数据。但具体数据量未明确提及总大小


上图中, 我们从一个预训练混合数据集开始,该数据集由我们自己的灵巧操作数据集和开源数据组成。我们使用这个混合数据集来训练我们的流匹配VLA模型,该模型由一个较大的视觉语言模型(VLM)主干和一个较小的动作专家组成,用于处理机器人状态和动作。VLM主干的权重从PaliGemma [5]初始化,提供从大规模互联网预训练中学习到的表示。由此得到的π0模型可用于控制具有不同动作空间的多种机器人实体,以完成各种各样的任务。

[6] 模型基于 PaliGemma 视觉语言模型 ,然后使用混合数据对其进一步训练。为了将基础的PaliGemma 视觉语言模型转变为 π0,我们添加了动作输出,该输出使用流匹配 [32, 28] 来生成连续动作分布。我们将在下一节详细介绍这种设计。我们使用 PaliGemma 是为了方便,并且由于它的规模相对较小(这对实时控制很有用),但我们的框架与任何基础预训练视觉语言模型都兼容。

[7] π₀模型主要由语言模型 Transformer 主干组成。遵循标准的晚期融合 VLM 方法,图像编码器将机器人的图像观察嵌入到与语言标记相同的嵌入空间中。我们进一步用特定于机器人的本体感受状态和机器人动作增强这个主干特征.

[8]π₀使用条件流匹配来建模动作的连续分布。流匹配为我们的模型提供了高精度和多模态建模能力,使其特别适合高频灵巧任务。我们的架构受到 Transfusion 的启发,Transfusion 使用多个目标训练单个 Transformer,其中标记对应于通过流匹配损失监督的连续输出,标记对应于通过交叉熵损失监督的离散输出。

[9] 在 Transfusion 的基础上,我们还发现对特定于机器人的(动作和状态)token使用单独的权重导致性能提升。这种设计类似于具有两个混合元素的专家混合,其中第一个元素用于图像和文本输入,第二个元素用于特定于机器人的输入和输出(动作和状态)。我们将第二组权重称为动作专家。

[10] 为了能够利用多种不同的机器人数据源,作者采用跨体态训练[10],即将多种类型机器人的数据合并到同一个模型中这些不同类型的机器人具有不同的构型空间和动作表示, 这里和RDT模型有相同的思想

[11] 灵巧的机器人操作需要π0以高频率(比如高达每秒 50 次)输出运动命令。为了提供这种级别的灵活性,他们通过流匹配为预训练的 VLM 提供连续动作输出

(一) 模型结构及细节

整体框架(摘抄博客:π0——用于通用机器人控制的VLA模型:一套框架控制7种机械臂(基于PaliGemma和流匹配的3B模型))



(二) 模型的改动


(二) 模型推理

参考:
[1] 多模态视觉层:CLIP、SigLIP谁更胜一筹
[2] 多模态PaliGemma 2(含1代):Google推出的基于SigLIP和Gemma 2的视觉语言模型(附SigLIP详解)
[3] Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model 回顾
[4] Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model(2024,8)

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容