适应多形态多任务，最强开源机器人学习系统「八爪鱼」诞生

2024-05-28 04:05

机器之心报道

编辑：Panda

一位优秀的相声演员需要吹拉弹唱样样在行，类似地，一个优秀的机器人模型也应能适应多样化的机器人形态和不同的任务，但目前大多数机器人模型都只能控制一种形态的机器人执行一类任务。现在 Octo（八爪鱼）来了！这个基于 Transformer 的模型堪称当前最强大的开源机器人学习系统，无需额外训练就能完成多样化的机器人操控任务并能在一定程度适应新机器人形态和新任务，就像肢体灵活的八爪鱼。

在机器人学习方面，一种常用方法是收集针对特定机器人和任务的数据集，然后用其来训练策略。但是，如果使用这种方法来从头开始学习，每一个任务都需要收集足够数据，并且所得策略的泛化能力通常也不佳。

原理上讲，从其它机器人和任务收集的经验能提供可能的解决方案，能让模型看到多种多样的机器人控制问题，而这些问题也许能提升机器人在下游任务上的泛化能力和性能。但是，即便现在已经出现了能处理多种自然语言和计算机视觉任务的通用模型，构建「通用机器人模型」依然困难重重。

要为机器人训练一个统一的控制策略非常困难，其中涉及诸多难点，包括操作不同的机器人机体、传感器配置、动作空间、任务规范、环境和计算预算。

为了实现这一目标，已经出现了一些「机器人基础模型」相关研究成果；它们的做法是直接将机器人观察映射成动作，然后通过零样本或少样本方式泛化至新领域或新机器人。这些模型通常被称为「通才机器人策略（generalist robot policy）」，简称 GRP，这强调了机器人跨多种任务、环境和机器人系统执行低阶视觉运动控制的能力。

举些例子：GNM（General Navigation Model，通用导航模型）适用于多种不同的机器人导航场景，RoboCat 可针对任务目标操作不同的机器人机体，RT-X 能通过语言操控五种不同的机器人机体。尽管这些模型确实是重要进展，但它们也存在多方面的局限：它们的输入观察通常是预定义的且通常很有限（比如单相机输入视频流）；它们难以有效地微调至新领域；这些模型中最大型的版本都没有提供人们使用（这一点很重要）。

近日，加州大学伯克利分校、斯坦福大学、卡内基梅隆大学和谷歌 DeepMind 的 18 位研究者组成的 Octo Model Team 发布了他们的开创性研究成果：Octo 模型。该项目有效地克服了上述局限。

论文标题：Octo: An Open-Source Generalist Robot Policy
论文地址：https://arxiv.org/pdf/2405.12213
开源项目：https://octo-models.github.io/

他们设计了一个系统，能让 GRP 更轻松地应对下游机器人应用的接口多样化问题。

该模型的核心是 Transformer 架构，其可将任意输入 token（根据观察和任务创建）映射成输出 token（然后编码成动作），而且该架构可使用多样化的机器人和任务数据集进行训练。该策略无需额外训练就能接受不同的相机配置，也能控制不同的机器人，还能通过语言命令或目标图像进行引导 —— 所有这些只需通过改变输入模型的 token 即可实现。

最重要的是，该模型还能适应传感器输入、动作空间或机器人形态不同的新机器人配置，所需的只是采用适当的适配器（adapter）并使用一个小的目标领域数据集和少量计算预算进行微调。

不仅如此，Octo 还已经在迄今为止最大的机器人操控数据集上完成了预训练 —— 该数据集包含来自 Open X-Embodiment 数据集的 80 万个机器人演示。Octo 不仅是首个可有效微调至新观察和动作空间的 GRP，也是首个完全开源（训练工作流程、模型检查点和数据）的通才机器人操控策略。该团队也在论文中强调了其组合 Octo 各组件的独特创新性。

Octo 模型

下面我们来看看 Octo 这个开源的通才机器人策略是如何构建的。总体而言，Octo 的设计目标是让其成为一个灵活且广泛适用的通才机器人策略，可被大量不同的下游机器人应用和研究项目使用。

架构

Octo 的核心是基于 Transformer 的策略 π。其包含三个关键部分：输入 token 化器、Transformer 骨干网络和读出头。

如图 2 所示，其中输入 token 化器的作用是将语言指令、目标和观察序列转换成 token，Transformer 骨干会把这些 token 处理成嵌入，读出头则是得出所需的输出，即动作。

任务和观察 token 化器

为了将任务定义（比如语言指令和目标图像）与观察（比如相机视频流）转换成常用的已 token 化的格式，该团队针对不同模态使用了不同的 token 化器：

对于语言输入，先 token 化，然后通过一个预训练的 Transformer 将其处理成一个语言嵌入 token 序列。具体而言，他们使用的模型是 t5-base (111M)。

对于图像观察和目标，则是通过一个较浅的卷积堆栈来处理，然后再拆分成平展后图块构成的序列。

最后，通过向任务和观察 token 添加可学习的位置嵌入并按一定顺序排列它们来构建 Transformer 的输入序列。

Transformer 骨干和读出头

将输入处理成一种统一化的 token 序列之后，就能交给 Transformer 处理了。这与之前的研究工作类似：基于观察和动作系列来训练基于 Transformer 的策略。

Octo 的注意力模式是逐块掩码式：观察 token 只能按照因果关系关注来自同一或之前时间步骤的 token 以及任务 token。对应于不存在观察的 token 会被完全掩蔽掉（比如没有语言指令的数据集）。这种模块化设计很方便，可在微调阶段添加或移除观察或任务。

除了这些输入 token 模块，该团队还插入了已学习完成的读出 token。读出 token 会关注其之前的观察和任务 token，但不会被任何观察或任务 token 关注。因此，读出 token 只能读取和处理内部嵌入，而无法影响内部嵌入。读出 token 的作用类似于 BERT 中的 [CLS] token，充当截至目前的观察序列的紧凑向量嵌入。针对读出 token 的嵌入，会使用一个轻量的实现扩散过程的「动作头」。这个动作头会预测多个连续动作构成的一个「块（chunk）」。

这样的设计可让用户在下游微调时向模型灵活地添加新的任务和观察输入或动作输出头。当在下游添加新的任务、观察或损失函数时，可以在整体上保留 Transformer 的预训练权重，仅添加新的位置嵌入、一个新的轻量编码器、或由于规范变化而必需的新头的参数。这不同于之前的架构 —— 对于之前的架构，如果添加或移除图像输入或改变任务规范，就需要重新初始化或重新训练预训练模型的大量组件。

要让 Octo 成为真正的「通才」模型，这种灵活性至关重要：由于我们不可能在预训练阶段覆盖所有可能的机器人传感器和动作配置，因此，如果能在微调阶段调整 Octo 的输入和输出，便能让其成为机器人社区的一种多功能工具。另外，之前使用标准 Transformer 骨干或融合使用视觉编码器与 MLP 输出头的模型设计固定了模型输入的类型和顺序。相较之下，切换 Octo 的观察或任务并不需要对大部分模型进行重新初始化。

训练数据

该团队从 Open X-Embodiment 中取用了包含 25 个数据集的混合数据集。图 3 给出了数据集的组成。

有关训练目标和训练硬件配置等更多细节请参阅原论文。

模型检查点和代码

重点来了！该团队不仅发了 Octo 的论文，还完全开源了所有资源，其中包括：

预训练完成的 Octo 检查点，包括 2700 万参数的 Octo-Small 和 9300 万参数的 Octo-Base。
用于 Octo 模型的微调脚本，基于 JAX。
用于在 Open X-Embodiment 数据集上预训练 Octo 的模型预训练工作流程，基于 JAX。用于 Open X-Embodiment 数据的数据加载器，兼容 JAX 和 PyTorch。

实验

该团队也通过实验对 Octo 进行了实证分析，在多个维度上评估了其作为机器人基础模型的性能：