0. 论文信息
标题：RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation
作者：Songming Liu, Lingxuan Wu, Bangguo Li, Hengkai Tan, Huayu Chen, Zhengyi Wang, Ke Xu, Hang Su, Jun Zhu
机构：Tsinghua University
原文链接：https://arxiv.org/abs/2410.07864
代码链接：https://github.com/thu-ml/RoboticsDiffusionTransformer
数据集链接：https://huggingface.co/datasets/robotics-diffusion-transformer/rdt-ft-data
官方主页：https://rdt-robotics.github.io/rdt-robotics/
1. 导读
双手操作在机器人技术中是必不可少的，然而由于协调两个机器人手臂的固有复杂性(导致多模态动作分布)和训练数据的稀缺，开发基础模型极具挑战性。在这篇论文中，我们提出了机器人扩散变压器(RDT)，一个开创性的双手操作的扩散基础模型。RDT建立在扩散模型的基础上，以有效地表示多模态，具有可扩展转换器的创新设计，以处理多模态输入的异质性，并捕捉机器人数据的非线性和高频率。为了解决数据缺乏的问题，我们进一步引入了一个物理可解释的统一动作空间，该空间可以统一各种机器人的动作表示，同时保留原始动作的物理意义，便于学习可迁移的物理知识。通过这些设计，我们成功地在迄今为止最大的多机器人数据集集合上对RDT进行了预训练，并将其扩展到1.2B参数，这是最大的基于扩散的机器人操纵基础模型。我们最终在一个自创的多任务双手数据集上对RDT进行了微调，该数据集超过6K集，以完善其操纵能力。在真实机器人上的实验表明，RDT明显优于现有的方法。它对看不见的物体和场景表现出零镜头概括，理解并遵循语言指令，只需1~5次演示就能学习新技能，并有效地处理复杂、灵巧的任务。
2. 引言
双手协调操作对于机器人完成现实世界的任务至关重要。在实际应用中，有效的操作策略应该能够泛化到未见过的场景，如未见过的物体和场景。然而，当前的方法要么依赖于特定任务的基元，要么仅限于小规模模型、数据和简单任务，因此仅表现出有限的泛化能力，在复杂任务中失败。继自然语言处理和计算机视觉领域取得成功后，实现可泛化行为的一个有前景的方向是通过在大型数据集上进行模仿学习来开发基础模型。
然而，开发双手协调操作的基础模型绝非易事。一个主要原因是，由于硬件成本高，针对特定双臂机器人的可用数据极为稀缺，这破坏了训练基础模型所需的数据密集型要求。受近期单手操作尝试的启发，我们寻求首先在广泛的多机器人数据集上进行预训练，然后在目标双臂机器人收集的小数据集上进行微调。这可以帮助我们将数据规模扩大三个数量级，有可能从其他机器人的数据集中学习到可迁移的物理知识。然而，存在两个关键的技术挑战。首先，一个可泛化的基础模型在表达能力和可扩展性方面需要高度强大的架构。因此，模型必须具有足够的表达能力来捕捉动作分布中的多模态性。以往的方法通常无法达到这一标准，导致性能不尽如人意。此外，该架构需要有效处理来自不同模态的输入，包括文本、图像和动作。它必须具有可扩展性，以便能够稳定地在大型机器人数据集上进行训练。其次，由于不同机器人的物理结构和动作空间定义存在差异，数据异质性可能导致在多机器人数据训练期间出现负迁移，阻碍策略泛化。现有方法要么丢弃具有不同动作空间的机器人，要么仅保留其结构在不同机器人之间保持恒定的数据部分，从而丢失了宝贵的数据。
在本文中，我们介绍了机器人扩散变换器（Robotics Diffusion Transformer，RDT），这是目前最大的具有强大泛化能力的双手协调操作基础模型。RDT采用扩散变换器（Diffusion Transformers，DiTs）作为其可扩展的主干网络，并针对基于语言的双臂协调操作与视觉进行了特殊设计。在表达能力方面，RDT通过利用扩散模型对复杂分布进行建模的能力，擅长从海量数据中捕捉双臂动作的全部模态。在可扩展性方面，我们利用Transformer主干网络，并精心设计多模态编码，以消除各种模态的异质性。为了表征机器人数据中固有的非线性动力学、高频变化和不稳定的数值范围，这些特性与具有时空连续性的图像和视频显著不同，我们对原始DiT结构进行了重要修改，包括多层感知器（MLP）解码、改进的正则化和交替条件注入。为了进一步实现在异构数据上训练RDT，我们提出了物理可解释统一动作空间，这是一种适用于各种带夹持臂的机器人的统一动作格式。这种创新的格式缓解了不同机器人之间的潜在冲突，同时保留了原始动作的物理意义，这可以促进模型在不同机器人数据集上学习可泛化的物理知识。
3. 效果展示
我们呈现机器人扩散变压器随着1.2B参数(RDT-1B)最大的基于扩散的机器人操作基础模型。它是预先训练好的最大的46个数据集的多机器人集合，超过100万集。为了提高它的双手操作能力，我们收集了6K+集(其中之一最大的迄今为止)对ALOHA双臂机器人进行微调。它在灵活性、零射击推广性和少射击学习方面树立了一个新的标杆。它支持对几乎所有现代机械手的控制(例如双臂、关节、eef，甚至轮式移动)，并准备好让社区对他们的机器人进行微调！
4. 主要贡献
我们成功地在迄今为止最大的多机器人数据集集合上对RDT模型进行了预训练，并将其扩展到12亿参数，这是用于机器人操作的最大基于扩散的预训练模型。为了进一步增强其双手协调操作能力，我们在一个包含超过6000多条轨迹的自收集多任务双臂数据集上对RDT进行了微调，这是最大规模的双臂数据集之一。在我们的实验中，我们全面评估了RDT在双手协调操作和机器人基础模型方面的强大基线。结果表明，RDT取得了最先进的性能，在广泛的挑战性任务中成功率提高了56%，优于基线。特别是，RDT在未见过的物体、场景、指令甚至技能方面具有出色的零样本和少样本（1~5样本）泛化能力。RDT还能够完成需要精细操作的任务，例如使用操纵杆控制机器狗。最后，消融研究表明，扩散建模、大模型规模和大数据量都对卓越性能做出了贡献。
5. 方法
(a)ALOHA双臂机器人示意图。(b)抓取立方体的玩具示例。与单手操作相比，双手操作有更多可能的动作模式，导致更强的多模态。从亮到暗的颜色表示时间在前进。双臂协调操作的动作空间维度是单手操作的两倍，这使得可行动作分布的模态性更高。
机器人操作有如下挑战：
挑战1：如何设计一个强大的架构？一个具有泛化能力的基础模型需要一个强大的架构。这一要求包含两个主要方面。首先，架构必须拥有足够的表达能力，以捕捉动作分布中的多模态性。图2b展示了一个玩具示例，其中机器人尝试抓取一个立方体。我们可以看到，完成这项任务有多种模式，与仅控制一个机械臂的单手操作相比有所不同。在收集演示数据时，人类操作者可能会随机选择其中一种模式，从而导致收集到的动作数据具有多模态性。其次，这样的架构必须具有可扩展性。作为基础模型，它应能够有效处理来自不同模态（文本、图像、动作等）的异构输入，同时能够在大型数据集上稳定训练并具有可扩展性。
挑战2：如何在异构数据上进行训练？在多机器人数据上进行训练面临数据异构性的独特挑战。不同机器人的物理结构和动作空间可能差异很大。以往的研究要么将自己限制在具有相似动作空间的机器人子集上，要么仅保留具有相同结构的输入子集，但代价是丢失了大量信息。如何在这样的异构数据上训练模型，目前仍在很大程度上没有得到解决。从表3的结果可以看出，机器人扩散变换器（RDT）始终优于其他基线方法。这是因为RDT采用强大的网络架构进行扩散，以准确建模多模态动作的分布，而离散化和变分自编码器（VAE）分别在准确性和表达能力上有所欠缺。此外，大规模预训练后的大量参数提供了大量先验知识，显著提高了模型的泛化能力。以下是详细分析：
我们提出如图3所示的机器人扩散变换器（RDT）。我们介绍了扩散模型及其相应的架构，以解决挑战1。我们通过提出物理可解释统一动作空间来解决挑战2，以统一各种机器人的动作空间，并实现多机器人预训练。我们还收集了一个全面的多任务双臂数据集用于微调，以提高RDT的双手协调操作能力。
