RDT——清华开源的双臂机器人扩散大模型：先预训练后微调，支持语言、图像、动作多种输入

第一部分清华开源全球最大双臂机器人扩散大模型RDT

2.1 什么是RDT

2.1.1 RDT推出的背景及其与以前工作的对比

受到最近在单手操作方面尝试的启发（Brohan等，2023；Kim等，2024），清华一研究团队推出了RDT，其对应的项目主页为：github.io/rdt-robotics，其对应的论文为《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

有意思的是，如此文所说，视频生成Sora的全面解析：从AI绘画、ViT到ViViT、TECO、DiT、VDT、NaViT等
“其实早在2022年9月，清华朱军团队(背后关联的公司为生数科技，后于24年3月上旬拿到一笔数亿元的融资)就发布了一篇名为《All are Worth Words: A ViT Backbone for Diffusion Models》的论文(其一作为Fan Bao，比 DiT早提出两个多月，后被 CVPR 2023 收录)，这篇论文提出了用「基于Transformer的架构U-ViT」替代基于卷积架构的U-Net ”

为何要提这点呢，如我之前预判的，sora这类视频生成模型一定会赋能机器人的训练，所以你在上文字节的GR2看到了视频生成模型在机器人领域的应用，且在本文你将再次看到
其首先寻求在广泛的多机器人数据集上进行预训练，然后在目标双臂机器人收集的小数据集上进行微调

2.1.2 RDT所要解决的问题及其面对的挑战

首先，作者团队选择了ALOHA双臂机器人作为目标机器人，因为他们认为其是最具代表性的双臂机器人之一，并且适合通过远程操作收集人类演示数据（Zhao et al.,2023; Fu et al.,2024; Aldaco etal.,2024）

下图图2a显示了目标机器人的示意图，其中包含两个带有夹持器的手臂和三个摄像头

当然，他们说他们的设置和基础模型对任何双臂夹持机器人都是通用的。且他们考虑视觉条件下的语言引导双手操作这一具体任务，这在机器人领域是基础性的，并且在家用等现实场景中具有重要价值（Stepputtis等，2020；Brohan等，2022；Zhao等，2023）

形式上，给定一个语言指令 $l$ ，策略在时间 $t \in \mathbb{N}^{+}$ 时接收到一个观测 $\boldsymbol{o}_{t}$ ；然后它生成一个动作 $\boldsymbol{a}_{t}$ 来控制两个机器人手臂，以实现 $l$ 指定的目标

观测表示为一个三元组
$\boldsymbol{o}_{t}:=\left(\boldsymbol{X}_{t-T_{\mathrm{img}}+1: t+1}, \boldsymbol{z}_{t}, c\right)$
其中
$\boldsymbol{X}_{t-T_{\text {img }}+1: t+1}:=\left(\boldsymbol{X}_{t-T_{\mathrm{img}}+1}, \ldots, \boldsymbol{X}_{t}\right)$ 是大小为 $T_{\mathrm{img}}$ 的RGB观测历史
$\boldsymbol{z}_{t}$ 是机器人的低维本体感知
$c$ 是控制频率
动作 $\boldsymbol{a}_{t}$ 通常是期望本体感知 $\boldsymbol{z}_{t+1}$ 的一个子集「即The action at is usually a subset of the desired proprioception zt+1」
双手操作中的特定任务通常由多个元素组成：技能（例如，动词如“拾取”或“擦拭”）、对象（例如，名词如“瓶子”或“桌子”）、场景（即任务发生的环境）以及描述技能如何执行的方式（例如，副词性短语如“用左手拿瓶子”）
面对新任务时，需要一种实用的策略来推广到训练数据中未出现的未知元素。这对以前基于规则的方法以及小模型/数据的学习方法来说都是一个挑战
作者旨在通过模仿学习训练一个基础模型策略，以实现通用性。然而，由于硬件成本高，对于特定双臂机器人的可用数据特别稀少（<10K条轨迹），远远达不到训练基础模型的常见要求

为了解决这个问题，他们建议采用预训练和微调流程（Radford等，2018），通过借鉴最近在单臂操作中的进展（Ghosh等，2023；Collaboration等，2023；Kim等，2024），利用多个机器人的数据

通过这种方式，将数据规模扩大三个数量级

具体而言，他们首先在一个大型多机器人数据集 $\mathcal{D}_{\text {pre }}$ （主要是单臂）上预训练模型，然后在目标机器人数据集 $\mathcal{D}_{\mathrm{ft}}$ 上进行微调

他们将数据集记为

$\mathcal{D} .=\left\{\left(\ell^{(i)}, \boldsymbol{o}_{t}^{(i)}, \boldsymbol{a}_{t}^{(i)}\right) \mid 0 \leq t<T^{(i)}, 1 \leq i \leq N\right\}$

其中 $T^{(i)}$ 是第i条轨迹的长度， $N$ 是轨迹的数量。此外，值得强调的是，他们的目标是利用多机器人数据来增强模型在双手操作中的普适性，而不是开发用于各种机器人的跨实体模型

而他们面临的挑战在于，一个在架构，一个在数据，具体而言，如下所示：

挑战一，在机器人领域，缺乏一个像GPT这样的通用、强大的模型架构
首先，架构必须具有足够的表达能力以捕捉动作分布中的多模态性。下图图2b展示了一个简单的例子，其中机器人尝试抓取一个立方体「其是一个抓取立方体的玩具示例。与单手操作相比，双手操作具有更多可能的动作模式，从而导致更强的多模态性。颜色从浅到深表示时间向前推进」
可以看到，完成这个任务有很多种模式，这与单臂操作形成对比，在单臂操作中只控制一个机器人手臂。但在收集演示时，人类操作员可能会随机选择其中一种方式，导致收集的动作数据中存在多模态性「说白了，在双手操作中，动作空间的维度是单手操作的两倍」

其次，这种架构应该能够有效处理来自各种模态（文本、图像、动作等）的异构输入。它必须具有可扩展性(比如增大模型参数性能可以随之提高)，以稳定地在大规模机器人数据上进行训练
挑战二，在之前的研究中，尚没有一个公认的可以在多种机器人数据上训练的方案
如何在异构数据上进行训练？在多机器人数据上进行训练带来了数据异构性的独特挑战。毕竟，不同的机器人在物理结构和动作空间上可能会有很大差异

先前的尝试要么局限于具有相似动作空间的机器人子集（Yang等，2023；Ghosh等，2023；Kim等，2024），要么只保留具有相同结构的输入子集（Collaboration等，2023；Yang等，2024），以牺牲大量信息为代价

如何在如此异构的数据上训练模型仍然很大程度上没有得到解决

2.1.3 扩散建模

由于多模态性，给定语言指令 $l$ 和观测 $\boldsymbol{o}_{t}$ ，可能存在多种可能的动作 $a_{t}$ 来执行任务。如果将策略建模为确定性映射 $\left(\ell, o_{t}\right) \mapsto a_{t}$ 并回归训练数据中的 $\left(\ell, \boldsymbol{o}_{t}, \boldsymbol{a}_{t}\right)$ 元组，策略将学习动作模式的“平均值”。这可能导致分布外的动作，例如多个模式的算术平均值，这可能完全不可行（Pearce等，2023），相反，作者选择建模连续条件分布 $p\left(\boldsymbol{a}_{t} \mid \ell, \boldsymbol{o}_{t}\right)$

由于在各种方法中，扩散模型在表现力和采样质量上均表现出色，但在对高维数据（例如图像）进行采样时可能速度较慢。幸运的是，对于作者的设置来说，这一缺点很小，因为 $\boldsymbol{a}_{t}$ 的维度比图像低得多，只需要最小的采样开销。这使得扩散模型成为策略的理想选择，如Chi等（2023）所示

然而，将扩散模型应用于机器人任务面临独特的挑战

因为机器人物理量（即动作和本体感受）的固有特性与图像/视频数据不同。图像和视频数据虽然是高维的，但通常表现出一定程度的时间和空间连续性（Chen et al., 2019; Liang et al.,2022），帧之间的变化通常是渐进的

相比之下，机器人物理量的特征是其非线性动态（de Wit et al., 2012）以及由于物理交互（如碰撞、约束和阻尼等材料特性）而可能出现的高频变化
此外，这些量还具有不稳定的数值范围，可能是由于不可靠的传感器导致的极端值。这凸显了调整当前扩散模型以有效捕捉机器人数据的不稳定性和非线性的必要性

接下来，将首先详细阐述扩散公式「如果还不了解何谓DDPM的，请看此文：图像生成发展起源：从VAE、VQ-VAE、扩散模型DDPM、DETR到ViT、Swin transformer」，然后介绍作者为解决这些挑战而设计的架构

在使用扩散策略进行决策时，首先从N(0,I)中采样一个完全噪声的动作 $\boldsymbol{a}_{t}^{K} \sim \mathcal{N}(\mathbf{0}, \boldsymbol{I})$ ，然后执行 $K \in \mathbb{N}^{+}$ 去噪步骤，将其去噪为一个干净的动作样本 $a_{t}^{0}$ ，来自 $p\left(\boldsymbol{a}_{t} \mid \ell, \boldsymbol{o}_{t}\right)$ 「When making a decision with diffusion policies, we first sample a totally noisy action aKt ∼N(0, I)and then perform K ∈N+ denoising steps to denoise it to a clean action sample a0t from p(at|ℓ, ot)」

$\boldsymbol{a}_{t}^{k-1}=\frac{\sqrt{\bar{\alpha}^{k-1}} \beta^{k}}{1-\bar{\alpha}^{k}} \boldsymbol{a}_{t}^{0}+\frac{\sqrt{\alpha^{k}}\left(1-\bar{\alpha}^{k-1}\right)}{1-\bar{\alpha}^{k}} \boldsymbol{a}_{t}^{k}+\sigma^{k} \boldsymbol{z}, \quad k=K, \ldots, 1$

其中 $\left\{\alpha^{k}\right\}_{k=1}^{K}$ 、 $\left\{\sigma^{k}\right\}_{k=1}^{K}$ 是由噪声计划预定义的标量系数（Nichol & Dhariwal，2021）
这里， $\beta^{k}:=1-\alpha^{k}$
且如果k >1， $\bar{\alpha}^{k-1}:=\prod_{i=1}^{k-1} \alpha^{i}$ ， $\boldsymbol{z} \sim \mathcal{N}(\mathbf{0}, \boldsymbol{I})$ ，否则 $\bar{\alpha}^{k-1}=1$ ， $\boldsymbol{z}=\mathbf{0}$

然而，在采样完成之前， $\boldsymbol{a}_{t}^{0}$ 是不可求解的。作者选择使用一个可学习的去噪网络 $f_{\boldsymbol{\theta}}$ ，参数为 $\theta$ ，从噪声样本中估计出干净样本： $\boldsymbol{a}_{t}^{0} \leftarrow f_{\boldsymbol{\theta}}\left(\ell, \boldsymbol{o}_{t}, \boldsymbol{a}_{t}^{k}, k\right)$

为了训练这样的网络，可最小化以下去噪的均方误差（MSE）

$\mathcal{L}(\boldsymbol{\theta}):=\operatorname{MSE}\left(\boldsymbol{a}_{t}, f_{\boldsymbol{\theta}}\left(\ell, \boldsymbol{o}_{t}, \sqrt{\bar{\alpha}^{k}} \boldsymbol{a}_{t}+\sqrt{1-\bar{\alpha}^{k}} \boldsymbol{\epsilon}, k\right)\right)$

其中 $k \sim \operatorname{Uniform}(\{1, \ldots, K\})$ ， $\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \boldsymbol{I})$ ，并且 $\left(\ell, \boldsymbol{o}_{t}, \boldsymbol{a}_{t}\right)$ 是他们的训练数据集中采样的

在本文后面，将用 $\tilde{\boldsymbol{a}}_{t}:=\sqrt{\bar{\alpha}^{k}} \boldsymbol{a}_{t}+\sqrt{1-\bar{\alpha}^{k}} \boldsymbol{\epsilon}$ 表示噪声动作输入，其中 $k$ 的上标为了简化而省略

这样，你可以和图像扩展中的这个去噪公式对比对比：

来自此文：图像生成发展起源：从VAE、VQ-VAE、扩散模型DDPM、DETR到ViT、Swin transformer

此外，在实践中，我们倾向于一次性预测一系列动作，即一个动作块，以促进时间一致性（Chi等，2023），并通过减少任务中的决策次数来缓解错误累积（Zhao等，2023）

具体来说，他们建模 $p\left(\boldsymbol{a}_{t: t+T_{a}} \mid \ell, \boldsymbol{o}_{t}\right)$ ，其中 $a_{t: t+T_{a}}:=\left(a_{t}, \ldots, a_{t+T_{a}-1}\right)$ 是一个动作块， $T_{a}$ 表示块的大小（Zhao等，2023）「其实，这里说的就是动作预测算法ACT，详见此文：ACT的原理解析：斯坦福炒虾机器人Moblie Aloha的动作分块算法ACT」

2.2 机器人扩散transformer RDT：解决架构与数据的双重挑战

2.2.1 RDT的模型架构：多模态输入的编码和网络结构

在异构多模态输入的编码上

多模态输入的异构性体现在其结构中；也就是说，每种模态的格式和维度数量显著不同。这给多模态训练带来了挑战

换言之，对于一个具体的机器人任务，模型在收到人类发出的语言指令后，需要结合自己的视觉观察，来预测完成任务所需的机械臂动作action，这里就涉及到了语言、图片和动作三种模态

为了解决这个问题，作者将这些不同的模态编码到一个统一的潜在空间中。以下是编码方法(下图是图3)：

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

除了结构，异质性还体现在不同输入中包含的信息量不同

首先，不同模态的数据包含不同的信息量。例如，图像通常比文本包含更多的信息，编码后会产生更多的token
其次，相同模态的不同输入可能包含非常不同的信息量。例如，机器人的外部摄像头具有更广阔的视野，比腕部摄像头包含更丰富的信息，如下图图3右上角所示
在这种情况下，模型可能会学习到一种捷径：只关注外部视图而忽略腕部视图，从而失去感知深度的能力
为了解决这个问题，作者在编码过程中以一定的概率随机独立地屏蔽每个多模态输入，以防止模型过度依赖特定的输入

说白了，就像人都喜欢看信息量大的图而不喜欢看信息量小的文字，其实模型也一样
故为了避免模型 “偷懒”，只看信息量大的模态，在训练中，研究者会以一定概率随机遮蔽（mask）各个模态

在网络结构上，作者选择Transformer作为可扩展的骨干网络（Bao等,2023；Peebles & Xie,2023），并根据他们机器人问题的特点对Diffusion Transformer（DiT）进行了以下三个关键修改「DiT结合了视觉 transformer 和 diffusion 模型的优点，即DiT = DDPM + ViT，相当于之前一般更多是DDPM + CNN。至于如还不了解什么是DiT的，请参见此文《视频生成Sora的全面解析：从AI绘画、ViT到ViViT、TECO、DiT、VDT、NaViT等》的第2.4节DiT(含U-ViT)：将扩散过程中的U-Net 换成ViT(2D图像生成，带文本条件融合)」：

QKNorm和RMSNorm
输入机器人物理量的不稳定数值范围可能导致梯度不稳定和数值溢出等问题，尤其是在训练大型基础模型时
为了解决这个问题，作者添加了QKNorm（Henry等,2020）以避免计算注意力时的数值不稳定

此外，他们还注意到，他们的问题可以被视为一个时间序列预测任务，原始DiT中的LayerNorm的中心化操作可能导致token偏移和注意力偏移，从而破坏时间序列的对称性（Huang等,2024）
因此，他们用没有中心化操作的RMSNorm（Zhang & Sennrich,2019）替换了LayerNorm。图4a显示，如果没有这种修改，大规模预训练往往会非常不稳定，甚至爆炸
说白了，就是由于传感器失灵等原因，机器人数据中往往会出现极端值。这种极端值可能导致梯度不稳定和数值溢出等问题。研究者采用更加先进的 QKNorm 和 RMSNorm 来进行缓解
MLP解码器
为了提高对非线性机器人动作的逼近能力，作者用非线性MLP解码器替换了最终的线性解码器，将其作为从潜在空间回到物理空间的投影「原因是机器人的动作往往符合非线性动力学的物理规律」

如下图图4b实证所示，没有这个设计，RDT无法有效捕捉非线性动态，因此失去了完成需要精细操作的灵巧任务的能力
交替条件注入（ACI）
在作者的模型中，图像和语言输入作为条件，这些条件是高维且长度可变的，与传统DiTs（Xie等，2020）中的类别标签条件形成对比。这些信息丰富的条件很难压缩成单一的标记，使得原始的自适应层规范方法不适用

因此，作者采用交叉注意力来适应不同长度的条件，避免在进一步压缩中信息丢失
此外，作者进一步分析，由于图像token通常比文本token多，同时注入两种模态往往会掩盖与文本相关的信息「说白了，图像的维度通常远高于文本的维度。同时将这两种模态注入到主干网络中时，往往图像会淹没文本」，从而削弱模型的指令遵循能力（见上图图4b的定量结果）

为此，作者在连续层的交叉注意力中策略性地交替注入图像和文本token，而不是在每一层中同时注入这两种token

2.2.2 数据：预训练数据与微调数据

首先，为了在异构多机器人数据上进行训练，他们需要一个在各种机器人之间共享的统一动作空间，这可以为多机器人动作提供统一的格式

机器人从原始动作空间到统一动作空间的映射应该具有物理可解释性，每个维度都应该有明确的物理意义。这可以鼓励模型从不同的机器人数据中学习共享的物理规律，从而提高从不同机器人数据中学习的效率（Shah et al., 2023a）——说白了，为了在多种机器人数据上进行预训练，故需要整合不同机器人现有的训练数据，故需要对数据格式进行统一

对此，作者构建了一个统一的动作空间，该空间的设计包括两个步骤

首先，对于每个机器人，可以使用一个单一的空间来容纳其本体感觉 $z_{t}$ 和动作 $a_{t}$ 。这是因为 $a_{t}$ 通常是所需 $\boldsymbol{z}_{t+1}$ 的一个子集（de Wit等，2012；Kouvaritakis & Cannon，2016），因此 $z_{t}$ 的空间自然包含 $a_{t}$ 的空间
其次，作者设计了一个统一的空间，涵盖了大多数带有夹持臂机器人的主要物理量，如下图图3左侧所示

通过根据其物理意义将原始动作向量的每个元素填充到统一动作空间向量的相应位置，将机器人的动作空间嵌入到这个统一空间中，剩余的位置则进行填充

空间的具体定义见附录C

其次，对于预训练数据，在这个统一的空间下，他们能够在几乎所有现代带夹持臂的机器人数据上对RDT进行预训练，并极大地扩展数据规模以满足基础模型的需求

具体来说，我们的预训练数据集集合包括46个各种机器人的数据集，总规模超过1M+条轨迹和21TB。更多细节和预处理推迟到附录D

最后，收集全面的多任务双手数据集

尽管已在大规模数据集上进行了预训练，但由于体现差距，RDT在零样本泛化到目标双臂机器人上仍可能需要帮助。为了弥合这一差距，他们需要在目标机器人上收集多任务双手数据集以进行微调

最近在大型语言模型方面的进展（Ziegler等，2019；Brown等，2020；Touvron等，2023）表明，高质量的微调数据集对于模型性能至关重要

他们从三个方面确保数据集的高质量：

在数量方面，已收集了6000多个轨迹，使我们的数据集成为当今最大的双手数据集之一；
在全面性方面，考虑了300多个具有挑战性的任务，涵盖了大多数操作任务类型，从拾取和放置到插入电缆，甚至包括书写数学方程式；
在多样性方面，准备了100多种具有不同尺寸和纹理的刚性和非刚性物体，以及15多个具有不同照明条件的不同房间
此外，他们进一步利用GPT-4-Turbo（Achiam等，2023）重写人工标注的指令以增加文本多样性

// 待更