一位优秀的相声演员需要吹拉弹唱样样在行,类似地,一个优秀的机器人模型也应能适应多样化的机器人形态和分歧的使命,但目前大多数机器人模型都只能控制一种形态的机器人执行一类使命。现在 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 进行了实证分析,在多个维度上评估了其作为机器人基础模型的性能:
能否直接应用 Octo 控制多台机器人机体并解决谈话和目的使命?
Octo 权重可否作为优质的初始化基础,支持面向新使命和机器人的数据高效型微调,以及它是否优于从头开始训练的方法和常用的预训练表征?
在构建通才机器人计谋方面,Octo 中的哪种设想决策最重要?
图 4 展示了评估 Octo 的 9 种使命。
直接应用 Octo 控制多台机器人
该团队比较了 Octo、RT-1-X、RT-2-X 的零样本操控能力,结果见图 5。
可以看到,Octo 的成功率比 RT-1-X(3500 万参数)高 29%。而在 WidowX 和 RT-1 Robot 评估上,Octo 与 550 亿参数的 RT-2-X 性能相当。
此外,RT-1-X 和 RT-2-X 仅支持谈话指令,而 Octo 还支持以目的图象为条件。该团队还发现,在 WidowX 使命上,如果应用目的图象为条件,成功率比应用谈话为条件高 25%。这可能是因为目的图象能提供更多有关使命完成的信息。
Octo 能高效地应用数据来适应新领域
表 1 给出了数据高效型微调的实验结果。
可以看到,相比于从头开始训练或应用预训练的 VC-1 权重进行预训练,微调 Octo 得到的结果更好。在 6 种评估设置上,Octo 相较于第二名基准的平均优势为 52%!
并且不得不提的是:针对所有这些评估使命,微调 Octo 时应用的配方和超参数全都一样,由此可见该团队找到了一个非常好的默认配置。
通才机器人计谋训练的设想决策
上面的结果表明 Octo 确实能作为零样本多机器人控制器,也能作为计谋微调的初始化基础。接下来,该团队分析了分歧设想决策对 Octo 计谋性能的影响。具体而言,他们关注的重点是以下方面:模型架构、训练数据、训练目的、模型规模。为此,他们进行了消融研究。
表 2 给出了在模型架构、训练数据和训练目的的消融研究结果。
图 6 则展现了模型规模对零样本成功率的影响,可以看出来更大的模型有更好的视觉场景感知能力。
整体而言,Octo 各组件的有效性得到了证明。