前言
大约有三年的时间没有更新过简书的文章了, 最近随着人型机器人比较火热并伴随着大模型的成熟, 具身智能技术也被业界关注着, 当然我也转向投入具身智能算法. 这段时间我也将自己看到的文章以及对应的代码进行进一步梳理, 一方面为了更好的巩固自己的模型理解, 同时也方便自己日后的复习.
今天我想写的文章是RDT-1B, 号称具身国产之光, 刘松铭博士所在的团队打造.老规矩, 这次还是先从论文进行梳理, 之后会详细分析代码. 这次我也贴上大佬介绍模型的两个视频, 让我对模型理解会深刻一点.
论文地址: https://arxiv.org/pdf/2410.07864
代码地址: https://github.com/thu-ml/RoboticsDiffusionTransformer
视频讲解: RDT-1B:解密全球最大的双臂机器人扩散大模型、 讲座 | 双臂机器人操作扩散大模型RDT——清华大学TSAIL团队博士生刘松铭
一、 论文阅读
文章还是说重点, 因为纯直白翻译直接用翻译软件即可, 且无法抓住稳重重点.
1. 引言
这篇文章才用了扩散模型的思路构建了模型, 且该模型数据有6k+个episode个数据, 该模型大小有1.2B参数量模型, 模型只需要学习1-5次的验收即可学习新的技能. 这个模型的优势是整个DiT骨干上进行扩散过程,相比pi0仅仅在action expert进行扩散,能够统一动作空间。劣势是没办法做多模态推理,对语言指令的语义泛化能力没有使用VLM的pi0好. 目前主要数据集还是单臂操作, 但是双臂场景的数据集较为稀缺, 因此作者考虑在广泛的多种机器人的数据集上进行预训练,之后再从目标的双臂机器人上收集的小数据集上进行微调, 有助于讲数据集量级别进行扩增, 因为我们知道机器人数据的异质化比较严重, 如何充足合理利用好这些数据进行模型训练是一个非常重要的过程.
这篇文章发明的模型叫做RDT, 全名是RoboticsDiffusionTransformer, 号称是最⼤的具有很强泛化性的双机械臂基础模型, 其主要使用了DiT(Diffusion Transformer)作为backbone(这里后面也会进行详细解释), 这使得其具备对复杂动作分布的建模能力,精心设计了多模态编码器,以消除模态间的异质性。以下是作者的设计思路:
-
表达性: 通过扩散模型建模捕获海量数据完整模态(请参考2.1部分) -
扩展性: 设计多模态编码解决模态之间的异构性 -
表征机器人数据的非线性动力、高频以及非稳定数值变化特性: 对原始的DiT结构进⾏修改, 包括MLP解码、改进归一化方式以及文本和视觉模态交替条件注入. -
数据异质性: 设计物理可解释的统⼀动作空间, 这是⼀种适⽤于各种带有抓⼿机器⼈的统⼀动作格式. 其创新减少了不同机器⼈之间的潜在冲突,同时保留了原始动作的物理意义, 促进模型在不同机器⼈数据集中学习可泛化的物理知识.
通过上述方面的改进, 使得模型在各种具有挑战性的任务中,提高了56%的成功率,超过了基准。在zero-shot, 甚至few-shot(1-5)有很强的能力, RDT还能够完成需要精细操作的任务,例如用操纵杆控制机器人狗等。 同时证明增加扩散模型的大小以及数据量都可以提高模型的效果.
2. 任务描述

这里还是一开始就放上了模型框架, 为了让读者可以第一眼就可以一目了然的知道整体模型的任务大概以及模型的大致思想. 这里为了在硬件上评估模型, 因此作者选择了在ALOHA双臂机机器人, 毕竟这种机器人还是具有很强的代表性.下面我们看下任务的整体pipeline. 可以看到上图左侧, 也是论文贡献处之一, 会将不同的异构空间的机器人会被编码到
统一的动作空间. 这里的输入包含感知信号首先我们会下发语言指令, 希望机器人可以完成什么动作. 希望该动作在t的时间范围内, 完成整体动作
来控制机械臂完成
的目标. 机器人不仅仅接受语言指令的输入, 同时也会输入观察信号, 这里我们称观察信号为
, 该信号是由三元素构成即
组成. 这里的
的视觉RGB图片观测值, 这里的时间段周期为
, 作者设定
, 即选择前后两帧视觉图片. 同时选择三个视角的图片作为输入分别为
、
、
, 分别代表外部相机图片, 右腕图片以及左腕图片. 可以看到上图, 共有三个视角的图片, 每个视角的图片共有2张图片. 作者这里说
往往是
的子集, 为什么这么说呢? 因为z_{t}代表当前时间段的机械臂的位置, a{t}则是模型预测出来的动作, 但是往往需要考虑机械臂会有一些机械扰动因此
往往是
的子集.
双手操作中的一项特定任务通常由多个要素组成:
一项技能(例如,“拿起” 或 “擦拭” 等动词)一个物体(例如,“瓶子” 或 “桌子” 等名词)一个场景(即任务发生的环境)描述技能如何执行的方式(例如,“用左手拿起瓶子” 等状语)
3. 问题与挑战

该任务就是我们分别输入语言指令(即告诉机器人我需要执行什么样的指令)以及对应的视觉输入(这里的输入一般来来自于两个机械别夹爪上的摄像头以及前段摄像头 这里作者不仅采用了一帧的数据而是采用了两帧的数据), 当前机械臂的状态以及对应的机械臂的运动控制频率, 模型需要输出机器执行的
一段动作. 这就好比我们去拿一个东西我们不仅要有意识我们要做什么(语言指令), 同时需要我们的双眼观察我们目前物体在什么地方(摄像头观测到的视觉信息), 同时很关键一点就是目前我们的双臂是什么一个状态, 即我们的双臂在空间什么位置(当前机械臂的状态), 才能更好的去预测我们下一步的双臂应该如何去移动完成我们目标的执行. 那我们为什么还需要输入机械臂的控制频率呢? 因为我们的机器人数据来自不同的机械臂, 这里我们要考虑到数据的异质化问题, 因此这里我认为输入机械臂的控制频率为了让模型更好的适应不同的机械别的属性, 以便更精准的进行机械臂的控制 . 下面有两个细节, 作者进行了回答:
3.1 为什么要加入2帧的输入信息?
我们可以看到下面的动作执行过程, 整体过程的意思是用右机械臂夹住物体, 并上下晃动, 接着将机械臂将物体放回原位. 但是在整个过程通过视觉发现动作初始状态和末尾状态是一致的, 如果仅用单帧的信息是无法让模型知道该动作是step1的抓取, 还是step3的放回夹爪. 同时在机器人运转过程中会有实效情况的出现, 因此加入历史信息使得机器人运动更加鲁棒
3.2 为什么不要使用回归模型(挑战一)?

但我们从上面的途中可以看出, 机器双臂进行物体抓取有很多种方法, 尤其是双臂的抓取比单臂的抓取复杂度上升好几个级别, 同时我们也可以看到抓取一个物体有很多种方式, 可以单臂抓取也可以双臂抓取, 可以正着抓取, 也可以侧着或者上面抓取. 这么多方式都是可以得到一个正确的结果, 这样的过程会使得模型非常难以学习, 如果仅仅用一个MSE的类似回归方法就会使得模型学到上图各个峰的平均值(每个峰代表的是一种方法), 但是这种平均值有可能是不可行的. 作者在他的视频中讲座 | 双臂机器人操作扩散大模型RDT——清华大学TSAIL团队博士生刘松铭 说了一个非常好的例子, 根据右上图我们可以看到, 我们在遇到之间灰色圆形障碍物的时候无论左边绕过去还是右边绕过去都可以很好的绕开障碍物进行前进, 但是如果模型学出来各50%的概率进行回归的话, 很有可能是的模型学出来仅仅就是直线前进导致最终路线垂直撞上对应的障碍物, 导致任务失败, 因此需要设计条件模型, 想到的是主流的扩散模型去解决上述问题.
3.3 如何处理缺乏的数据(挑战二)
因为基于模仿学习收集数据但是由于高费用的硬件,缺乏共同需求的底座模型, 所以作者采用pretrain+finetune的方式获取大量单臂的数据集经验, 在进行下游任务fintune, 这得益于单臂机械臂的数量多.作者设计了统⼀各种机器⼈动作空间, 并实现多机器⼈预训练. 作者收集了⼀个全⾯的多任务双⼿数据集进⾏微调,以提⾼RDT的双⼿作能⼒.
针对挑战一作者设计了扩散相关模型进行解决, 针对挑战二作者设计了机器人统一运动空间进行解决.
4. 模型设计
*因为这里预测的维度特别小, 因此基于扩散模型不想图像特征那样高纬度, 所以效率会很快. 说到这里的模型我们需要提前了解一些模型的先验知识, 这里我们分别讲解一下
扩散模型以及DIT.
4.1 扩散模型
这里提供B站唐宇迪讲解的扩散模型DDPM, 说白了扩散模型在训练的时候就是不断预测噪音的过程, 推理的时候就是用预测的噪音信号进行去噪的过程. 下面给出我在NOTION添加的笔记(非常推荐这个笔记软件, 打开需要梯子): Diffusion原理, 下面几张图也可以让大家很简单理解扩散模型的原理(讲解来源B站【大白话01】一文理清 Diffusion Model 扩散模型 | 原理图解+公式推导).

通过上图可以看到上面的过程就是不断加噪音的过程随着时间越来越靠后加入的噪音也越来越多, 输入到模型为加完噪音的图片, 但是模型需要预测出在t时刻对应的纯噪音. 这样在模型推理的时候只要预测出t时刻加了噪音的图片中的纯噪音, 只要将该输入的纯噪音剥离出来这样就可以还原原图了. 思想基本非常简单. 下面的两张图分别也解释了这个过程, 一个是训练一个是推理, 两个过程其实是相反的两个思路.


最后用一个公式来解释上面的整个思路. 将下面的公式与上面的过程图片一结合就非常直观了解整体思路了. 需要注意的是, 最后在最后一步没有加入这个扰动项了(实验表明,省略最后一步的噪声能提高生成样本的保真度。例如,DDPM论文中提到,最终步骤的噪声可以忽略,因为其贡献极小)

个人觉得扩散模型将时间信息引入到模型中通过扩散一步一步去解决问题, 相比较基于回归进行问题解决更符合认知.
4.2 DIT模型
可以参考这篇博客DiT(Diffusion Transformer)详解——AIGC时代的新宠儿, 整体框架可以看到如下的图以及公式则一目了然



当然可以直接看B站这篇博主的视频【研1基本功 (真的很简单)Diffusion Vision Transformer (DiT)】构建DiT核心代码以及【Sora重要技术】复现DiT(Diffusion Transformer)模型
通过条件控制的方式解决了交叉注意力算力耗费多的问题
4.3 RDT模型细节
4.3.1 扩撒模型细节

根据上图扩散公式可以看出和4.1部分推理的公式一摸一样, 损失函数仍然是MSE损失函数.
4.3.2 异质多模态输入编码
异质性体现在模态的格式、维度数量和信息量(编码后的token)方面。为了解决这个问题,需要将不同的模态编码到一个统一的隐空间中,使用的编码器如下:
- 低维输入:包括机器人诶本体感受信息,动作块、控制频率,使用带有傅立叶特征的 MLP 编码捕获低维空间中的高频变化
- 图像输入:使用 SigLIP 从图像中提取高维的、包含丰富的空间和语义信息.
为了提取紧凑的表⽰,使⽤视觉编码器SigLIP。在训练期间固定其权重以节省GPU内存 - 语言输入:使用 T5-XXL 从语言中提取长度变化、高度抽象、模糊和复杂的信息. 在预训练期间固定其权重以节省GPU内存, 在训练期间, 来 ⾃各种模态的每个输⼊都以10%的概率独立mask.
除了结构之外,异构性还意味着不同输⼊中包含的信息量不同。⾸先, 不同模态的数据包含的信息量不同。例如,图像通常⽐⽂本包含更多的信息, 同⼀模态的不同输⼊可能包含⾮常不同的信息量。 例如, 与腕式摄像头相⽐,机器⼈的外部摄像头具有更全局的视野,包含的信息也更丰富,在 这种情况下,模型可能会学习⼀条捷径:只关注外部视图⽽忽略⼿腕视图,从⽽失去感知深度的能⼒。为了解决这个问题,在编码过程中以⼀定的概率随机独⽴地屏蔽每个多模态输⼊,以防⽌模型过度依赖特定的输⼊
4.3 DIT模型改进
- QKNorm 和 RMSNorm:为了解决机器人物理量的数值不稳定导致的梯度不稳定和数值溢出问题,在计算attention时,增加了 QKNorm。action prediction本质上可以看作一个时间序列预测问题,而原始DiT 中的LayerNorm中的centering 操作会导致 token shift 和 attention shift 从而破坏时间序列的对称性。因此,替换为了 RMSNorm。
关于QKNorm可以理解一种修改注意力机制的规范化技术,使得softmax函数不易受到任意饱和的影响,同时不牺牲表达能力。具体来说,作者在将查询和键矩阵相乘之前,沿着头部维度对它们应用ℓ2规范化,然后用一个可学习的参数进行放大,而不是除以嵌入维度的平方根。具体可以参考这篇博客Transformer升级之路:15、Key归一化助力长度外推
关于RMSNorm与LayerNorm对比可以参考这篇博客RMSNorm(RMS归一化)介绍, layerNorm的问题, 下面是deepseek的解释:

如果没有上述这种修改,⼤规模的预训练往往⾮常不稳定,甚⾄会爆炸, 如下图所示。

-
MLP解码器:为了提高对非线性的机器人动作的建模能力,将现行解码器替换为非线性MLP解码器,作为从隐空间到物理空间的投影。从下图可以看出,如果没有这个改进,RDT无法有效地捕获非线性动态,从而导致无法完成需要精细操作的灵巧操作任务。
Task
- Alternating Condition Injection (ACI):RDT中的图像和语言输入作为条件,是高维的,并且长度可变,这与传统 DiTs 中的类别标签条件不同。这些信息丰富的条件难以压缩成单个 token,使得原始的 adapter layer 的 norm 方法不适用。因此,采用交叉注意力来适应不同长度的条件,避免在进一步压缩时造成信息丢失。此外,
由于图像通常远多于文本token,同时注入两种模态往往会掩盖文本相关信息,从而削弱指令遵从能力(如上图所示)。 为了缓解这个问题,在successive layer 的交叉注意力中交替注入图像和文本token,而不是在每一层都注入, 后面代码也会详细介绍。
这里还有一处可以提下模型结构中的频率$c$类似于Dit模型作为条件输入, 其实在文生图中一般为文本, 这里改为频率.
2. 数据构建
2.1 如何训练多种不同机器人的异构数据
为了启用有关异质多机器人数据的训练,我们需要在各种机器人之间共享一个统一的动作空间,这可以为多机器人动作提供统一的格式。从机器人的原始动作空间到统一动作空间的映射应在物理上可以解释,每个尺寸的尺寸都应具有明确的物理含义。这可以鼓励模型从不同的机器人数据中学习共享的物理定律,从而提高从不同机器人的数据中学习的效率
通过将原始动作向量中的每个元素根据其物理意义填充到统一动作空间向量的对应位置来将机器人的动作空间嵌入到这个统一空间中,剩余位置则用填充值填充(表 4 列出了具体的含义)。在这个统一空间下,能够利用几乎所有带夹爪的现代机器人的数据上预训练 RDT,并极大地扩展数据规模,以满足基础模型的要求。 预训练数据包含46种不同机器人的数据集(表 6),总规模为 1M+条轨迹,共计21TB


2.2 双臂操作任务数据收集
高质量的微调数据集对于模型性能至关重要。从三个方面确保数据集的高质量:
- 在数量方面,收集了 6K+ 条轨迹,是目前最大的双臂数据集之一;
- 在全面性方面,考虑了 300+ 个具有挑战性的任务,涵盖了从抓取和放置到连接电缆等大多数操作任务类型,甚至包括书写数学方程;
- 在多样性方面,准备了 100+ 个物体,包括不同尺寸和纹理的刚性和非刚性物体,以及 15+ 个不同照明条件的不同房间
此外,进一步利用 GPT-4-Turbo 来重写人工标注的指令,以增加文本多样性。
3. 实验
3.1 实验设置
旨在通过真实机器人实验回答以下问题:
- 问题1: RDT 的零样本泛化能力是否可以泛化到未见的物体和场景?
- 问题2: RDT 的零样本指令遵循能力对未见的模态的有效性?
- 问题3: RDT 能否促进对未见的技能的少样本学习?
- 问题4: RDT 是否能够完成需要精细操作的任务?
-
问题5: 大模型参数、大数据量和扩散模型对 RDT 的性能有帮助吗?
任务定义和可视化。对于7个具有挑战性的任务,我们描述了他们的语言指导,随机化和每个子任务的定义。对于倒入Water-L-1/3并倒入Water-R2/3,我们在两张图像中显示了所得的水位。
3.2 模型训练及推理消耗资源
-
模型训练
- 模型大小: 1.2B参数量
- 训练:48 个 H100 GPU(80GB显寸),一个月,1M steps
- 微调:48 个 H100 GPU(80GB显寸),三天,130K steps
模型推理:
- 框架:
DPM-Slover++(参考 【组会粗读】DPM-solver 求ODE的解析解来加速diffusion model采样) - RTX 4090 24GB GPU 推理频率 6Hz Action Trunk,381 Hz Actions
-
测试配置
image.png -
实验结论
实验结论- 扩散模型好于回归模型
- 大的模型参数好于小的模型参数
- finetune效果好于直接训练的模型
从表 3 的结果可以看出,RDT 始终优于其他基线方法。 这是因为 RDT 使用强大的网络架构与扩散相结合,能够准确地对多模态动作的分布进行建模,而离散化和 VAE 分别缺乏准确性和表达能力。 此外,大规模预训练后的大量参数提供了大量的先验知识,这极大地提高了泛化能力。 下面是详细分析:
问题1和2: RDT可以将zero-shot泛化到到看不见的物体、场景和模态。在Wash Cup和PourWater中,RDT在看不见的场景下仍然可以达到很⾼的成功率,其性能与看不见的场景没有太⼤区别。相⽐之下,其他基线甚⾄⽆法完成整个任务。在倒⽔-L-1/3和倒⽔-R-2/3中,从图5或图10的第三⾏开始,我们可以发现RDT精确地理解要纵哪只⼿以及要倒多少⽔,并通过其动作密切遵循指令,即使它从未见过“三分之⼀”或“三分之⼆”这样的词。正是因为⼤规模的预训RDT看到了⼤量不同的对象、场景和指令,导致了如此强⼤的零镜头泛化。
问题3:RDT仅通过少量样本就能学习新技能。在物品交接和折叠短裤任务中,RDT通过少样本学习掌握了物品交接和折叠等全新且复杂的技能,其动作模式与已知技能差异很大,而其他方法的成功率几乎为零。这种提升同样得益于大规模预训练。少样本学习能帮助RDT快速适应新的工作环境,这对于实际应用具有重要意义。
问题4:RDT能够处理灵巧任务。在机器狗任务中,RDT在推动操纵杆时能精确控制角度,而其他方法则导致机器狗偏离。这是因为基于我们强大的网络架构的扩散模型可以对多模态和非线性动作的分布进行建模,从而使动作精度能够满足灵巧任务的要求。我们还注意到操纵杆和遥控器都是黑色的,这使得操纵杆在视觉上不明显。这可能导致ACT容易失败。相比之下,大规模预训练使RDT对操纵杆概念学习到了更好的视觉 - 语言表征,提高了识别能力。
问题5:大模型规模、海量数据和扩散模型都是我们表现出色的关键因素。在表2中,缺少任何一个因素都会导致性能严重下降,这证明了我们这些贡献的必要性。特别是从头开始训练的RDT在未见物体和场景上表现不佳,这表明预训练获得的知识对于模型泛化至关重要。
4. 实验结论
实验结论
本文提出了 RDT 来解决可泛化双臂操作中数据稀缺和操作复杂性增加的挑战,RDT 是一种基于扩散的语言条件视觉运动模仿学习基础模型(diffusion-based foundation model for language-conditioned visuomotor imitation learning)。 RDT 是在一个广泛的多机器人数据集上预训练的,并在一个自收集的双臂数据集上微调。引入了可解释的统一动作空间,以统一不同机器人之间的动作表示,增强了鲁棒性和可迁移性。RDT 超越了现有方法,不仅在灵巧的双臂能力和指令遵循方面表现出显著改进,而且在少样本学习和对未见物体和场景的零样本泛化方面也取得了卓越的性能。
5 实验细节
5.1 微调数据集

微调数据集包含以下关键特征:
(1)多样的物体和场景。我们的数据集包含在不同场景和条件下被操控的具有不同属性的物体。
(2)具有挑战性的任务。我们的数据集纳入了灵巧操作、语言和视觉理解以及双手操作任务。
(3)多模态特征。我们的数据集使用丰富的多模态数据进行标注,包括三视角 RGB 相机数据、关节信息和增强指令。
5.2 训练细节
平台
使用Pytorch和DeepSpeed来促进并行训练,并使用Tensorflow数据集作为生产者消费者框架来快速数据加载。由于开放X-Embodiment中的大多数数据集都以Tfrecord的形式存储,因此我们将所有预训练数据集转换为Tfrecord以进行存储。在预训练中,我们使用生产者进程将Tfrecord的数据解压缩并将其存储在硬盘上的缓冲区中。同时,我们使用消费者流程以无序的顺序从缓冲区中读取数据,并将其送入模型训练。这不仅解耦了tensorflow和Pytorch环境,而且还减轻了由内存中的shuffleing buffer的小尺寸造成的训练性能损失。在微调阶段,由于数据集相对较小,因此我们还使用HDF5数据集来实现数据读取管道进行存储.动作padding以及本体知觉
为了将特定的机器人动作嵌入128维统一动作空间,我们需要填充不可用的动作元素。常见做法是用0值或特定值进行填充。但“0”实际上有物理意义。例如,速度“0”通常表示相对于地面静止。这可能会使模型困惑:“0”代表静止还是填充值?为解决此问题,我们在将动作和本体感受编码到令牌空间之前,用一个0 - 1向量连接动作和本体感受,该向量指示每个维度是否被填充,从而得到一个256维向量。这样可以补充缺失的可用性信息并消除混淆。上述可以理解为通过mask信息告诉模型那些位置是被填充了, 具体代码会做介绍训练过程的检查
训练期间,每个固定时期,我们进行扩散采样,并将采样动作与训练数据集的基础真相进行比较。从经验上讲,我们发现了两者的平均平方误差(MSE)与机器人部署性能之间的一般正相关。该观察结果使我们能够轻松监视模型的培训进度。当这种MSE收敛时,我们通常可以停止培训。我们注意到,过度低的MSE也可能意味着过度拟合。数据增广
数据增强。过度拟合是训练大型神经模型的普遍挑战,尤其是在微调阶段。我们利用数据增强技术来解决它。我们执行图像增强,包括颜色抖动和图像损坏,并以40dB的信噪比(SNR)(SNR)为输入本体感受添加高斯噪声。我们还使用GPT-4-Turbo来增强和扩展语言说明其他finetune细节
在微调过程中,我们删除了每个情节开头的一个静态部分(代码会详细介绍),这可能是由于记录开始后操作员没有做出反应造成的。我们的语言指令是从原始手动标注的指令、扩展指令和简化指令中以三分之一的概率采样得到的。当抽取到扩展指令时,我们从与任务对应的 100 条扩展指令中均匀采样一条。我们没有应用Classifier-Free Guidance(CFG),因为我们发现这并没有提高模型的性能,反而导致了机械臂行为不稳定。
-
RDT超参数
下面表 8 中列出了多模态编码器的详细信息,在表 9 中列出了模型参数。图像历史大小为 Timg = 2,动作块大小为 Ta = 64,语言token空间维度为 4096,图像token空间维度为 1152,RDT 的标记空间维度为 2048。我们使用适配器将每个模态的标记维度对齐到 2048。并且多模态编码器的所有激活函数都采用 GeLU 激活函数。在预训练和微调阶使用 AdamW 优化器,搭配常数学习率调度器以及表 10 中的超参数。该模型分别在 48 个 H100 80GB GPU 上进行 100 万步的预训练和 13 万步的微调。由于调度原因,我们没有从 100 万步预训练的checkpoint开始微调,而是选择了 50 万步的checkpoint。在训练阶段,我们使用带有 glide 余弦调度器(即 squaredcos cap v2)的 DDPM 调度器,步数为 1000。在采样阶段,我们利用带有 glide cosine scheduler的 DPM-Solver++,采样步数为 5。在微调过程中,我们还会过滤掉长度低于 32 的episode,并将长度高于 2048 的episode下采样到 2048
image.png
二、代码阅读
未完待续, 参考RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation 论文及代码总结(二)




