当前位置：首页 > news >正文

清华双臂机器人扩散大模型RDT：先预训练后微调，支持语言、图像、动作多种输入(1B参数)

news 2024/11/5 11:32:41

前言

通过上文介绍的GR2，我们看到了视频生成模型在机器人训练中的应用

无独有偶，和GR2差不多一个时期出来的清华RDT，其模型架构便基于视频生成架构DiT改造而成(当然，该清华团队其实也在DiT之前推出了U-ViT，具体下文会详述，总之，该清华团队的这两个工作——U-ViT和RDT确实不错，也把我之前研究过的视频生成结合起来了)

且RDT也是先预训练后微调，且它们都支持语言指令、图像序列以及机器人状态序列等多种输入

另，值得一提的是，我个人认为，这两个模型算是真正开启了「机器人大模型」的路线，比如

GR2的参数规模为230M(当然了，严格意义上讲，也没有多大，和BERT的两个版本的大小——110M/1.1亿、340M/3.4亿差不太多)，相当于2.3亿参数
而RDT则是1B——10亿参数，而上了1B/10亿规模的，在机器人领域确实是可以称之为大模型了
顺带做下预告，本博客的下一篇是一个3B的机器人大模型打通7种不同品牌的机械臂

总之，这几个工作宣告机器人真正进入大模型时代：一个大模型大一统的时代

最后，友情提醒一下本博客的读者

我的解读一般不会严格按照原论文的行文顺序、结构编排，而是从最大化提高解读的逻辑性、可读性、清晰性而出发
且行文标题上也不会直接照搬原论文的各节标题，而是为了更精准概括每一部分的核心痛点，比如到底要解决什么问题、如何解决对应问题等等..

这一切的目的是，让解读相比读原论文，更通俗易懂、能更快理解，这也是解读的价值与意义所在

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

1.1 什么是RDT

1.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看到了视频生成模型在机器人领域的应用，且在本文你将再次看到
其首先寻求在广泛的多机器人数据集上进行预训练，然后在目标双臂机器人收集的小数据集上进行微调

1.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)$
注意
$\rightarrow$    $\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观测历史
$\rightarrow$    $\boldsymbol{z}_{t}$ 是机器人的低维本体感知
$\rightarrow$    $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），以牺牲大量信息为代价

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

1.1.3 对模型架构的探索：扩散建模(但数据不稳定且非线性)且ACT预测动作序列

由于多模态性，给定语言指令 $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)$ 「The policy will learn the “average”of action modes if we model it as a deterministic mapping (ℓ, ot) �→at and regress the tuples of(ℓ, ot, at) in the training data. This may result in out-of-distribution actions, such as the arithmeticmean of multiple modes, which can be completely infeasible (Pearce et al., 2023). Instead, we chooseto model the continuous conditional distribution p(at|ℓ, ot)」
由于在各种方法中，扩散模型在表现力和采样质量上均表现出色，但在对高维数据（例如图像）进行采样时可能速度较慢。幸运的是，对于机器人的动作训练来说，这一缺点很小，因为 $\boldsymbol{a}_{t}$ 的维度比图像低得多，只需要最小的采样开销。这使得扩散模型成为策略的理想选择，如Chi等（即2023发表的： Diffusion policy: Visuomotor policy learning via action diffusion）所示
among various approaches, diffusion models excel in both expressiveness and sampling quality, but can beslow to sample a high-dimensional data (e.g., images). Luckily, for our settings, the drawback is minor since that a_t has a much lower dimension than images, which requires only minimal samplingoverhead.

至于扩散策略的介绍，详见此文《UMI——斯坦福刷盘机器人：从手持夹持器到动作预测Diffusion Policy(含代码解读)》的第三部分 Diffusion Policy：基于CNN或Transformer

然而，将扩散模型应用于机器人任务面临独特的挑战，一个是机器人数据的不稳定性、一个是机器人数据的非线性特征，而这也是如果用扩散模型作为RDT架构所需要解决的两大问题

具体而言

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

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

接下来，将首先详细阐述扩散公式「如果还不了解何谓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）「Besides, in practice, we prefer to predict a sequence of actions, i.e., an action chunk, in one shot to encourage temporal consistency (Chi et al., 2023) 」

当然了，该Diffusion policy(Chi et al., 2023)论文中确实有说道：“一个有效的动作制定应该鼓励在长期规划中的时间一致性和平滑性，同时允许对意外观察做出迅速反应
为了实现这一目标，应该在重新规划之前，采用扩散模型生成的固定持续时间的行动序列预测 ”
详见上文提到过的此文第三部分：《UMI——斯坦福刷盘机器人：从手持夹持器到动作预测Diffusion Policy(含代码解读)》
并通过减少任务中的决策次数来缓解错误累积（Zhao等，2023）「 and to alleviate error accumulation over time by reducing number of decisions in a task (Zhao et al., 2023)」
其中，「Zhao等，2023」这篇论文即值得是动作分块预测算法ACT 「关于动作分块预测算法ACT，详见此文：ACT的原理解析：斯坦福炒虾机器人Moblie Aloha的动作分块算法ACT」

具体来说，他们建模 $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）

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

1.2.1 解决架构的问题：先解决多模态输入的编码、再解决网络结构(解决不稳定性和非线性特征)

首先，一个网络得先确定其输入，故先考虑异构多模态输入的编码问题

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

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

为了解决这个问题，作者将这些不同的模态编码到一个统一的潜在空间中(他们使用适配器将每种模态的token维度对齐到2048，且所有多模态编码器的适配器均使用GeLU激活)

具体而言，以下是编码方法(下图是图3)：

第一种输入模态，低维输入是表示机器人物理量的低维向量，包括本体感知、动作块和控制频率
为了对它们进行编码，使用「带有傅里叶特征（Tancik等人，2020））」的MLP，可以有效捕捉低维空间中的高频变化

具体而言
首先，将本体感知 $\boldsymbol{z}_{t}$ ，和噪声动作块 $\tilde{\boldsymbol{a}}_{t: t+T_{a}}$ 嵌入到统一的动作空间。该空间用于统一不同机器人中 $\boldsymbol{z}_{t}$ 和 $\tilde{\boldsymbol{a}}_{t: t+T_{a}}$ 的表示「在原论文第19页附录B中，对于多模态输入的编码中，其中在低维输入——Low-Dimensional Inputs上，The proprioception zt and the noisy action chunk ˜at:t+Ta are firstembedded into the unified action space. This space is used to unify the representation of zt and˜at:t+Ta across various robots, which is elaborated in Sec. 4.2.」
1，为了在异构多机器人数据上进行训练，需要一个在各种机器人之间共享的统一动作空间，这可以为多机器人动作提供统一的格式

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

2，具体怎么做呢？分为两个步骤：
步骤1 对于每个机器人，可以使用一个单一的空间来容纳其本体感觉 $z_{t}$ 和动作 $a_{t}$ 。这是因为 $a_{t}$ 通常是所需 $\boldsymbol{z}_{t+1}$ 的一个子集（de Wit等，2012；Kouvaritakis & Cannon，2016），因此 $z_{t}$ 的空间自然包含 $a_{t}$ 的空间
步骤2 设计一个统一的空间，以涵盖大多数带有夹持臂机器人的主要物理量。怎么做到呢？很简单，通过根据其物理意义将原始动作向量的每个元素填充到统一动作空间向量的相应位置，将机器人的动作空间嵌入到这个统一空间中，剩余的位置则进行填充

3，补充一下对这个空间的定义
为了将各种机器人的动作嵌入到一个包含所有主要物理量的统一空间中，先将这个统一动作空间的维度是128
下表表4描述了该统一动作空间中向量的每个元素。对于一个特定的机器人，原始动作向量的每个元素根据其物理含义填充到统一动作向量的相应位置，剩余位置则填充「对于单臂机器人，其手臂映射为“右”臂。对于只有6自由度的机器人手臂，其关节位置将填入对应的10个位置中的前6个。其他物理量也是如此」

如此，构建了一个统一的动作空间之后，便可以在异构多机器人数据上进行训练

然后，由于它们具有相似的物理意义，它们通过一个共享的MLP被编码到token空间中。这种连续编码可以避免离散编码所导致的精度损失（Brohan等，2022；2023；Kim等，2024）
对于频率 $c$ 「之所以将控制频率输入模型，是为了在解决数据集中不同控制频率带来的挑战时，使模型在做出决策时能够考虑到这种变化」，以及扩散时间步 $k$ ，作者分别通过两个MLP将它们编码到token空间中
之后，所有这些在长度方向上被拼接在一起，以实现上下文条件化（Peebles & Xie,2023；Bao等，2023），最终得到长度为 $1+T_{a}+1+1$ 的输入token序列
最后，添加位置嵌入以区分不同的模态并注入 $\tilde{\boldsymbol{a}}_{t: t+T_{a}}$ 中的时间信息

在预处理上，通过统一物理量的单位（米、弧度、米/秒、弧度/秒等），而不是像之前的工作（Chi et al., 2023; Ghosh et al., 2023）那样严格归一化到[−1,1]或N(0,1)，来大致对齐各种数据集的尺度。例如，不同数据集中的“1（米）”对应相同的现实世界长度。重新缩放物理量会破坏这种共享属性，从而削弱模型在不同机器人间的迁移能力
且他们还采用6D表示法（Zhou et al., 2019）来解决EEF旋转中的万向节锁问题。在选择 $T_{a}=64$ (相当于动作块大小)之前，参考了Zhao等人（2023）的先前消融研究，并在性能和计算开销之间取得了平衡。此外，历史本体感知 $\boldsymbol{z}_{i}, i<t$ 被排除在外，以防止模型仅通过低维输入学习捷径，从而固定在固定的运动模式上。相反，他们鼓励模型从高维图像特征中学习可推广的决策结构
第二种输入模态，图像输入是高维的(图像token空间维度为1152，RDT的token空间维度为2048)，包含丰富的空间和语义信息
为了提取紧凑的表示，作者使用了一个图像-文本对齐的预训练视觉编码器：SigLIP（Zhai等，2023），且作者在训练过程中固定其权重以节省GPU内存
具体而言，使用冻结的 SigLIP (Zhai et al., 2023) 对 RGB 图像进行编码，并利用额外的 MLP 将输出投射到token空间
为了增强模型根据视角和时间步区分图像的能力，作者将传统的正弦位置嵌入扩展到多维网格，如上图图 3 右侧所示

此修改集成了时空信息，使模型能够捕捉输入图像之间的关系。具体来说，作者采用 Liu et al. (2022) 的实现，使用网格维度 $\left(T_{\mathrm{img}}, N_{\mathrm{cam}}, N_{\text {patch }}, D\right)$
这里，Ncam 代表相机的数量，在他们的配置中设置为三个，Npatch 表示每幅图像被 ViT 基于图像编码器划分的补丁数量，D 表示嵌入维度

在预处理上，他们采用固定长度的图像输入策略。为所有机器人固定图像输入顺序和格式，共有三个视图：一个静态外部视图，一个右手腕视图和一个左手腕视图，这被认为足以满足大多数双手任务的要求
他们将单臂机器人的手腕摄像头视为右手腕摄像头，并用背景色填充不可用的视图。当输入模型时，每张图像被填充成一个正方形并调整为384×384的大小，保持其原始纵横比
此外，选择 $T_{\mathrm{img}}=2$ (相当于图像历史大小)，因为两帧的历史长度对于大多数情况来说已经足够，在效率和性能之间达到了平衡（Ghosh et al.,2023;Wu et al.,2024）
最后，可以将图像输入写为 $\boldsymbol{X}_{t-1: t+1}:=\left(\left\{\boldsymbol{X}_{t-1}^{1}, \boldsymbol{X}_{t-1}^{2}, \boldsymbol{X}_{t-1}^{3}\right\},\left\{\boldsymbol{X}_{t}^{1}, \boldsymbol{X}_{t}^{2}, \boldsymbol{X}_{t}^{3}\right\}\right)$
第三种输入模态，语言输入(语言token的空间维度为4096)的长度各异且高度抽象，由于其复杂性和模糊性，带来了集成挑战
为了对其进行编码，作者使用了基于Transformer的预训练语言模型T5-XXL（Raffel等，2020），且作者还在训练过程中固定其权重以节省GPU内存

具体而言，语言指令由一个冻结的T5-XXL（Raffel等，2020）进行编码，并使用一个MLP将输出投射到词元空间。在计算语言词元的注意力时，作者应用语言注意力掩码以掩盖在批处理过程中附加的填充词元

在预处理上，会针对这部分对原始文本进行简单清理，例如去除非法字符和多余空格，将句子开头大写，并在句末添加句号，且保留文本为可变长度

最终，在LDiT块之后，再对输出进行归一化——即Norm，并通过MLP解码器将其投射回动作空间(如下图图3中部下部分所示)

此外，异质性还体现在不同输入中包含的信息量不同

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

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

其次，在网络结构上，作者选择Transformer作为可扩展的骨干网络（Bao等,2023；Peebles & Xie,2023），并根据他们机器人问题的特点对Diffusion Transformer（DiT）进行了三个关键修改

继续行文之前，有两个重要的问题需要先解释下

首先，什么是DiT呢？DiT结合了视觉 transformer 和 diffusion 模型的优点，即DiT = DDPM + ViT，相当于之前一般更多是DDPM + CNN

更具体的请参见此文《视频生成Sora的全面解析：从AI绘画、ViT到ViViT、TECO、DiT、VDT、NaViT等》的第2.4节DiT(含U-ViT)：将扩散过程中的U-Net 换成ViT(2D图像生成，带文本条件融合)
为何要针对DiT做三个关键修改呢，原因很简单，在上文「2.1.3 对模型架构的探索：扩散建模(但数据不稳定且非线性)且ACT预测动作序列」中早已说了：
“将扩散模型应用于机器人任务面临独特的挑战，一个是机器人数据的不稳定性、一个是机器人数据的非线性特征，而这也是如果用扩散模型作为RDT架构所需要解决的两大问题”

那具体针对DiT做了哪三大修改呢？具体如下

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

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

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

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

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

1.2.2 解决数据的问题：有了数据后则先预训练然后微调

首先，对于预训练数据，通过上文1.2.1节开头构建的机器人统一动作空间之下，他们能够在几乎所有现代带夹持臂的机器人数据上对RDT进行预训练，并极大地扩展数据规模以满足基础模型的需求

具体来说，他们的预训练数据集集合包括46个各种机器人的数据集，总规模超过1M+条轨迹和21TB

他们的预训练数据集集合包括46个数据集，总规模超过1M+的轨迹和21TB，是迄今为止最大的机器人数据集预训练集合

下表表5展示了他们预训练数据集的完整列表及其采样权重

他们分配了一个初始权重为 p为每个数据集分配权重，并根据每个数据集的多样性和质量进行调整。与线性加权相比，这种方法防止了对大数据集的过度采样，同时确保小数据集得到充分采样，从而提高了每个小批次预训练样本的多样性

在预训练阶段，进一步观察并调整了不同数据集的权重，基于其中间损失结果。且增加了那些收敛较慢的数据集的权重

具体而言，他们所用到的主要的数据集如下：

RT-1数据集（Brohan等，2022）是一个大型多样化数据集，包括130K条轨迹，涉及多个任务、对象和环境。它是从13种不同的实施体中收集的，每种实施体都配备一个外部RGB摄像头
动作空间包括6D末端执行器（EEF）、夹爪打开和基座位移，控制频率为3Hz
DROID(Khazatsky et al.,2024) 是一个大规模多任务数据集，包含76K条轨迹和564个场景。它通过远程操作一个Franka Panda7-DoF机器人手臂收集，配备了腕部和外部RGB-D摄像头
动作空间包括7-DoF关节位置和一个夹爪宽度，而本体感知还包括6D EEF，控制频率为15Hz
RH20T（Fang等，2023）是一个全面的数据集，涵盖了110K的轨迹和140个任务。它包括四种不同的机器人化身和三种不同的相机视角，以10Hz的频率采样。它还包括长短任务
其状态空间是6自由度和7自由度关节位置的混合，并且具有第三人称视角的RGB-D相机
移动ALOHA数据集（Fu等人，2024）是一个双手数据集，包含由移动ALOHA机器人收集的1000多个轨迹。其状态空间包括基座运动和双手的14维关节位置，以及三到四个第一人称视角摄像头
其部分数据包括广泛的视角变化和基座运动，最初适合用于模仿学习
其他数据集。其他数据来自RH20T（Fang等，2023），RoboSet（Kumar等，2024），BridgeData V2（Walke等，2023），以及Open X-Embodiment（Collaboration等，2023）
它们大多数具有不同的机器人形态和摄像头观测，增强了他们预训练数据集的异质性和多样性

在数据清洗上，其中重复的情节和失败的情节被排除，以确保预训练数据集的质量。他们移除空白图像，排除错误记录的速度，并过滤掉过短的轨迹。对于过长的轨迹，将进行下采样以避免不公平性

其次，收集全面的多任务双手数据集好做微调

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

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

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

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

微调数据集是使用Mobile ALOHA机器人（Fu等，2024）创建的，包括300多个任务、6K+轨迹和3M+帧。它也是迄今为止最大的开源多任务双手机器人数据集之一

下图图6总结了该数据集「简言之，该数据集包含这几大特性：（1）多样化的物体和场景。该数据集包含在不同场景和条件下操控的具有不同属性的物体。（2）挑战性任务。该的数据集结合了灵巧的操作、语言和视觉理解以及双手任务。（3）多模态特性。该数据集标注了丰富的多模态数据，包括三视图RGB相机、关节信息和增强指令」

且RDT的作者从开源的Songling数据集（Wang等，2024）中借用了3个任务（总共140集）

最终，他们的微调数据集具备以下几个特征

多模态特征
他们使用三个RGB摄像头收集数据集，摄像头分别位于前方以及左右夹持器上。我们记录了双臂6自由度关节的位置和速度，以及夹持器的角度

且为每个任务手动标注了指令。为了进一步扩充他们的指令并使其与预训练数据集对齐，他们使用GPT-4-Turbo（Achiam等，2023）为每个任务生成了100条扩展指令和一条简化指令。这些多模态信息进一步增强了数据集的丰富性和质量
多样化的物体和场景
他们的数据集包含多样化的任务和场景，涵盖了300多个任务，包括拾取、插入、书写、推动和拉动等技能。它具有100多个具有不同尺寸和纹理的刚性和非刚性物体
且他们在15个以上的场景中收集数据集，并在每个任务的数据收集过程中引入随机性，例如改变物体和机器人的初始位置

更为了进一步增加多样性，还添加了随机光照条件。例如，在正常光照和变色条件下都进行了倒水操作。这些措施进一步增强了该数据集的多样性
挑战性任务
他们还考虑了各种挑战性任务，包括灵巧操作，例如从塑料瓶上拧开瓶盖，以及理解任务，例如用字母积木拼写“爱”

此外，数据集还包括结合了灵巧性和理解力的任务，例如在白板上解数学方程。此外，他们的数据集还包含双手任务，例如将充电线插入手机
这些复杂且高质量的任务进一步增强了模型的下游理解能力和通用性

1.3 实验及与其他基线的对比

1.3.1 实验设置：模型训练和推理

在模型训练和推理上，作者将RDT的规模扩大到12亿参数，确立其为目前最大的基于扩散的机器人基础模型。该模型在48台H100 80GB GPU上进行了一个月的预训练，总共进行了100万次训练迭代步骤。使用相同的GPU进行13万步的微调(历时三天)

对于实时推理，作者采用了DPM-Solver++（Lu等，2022），这是一种最近的扩散模型采样加速器。它可以将采样一个动作块所需的扩散步骤从100步减少到5步，实现了6Hz（每秒动作块）的动作块推理频率和381Hz（每秒动作）的平均动作推理频率，在目标机器人的机载RTX4090 24GB GPU上

且他们在预训练和微调阶段使用AdamW优化器（Adam等，2019），采用恒定学习率调度器和下表表10中的超参数

模型在48个H100 80GB GPU上分别进行1M步和130K步的预训练和微调。由于调度原因，并没有从1M预训练检查点开始微调，而是选择了500K检查点

在训练阶段，使用DDPM调度器和滑动余弦调度器（即，squaredcos cap v2）和1000步。在采样阶段，使用DPM-Solver++（Lu等，2022）与滑动余弦调度器，采样步数为5
在微调期间，还过滤掉长度低于32的片段，并将长度超过2048的片段下采样至2048

1.3.2 与基线的对比

在与基线的对比上，作者为了全面评估RDT，他们考虑了机器人基础模型和双手操作中最先进的基线，包括使用transformer的动作分块（ACT）（Zhao et al.,2023）、OpenVLA（Kimet al.,2024）和Octo（Ghosh et al.,2023）

ACT如上文所述，是一种在双手操作中使用VAE建模动作分布的最先进方法
直接采用了与原始论文（Fu et al.,2024）中相同的ACT架构和超参数，除了下表表11中的超参数
他们用90%的6K微调集训练了ACT共8000个epoch（总共约8天），其余10%作为验证集。根据验证集中的最佳表现，将epoch 5413时的检查点作为最终结果
OpenVLA是最大的开源基础模型（7B），采用离散化建模
Octo是基于扩散的基础模型，其最大版本只有93M参数

在消融研究上，作者进行了关于模型大小、预训练和建模方法的消融研究，以了解它们的重要性

他们考虑了以下变体：

RDT（我们的）：原始RDT
RDT（回归）：没有扩散建模的RDT，它对确定性映射 $\left(\ell, o_{t}\right) \mapsto a_{t}$ 进行建模
RDT（小型）：没有大参数的RDT，它只有166M参数
RDT（从头开始）：没有预训练的RDT。在微调期间从头开始训练

下表表2中，他们从泛化能力的三个维度评估这些变体

下表表7则提供了不同RDT变体及基准的比较

从下表表3的结果中可以看出来

RDT始终优于其他基线，这是因为

RDT采用了强大的网络架构进行扩散，以准确建模多模态动作的分布，而离散化和VAE分别缺乏准确性和表现力
此外，大规模预训练后大量的参数提供了丰富的先验知识，这显著提高了泛化能力

第二部分 RDT的训练细节、硬件配置、实验详情

2.1 RDT的更多训练细节：软件平台、数据增强、微调细节

2.1.1 平台

他们使用 Pytorch（Paszke 等，2019）和 DeepSpeed（Rasley 等，2020）来促进并行训练，并使用 TensorFlow Dataset（TFD）的生产者-消费者框架来实现快速数据加载

由于 Open X-Embodiment（Collaboration 等，2023）中的大多数数据集是以TFRecord 的形式存储的，他们将所有预训练数据集转换为 TFRecord 进行存储
在预训练中，使用生产者进程从 TFRecord 解压数据并将其存储在硬盘的缓冲区中。同时，使用消费者进程以无序方式从缓冲区读取数据并将其提供给模型训练。这不仅解耦了 TensorFlow（Abadi 等，2015）和 PyTorch 环境，还缓解了由于内存中混洗缓冲区尺寸较小而导致的训练性能损失
在微调阶段，由于数据集相对较小，他们还实现了使用 HDF5 数据集进行存储的数据读取管道

2.1.2 填充动作和本体感受

为了将一个特定的机器人动作嵌入128维的统一动作空间，作者需要填充不可用的动作元素。通常的做法是用0值或特定值进行填充。但“0”实际上具有物理意义。例如，速度为“0”通常表示相对于地面静止。这可能会让模型感到困惑：“0”是代表静止还是填充值？

为了解决这个问题，他们将动作和本体感受与一个0-1向量连接，指示每个维度在编码到token空间之前是否被填充，最终形成一个256维的向量。这可以补充缺失的可用性信息并消除混淆

2.1.3 检查训练过程

在训练期间，每隔固定时间段，进行一次扩散采样，并将采样的动作与训练数据集的真实值进行比较。根据经验，发现两者的均方误差（MSE）与机器人部署性能之间存在普遍的正相关关系

这个观察使他们能够轻松监控模型的训练进度。当这个MSE收敛时，通常可以停止训练。且注意到，过低的MSE也可能意味着过拟合

2.1.4 数据增强与一些微调细节

过拟合是训练大型神经模型时的一个常见挑战，特别是在微调阶段。他们利用数据增强技术来解决这个问题

进行图像增强，包括颜色抖动和图像损坏，并在输入本体感受中添加高斯噪声，信噪比（SNR）为40dB
如上文所述，他们还使用了GPT-4-Turbo来增强和扩展语言指令

以下是一些微调细节

在微调期间，他们移除了每个片段开头的一个静态部分，这可能是由于操作员在录制开始后没有反应造成的
他们的语言指令是从原始手动注释的指令、扩展指令和简化指令中以三分之一的概率抽取的。当抽取到扩展指令时，从对应任务的100个扩展指令中均匀抽取一个

他们没有应用无分类器指导（CFG），因为他们发现这并没有提高模型的性能，反而带来了不稳定的机器人手臂行为