[EAI-027] RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation

Paper Card

论文标题：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
论文链接：https://arxiv.org/abs/2410.07864v1
项目主页：https://rdt-robotics.github.io/rdt-robotics/
论文出处：/
论文被引：/

Abstract

双臂操作由于协调两只机械臂的复杂性（导致多模态动作分布）以及训练数据的稀缺性，开发基础模型极具挑战性。本文提出了用于双臂操作的扩散基础模型 Robotics Diffusion Transformer (RDT)。 RDT 构建在扩散模型的基础上，有效地表征多模态动作分布，并创新性地设计了可扩展的Transformer来处理异质的多模态输入，以学习非线性和高频的机器人数据。为了解决数据稀缺问题，引入了物理可解释的统一动作空间，该空间可以统一各种机器人的动作表示，同时保留原始动作的物理意义，便于学习可迁移的物理知识。凭借这些设计，成功地对迄今为止最大的多机器人数据集完成了参数量最大的基于扩散的机器人操作基础模型 RDT-1.2B 的训练。最后，在一个自建的多任务双臂数据集上对RDT进行了微调，该数据集包含超过 6K+ episode，以改进其操作能力。真实机器人上的实验表明，RDT 显著优于现有方法。展示了对未见的物体和场景的零样本泛化，理解并遵循语言指令，只需 1∼5 次演示数据就能学习新技能，并有效地处理复杂、灵巧的任务。

Summary

用于机器人双臂操作的基础模型 1.2B，跟pi0是同期发布的工作，应该不是参数规模最大的了。这个模型的优势是整个DiT骨干上进行扩散过程，相比pi0仅仅在action expert进行扩散，能够统一动作空间。劣势是没办法做多模态推理，对语言指令的语义泛化能力没有使用VLM的pi0好。开源了 6k+ 条的双臂操作数据，功德无量～
在这里插入图片描述

研究背景

当前缺少用于机器人双臂操作的基础模型，主要原因是包含高维动作空间的数据稀缺（可用数据 < 10K episode），硬件成本高，数据采集困难。为了解决这个问题，首先能够使用已有的机器人数据，这些数据可以是单臂机器人的数据，提供一个比较好的start up之后，再使用专门的双臂数据进行微调，并且要具备一定的泛化能力。要达成这个目标，需要解决下述模型层面问题：1）需要一个在表达能力和可扩展性方面非常强大的架构，尤其是对机器人动作空间多模态分布特性的建模能力；2）能够有效地处理多模态的异质数据，包括文本、图像和动作（不同机器人本体、不同机器人观测空间等等）；3）能够高效稳定地在大规模数据上进行训练。

方法介绍

本文提出了一个具备强大泛化能力的双臂操作基础模型 Robotics Diffusion Transformer (RDT)，其以 DiT 作为骨干网络，这使得其具备对复杂动作分布的建模能力，精心设计了多模态编码器，以消除模态间的异质性。

为了适配机器人动作的高频变化和非线性动力学，针对DiT做了优化，主要包括MLP解码，改进的归一化方法和交替条件注入（alternating condition injecting）。
为了能够在异质数据上训练，设计了可解释的统一动作空间，这是一个用于具有夹爪机器人的统一动作格式，可以在保留原始操作物理语义的前提下，缓解不同机器人动作空间配置的冲突，从而学习跨机器人本体的课泛化的物理世界知识。

问题描述

机器人本体：松灵 ALOHA 双臂机器人，通过遥操作收集数据，包括两个二指夹爪，两个腕部相机和第三视角相机
在这里插入图片描述

解决的问题：以语言为条件的双臂视觉操作任务，在家居场景具有重要价值。给定语言指令 l，和 t 时刻的观测信息 ot（包括：相机的历史序列长度为T的RGB输入x，机器人低维的本体感受信息 zt，以及控制频率 c），模型需要生成动作 at 来控制双臂实现 l 指定的目标。注意，动作 at 通常是本体感受信息 zt+1的子集。

一个典型的双臂操作任务涉及的元素：

一个技能：例如，抓取，擦拭等
一个物体：例如，瓶子，桌子等
一个场景：任务发生的场景
一个关于如何执行技能的模态描述：例如用左手臂拿起瓶子

对于一个新任务，需要模型能够泛化到训练数据未出现过的元素。本文的的目标是训练一个适用于双臂操作机器人的基础模型 RDT 实现上述元素的泛化。注意，是利用多种机器人数据增强双臂操作的泛化能力，而非训练一个跨机器人本体泛化的基础模型。

需要解决两个问题：

如何设计强大的模型架构？
如何针对异质的数据训练？

模型架构

什么是机器人动作的多模态特性？给定语言指令和观测信息，可能存在多条可以执行的机器人动作轨迹来执行任务。如果是确定性映射的方式，模型会学习动作分布的平均值，从而导致无法完成任务。而如果使用连续条件分布建模，比如基于扩散模型，则可以解决这个问题。但随之而来的是原生的DiT架构建模的是图像和视频，并且在时间和空间上是连续的，帧与帧之间的变化是增量的。然而，机器人相关的低维动作模态等物理量是非线性的并且是高频变化的，这些特性源于注入碰撞、约束和材料阻尼等物理交互产生。此外，还可能由于传感器不可靠导致存在数值不稳定的问题。

基于扩散的动作建模

两点需要注意：去噪过程在整个网络进行，可以学习到统一的动作空间；预测的是 action trunk 而非单点的action。
在这里插入图片描述

异质多模态输入编码

异质性体现在模态的格式、维度数量和信息量（编码后的token）方面。为了解决这个问题，需要将不同的模态编码到一个统一的隐空间中，使用的编码器如下：

低维输入：包括机器人诶本体感受信息，动作块、控制频率，使用带有傅立叶特征的 MLP 编码捕获低维空间中的高频变化
图像输入：使用 SigLIP 从图像中提取高维的、包含丰富的空间和语义信息
语言输入：使用 T5-XXL 从语言中提取长度变化、高度抽象、模糊和复杂的信息

小的trick：不同模态输入可能包含非常不同的信息量，比如第三视角相机捕获的全局视图的图像包含的信息量多于腕部相机捕获的图像，为了避免模型只关注全局视图走捷径而忽略腕部相机视图，从而失去感知深度的能力，在编码过程中以一定概率随机地对每个多模态输入进行mask，以防止模型过度依赖特定输入。

DiT_61">DiT骨干网络改进

对 DiT 的三个改进：

QKNorm 和 RMSNorm：为了解决机器人物理量的数值不稳定导致的梯度不稳定和数值溢出问题，在计算attention时，增加了 QKNorm。action prediction本质上可以看作一个时间序列预测问题，而原始DiT 中的LayerNorm中的centering 操作会导致 token shift 和 attention shift 从而破坏时间序列的对称性。因此，替换为了 RMSNorm。图 4 可以看出，如果没有这种修改，预训练会非常不稳定，并且可能导致梯度爆炸。
MLP解码器：为了提高对非线性的机器人动作的建模能力，将现行解码器替换为非线性MLP解码器，作为从隐空间到物理空间的投影。图 4 可以看出，如果没有这个改进，RDt 无法有效地捕获非线性动态，从而导致无法完成需要精细操作的灵巧操作任务。
Alternating Condition Injection (ACI)：RDT中的图像和语言输入作为条件，是高维的，并且长度可变，这与传统 DiTs 中的类别标签条件不同。这些信息丰富的条件难以压缩成单个 token，使得原始的 adapter layer 的 norm 方法不适用。因此，采用交叉注意力来适应不同长度的条件，避免在进一步压缩时造成信息丢失。此外，由于图像通常远多于文本token，同时注入两种模态往往会掩盖文本相关信息，从而削弱指令遵从能力（图 4）。为了缓解这个问题，在 successive layer 的交叉注意力中交替注入图像和文本token，而不是在每一层都注入。

数据构建

统一的动作空间

为了能够在异构多机器人数据上进行训练，需要一个在各种机器人之间共享的统一动作空间，它可以为多机器人动作提供统一的格式。从机器人的原始动作空间到统一动作空间的映射应该是物理上可解释的，其每个维度都应该具有明确的物理意义。这可以使得模型从不同的机器人数据中学习共享的物理规律，从而提高从不同机器人数据中学习的效率。

如图 3 左侧所示，通过将原始动作向量中的每个元素根据其物理意义填充到统一动作空间向量的对应位置来将机器人的动作空间嵌入到这个统一空间中，剩余位置则用填充值填充（表 4 列出了具体的含义）。在这个统一空间下，能够利用几乎所有带夹爪的现代机器人的数据上预训练 RDT，并极大地扩展数据规模，以满足基础模型的要求。预训练数据包含 46 种不同机器人的数据集（表 6），总规模为 1M+ 条轨迹，共计 21TB。
在这里插入图片描述

双臂操作任务数据收集

高质量的微调数据集对于模型性能至关重要。从三个方面确保数据集的高质量：

在数量方面，收集了 6K+ 条轨迹，是目前最大的双臂数据集之一；
在全面性方面，考虑了 300+ 个具有挑战性的任务，涵盖了从抓取和放置到连接电缆等大多数操作任务类型，甚至包括书写数学方程；
在多样性方面，准备了 100+ 个物体，包括不同尺寸和纹理的刚性和非刚性物体，以及 15+ 个不同照明条件的不同房间
此外，进一步利用 GPT-4-Turbo 来重写人工标注的指令，以增加文本多样性。

消融实验

实验设置

旨在通过真实机器人实验回答以下问题：

Q1: RDT 的零样本泛化能力是否可以泛化到未见的物体和场景？

Q2: RDT 的零样本指令遵循能力对未见的模态的有效性？

Q3: RDT 能否促进对未见的技能的少样本学习？

Q4: RDT 是否能够完成需要精细操作的任务？

Q5: 大模型参数、大数据量和扩散模型对 RDT 的性能有帮助吗？

RDT (ours): the original RDT.
RDT (regress): RDT without diffusion modeling. It models the deterministic mapping (ℓ,𝒐t)↦𝒂t.
RDT (small): RDT without large parameters. It has only 166M parameters.
RDT (scratch): RDT without pre-training. It is trained from scratch during fine-tuning.

研究任务：一共7个，来评估 RDT 从不同维度的泛化能力和功能，包括模型在现实世界任务中可能遇到的复杂场景，例如各种未见的元素和灵巧的操作。如图5和表1所示。
在这里插入图片描述

数据量：

Wash Cup: 133 demos for seen cups combined and 0 demos for unseen cups;
Pour Water: 350 demos for seen rooms combined and 0 demos for unseen rooms;
Pour Water-L-1/3 & Pour Water-R-2/3: 18 demos for the water level of little, 19 demos for half, and 19 demos for full;
Handover: 5 demos; Fold Shorts: 1 demo;
Robot Dog: 68 demos.

模型训推：

训练：48 个 H100 GPU，一个月，1M steps
微调：48 个 H100 GPU，三天，130K steps

模型推理：

框架：DPM-Slover++
RTX 4090 24GB GPU 推理频率 6Hz Action Trunk，381 Hz Actions

测试配置：

Wash Cup is tested with 8 trials for each cup (one seen cup, two unseen cups, 24 trials in total).
Pour Water is tested with 8 trials for each room (three unseen rooms, 24 trials in total).
Pour Water-L-1/3 and Pour Water-R-2/3 are tested with 8 trials each.
Handover, Fold Shorts, and Robot Dog are tested with 25 trials each.

实验分析

在这里插入图片描述

从表 3 的结果可以看出，RDT 始终优于其他基线方法。这是因为 RDT 使用强大的网络架构与扩散相结合，能够准确地对多模态动作的分布进行建模，而离散化和 VAE 分别缺乏准确性和表达能力。此外，大规模预训练后的大量参数提供了大量的先验知识，这极大地提高了泛化能力。下面是详细分析：

Q1 & Q2: RDT 可以零样本泛化到未见的物体、场景和模态。在 Wash Cup 和 Pour Water 任务中，RDT 仍然可以在未见的场景中取得很高的成功率，其性能与可见场景中的性能相差无几。相反，其他基线方法甚至无法完成整个任务。在 Pour Water-L-1/3 和 Pour Water-R-2/3 中，从图 5 或图 10 可以发现，RDT 精确地理解了要操纵哪只手以及要倒多少水，并通过其动作严格遵循指令，即使它从未见过像“三分之一”或“三分之二”这样的词。正是因为大规模预训练，RDT 已经看到了大量的不同物体、场景和指令，从而导致了如此强大的零样本泛化。

Q3: RDT 仅使用少量样本就可以学习新技能。在 Handover 和 Fold Shorts 中，RDT 通过少样本学习学习了新的复杂技能，即移交和折叠，其动作模式与已知技能非常不同，而其他模型的成功率几乎为零。这种改进也归因于大规模预训练。 少样本学习可以帮助 RDT 快速适应新的工作环境，这对实际应用具有重大意义。

Q4: RDT 可以处理灵巧的任务。在 Robot Dog 中，RDT 在推动操纵杆时准确地控制角度，而其他机器人则会导致机器狗偏离。这是因为扩散，可以对多模态和非线性动作的分布进行建模，从而使动作精度能够满足灵巧任务的要求。操纵杆和遥控器都是黑色的，使得操纵杆在视觉上不明显。这可能会导致 ACT 容易发生故障。相比之下，大规模预训练使 RDT 学习了更好的操纵杆概念的视觉语言表示，提高了识别能力。

Q5: 大型模型尺寸、大量数据和扩散都是我们取得优异成绩的重要因素。在表 2 中，如果没有这些因素，性能会严重下降。RDT (scratch) 在未见的物体和场景上的表现很差，这表明预训练的知识对于泛化至关重要。

实验结论

本文提出了 RDT 来解决可泛化双臂操作中数据稀缺和操作复杂性增加的挑战，RDT 是一种基于扩散的语言条件视觉运动模仿学习基础模型（diffusion-based foundation model for language-conditioned visuomotor imitation learning）。 RDT 是在一个广泛的多机器人数据集上预训练的，并在一个自收集的双臂数据集上微调。引入了可解释的统一动作空间，以统一不同机器人之间的动作表示，增强了鲁棒性和可迁移性。RDT 超越了现有方法，不仅在灵巧的双臂能力和指令遵循方面表现出显著改进，而且在少样本学习和对未见物体和场景的零样本泛化方面也取得了卓越的性能。