UnifoLM-VLA 预训练与 Action Head 训练总结

1. 总体结论

UnifoLM-VLA 的训练流程更适合理解为两阶段，而不是把所有通用预训练数据与机器人数据从头到尾混在一个阶段中联合训练。

• 第一阶段：VLM continued pre-training
  以 Qwen2.5-VL-7B 为基础，在覆盖机器人场景与通用场景的多任务数据上继续预训练，得到更强的视觉-语言 backbone，即 UnifoLM-VLM-0 / UnifoLM-VLM-Base。
• 第二阶段：VLA training
  以第一阶段得到的 VLM 为初始化，接入 action head，并在机器人操作数据上训练动作预测能力，最终得到 UnifoLM-VLA-0。

因此，它不是简单地把 pretrain 数据集和宇树自己的数据集在同一个训练阶段统一混合训练，而是先完成面向空间语义与轨迹理解的 VLM 继续预训练，再进行 VLA/action 训练。

---

2. 第一阶段：UnifoLM-VLM 的继续预训练

根据项目页面描述，UnifoLM-VLM-0 以 Qwen2.5-VL-7B 为基础，通过 continued pre-training 强化其在机器人相关空间理解任务上的能力。

2.1 预训练目标

这一阶段并不只是普通图文对齐，而是显式加入了与机器人理解更相关的监督任务，包括：

• 2D detection
• segmentation
• hierarchical task decomposition
• 3D object detection
• spatial reasoning
• trajectory prediction

这说明该阶段的目标是把通用 VLM 强化为一个更擅长：

1. 空间感知：理解目标、位置、遮挡与可操作区域；
2. 语言-空间对齐：将任务指令与场景结构对应起来；
3. 轨迹与时序理解：具备对运动趋势与任务执行顺序的建模能力。

2.2 这一阶段的作用

这一阶段得到的并不是最终机器人策略，而是一个更适合作为 VLA backbone 的视觉-语言模型。换言之，它先解决的是：

• 看懂场景；
• 理解操作对象与空间关系；
• 建立高层任务语义与动作意图之间的联系。

---

3. VLM 能不能输出轨迹运动

从原理上说，VLM 可以输出轨迹表示，但通常并不是直接输出高精度底层控制量，而更适合输出中高层的轨迹或动作表示。

3.1 可以输出的典型形式

1. Waypoint / 关键点序列
   例如输出若干个 2D/3D waypoint、末端执行器目标位姿。

2. Action chunk / 动作片段
   即一次预测一段时间窗口内的连续动作。

3. 语言化或符号化运动计划
   例如“先靠近把手，再抬起，再向左移动”。

3.2 为什么纯 VLM 通常不够

纯 VLM 虽然可以把轨迹视为序列预测问题，但在机器人控制中仍存在不足：

• 连续精度有限；
• 闭环修正能力不足；
• 动力学约束较弱；
• 长时间序列容易累积误差。

因此，实际系统通常采用：

• VLM 负责高层感知与语义理解；
• Action head / policy head 负责输出动作表示；
• 低层控制器 再将其变成可执行轨迹。

---

4. 第二阶段：VLA 中的 Action Head 是怎么训练的

4.1 基本方式

在 VLA 阶段，模型先加载第一阶段得到的 UnifoLM-VLM-Base 作为 backbone，然后在此基础上接入 action head。训练时，输入通常包括：

• 图像（例如主视角图像、腕部图像）
• 本体状态（proprioception / state）
• 对应的动作标签（actions）

训练的目标是让模型根据视觉观测、语言条件与机器人状态，预测正确的动作序列。

4.2 训练信号

公开代码中的训练循环表明，VLA 阶段的核心监督信号是：

• action_loss

即：

1. 模型前向传播；
2. 输出动作预测结果；
3. 根据真实动作标签计算 action_loss；
4. 以 action_loss 作为主要损失进行反向传播与参数更新。

因此，action head 的训练本质上就是动作监督学习，并不是一个与 VLA 完全分离的独立训练流程。

4.3 是否只训练 action head

从公开脚本来看，模型支持通过 freeze_modules 冻结部分模块，但默认训练脚本中并没有显式冻结所有 backbone 模块。

这意味着默认设置下更接近于：

• 用预训练好的 VLM 初始化；
• 在 action loss 驱动下，对 action head 以及部分可训练 backbone 一起进行优化。

因此，严格地说它不一定是“只训练 action head”，而更像是：

以预训练 VLM 为初始化，在动作监督下对 VLA 模型整体的可训练部分进行联合优化。

---

5. 宇树是不是把所有 mix 数据集一起训练

结论是：官方主模型的思路是多任务混合训练单一策略，但训练框架本身支持单任务或自定义组合。

也就是说：

• 如果配置 data_mix=Unitree_all_task，那么会把宇树公开的多个任务混在一起训练；
• 如果配置为某个单任务数据集，如 g1_stack_block，则可以进行单任务训练；
• 因此框架本身不是只能 all-mix，而是由 data_mix 配置控制。

官方展示的核心能力是：

• 单一策略完成多类复杂任务，

这说明其主要展示结果对应的是多任务混合训练策略。

---

6. OXE 数据集和宇树数据集是不是一起训练 Action Head

这是最关键的问题。结论是：

公开仓库默认不是 OXE 数据集 + 宇树数据集一起训练 action head。

6.1 为什么这么说

在 mixtures.py 中，数据混合被分成了不同类别：

1. OXE / RT-X / magic soup 等 mixture
   包含 Open-X Embodiment 相关数据集，如：

   • fractal20220817_data
   • bridge_orig
   • roboturk
   • 等等

2. Unitree Fine-Tuning Datasets
   单独定义了宇树自己的数据混合，如：

   • Unitree_all_task
   • g1_stack_block
   • g1_pack_pencilbox
   • g1_wipe_table
   • 等多个 g1_* 任务

这说明：

• OXE mixture 和 Unitree mixture 是分别定义的；
• 训练时只会读取你在 data_mix 中指定的那个 mixture；
• 默认并不会自动把 OXE 与 Unitree 拼接起来共同训练。

6.2 训练入口如何使用 mixture

训练脚本中，data_mix 会作为参数传入 dataloader，类似于：

• --datasets.vla_data.data_mix=${data_mix}

然后在数据准备阶段，用这个 data_mix 构建 RLDSDataset(...)。

因此训练时真正加载哪些数据集，完全取决于你传入的 mixture 名称。

6.3 具体意味着什么

• 如果你设置：data_mix=Unitree_all_task
  那么训练的是宇树自己的 12 个任务混合数据；
• 如果你设置：data_mix=rtx 或 oxe_magic_soup
  那么训练的是对应的 OXE mixture；
• 如果你希望 OXE + Unitree 一起训练，则需要你自己在 mixtures.py 中新建一个新的 mixture，把两边的数据集条目都写进去。

因此，公开版本的默认做法更接近：

先用上游 VLM 预训练获得强 backbone，再在某一个指定的机器人数据 mixture 上进行 VLA/action 训练。

而不是：

默认把 OXE 与 Unitree 数据自动合并到同一个 action head 训练阶段中。

---

7. 对整个训练流程的简化理解

可以将 UnifoLM-VLA 的整体训练流程概括为：

阶段一：VLM continued pre-training

• 基础模型：Qwen2.5-VL-7B
• 数据：通用 + 机器人相关多任务视觉语言数据
• 目标：提升空间理解、几何理解与轨迹建模能力
• 输出：UnifoLM-VLM-0 / UnifoLM-VLM-Base

阶段二：VLA / Action Head Training

• 初始化：加载 UnifoLM-VLM-Base
• 输入：图像 + 机器人状态 + 动作标签
• 数据：由 data_mix 指定的数据集混合（例如 Unitree_all_task）
• 目标：通过 action_loss 学习动作预测
• 输出：UnifoLM-VLA-0

---

8. 一句话总结

关于预训练

UnifoLM-VLA 并不是把通用预训练数据与宇树机器人数据从头到尾放在同一阶段混合训练，而是先完成 VLM 的 continued pre-training，再进行 VLA/action 训练。

关于 action head

Action head 是在 VLA 阶段结合图像、状态和动作标签，通过 action loss 进行监督训练的。

关于 OXE + Unitree

公开仓库默认不是 OXE 数据集和宇树数据集一起训练 action head；它们在 mixtures.py 中是分开定义的，实际训练使用哪一类数据，取决于 data_mix 配置。若想联合训练，需要手动自定义新的 mixture。