【VLAs篇】05:RDT模型结构和流程分析
文章目录
- RDT 模型架构分析:
- 模型组件 (RDTBlock) 深入分析:
- RDTRunner 总结:
- 项目整体梳理报告
- 1. 模型结构 (Model Architecture)
- 2. 训练流程 (Training Process)
- 3. 数据格式 (Data Format)
- 4. 数据处理 (Data Processing)
- Episode 在训练中的体现方式
- 为什么采用这种方式?
- Episode 处理流程图
RDT 模型架构分析:
- RDT(nn.Module) 类: 这是模型的主体。
- 初始化 (init):
- 接收 hidden_size (隐藏层维度), depth (Transformer层数), num_heads (注意力头数) 等关键超参数。
- 时间步和频率编码器 (TimestepEmbedder): 模型创建了两个独立的编码器,一个用于编码扩散模型的时间步 t,另一个用于编码机器人的控制频率 freq。这意味着模型不仅知道当前输入的噪声水平,还能适应不同机器人的操作频率,这是一个非常精巧的设计。
- 位置编码 (pos_embed): 模型为三种不同的输入序列定义了独立的位置编码:
- x_pos_embed: 用于主输入序列(包含机器人状态和动作序列)。
- lang_cond_pos_embed: 用于文本指令条件。
- img_cond_pos_embed: 用于图像观测条件。
- Transformer 主体 (RDTBlock): 模型的“躯干”是由 depth 个 RDTBlock 堆叠而成的列表 (nn.ModuleList)。
- 输出层 (FinalLayer): 最后是一个 FinalLayer,它负责将 Transformer 的输出处理成最终的预测结果(即噪声)。
- 权重初始化 (initialize_weights):
- 采用标准的 xavier_uniform_ 初始化线性层。
- 使用经典的正弦/余弦 (sin-cos) 函数来初始化位置编码,这是一种无需学习却非常有效的位置编码方式。
- 一个关键细节是:它将 输出层最后一个线性层的权重初始化为零。这种做法借鉴于 DiT (Diffusion Transformer) 等先进的生成模型,可以极大地稳定训练初期。模型在开始训练时会预测“零噪声”,然后在此基础上逐渐学习。
- 前向传播 (forward):
- 输入: 接收 x (带噪声的动作序列), freq (控制频率), t (扩散时间步), lang_c (语言条件), 和 img_c (图像条件)。
- 输入构建: 它首先将 t 和 freq 通过各自的编码器转换为高维向量,然后将它们与输入 x(状态+动作)拼接在一起,形成一个完整的序列:[时间步向量, 频率向量, 状态向量, 动作向量1, …, 动作向量H]。
- 注入位置信息: 将预先定义好的位置编码分别加到主序列、语言序列和图像序列上。
- 核心循环: 接着,数据流经一系列的 RDTBlock。这里有一个非常有趣的设计:它在不同的层交替使用语言条件和图像条件 (conds[i%2])进行交叉注意力计算。这意味着模型在第1层关注语言,第2层关注图像,第3层再关注语言,以此类推。这种交错注意力机制让模型可以深度融合两种模态的信息。
- 输出处理: 最后,序列经过 FinalLayer 处理,并只返回对应于动作部分的输出。因为模型的任务是预测添加到动作序列上的噪声。
总的来说,这是一个设计精良的 Diffusion Transformer 模型,它通过独特的条件注入方式(时间步、控制频率)和模态融合策略(交错交叉注意力)来处理复杂的机器人控制任务。
模型组件 (RDTBlock) 深入分析:
-
TimestepEmbedder: 这个模块的作用是将标量值(如扩散时间步 t 或控制频率 freq)映射成高维向量。它采用了经典的 正弦/余弦位置编码,并通过一个小型MLP(多层感知器)进行变换。这是将连续的标量数值有效融入模型的标准方法。
-
CrossAttention (交叉注意力): 这是模型融合多模态信息的核心。它接收来自主干网络序列的查询 x,以及来自条件(如语言或图像)的键和值 c。通过这个模块,主干序列中的每个元素都可以“关注”并提取来自外部条件信息。代码中还集成了 flash-attention 优化,以提高计算效率。
-
RDTBlock (RDT 核心模块): 这是 Transformer 的基本重复单元,其内部处理流程非常清晰,遵循了 [输入 -> 自注意力 -> 交叉注意力 -> 前馈网络 -> 输出] 的经典结构:
- 自注意力 (Self-Attention): 序列 x 首先进行自注意力计算,这使得序列内部的元素(例如,不同时间步的动作)能够相互通信和关联。
- 交叉注意力 (Cross-Attention): 自注意力的输出,会作为查询(Query),去“关注”外部的条件输入 c(即文本或图像特征)。这是模型注入条件信息的关键步骤。
- 前馈网络 (Feed-Forward Network): 结果最后通过一个MLP网络,进行非线性变换,增加模型的表达能力。
- 归一化 (Normalization): 在每个关键操作(自注意力、交叉注意力、前馈网络)之前都使用了 RmsNorm (Root Mean Square Normalization)。这是一种计算效率更高的层归一化变体,有助于稳定训练过程。
- 残差连接 (Residual Connection): 每个模块的输出都与输入相加(x = module(x) + x),这是稳定深度网络训练的标准技术。
-
FinalLayer: 一个简单的收尾模块,在最后进行一次归一化和MLP变换,将Transformer的输出维度映射到最终所需的维度(out_channels)。
RDTRunner 总结:
- 初始化 (init):
- 实例化 RDT 模型: 读取配置,创建 RDT 核心网络模型。
- 创建数据适配器 (Adaptors): 它为语言 (lang_adaptor)、图像 (img_adaptor) 和状态 (state_adaptor) 分别创建了适配器。这些适配器本质上是简单的MLP网络,其唯一作用是将各种不同维度的数据(如T5语言编码器的输出、SigLIP视觉编码器的输出、机器人自身状态向量)投影到 RDT 模型所要求的统一隐藏维度 (hidden_size)。这是实现多模态信息融合前至关重要的一步,确保所有输入都在同一个“频道”上对话。
- 初始化噪声调度器 (Noise Schedulers): 它利用 diffusers 库配置了两个调度器:
- DDPMScheduler: 用于训练阶段。负责在前向扩散过程中,将随机噪声添加到真实的动作序列上。
- DPMSolverMultistepScheduler: 用于推理阶段。负责在反向扩散过程中,从纯噪声中逐步生成清晰的动作。它是一种比 DDPM 更高效的采样器,可以用更少的步数生成高质量结果。
-
损失计算 (compute_loss): 此方法定义了模型的训练目标,完整地再现了扩散模型的训练步骤:
a. 取一批真实的动作序列 action_gt。
b. 随机采样高斯噪声和扩散时间步 t。
c. 前向扩散: 使用 noise_scheduler.add_noise,根据时间步 t 将噪声添加到 action_gt 上,得到带噪的动作 noisy_action。
d. 将语言、图像、状态和带噪的动作序列通过各自的 adaptor 适配维度。
e. 将所有处理好的数据送入 RDT 模型,让模型预测被添加的噪声 epsilon。
f. 计算损失: 使用均方误差损失 (MSE Loss) 来衡量模型预测的噪声和步骤2中真实采样的噪声之间的差距。这是 epsilon-predicting 扩散模型的标准损失函数。 -
推理与采样 (predict_action & conditional_sample): 此方法展示了模型如何生成动作序列:
a. 从一个与输出形状相同的纯高斯噪声张量开始。
b. 启动一个去噪循环,从最后一个时间步 T 向 0 迭代。在循环的每一步:- 将当前的噪声动作、语言、图像等条件送入 RDT 模型,预测当前步的噪声。
- 调用 noise_scheduler_sample.step 方法执行一步反向扩散,即根据预测的噪声,对当前噪声动作进行一次“净化”,得到一个稍微清晰一点的动作。
c. 循环结束后,初始的纯噪声张量就被完全“净化”成了模型预测的最终动作序列。
根据 README , data/hdf5_vla_dataset.py 文件是理解数据处理流程的关键。这个文件完整地揭示了数据处理的全貌。它是一个为用户定制而设计的模板,通过其中的 [Modify] 和 [Optional] 注释可以清晰地看出。
以下是基于 HDF5VLADataset 类对数据管道的详细解析:
-
数据存储:
- 项目假定数据集是由大量的 HDF5 (.hdf5) 文件组成的,每一个HDF5文件代表一个完整的机器人操作回合(episode)。代码会自动扫描指定目录 (HDF5_DIR) 下的所有 .hdf5 文件。
-
数据采样策略:
- 在初始化 (init) 时,代码会预先计算每个 episode 的长度(总步数),并以此为权重创建一个采样概率分布。
- 在请求数据 (get_item) 时,系统会根据这个权重随机选择一个 episode。这意味着更长的 episode 有更高的概率被选中,这是一种非常好的实践,可以有效平衡不同长度数据在训练中的贡献。
-
核心数据解析 (parse_hdf5_file):
这是从单个HDF5文件中提取一个训练样本的核心方法。
a. 加载数据: 打开 HDF5 文件,读取 observations/qpos (机器人关节位置) 和 action (目标动作) 数据。
b. 数据清洗: 包含一些可选的逻辑,例如可以丢弃过短的 episode,或跳过每个 episode 开头机器人静止的部分。
c. 随机时间点采样: 在 episode 的有效时间范围内,随机抽取一个时间点 t。这个 t 就构成了训练样本的“当前”时刻。
d. 加载文本指令: 从对应的 JSON 文件中加载文本指令。一个巧妙的设计是,它会随机选择使用原始指令、简化指令或扩展指令,这相当于对语言输入进行了一种数据增强。
e. 提取当前状态: 状态 state 就是机器人在 t 时刻的 qpos。
f. 提取未来动作序列: 模型的学习目标不是单个动作,而是一个未来的动作序列,即从 t 时刻开始,长度为 CHUNK_SIZE (预测时域) 的动作块 action[t : t + CHUNK_SIZE]。如果 episode 在这个窗口结束前就完结了,代码会用最后一个有效动作进行填充。
g. 构建“统一动作向量” (Unified Action Vector): 这是整个数据处理流程的精髓所在。- 它利用 configs/state_vec.py 中定义的 STATE_VEC_IDX_MAPPING 映射关系。
- 将当前机器人 specific 的关节值(例如,这个机器人只有左臂6个关节+1个夹爪)填充到一个预定义的、非常大的、标准化的全零向量 (uni_vec) 中的指定位置。
- 无论是当前状态、未来动作序列,还是 state_indicator(一个指示该机器人在统一向量中实际使用了哪些维度的掩码),都经过了这一处理。
- 正是这种“统一向量”的设计,使得模型可以同时在来自不同机器人(拥有不同关节数、不同结构)的数据上进行训练。模型看到的输入维度始终是固定的,未被使用的维度始终为零。
h. 解析图像: - 代码会读取最多三个摄像头的图像:cam_high (外部固定视角), cam_left_wrist (左手腕), cam_right_wrist (右手腕)。
- 它不仅获取 t 时刻的图像,还会提取包括 t 在内过去 IMG_HISORY_SIZE 帧的图像历史。
- 如果在 episode 的开头,历史图像不足,它会用第一帧图像进行填充。同时会生成一个掩码 cam_*_mask 来标记哪些是真实图像,哪些是填充图像。
i. 返回样本: 最后,方法返回一个包含所有处理好信息的字典,包括元数据、状态、动作、图像、掩码和文本指令。
configs/state_vec.py文件定义了“统一动作向量”的骨架,是理解数据格式的最后一块拼图。
STATE_VEC_IDX_MAPPING 是一个巨大的字典,它将人类可读的名称(如 ‘right_arm_joint_0_pos’, ‘left_eef_pos_x’)映射到这个128维向量中的具体索引位置。
这个结构的要点如下:
- 固定维度: 状态/动作向量的总维度被固定为 128。
- 双臂+移动平台设计: 向量中为双臂和移动底盘预留了专门的“槽位”:
- 右臂 (索引 0-49): 包括关节位置、夹爪位置、关节速度、夹爪速度、末端执行器位姿(XYZ位置 + 6D旋转)和速度。
- 左臂 (索引 50-99): 结构与右臂相同。
- 移动底盘 (索引 100-102): 为带轮子的移动底盘预留了线速度和角速度的槽位。
- 兼容多种控制模式: 该向量同时支持关节空间控制 (_pos, vel) 和末端笛卡尔空间控制 (eef_pos*)。
- 预留与扩展性: 向量中包含大量预留(reserved)槽位,这为未来支持更复杂的机器人提供了良好的扩展性,而不会破坏现有结构。
- 别名: 它为常用的控制维度提供了更方便的别名,例如 gripper_open 就是 right_gripper_joint_0_pos 的别名。这也与 README 中提到的“单臂机器人请使用右臂槽位”的建议相符。
项目整体梳理报告
1. 模型结构 (Model Architecture)
整个模型的结构可以分为三个层次:多模态编码器 (不在本项目代码中,需预先下载)、数据适配器 (Adaptor) 和 核心的RDT模型。
-
外部多模态编码器:
- 语言编码器: 使用 Google 的 T5-v1.1-XXL 模型,将自然语言指令(如“把桌上的苹果递给我”)编码成高维度的文本特征向量。
- 视觉编码器: 使用 Google 的 SigLIP 模型,将多视角的摄像头图像(外部、左手腕、右手腕)编码成图像特征向量。
-
数据适配器 (Adaptor):
- 这些是在 RDTRunner 中定义的简单MLP网络。
- 它们的作用是“桥梁”,将来自外部编码器的不同维度的输出(文本特征、图像特征)以及机器人自身的状态向量,全部投影到 RDT 模型内部要求的统一维度(例如1152维)。这确保了所有信息可以在同一个特征空间内被处理和融合。
-
核心 RDT 模型 (RDT):
- 这是一个 Diffusion Transformer 模型,其设计深受 DiT (Diffusion Transformer) 模型的启发。
- 输入: 它的输入是一个拼接起来的序列,包含了 [扩散时间步, 控制频率, 机器人当前状态, 带噪声的未来动作序列]。
- 主体: 模型的主体是堆叠了多层(例如28层)的 RDTBlock。
- RDTBlock 内部: 每个 RDTBlock 都包含三个关键部分:
- 自注意力 (Self-Attention): 允许输入序列内部的元素(例如,不同未来时间步的动作之间)进行信息交换和依赖建模。
- 交叉注意力 (Cross-Attention): 这是融合指导信息的关键。它让序列去“关注”来自语言和图像的条件特征。一个非常独特的设计是,它在不同层之间交替地对语言和图像进行交叉注意力计算,从而实现两种模态的深度、交错融合。
- 前馈网络 (FFN): 标准的MLP网络,用于增加模型的非线性计算能力。
- 输出: 模型的最终输出是它预测的“噪声”,其维度与输入的未来动作序列完全相同。
2. 训练流程 (Training Process)
训练流程遵循标准的扩散模型训练范式,在 RDTRunner 的 compute_loss 方法中定义:
- 数据准备: 从数据集中取出一个样本,包含:真实的未来动作序列 (action_gt)、机器人当前状态、文本指令、历史图像。
- 前向扩散 (加噪):
- 随机采样一个高斯噪声,其形状与 action_gt 相同。
- 随机采样一个扩散时间步 t (例如,从0到1000之间)。
- 调用 DDPMScheduler.add_noise 方法,根据时间步 t 将噪声“注入”到 action_gt 中,得到一个带噪的动作序列 noisy_action。
- 模型预测:
- 将 noisy_action、机器人状态、以及经过编码器和适配器处理后的语言/图像条件,全部送入 RDT 模型。
- 模型对输入的 noisy_action 进行处理,并预测出它认为当初被加入的噪声 predicted_noise。
- 损失计算:
- 计算模型预测的 predicted_noise 和步骤2中真实采样的噪声之间的 均方误差 (MSE Loss)。
- 反向传播: 使用此 Loss 进行反向传播,更新整个模型(包括适配器和RDT模型)的权重。训练框架使用了 DeepSpeed 来支持大规模分布式训练。
3. 数据格式 (Data Format)
数据被组织成 .hdf5 文件,每个文件是一个 episode。当从中采样一个训练数据点时,其格式是一个包含丰富信息的字典:
- actions: (CHUNK_SIZE, STATE_DIM) 的 ndarray。这是模型需要学习预测的目标,即未来 CHUNK_SIZE 步的动作序列。STATE_DIM 是统一向量的维度(128)。
- state: (1, STATE_DIM) 的 ndarray。机器人“当前”时刻的状态。
- instruction: 字符串。该 episode 对应的自然语言指令。
- cam_high, cam_left_wrist, cam_right_wrist: (IMG_HISORY_SIZE, H, W, 3) 的 ndarray。分别对应外部、左手腕、右手腕相机在当前时刻及之前的历史图像。
- state_indicator: (STATE_DIM,) 的 ndarray。一个0-1掩码,标记了在该向量的128个维度中,哪些是当前机器人实际使用的。
- 其他: 还包括用于归一化的统计数据 state_std, state_mean 等。
4. 数据处理 (Data Processing)
数据处理的核心思想是标准化和泛化,关键在于 HDF5VLADataset 类和 state_vec.py 配置文件。
- 统一动作向量 (Unified Action Vector): 这是整个框架的基石。
- configs/state_vec.py 中定义了一个长度为 128 的标准向量模板。
- 这个模板为各种可能的机器人自由度(双臂各10个关节、夹爪、末端位姿、移动底盘等)都预留了固定的索引位置。
- 当处理来自某个特定机器人的数据时(例如,一个只有6自由度单臂的机器人),HDF5VLADataset 会创建一个128维的全零向量,然后根据 STATE_VEC_IDX_MAPPING,仅将该机器人拥有的6个关节数据填入预留好的特定位置。
- 通过这种方式,无论原始数据来自多么不同的机器人,输入到模型中的状态/动作向量始终是128维,这使得模型可以无缝地处理来自异构机器人的数据。
- 数据增强:
- 语言: 训练时会从原始指令、简化指令、扩展指令中随机选择一种,增强模型对语言变化的鲁棒性。
- 图像: 训练脚本中可以通过 --image_aug 参数开启图像增强,对输入的图像进行随机的视觉变换。
- 时序处理:
- 动作: 模型预测的是未来一个时间窗口 (CHUNK_SIZE) 的动作,而非单个动作,这有助于生成更连贯的行为。
- 图像: 模型接收的是过去一个时间窗口 (IMG_HISORY_SIZE) 的图像,这为模型提供了动态变化的视觉信息。
Episode 在训练中的体现方式
想象一下,您的数据集里有很多部电影(每一个 .hdf5 文件就是一个 Episode),这些电影长短不一。训练模型并不像让它把每一部电影从头到尾看完,而是像下面这样:
- 第一步:加权选择电影 (Episode Sampling)
- 系统首先会看一下所有电影的片长。长的电影(包含更多有效操作的 Episode)会获得更高的权重。
- 然后,系统根据这个权重随机选择一部电影(一个 Episode 文件)来观看。这意味着,内容更丰富的长电影有更高的几率被选中。
- 这一步是在 HDF5VLADataset 的 init 和 get_item 方法中完成的。
- 第二步:在电影中随机定位 (Timestep Sampling)
- 电影被选中后,系统并不会从头播放,而是在电影的进度条上随机选择一个时间点 t。
- 这个时间点 t 就成了我们这个训练样本的“现在”。
- 这一步是在 parse_hdf5_file 方法中,通过 np.random.randint(…) 来实现的。
- 第三步:截取一个片段 (Window Slicing)
- 以这个随机时间点 t 为中心,系统会截取一个标准长度的“片段”:
- 回顾过去: 提取 t 时刻以及它之前的几帧图像 (由 IMG_HISORY_SIZE 决定)。
- 观察现在: 提取 t 时刻的机器人状态 state。
- 预测未来: 提取从 t 时刻开始,未来 CHUNK_SIZE 步的真实动作 actions。
- 这个包含“过去-现在-未来”的短小片段,就构成了一个独立的、标准化的训练样本。
- 第四步:组成一个批次 (Batching)
- 训练时,PyTorch 的 DataLoader 会重复执行前三步,收集大量的这种“小片段”。
- 重要的是,一个训练批次 (batch) 中可能包含了来自几十个不同电影(Episode)、在各自完全不同的时间点上截取的片段。
- 模型在训练时看到的,就是这样一个由各种不同场景的“小片段”混合而成的批次。
为什么采用这种方式?
- 计算效率: GPU 擅长处理形状规整、大小统一的张量。将长短不一的 Episode 切割成固定大小的样本,是实现高效并行计算的基础。
- 数据利用率最大化: 一个长达500步的 Episode,理论上可以产生近500个不同的训练“片段”。这种方式极大地丰富了训练数据,让模型从各个角度学习同一个 Episode 中的技能。
- 学习通用策略: 这种方法迫使模型学习一个更通用的策略:“在任意给定情况下,接下来该怎么做”,而不是去死记硬背某一个 Episode 的完整流程。这大大增强了模型的泛化能力。