【VLAs篇】05：RDT模型结构和流程分析

RDT 模型架构分析:

1. 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 的输出处理成最终的预测结果（即噪声）。

2. 权重初始化 (initialize_weights):

• 采用标准的 xavier_uniform_ 初始化线性层。
• 使用经典的正弦/余弦 (sin-cos) 函数来初始化位置编码，这是一种无需学习却非常有效的位置编码方式。
• 一个关键细节是：它将 输出层最后一个线性层的权重初始化为零。这种做法借鉴于 DiT (Diffusion Transformer) 等先进的生成模型，可以极大地稳定训练初期。模型在开始训练时会预测“零噪声”，然后在此基础上逐渐学习。

3. 前向传播 (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 的基本重复单元，其内部处理流程非常清晰，遵循了 [输入 -> 自注意力 -> 交叉注意力 -> 前馈网络 -> 输出] 的经典结构：

  1. 自注意力 (Self-Attention): 序列 x 首先进行自注意力计算，这使得序列内部的元素（例如，不同时间步的动作）能够相互通信和关联。
  2. 交叉注意力 (Cross-Attention): 自注意力的输出，会作为查询（Query），去“关注”外部的条件输入 c（即文本或图像特征）。这是模型注入条件信息的关键步骤。
  3. 前馈网络 (Feed-Forward Network): 结果最后通过一个MLP网络，进行非线性变换，增加模型的表达能力。
  • 归一化 (Normalization): 在每个关键操作（自注意力、交叉注意力、前馈网络）之前都使用了 RmsNorm (Root Mean Square Normalization)。这是一种计算效率更高的层归一化变体，有助于稳定训练过程。
  • 残差连接 (Residual Connection): 每个模块的输出都与输入相加（x = module(x) + x），这是稳定深度网络训练的标准技术。

• FinalLayer: 一个简单的收尾模块，在最后进行一次归一化和MLP变换，将Transformer的输出维度映射到最终所需的维度（out_channels）。

RDTRunner 总结:

1. 初始化 (init):

• 实例化 RDT 模型: 读取配置，创建 RDT 核心网络模型。
• 创建数据适配器 (Adaptors): 它为语言 (lang_adaptor)、图像 (img_adaptor) 和状态 (state_adaptor) 分别创建了适配器。这些适配器本质上是简单的MLP网络，其唯一作用是将各种不同维度的数据（如T5语言编码器的输出、SigLIP视觉编码器的输出、机器人自身状态向量）投影到 RDT 模型所要求的统一隐藏维度 (hidden_size)。这是实现多模态信息融合前至关重要的一步，确保所有输入都在同一个“频道”上对话。
• 初始化噪声调度器 (Noise Schedulers): 它利用 diffusers 库配置了两个调度器：
  • DDPMScheduler: 用于训练阶段。负责在前向扩散过程中，将随机噪声添加到真实的动作序列上。
  • DPMSolverMultistepScheduler: 用于推理阶段。负责在反向扩散过程中，从纯噪声中逐步生成清晰的动作。它是一种比 DDPM 更高效的采样器，可以用更少的步数生成高质量结果。

2. 损失计算 (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 扩散模型的标准损失函数。

3. 推理与采样 (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 都包含三个关键部分：
    1. 自注意力 (Self-Attention): 允许输入序列内部的元素（例如，不同未来时间步的动作之间）进行信息交换和依赖建模。
    2. 交叉注意力 (Cross-Attention): 这是融合指导信息的关键。它让序列去“关注”来自语言和图像的条件特征。一个非常独特的设计是，它在不同层之间交替地对语言和图像进行交叉注意力计算，从而实现两种模态的深度、交错融合。
    3. 前馈网络 (FFN): 标准的MLP网络，用于增加模型的非线性计算能力。
  • 输出: 模型的最终输出是它预测的“噪声”，其维度与输入的未来动作序列完全相同。

2. 训练流程 (Training Process)

训练流程遵循标准的扩散模型训练范式，在 RDTRunner 的 compute_loss 方法中定义：

1. 数据准备: 从数据集中取出一个样本，包含：真实的未来动作序列 (action_gt)、机器人当前状态、文本指令、历史图像。
2. 前向扩散 (加噪):

• 随机采样一个高斯噪声，其形状与 action_gt 相同。
• 随机采样一个扩散时间步 t (例如，从0到1000之间)。
• 调用 DDPMScheduler.add_noise 方法，根据时间步 t 将噪声“注入”到 action_gt 中，得到一个带噪的动作序列 noisy_action。

3. 模型预测:

• 将 noisy_action、机器人状态、以及经过编码器和适配器处理后的语言/图像条件，全部送入 RDT 模型。
• 模型对输入的 noisy_action 进行处理，并预测出它认为当初被加入的噪声 predicted_noise。
  

4. 损失计算:

• 计算模型预测的 predicted_noise 和步骤2中真实采样的噪声之间的 均方误差 (MSE Loss)。

5. 反向传播: 使用此 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 配置文件。

1. 统一动作向量 (Unified Action Vector): 这是整个框架的基石。

• configs/state_vec.py 中定义了一个长度为 128 的标准向量模板。
• 这个模板为各种可能的机器人自由度（双臂各10个关节、夹爪、末端位姿、移动底盘等）都预留了固定的索引位置。
• 当处理来自某个特定机器人的数据时（例如，一个只有6自由度单臂的机器人），HDF5VLADataset 会创建一个128维的全零向量，然后根据 STATE_VEC_IDX_MAPPING，仅将该机器人拥有的6个关节数据填入预留好的特定位置。
• 通过这种方式，无论原始数据来自多么不同的机器人，输入到模型中的状态/动作向量始终是128维，这使得模型可以无缝地处理来自异构机器人的数据。

2. 数据增强:

• 语言: 训练时会从原始指令、简化指令、扩展指令中随机选择一种，增强模型对语言变化的鲁棒性。
• 图像: 训练脚本中可以通过 --image_aug 参数开启图像增强，对输入的图像进行随机的视觉变换。

3. 时序处理:

• 动作: 模型预测的是未来一个时间窗口 (CHUNK_SIZE) 的动作，而非单个动作，这有助于生成更连贯的行为。
• 图像: 模型接收的是过去一个时间窗口 (IMG_HISORY_SIZE) 的图像，这为模型提供了动态变化的视觉信息。

Episode 在训练中的体现方式

想象一下，您的数据集里有很多部电影（每一个 .hdf5 文件就是一个 Episode），这些电影长短不一。训练模型并不像让它把每一部电影从头到尾看完，而是像下面这样：

1. 第一步：加权选择电影 (Episode Sampling)

• 系统首先会看一下所有电影的片长。长的电影（包含更多有效操作的 Episode）会获得更高的权重。
• 然后，系统根据这个权重随机选择一部电影（一个 Episode 文件）来观看。这意味着，内容更丰富的长电影有更高的几率被选中。
• 这一步是在 HDF5VLADataset 的 init 和 get_item 方法中完成的。

2. 第二步：在电影中随机定位 (Timestep Sampling)

• 电影被选中后，系统并不会从头播放，而是在电影的进度条上随机选择一个时间点 t。
• 这个时间点 t 就成了我们这个训练样本的“现在”。
• 这一步是在 parse_hdf5_file 方法中，通过 np.random.randint(…) 来实现的。

3. 第三步：截取一个片段 (Window Slicing)

• 以这个随机时间点 t 为中心，系统会截取一个标准长度的“片段”：
  • 回顾过去: 提取 t 时刻以及它之前的几帧图像 (由 IMG_HISORY_SIZE 决定)。
  • 观察现在: 提取 t 时刻的机器人状态 state。
  • 预测未来: 提取从 t 时刻开始，未来 CHUNK_SIZE 步的真实动作 actions。
• 这个包含“过去-现在-未来”的短小片段，就构成了一个独立的、标准化的训练样本。

4. 第四步：组成一个批次 (Batching)

• 训练时，PyTorch 的 DataLoader 会重复执行前三步，收集大量的这种“小片段”。
• 重要的是，一个训练批次 (batch) 中可能包含了来自几十个不同电影（Episode）、在各自完全不同的时间点上截取的片段。
• 模型在训练时看到的，就是这样一个由各种不同场景的“小片段”混合而成的批次。

为什么采用这种方式？

• 计算效率: GPU 擅长处理形状规整、大小统一的张量。将长短不一的 Episode 切割成固定大小的样本，是实现高效并行计算的基础。
• 数据利用率最大化: 一个长达500步的 Episode，理论上可以产生近500个不同的训练“片段”。这种方式极大地丰富了训练数据，让模型从各个角度学习同一个 Episode 中的技能。
• 学习通用策略: 这种方法迫使模型学习一个更通用的策略：“在任意给定情况下，接下来该怎么做”，而不是去死记硬背某一个 Episode 的完整流程。这大大增强了模型的泛化能力。

Episode 处理流程图