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

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

文章目录

    • RDT 模型架构分析:
    • 模型组件 (RDTBlock) 深入分析:
    • RDTRunner 总结:
    • 项目整体梳理报告
      • 1. 模型结构 (Model Architecture)
      • 2. 训练流程 (Training Process)
      • 3. 数据格式 (Data Format)
      • 4. 数据处理 (Data Processing)
    • Episode 在训练中的体现方式
      • 为什么采用这种方式?
      • Episode 处理流程图

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 的输出处理成最终的预测结果(即噪声)。
  1. 权重初始化 (initialize_weights):
  • 采用标准的 xavier_uniform_ 初始化线性层。
  • 使用经典的正弦/余弦 (sin-cos) 函数来初始化位置编码,这是一种无需学习却非常有效的位置编码方式。
  • 一个关键细节是:它将 输出层最后一个线性层的权重初始化为零。这种做法借鉴于 DiT (Diffusion Transformer) 等先进的生成模型,可以极大地稳定训练初期。模型在开始训练时会预测“零噪声”,然后在此基础上逐渐学习。
  1. 前向传播 (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 更高效的采样器,可以用更少的步数生成高质量结果。
  1. 损失计算 (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 扩散模型的标准损失函数。

  2. 推理与采样 (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。
  1. 模型预测:
  • 将 noisy_action、机器人状态、以及经过编码器和适配器处理后的语言/图像条件,全部送入 RDT 模型。
  • 模型对输入的 noisy_action 进行处理,并预测出它认为当初被加入的噪声 predicted_noise。
    在这里插入图片描述
  1. 损失计算:
  • 计算模型预测的 predicted_noise 和步骤2中真实采样的噪声之间的 均方误差 (MSE Loss)。
  1. 反向传播: 使用此 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维,这使得模型可以无缝地处理来自异构机器人的数据。
  1. 数据增强:
  • 语言: 训练时会从原始指令、简化指令、扩展指令中随机选择一种,增强模型对语言变化的鲁棒性。
  • 图像: 训练脚本中可以通过 --image_aug 参数开启图像增强,对输入的图像进行随机的视觉变换。
  1. 时序处理:
  • 动作: 模型预测的是未来一个时间窗口 (CHUNK_SIZE) 的动作,而非单个动作,这有助于生成更连贯的行为。
  • 图像: 模型接收的是过去一个时间窗口 (IMG_HISORY_SIZE) 的图像,这为模型提供了动态变化的视觉信息。

Episode 在训练中的体现方式

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

  1. 第一步:加权选择电影 (Episode Sampling)
  • 系统首先会看一下所有电影的片长。长的电影(包含更多有效操作的 Episode)会获得更高的权重。
  • 然后,系统根据这个权重随机选择一部电影(一个 Episode 文件)来观看。这意味着,内容更丰富的长电影有更高的几率被选中。
  • 这一步是在 HDF5VLADataset 的 init 和 get_item 方法中完成的。
  1. 第二步:在电影中随机定位 (Timestep Sampling)
  • 电影被选中后,系统并不会从头播放,而是在电影的进度条上随机选择一个时间点 t。
  • 这个时间点 t 就成了我们这个训练样本的“现在”。
  • 这一步是在 parse_hdf5_file 方法中,通过 np.random.randint(…) 来实现的。
  1. 第三步:截取一个片段 (Window Slicing)
  • 以这个随机时间点 t 为中心,系统会截取一个标准长度的“片段”:
    • 回顾过去: 提取 t 时刻以及它之前的几帧图像 (由 IMG_HISORY_SIZE 决定)。
    • 观察现在: 提取 t 时刻的机器人状态 state。
    • 预测未来: 提取从 t 时刻开始,未来 CHUNK_SIZE 步的真实动作 actions。
  • 这个包含“过去-现在-未来”的短小片段,就构成了一个独立的、标准化的训练样本。
  1. 第四步:组成一个批次 (Batching)
  • 训练时,PyTorch 的 DataLoader 会重复执行前三步,收集大量的这种“小片段”。
  • 重要的是,一个训练批次 (batch) 中可能包含了来自几十个不同电影(Episode)、在各自完全不同的时间点上截取的片段。
  • 模型在训练时看到的,就是这样一个由各种不同场景的“小片段”混合而成的批次。

为什么采用这种方式?

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

Episode 处理流程图

在这里插入图片描述

http://www.dtcms.com/a/272208.html

相关文章:

  • HTML颜色定义
  • 深入了解Modbus TCP:工业通信的“通用语言”
  • Docker-构建镜像并实现LNMP架构
  • C语言 | 函数核心机制深度解构:从底层架构到工程化实践
  • 西电考研录取:哪些省份考研上岸西电更容易?
  • PyTorch Tensor 的创建与操作入门
  • 低版本hive(1.2.1)UDF实现清除历史分区数据
  • 1.1.1数据类型与变量——AI教你学Django
  • 基于Uniapp+MySQL+PHP的景区多商户小程序源码系统 带完整的搭建指南
  • vue引入应用通义AI大模型-(一)前期准备整理思路
  • idea如何打开extract surround
  • 【CPU】不同核数下的翻译
  • hbuilder开发app记录
  • ReaLTaiizor:WinForms 界面设计利器
  • 《未来已来:当人类智慧遇上AI智能体》
  • 【WPF实战】MVVM中如何从数据模型反查自定义控件实例(ImageView + Halcon)
  • 学习开发之条件函数
  • 如何用 LangChain 自定义 Chat Model —— 测试工程师实践指南
  • Maven生命周期:构建流程深度解析
  • eVTOL动力测试台架气动干扰分析与应对措施
  • TestCafe ➜ Playwright fixture 架构迁移指南
  • 上位机与Modbus的32位数据交互
  • 嘿嘿嘿嘿嘿
  • C++---多态
  • Camera2API笔记
  • Unity WebGL文本输入
  • centos7 安装jenkins
  • jenkins部署springboot项目
  • 抽象类与接口:Java面向对象设计的两大支柱
  • 表达式索引海外云持久化实践:关键技术解析与性能优化