Octo：伯克利开源机器人开发框架

【摘要】在各种机器人数据集上预先训练的大型策略有可能改变机器人学习：这种通用机器人策略无需从头开始训练新策略，只需使用少量领域内数据即可进行微调，但具有广泛的泛化能力。然而，为了广泛应用于各种机器人学习场景、环境和任务，这些策略需要处理不同的传感器和动作空间，适应各种常用的机器人平台，并轻松高效地微调到新领域。在这项工作中，我们旨在为开发开源、广泛适用的通用机器人操作策略奠定基础。作为第一步，我们引入了 Octo，这是一种基于大型 Transformer 的策略，在 Open X-Embodiment 数据集（迄今为止最大的机器人操作数据集）上的 800k 条轨迹上进行训练。它可以通过语言命令或目标图像进行指示，并且可以在标准消费级 GPU 上在几个小时内通过新的感官输入和动作空间有效地微调到机器人设置。在 9 个机器人平台上的实验中，我们证明了 Octo 是一种多功能策略初始化，可以有效地微调到新的观察和动作空间。我们还对 Octo 模型的设计决策进行了详细的讨论，从架构到训练数据，以指导未来构建通用机器人模型的研究。

原文：Octo: An Open-Source Generalist Robot Policy
地址：https://arxiv.org/abs/2405.12213v2
代码：https://octo-models.github.io
出版：未知
机构: UC Berkeley, Stanford

写的这么辛苦，麻烦关注微信公众号“码农的科研笔记”！

1 研究问题

本文研究的核心问题是: 如何设计一个开源的、通用的、可适配不同机器人和任务的机器人操作策略模型。

想象一下在一个机器人实验室里,有各种不同型号和功能的机器人,比如WidowX机械臂、UR5工业机器人、CMU厨师机器人等。研究人员希望能有一个统一的模型框架,通过少量的适配和微调,就能让不同的机器人完成多样化的任务,比如抓取物体、开关抽屉、操作家电等。同时这个框架最好是开源的,能够方便地在学术界和工业界推广应用。

本文研究问题的特点和现有方法面临的挑战主要体现在以下几个方面:

现有的机器人学习方法大多针对特定的机器人和任务,泛化能力有限。训练一个新的机器人完成一个新的任务,往往需要从头收集数据、训练模型,代价很高。
不同的机器人在硬件构型、传感器类型、动作空间等方面存在很大差异。很难设计一个统一的模型架构,能够灵活适配不同机器人的观察和动作接口。
大规模机器人操作数据的缺乏。与计算机视觉和自然语言处理不同,获取大量多样化的机器人互动数据需要巨大的硬件投入和人力成本。缺乏高质量的数据限制了通用机器人模型的训练。
模型的实用性和可访问性有待提高。很多先进的机器人学习模型要么是私有的,要么对计算资源要求很高,在实际场景中的应用受到限制。

针对这些挑战,本文提出了一种灵活且可扩展的"八爪鱼(Octo)"机器人通用策略:

Octo的核心是一个基于Transformer的神经网络。它将任务定义(如语言指令、目标图像)和机器人观察(如相机图像、力传感器读数)统一编码为一个token序列,然后用一个通用的Transformer骨干网络提取特征,最后用轻量级的输出头解码成机器人动作。这种模块化的设计赋予了Octo极大的灵活性:通过简单地增删输入输出token和头,它可以适配不同的机器人和任务,而无需修改预训练的骨干网络参数。同时Octo在目前最大规模的机器人操作数据集Open X-Embodiment(包含800k个机器人轨迹)上训练,习得了一个强大的跨任务、跨机型的通用策略表征。在9个不同的实验平台上,Octo展现出了卓越的零样本迁移和少样本微调能力。值得一提的是,本文还细致地分析了模型设计的关键要素,如数据分布、模型架构、策略形式等,为未来机器人通用模型的研究提供了重要参考。最后,Octo的代码和模型权重被完全开源,使得这一进展能够惠及整个机器人研究社区。

2 研究方法

本节详细介绍论文提出的Octo模型，一个用于机器人操控的通用Transformer策略网络。内容组织如下：首先介绍模型的整体架构，然后说明模型训练使用的数据来源和预处理方法，接着阐述模型训练时的目标函数和一些训练细节，最后介绍本工作开源的代码和模型权重。

2.1 模型架构

Octo模型主要由三个部分组成：输入编码器，用于将不同形式的输入统一编码为token序列；Transformer主干网络，用于处理输入token序列并输出对应的embedding；输出头，将Transformer输出解码为机器人动作。模型的整体架构如图2所示。

输入编码器将语言指令、目标图像、机器人传感器观察等不同modality的输入分别编码为统一格式的token。其中，语言指令通过预训练的语言模型编码为定长的token embedding序列；图像则先经过一个较浅的CNN提取特征，然后划分为多个patches并展平，从而得到图像token序列。最后将这些不同来源的输入token拼接为一个统一的序列送入Transformer主干网络。

Transformer接收这些输入token后，以块状注意力(block-wise attention)的方式处理它们。如图2顶部所示，不同来源的输入token序列先分别与对应的可学习的位置编码(position embedding)相加，然后依次输入到Transformer的每一层。每个输入token只能关注当前时刻之前的输入，而任务指令token(绿色方块)可以被所有其他token关注。此外，还引入了一些可学习的输出token(readout token，图中紫色方块)。每个输出token都关注它之前的输入token，但反过来输入token不会关注输出token。这些输出token就像是BERT中的[CLS]标记，作为截至当前步的观察序列的紧凑embedding表示。最后将输出token的embedding传入输出头，解码为机器人动作。

模型输出头采用diffusion过程对机器人动作进行建模。它将一个初始的高斯噪声向量解码为连续的机器人动作，解码过程通过多步的逐步去噪来实现。每一步去噪过程由浅层MLP实现，它以当前步的输出、Transformer输出的embedding、以及步骤编号作为输入。通过这种方式，模型能输出一段连续的未来机器人动作序列，即所谓的action chunk。

本文提出的Octo模型最大的特点在于其灵活性。得益于块状注意力和输出head的设计，在迁移学习到新机器人或任务时，可以很方便地增删输入输出modalities，而无需重新初始化或训练模型的大部分参数。如图2底部所示，在下游任务中需要新的观察信息时，只需引入新的对应输入token即可(虚线蓝色方块)；类似地，需要输出新的动作空间时，只需引入新的可学习输出token(虚线紫色方块)并搭配新的输出头即可。预训练好的Transformer参数都可以原封不动地继承使用。这种灵活的设计使得Octo成为一个通用的机器人控制策略，能适应多种机器人平台和任务。

2.2 训练数据

Octo模型在目前最大规模的机器人操控数据集Open X-Embodiment dataset上训练，该数据集包含了约150万个机器人操控的轨迹。本文从中精选了25个子数据集共计800k个轨迹进行训练，涵盖了多个机器人平台和任务环境，数据来源如图3所示。

为了进行跨数据集的训练，本文对原始数据进行了必要的预处理和对齐。以统一夹爪动作空间为例，将所有数据集的夹爪指令对齐为: +1表示张开，0表示闭合。另外，为了防止过长的单个轨迹主导整个训练过程，本文对每条轨迹随机下采样至多100个时间步。

在训练时需要从这25个数据集中采样数据，一个简单的策略是按数据集大小等比例采样。但是其中一些更丰富、更多样化的数据集可能对提高模型泛化性更有帮助。因此，本文依据数据集大小设置基础采样概率，并对一些表现更佳的数据集的概率进行少量提升，最终各数据集的采样概率如图3所示。这种经过调整的采样策略在实验中取得了更好的效果。

2.3 训练目标与训练细节

与许多先前工作使用离散化动作空间或MSE回归动作不同，本文采用条件扩散模型(conditional diffusion model)作为动作输出头，以期对连续动作空间进行更好的建模。具体来说，每次训练时，先将专家动作序列加入高斯噪声，然后训练模型去除这些噪声以还原原始动作。通过这种方式，模型能学会输出与专家演示动作更相似的动作序列。

另一个细节是，本文模型每次输出一段连续的未来动作序列(action chunking)，即预测从当前时刻开始的若干个未来时间步的动作，而不是单步预测。这样的序列化输出能得到更连贯平滑的动作轨迹。在测试时则采用滚动时域预测的方式，即预测未来若干步的动作，执行前几步，观察最新状态，再重新预测，以此类推。

本文还采用了一些常见的数据增强和正则化技巧，包括对训练图像进行随机裁剪、颜色变换等，以及对Transformer模型使用dropout和LayerNorm。模型训练时使用AdamW优化器，并设置初始学习率warmup和余弦衰减的学习率调度。表IV列出了主要的训练超参数。

最后，本文还训练了不同规模的Octo变体，如表V所示，包括参数量为2700万的Octo-Small和9300万的Octo-Base等。在实验中，更大规模的模型在开箱即用的零样本机器人控制上展现出了更强的能力。

2.4 开源代码与模型

为了方便研究人员的进一步研究和使用，本文将Octo训练和测试的全部代码开源，并提供了预训练好的模型权重。这包括：

Octo-Small和Octo-Base的预训练权重，可直接用于测试或在下游任务上进行微调
模型微调的示例代码，可适配到新的观察和动作空间
完整的模型训练流水线代码，包括高效的Open X-Embodiment数据加载器
方便推理部署的模型前向代码

有了这些开源代码和模型，研究人员只需几行代码就能调用Octo模型对机器人进行控制。例如，加载一个预训练权重，给定语言指令和传感器观察，就能输出对应的机器人动作。同时，在新的机器人和任务上对模型进行微调也变得简单易行。清单1展示了一个最简单的模型推理代码例子。

综上，本文提出的Octo模型兼顾了通用性和灵活性，能以开箱即用的方式对多个机器人执行语言指令，也能高效地迁移到新的观察空间和动作空间。模型的训练和推理代码全部开源，为今后机器人学习领域的研究提供了一个很好的基础模型。

4 实验

4.1 实验场景介绍

该论文提出了一个开源的通用机器人操作策略Octo,论文实验主要评估Octo在零样本多机器人控制和few-shot策略微调中的性能,以及不同设计决策的影响。

4.2 实验设置

实验平台:在4个机构的9个真实机器人设置上进行评估,涵盖不同机器人、传感器配置和任务类型。
Datasets:Octo在Open X-Embodiment数据集的800k机器人轨迹上预训练。few-shot微调使用~100个目标域演示。
Baseline:零样本对比RT-1-X和RT-2-X。few-shot微调对比从头训练和VC-1预训练视觉表征。
metric:任务成功率

4.3 实验结果

4.3.1 实验一、Octo的零样本多机器人控制性能

目的:评估Octo在来自预训练数据的环境中对多个机器人的开箱即用控制能力

涉及图表:图5

结果:

Octo在WidowX、UR5和RT-1机器人上的语言指定任务上优于RT-1-X。
在WidowX任务上,Octo与更大的RT-2-X模型表现相似。
Octo还支持目标图像条件,在WidowX上比语言条件高出25%的成功率。

4.3.2 实验二、Octo在新域上few-shot微调的性能

目的:评估Octo作为新任务和新机器人策略初始化的few-shot微调性能

涉及图表:表I

实验细节概述：在6个包含新观测(力矩)、新动作空间(关节位置控制)和新机器人的评估设置上,使用~100个目标域演示和相同超参数进行微调,每个域评估20次。

结果:

在所有设置中,微调Octo优于从头训练和VC-1预训练视觉表征,平均高出52%。
结果凸显了Octo适应新观测、动作空间和机器人的能力,使其广泛适用于单臂双臂操作问题。

4.3.3 实验三、设计决策消融实验

目的:评估不同设计决策(训练数据、模型架构、训练目标、模型尺度)对Octo性能的影响

涉及图表:表II,表VI,图6

实验细节概述：除模型尺度外,所有消融在Octo-Small上进行。每项评估40次,跨两个语言指定和两个目标图像指定任务。还分析了Octo在新对象、新环境、新技能上的泛化能力。

结果:

模型架构:ViT比常见的ResNet编码器表现更好。
训练数据:在更多样化的数据集上训练可提高性能。
训练目标:扩散解码优于MSE和离散化动作预测。
模型尺度:性能随模型尺寸增加而提高。更大模型对初始场景更稳健,不易过早尝试抓取。
泛化性:Octo在新对象和环境上泛化较好,但在未见技能上表现较差。

4 总结后记

本论文针对如何训练一个通用的机器人控制策略这一问题,提出了一个名为Octo的Transformer模型。通过在大规模多机器人数据集上预训练,再在小规模目标领域数据上微调的方式,Octo可以灵活适应新的观察和动作空间,实现了跨机器人和跨任务的强大泛化能力。实验结果表明,Octo在零样本和少样本场景下都取得了优异的表现,为构建通用机器人智能迈出了重要一步。

疑惑和想法：