清华重磅开源RDT-1B：全球最大的双臂机器人扩散模型！

来源：投稿  作者：橡皮
编辑：学姐

论文链接：https://arxiv.org/abs/2410.07864

项目主页：https://rdt-robotics.github.io/rdt-robotics/

摘要：

双手操作在机器人技术中至关重要，但由于协调两个机械臂的固有复杂性（导致多模态动作分布）和训练数据的稀缺性，开发基础模型极具挑战性。在本文中，我们提出了机器人扩散变换器（RDT），这是一种用于双手操作的开创性扩散基础模型。RDT 以扩散模型为基础，有效地表示多模态，并创新性地设计了可扩展的变换器来处理多模态输入的异质性并捕获机器人数据的非线性和高频性。为了解决数据稀缺问题，我们进一步引入了物理上可解释的统一动作空间，它可以统一各种机器人的动作表示，同时保留原始动作的物理意义，促进学习可转移的物理知识。通过这些设计，我们成功地在迄今为止最大的多机器人数据集上对 RDT 进行了预训练，并将其扩展到 1.2B 参数，这是最大的基于扩散的机器人操作基础模型。我们最终在自创的多任务双手数据集上对 RDT 进行了微调，该数据集包含超过 6K+ 集，以改进其操作能力。在真实机器人上进行的实验表明，RDT 的表现明显优于现有方法。它对未见过的物体和场景表现出零样本泛化能力，能够理解和遵循语言指令，仅通过 1∼5 次演示即可学习新技能，并能有效处理复杂、灵巧的任务。

1 引言

双手操作对于机器人完成现实世界的任务至关重要。对于实际应用，有用的操作策略应该能够推广到未见过的场景，例如未见过的物体和场景。然而，当前的方法要么依赖于特定于任务的基元，要么仅限于小规模模型、数据和简单任务，因此只表现出狭窄的泛化能力，无法完成复杂的任务。继自然语言处理和计算机视觉取得成功之后，实现可泛化行为的一个有希望的方向是通过对大规模数据集进行模仿学习来开发基础模型。

然而，开发双手操作基础模型并非易事。一个主要原因是由于硬件成本高昂，特定双臂机器人的可访问数据非常稀缺，从而削弱了训练基础模型的数据密集型要求。受到最近单手操作尝试的启发，我们试图首先在广泛的多机器人数据集上进行预训练，然后在目标双臂机器人上收集的小数据集上进行微调。这可以帮助我们将数据大小扩展到三个数量级，并有可能从其他机器人的数据集中学习可迁移的物理知识。然而，有两个关键的技术挑战。首先，可通用的基础模型需要具有高度表达能力和可扩展性的架构。双手操作中动作空间的维度是单手操作的两倍，从而在可行动作的分布中带来更高程度的多模态性，如图 2b 所示。因此，模型必须具有足够的表达能力才能捕捉动作分布中的多模态性。以前的方法通常无法满足这一标准，导致性能不尽人意。此外，架构需要有效地处理来自不同模态的输入，包括文本、图像和动作。它必须是可扩展的，才能在大规模机器人数据上稳定地进行训练。其次，日期异质性是由不同机器人的物理结构和动作空间定义的变化引起的，在多机器人数据训练期间会导致负迁移并阻碍策略泛化。现有的方法要么丢弃具有不同动作空间的机器人，要么只保留结构在整个机器人中保持不变的数据部分，但代价是丢失宝贵的数据。

在本文中，我们介绍了机器人扩散Transformer (RDT)，它是最大的双手操作基础模型，具有很强的泛化能力。RDT 采用扩散Transformer (DiT) 作为其可扩展的主干网络，并针对语言条件下的视觉双手操作进行了特殊设计。在表现力方面，RDT 利用扩散模型对复杂分布进行建模的能力，擅长从海量数据中捕获双手动作的完整模态。在可扩展性方面，我们利用 Transformer 主干并精心设计多模态编码以消除各种模态的异质性。为了表征机器人数据固有的非线性动力学、高频变化和不稳定的数值范围，与具有时间和空间连续性的图像和视频有显著不同，我们对原始 DiT 结构进行了重要修改，包括 MLP 解码、改进的规范化和交替条件注入（参见图 4 了解它们的重要性）。为了进一步实现在异构数据上训练 RDT，我们提出了物理可解释统一动作空间，这是针对各种带有夹臂的机器人的统一动作格式。这种创新格式可以缓解不同机器人之间的潜在冲突，同时保留原始动作的物理含义，从而促进模型学习跨各种机器人数据集的可泛化物理知识。

通过上述设计，我们成功地在迄今为止最大的多机器人数据集上对 RDT 模型进行了预训练，并将其扩展到 1.2B 参数，这是最大的基于扩散的机器人操作预训练模型。为了进一步增强其双手操作能力，我们在自收集的多任务双手数据集上对 RDT 进行了微调，该数据集包含超过 6K+ 条轨迹，这是最广泛的双手数据集之一。在我们的实验中，我们根据双手操作和机器人基础模型中的强基线全面评估了 RDT。结果表明，RDT 达到了最先进的性能，在广泛的具有挑战性的任务中成功率提高了 56%，超越了基线。特别是，RDT 具有出色的零样本和小样本（1 ∼ 5 次样本）泛化能力，可推广到未见过的物体、场景、指令甚至技能。RDT 还能够完成需要精细操作的任务，例如用操纵杆控制机器狗。最后，消融研究表明，扩散建模、大型模型大小和大数据量都有助于实现卓越的性能。

2 相关工作

基于学习的双手操作。 学习双手操作策略的一大挑战是动作空间的高维性，这加剧了数据稀缺性和多模态行为。一些研究已经开发出更具成本效益的数据收集接口，但它们仅限于特定的硬件配置，仍然不足以弥补可推广策略的数据差距。其他人则试图通过引入归纳偏差来减少数据要求，例如区分用于稳定和功能的两个臂、参数化运动原语，或使用体素表示。这些方法使用强先验或简化建模，成功地减少了动作空间，但代价是应用范围缩小，无法表达双手行为的多模态性。

机器人基础模型。 基础模型通过在大型多任务机器人数据集上训练多任务“通才”模型 ，在实现可泛化行为方面表现出巨大的潜力。大多数研究采用大型视觉语言模型来直接预测动作。虽然它们展示了对新对象和任务的泛化能力，但它们在应用于双手操作时面临着量化误差和不协调行为的问题，这主要是由于它们对动作空间进行了离散化。为了提高精度，扩散模型已用于连续控制。Ghosh 等人在 Open X-Embodiment数据集（25 个数据集）的子集上预训练了基于 Transformer 的扩散策略，最多有 93M 个参数。

3 问题表述和挑战

我们首先制定任务并阐述挑战。为了在硬件上评估模型，我们选择 ALOHA 双臂机器人作为目标机器人，因为它是最具代表性的双臂机器人之一，适合通过远程操作收集人类演示数据。图 2a 显示了目标机器人的示意图，该机器人由两个带夹持器的手臂和三个摄像头组成。请注意，我们的设置和基础模型适用于任何双臂夹持器机器人。

我们考虑语言条件下的视觉双手操作的具体任务，这项任务是机器人技术的基础，在家庭等现实场景中具有重要价值。正式地，给定一个语言指令 ℓ，策略在时间 处呈现一个观察 ；然后它产生一个动作 来控制两个机械臂来实现 ℓ 指定的目标。观察表示为三元组 ，其中 是大小为 的 RGB 观察历史， 是机器人的低维本体感受， 是控制频率。动作 通常是所需本体感受 的子集。

双手操作中的一项特定任务通常由多个元素组成：技能（例如动词“挑选”或“擦拭”）、对象（例如名词“瓶子”或“桌子”）、场景（即执行任务的环境）和描述如何执行技能的模态（例如副词“用左手挑选瓶子”）。给定一个新任务，需要一个实用的策略来推广到训练数据中未出现的看不见的元素。这对于以前的基于规则的方法以及使用小模型/数据的基于学习的方法来说都是具有挑战性的，如第 2 节所述。

我们旨在通过模仿学习来训练基础模型策略，以实现普遍性。然而，由于硬件成本高昂，特定双臂机器人的可用数据特别稀少（< 10K 条轨迹），远远达不到训练基础模型的普遍要求。为了解决这个问题，我们建议采用预训练和微调流程，以利用来自多个机器人的数据，灵感来自单手操作的最新进展。通过这种方式，我们将数据量扩大三个数量级。具体来说，我们首先在大规模多机器人数据集 （主要是单臂）上对模型进行预训练，然后在目标机器人的数据集 上进行微调。我们将数据集表示为$\mathcal{D} .=\left\{\left(\ell^{(i)}, \boldsymbol{o}_{t}^{(i)}, \boldsymbol{a}_{t}^{(i)}\right) \mid 0 \leq t

挑战1：如何设计一个强大的架构？ 可泛化的基础模型需要强大的架构。这一要求包括两个主要方面。首先，架构必须具有足够的表现力，以捕捉动作分布中的多模态性。图 2b 展示了一个玩具示例，其中机器人试图抓住一个立方体。我们可以看到，有很多种模式可以完成这项任务，而单手操作则只能控制一个机械臂。在收集演示时，人类操作员可能会随机选择其中一种，从而导致收集到的动作数据具有多模态性。其次，这种架构必须具有可扩展性。作为基础模型，它应该能够有效地处理来自各种模态（文本、图像、动作等）的异构输入，同时还具有可扩展性，可以在大型数据集上进行稳定的训练。

挑战2：如何在异构数据上进行训练？ 在多机器人数据上进行训练提出了数据异质性的独特挑战。不同机器人的物理结构和动作空间可能有很大差异。以前的尝试要么将自己限制在具有相似动作空间的机器人子集上，要么仅保留具有相同结构的输入子集，但代价是丢失大量信息。如何在这种异构数据上训练模型仍然没有得到充分解决。

4 机器人扩散Transformer

我们现在介绍机器人扩散Transformer (RDT)，如图 3 所示。在第 4.1 节中，我们介绍了扩散模型和相应的架构以解决挑战1。在第 4.2 节中，我们通过提出一个物理上可解释的统一动作空间来解决挑战 2，以统一各种机器人动作空间并实现多机器人预训练。我们还收集了一个全面的多任务双手数据集进行微调，以提高 RDT 的双手操作能力。

4.1 RDT 模型

扩散建模。由于多模态性，给定语言指令 ℓ 和观察 ，可能有许多可能的动作 来继续执行任务。如果我们将其建模为确定性映射 并回归训练数据中的 元组，策略将学习动作模式的“平均值”。这可能会导致分布外的动作，例如多个模式的算术平均值，这可能是完全不可行的。相反，我们选择对连续条件分布 进行建模。如第 2 节所述，在各种方法中，扩散模型在表达能力和采样质量方面都表现出色，但对高维数据（例如图像）进行采样可能会很慢。幸运的是，对于我们的设置，缺点很小，因为 的维度比图像低得多，这只需要最小的采样开销。这使得扩散模型成为策略的理想选择，就像 Chi 等人。

然而，将扩散模型应用于机器人任务面临着独特的挑战，因为机器人物理量的固有属性（即动作和本体感受）与图像/视频数据不同。图像和视频数据虽然是高维的，但通常表现出一定程度的时间和空间连续性，帧之间的变化通常是增量的。相比之下，机器人物理量的特点是其非线性动力学和由物理相互作用（例如碰撞、约束和阻尼等材料特性）引起的高频变化的可能性。此外，这些量还具有不稳定的数值范围，可能是由于不可靠的传感器导致的极端值。这强调了调整当前扩散模型以有效捕捉机器人数据的不稳定性和非线性的必要性。接下来，我们将首先详细阐述扩散公式，然后介绍我们为解决这些挑战而设计的架构。

在使用扩散策略进行决策时，我们首先对一个完全嘈杂的动作 进行采样，然后执行 个去噪步骤，从 中将其去噪为一个干净的动作样本 ：

其中 是根据噪声计划预先定义的标量系数。这里， ，且 如果 k > 1，否则 。然而，在采样完成之前， 是难以处理的。我们选择使用具有参数 θ 的可学习去噪网络 fθ 来从噪声样本中估计干净样本: 。为了训练这样的网络，我们将最小化以下去噪均方误差 (MSE)：

其中 ，且 从我们的训练数据集中采样。在本文后面，我们将用 表示嘈杂的动作输入，其中为简单起见，删除了 k 的上标。此外，在实践中，我们更喜欢一次性预测一系列动作，即一个动作块，以鼓励时间一致性并通过减少任务中的决策数量来减轻随时间的误差积累。具体来说，我们建模 ，其中 是一个动作块， 表示块大小。

我们现在介绍架构的设计，包括多模态输入的编码和 fθ 的网络结构，详细信息请参阅附录。

异构多模态输入的编码。 多模态输入的异质性体现在结构上，即每个模态的格式和维数有显著差异。这给多模态训练带来了挑战。为了解决这个问题，我们将这些不同的模态编码到统一的潜在空间中。以下是编码方法：

• 低维输入是代表机器人物理量的低维向量，包括本体感受、动作块和控制频率。为了对它们进行编码，我们使用 MLP（具有傅里叶特征），它可以有效捕捉低维空间中的高频变化。
• 图像输入是高维的，包含丰富的空间和语义信息。为了提取紧凑的表示，我们使用图像文本对齐的预训练视觉编码器 SigLIP。我们在训练期间固定其权重以节省 GPU 内存。
• 语言输入长度不一且高度抽象，由于其复杂性和模糊性，带来了集成挑战。为了对它们进行编码，我们使用预先训练的基于 Transformer 的语言模型 T5-XXL。我们还在训练期间固定其权重以节省 GPU 内存。

异质性除了结构外，还体现在不同的输入中蕴含的信息量不同。首先，不同模态的数据包含的信息量不同，比如图像通常比文本包含更多的信息，经过编码后，产生的token更多。其次，同一模态的不同输入可能包含的信息量差别很大。比如，机器人的外部摄像头比手腕摄像头具有更全局的视野，包含的信息更丰富，如图3右上所示。这种情况下，模型可能会学一条捷径：只关注外部视野，忽略手腕视野，从而失去感知深度的能力。为了解决这个问题，我们在编码时以一定概率随机独立地mask每一个多模态输入，以防止模型过度依赖某个特定的输入。

fθ 的网络结构。 我们选择 Transformer 作为可扩展的骨干网络，并根据机器人问题的特点对 Diffusion Transfomer (DiT) 进行了以下三个关键修改：

• QKNorm 和 RMSNorm。 输入的机器人物理量的数值范围不稳定会导致梯度不稳定、数值溢出等问题，尤其是在训练大型基础模型时。为了解决这个问题，我们添加了 QKNorm 以避免计算注意力时出现数值不稳定的情况。此外，我们还注意到我们的问题可以看作是一个时间序列预测任务，而原始 DiTs 的 LayerNorm 中的居中操作会导致 token 偏移和注意力偏移，从而破坏时间序列的对称性。因此，我们将 LayerNorm 替换为没有居中操作的 RMSNorm。图 4a 显示，如果不进行这种修改，大规模预训练往往会非常不稳定，甚至爆炸。

• MLP 解码器。 为了提高非线性机器人动作的近似能力，我们将最终的线性解码器替换为非线性 MLP 解码器，作为从潜在空间到物理空间的投影。如图 4b 中经验所示，如果没有这种设计，RDT 就无法有效捕捉非线性动态，从而失去完成需要精细操作的灵巧任务的能力。

• 交替条件注入(ACI)。 在我们的模型中，图像和语言输入作为条件，它们是高维的并且长度可变，与传统 DiT 中的类标签条件形成对比。这些信息条件很难压缩成单个标记，这使得原始的自适应层规范方法不适合。因此，我们采用交叉注意来适应不同长度的条件，避免进一步压缩时的信息丢失。此外，我们进一步分析，鉴于图像标记通常比文本标记多得多，同时注入两种模态往往会掩盖与文本相关的信息，从而削弱指令跟踪的能力（定量结果见图 4b）。为了缓解这个问题，我们在连续层的交叉注意中策略性地交替注入图像和文本标记，而不是在每一层都注入两者。