目标
这篇文章介绍了名为 π0 的模型,旨在解决机器人学习中的主要挑战,包括数据稀缺性、泛化能力和鲁棒性。目标是开发一种通用的机器人控制策略,通过结合预训练视觉-语言模型(VLM)和机器人动作生成技术,实现机器人在物理世界中的多任务、灵活操作能力。这种模型被设计为适用于各种机器人平台,并能够完成从简单到复杂的任务,如折叠衣物、清洁桌面和组装盒子。
方法
模型架构:
- 预训练视觉-语言模型(VLM):利用互联网规模的视觉和语义知识,作为模型的基础。
- 动作生成模块:通过一种基于流匹配(flow matching)的技术生成连续的机器人动作,以支持高频率和复杂的物理操作。
- 跨平台数据训练:融合来自多种机器人类型(如单臂、双臂、移动机器人等)的大量数据,建立通用的表示和策略。
- 分阶段训练:
- 预训练:在多样化的机器人数据集上训练,赋予模型基本的通用能力。
- 微调(后训练):在高质量的任务特定数据集上进行优化,以提高任务执行的精度和流畅性。
动作专家模块
动作专家模块 (Action Expert) 是 π0 模型架构中的关键组成部分,它专注于处理与机器人状态和动作生成相关的输入与输出,为实现高频率、连续性和多样化的动作生成提供支持。以下是对动作专家模块的详细展开:
专用参数和独立权重
-
独立的动作权重:
- 动作专家模块使用了与 VLM 背景分离的一组专用参数。
- VLM 的主要权重用于处理视觉和语言信息,而动作专家的专用权重则处理与机器人状态和动作生成相关的输入和输出。
-
模块化设计:
- 动作专家模块类似于 “混合专家网络 (Mixture of Experts)” 的设计。整个 π0 模型架构可以看作由两个子模块组成:
- VLM 专家:负责处理图像和语言输入。{width=2048, depth=18, mlp dim=16,384, num heads=18, num kv heads=1, head dim=256}。
- 动作专家:处理机器人关节状态(proprioceptive state)和生成动作。 {width=1024, mlp dim=4096}, resulting in a parameter count of ∼300M。
- transformer层中的self attention层用来混合VLM专家和动作专家;
- 动作专家模块类似于 “混合专家网络 (Mixture of Experts)” 的设计。整个 π0 模型架构可以看作由两个子模块组成:
输入和输出特征
-
输入特征:
- 多模态输入,包含:
- 视觉信息:来自摄像头的 RGB 图像。
- 语言指令:任务相关的文本提示或高层策略生成的中间指令。
- 机器人状态:关节角度(joint angles)、速度等传感器信息。
- 所有输入通过嵌入层(embedding layer)映射到与 VLM 输出一致的嵌入空间。
- 多模态输入,包含:
-
输出特征:
- 动作输出是 连续的多步动作序列(action chunks),由动作专家生成,包含未来的多个高频率动作。只输出未来预测的action;
- 输出动作的形式支持多种机器人配置(如单臂、双臂和移动基座):
- 单臂:6自由度(6-DoF)的动作序列。
- 双臂:14自由度(两个6-DoF机械臂+抓手动作)。
- 移动机器人:包括移动基座的额外自由度(如位移和旋转)。
动作生成机制
流匹配 (Flow Matching)
动作专家模块利用 流匹配技术 来生成复杂的连续动作。这种技术是 π0 模型能够高效生成高频率动作的核心。
-
流匹配的基本原理:
- 学习动作分布 $ p(A_t|o_t) $,其中 $ A_t $ 是动作块,$ o_t $是机器人观测(如图像、语言、关节状态)。
- 通过条件流匹配损失函数 $ L_\tau(\theta) $:
$$
L_\tau(\theta) = \mathbb{E}{q(A^\tau_t|A_t), p(o_t)} \left[ \left| v\theta(A^\tau_t, o_t) - u(A^\tau_t | A_t) \right|^2 \right]
$$
其中:
- $ q(A^\tau_t | A_t) $ 是带噪声的动作分布。
- $ v_\theta(A^\tau_t, o_t) $ 是模型生成的动作。
- $ u(A^\tau_t | A_t) $ 是目标动作分布。
- 动作生成过程通过从随机噪声开始,逐步移除噪声生成最终动作。
动作分块 (Action Chunking)
- 模块生成动作块 $ A_t = [a_t, a_{t+1}, ..., a_{t+H-1}] $,其中 H 是动作块长度(如50)。
- 动作分块的优势:
- 高频控制:通过生成多个连续动作块,模型能够以 50Hz 的频率控制机器人。
- 复杂性支持:分块机制使得模型能够应对复杂任务中连续性和多模态需求。
** 动作推理**
- 动作生成在推理时通过迭代的积分过程完成:
- 从随机初始动作 $ A_0 \sim \mathcal{N}(0, I) $ 开始。
- 使用欧拉积分公式:
$$
A^{\tau+\delta}t = A^\tau_t + \delta v\theta(A^\tau_t, o_t)
$$
其中 \( \delta \) 是步长(如 0.1)。
- 整个过程通过缓存中间结果(如注意力键值)优化计算效率。
数据收集与使用:
- 总计10,000小时的机器人操作数据,包括来自68种任务和7种机器人配置的示例。
- 数据分为预训练数据(覆盖广泛的行为)和微调数据(聚焦于高质量、特定任务行为)。
- 包括:OXE dataset(Open X-Embodiment: Robotic learning datasets and RT-X models [2023.10]);Bridge v2(BridgeData v2: A dataset for robot learning at scale [2023]);DROID(DROID: A large-scale in-the-wild robot manipulation dataset [2024.03])
任务表现
-
基础模型评估:
- 零样本能力:预训练后的模型在没有微调的情况下即可完成多种任务,如折叠衣服和简单的桌面清洁,性能优于基准模型(如 OpenVLA 和 Octo)。
- 评估显示,π0在各种任务中的表现大幅领先于其他机器人基础模型。
-
语言指令跟随:
- π0 在理解和执行复杂语言指令方面表现显著优于未使用预训练语言模型初始化的小型版本(π0-small)。
- 添加高级语义策略(如高层语言规划)进一步提高了任务完成率。
-
微调表现:
- π0 在复杂任务上的微调性能显著高于从头开始训练的模型,尤其是在需要高精度操作的任务中(如折叠毛巾和替换纸巾卷)。
- 对较少的数据量(如1小时的微调数据)也能表现出显著优势。
-
复杂任务解决能力:
- 在高度复杂、多阶段的任务(如组装盒子、分类物品)中,π0结合预训练和微调策略表现出强大的通用性和精确性。
- 在所有任务中,π0的综合得分均超过50%,并在最困难的任务中显示出预训练的显著益处。