WAM-RL: World-Action Model Reinforcement Learning with Reconstruction Rewards and Online Video SFT

5,760 字 · English 395 词 · 约 18 分钟阅读

论文信息 - 作者:Zezhong Qian(北京大学), Xiaowei Chi(北京大学,Project Leader), Yu Qi(东北大学), Haozhan Li(清华大学), Zhi Yang Chen(北京大学), Shanghang Zhang(北京大学,通讯作者) - 发表:arXiv 2026 - arXiv ID:2606.17906 - 代码:未提供

本文基于以下本地材料整理:arXiv TeX 源码(notes/2026-06/2026-06-24/arXiv-2606.17906v1/)


一、核心问题

World-Action(WA)模型在机器人策略学习中展现了强大的泛化能力和数据效率,但它们为什么无法持续进化?

以 Genie Envisioner、Cosmos Policy、DreamZero 为代表的 WA 模型通过视频生成联合建模未来观测和动作,利用视觉预测结构为长时域决策提供归纳偏置,在仿真和真实场景中均表现优异。然而,现有 WA 模型主要依赖专家轨迹的监督学习进行训练,这带来了两个根本性限制:

  1. 技能受限于演示分布:学到的策略被训练数据的支持集所约束,难以在演示分布之外习得精细操控技能。
  2. 缺乏持续改进能力:模型无法通过与环境的在线交互持续改进,而这对于适应新场景和纠正执行错误至关重要。

一个自然的方向是将强化学习(RL)引入 WA 范式。但直接套用现有 VLA 模型的 RL 方法行不通:VLA 的 RL 方法主要优化动作策略而固定视觉表征;而 WA 模型中,动作模型深度依赖世界模型的潜空间,naively 施加 RL 或在线微调会导致潜表征分布漂移(distribution shift),引发不稳定和性能退化。

图1:WAM-RL 整体框架

图 1:WAM-RL 框架总览。通过在线交互联合优化世界模型和动作模型(actor)。世界模型生成想象的未来观测,由 actor 翻译为动作并在真实环境中执行。产生的观测用于更新两个组件:世界模型通过成功轨迹的在线视频自监督微调加 KL 正则化来精炼,actor 则通过衡量想象结果与执行结果一致性的重建稠密奖励来优化。

为什么 WA 模型的 RL 比 VLA 更难?

WA 模型由两个紧耦合组件构成:

论文的核心观察是:WA 模型的主要能力来自世界模型,而动作模型只是将潜预测转换为动作的翻译器。因此,naive 地只优化 actor(如 VLA 的常见做法)是不够的——世界模型的预测精度从根本上约束了 actor 的表现上限。


二、核心思路与方法

2.1 背景:Flow Matching 与 Flow-SDE(逐公式解读)

在讲 WAM-RL 的方法之前,先理解其技术基础。本节把每个公式的原理和直觉讲透。

2.1.1 什么是 Flow Matching?一个直观类比

想象你有一大盆完全浑浊的水(高斯噪声),你想让它变成一杯分层漂亮的鸡尾酒(真实视频数据)。Flow Matching 就是学习"把每一滴水分子从浑浊位置推到鸡尾酒位置"的速度场。

形式化:你有两个分布——

模型要学习一个时变向量场 $v_\theta(x_t, t)$,告诉你:在时刻 $t$、当前位置 $x_t$,应该往哪个方向、以多大速度移动。

$$\frac{d x_t}{d t} = v_\theta(x_t, t)$$

逐符号解读:

怎么训练? 需要告诉模型"正确答案"是什么:

$$\mathcal{L}_{\text{FM}} = \mathbb{E}_{x_0, x_1, t} \left[ \left\| v_\theta(x_t, t) - v^*(x_t, t) \right\|^2 \right]$$

逐符号解读:

为什么叫 Flow Matching? 因为模型在"匹配"(matching)真实的"流"(flow)——学习一条从噪声通向数据的运输线。

2.1.2 Flow-SDE:为什么要在 ODE 里加噪声?

上面的 Flow Matching 是确定性的 ODE(常微分方程)——给定起点 $x_0$,终点 $x_1$ 唯一确定,没有任何随机性。但 RL 需要两样东西:

  1. 探索(Exploration):不能每次都走同样的路——需要偶尔随机偏离,发现更好的策略
  2. 可计算的动作概率(Action Likelihood):RL 的策略梯度公式需要 $\log \pi_\theta(a|s)$——"在当前状态 $s$ 下选择动作 $a$ 的概率有多大"。确定性过程没有概率可言

Flow-SDE 的解决方案:在原本的确定性 ODE 中注入布朗运动(白噪声):

$$d x_t = v_\theta(x_t, t)\, dt + \sigma\, dW_t$$

逐符号解读:

关键效果——动作变得可计算概率了:

加入噪声后,每一步的去噪转移从一个"确定的跳跃"变成了一个高斯分布:

$$p(x_{t-1} \mid x_t) = \mathcal{N}\big(x_{t-1}; \mu_\theta(x_t, t), \sigma^2 I \big)$$

逐符号解读:

通俗理解:确定性的 ODE 说"从 A 点你一定会走到 B 点";SDE 说"从 A 点你可能走到 B 点附近的一个范围,越靠近 B 越可能,但也可能偏一点"。

2.1.3 从 SDE 到 REINFORCE 策略梯度:完整推导链

这一步是整个 WAM-RL 数学链路的核心:如何把一个视频生成模型变成一个可以用 RL 优化的策略?

第一步:把去噪轨迹解释为 MDP

第二步:计算整个动作序列的概率

一个去噪过程走 $T$ 步(比如从纯噪声 $x_0$ 到干净视频 $x_1$),整个轨迹的对数概率是各步对数概率之和:

$$\log \pi_\theta(a \mid s) = \sum_{t=1}^{T} \log p(x_{t-1} \mid x_t)$$

逐符号解读:

第三步:REINFORCE 策略梯度

$$\nabla_\theta J = \mathbb{E} \left[ \nabla_\theta \log \pi_\theta(a \mid s)\, A(s,a) \right]$$

逐符号解读——每个部分在做什么:

为什么不用 PPO/GRPO? PPO 需要 clip + 重要性采样(多次 rollout),GRPO 需要同一 prompt 的多条采样做组内对比——都与单步视频生成的高昂推理成本冲突。Flow-SDE 已提供 $\log \pi_\phi$,重建奖励直接给 $A_t$(无需 critic bootstrap),REINFORCE 足够。

2.2 整体框架:联合优化两大组件

WAM-RL 的核心设计理念是:让世界模型和 actor 协同进化。这通过两个协调机制实现:

组件 优化方式 核心约束
世界模型 在线视频 SFT(成功轨迹) KL 正则化稳定潜空间
动作模型(Actor) RL + 重建稠密奖励 奖励信号对齐世界模型的预测目标

2.3 在线视频 SFT + KL 正则化:稳定地改进世界模型

为什么需要 KL 正则化?

直觉上,用 RL 交互中收集的成功轨迹来微调世界模型是自然的——让模型更好地预测任务相关的动态。但直接微调会导致严重不稳定:

Actor 深度依赖世界模型的潜特征分布。在线更新会显著偏移这一分布,导致 actor 学到的策略与演化的潜空间之间出现不匹配——actor 迅速失效。

KL 正则化的具体做法(逐公式解读)

第一步:把确定性特征变成概率分布

DiT 架构的世界模型输出的中间特征 $z_t = f_\theta(x_{<t})$ 是确定性的——同样的输入永远得到同样的 $z_t$。但 KL 散度只能衡量两个概率分布之间的差异。所以论文做了一个巧妙的高斯近似:

$$p_\theta(z_t \mid x_{<t}) = \mathcal{N}\big(z_t, \Sigma_\theta\big), \quad p_{\text{old}}(z_t \mid x_{<t}) = \mathcal{N}\big(z_t^{\text{old}}, \Sigma_{\text{old}}\big)$$

逐符号解读:

第二步:KL 散度——衡量分布漂移的尺子

$$\mathcal{L}_{\text{KL}} = \mathbb{E}_{t} \left[ D_{\text{KL}}\Big( \mathcal{N}(z_t, \Sigma_\theta) \;\|\; \mathcal{N}(z_t^{\text{old}}, \Sigma_{\text{old}}) \Big) \right]$$

逐符号解读:

第三步:最终损失——在学习和稳定之间折中

$$\mathcal{L}_{\text{WM}} = \mathcal{L}_{\text{video}} + \lambda_{\text{KL}} \mathcal{L}_{\text{KL}}$$

直观理解:KL 项约束更新后的世界模型保持预训练模型的潜特征几何结构,防止 abrupt distribution shift,同时允许渐进式适应。论文观察到这种方法稳定了联合训练,但改进是温和的——反映了稳定性和适应性之间的权衡。

2.4 基于重建奖励的 Actor RL(逐公式解读)

核心直觉:为什么用"想象 vs 现实"作为奖励?

Actor 的职责是将世界模型的"想象"翻译为真实动作。传统 RL 奖励是任务级的("抓到了 +1,没抓到 0")——稀疏、延迟、难以优化。WAM-RL 换了一个思路:既然世界模型擅长预测未来,那就让 actor 的目标是"让真实执行结果尽可能接近世界模型的预测"。

这就像一个导演(世界模型)画了分镜(未来视频预测),演员(actor)照着演。演员演得好不好,不看票房(任务奖励),而看表演和分镜有多像(重建一致性)。

第一步:重建奖励的定义

$$r_t = \mathrm{sim}(\hat{x}_{t+1:t+H}, x_{t+1:t+H})$$

逐符号解读:

与常规 RL 奖励的区别:

第二步:四种相似度函数的原理与取舍

函数 计算方式 测量什么 直觉
Pixel MSE $\|\hat{x} - x\|^2$ 逐像素平均 低级外观匹配 "两张图叠在一起看,每个像素差多少"
Optical Flow MSE 先提取光流场,再算 MSE 运动模式一致性 "物体移动的方向和速度是否一样"
DINOv2 MSE 用 DINOv2 提取语义特征,再算 MSE 高层语义对齐 "物体类别、结构、关系是否一致"
V-JEPA2 MSE 用 V-JEPA2 提取视频特征,再算 MSE 视频级时空语义 "整个动作序列的语义是否合理"

关键实验发现(反直觉):

Optical Flow 在奖励空间中区分成功/失败的 gap 最大(见第四节图),但下游任务成功率反而最低。Pixel MSE 区分度弱但效果最好。

原理解释:

第三步:策略梯度优化(REINFORCE)

$$\nabla_\phi J = \mathbb{E} \left[ \nabla_\phi \log \pi_\phi(a_t \mid s_t) A_t \right]$$

逐符号解读(与 2.1.3 节的链路对接):

注意:$A_t$ 这里不是 PPO 的 GAE(需要 critic 网络做 bootstrap),而是直接从重建奖励计算。论文没有详细说明 $A_t$ 的具体计算方式(可能是简单的 reward-to-go 或带 baseline 的 REINFORCE),但核心不依赖价值网络。


三、实验与结果

3.1 设置

3.2 主要结果

方法 LIBERO-Object RLBench (Water Plants)
Base 68% 19%
$\pi_{\text{RL}}$(actor-only) 78% 18%
WAM-RL (Ours) 82% 22%

关键发现:

  1. 短时域任务(LIBERO-Object):Actor-only RL($\pi_{\text{RL}}$)从 68% 提升到 78%(+10%),说明优化 actor 本身对短时域任务是有效的。但 WAM-RL 的联合优化进一步提升到 82%(+4%),说明世界模型的改进即使在短时域也有增益。

  2. 长时域任务(RLBench Water Plants):Actor-only RL 完全无效——$\pi_{\text{RL}}$ 的 18% 甚至略低于 Base 的 19%。这是因为在长时域下,actor 受世界模型预测精度的约束,累积预测误差无法仅靠优化 actor 来弥补。只有联合优化世界模型才能将成功率从 19% 提升到 22%。

  3. 恢复行为的涌现:在线视频 SFT 后的世界模型更频繁地生成恢复行为——例如抓取失败后预测重新调整夹爪、重新尝试抓取的动作。这种恢复动态的预测使得 actor 能执行更鲁棒的策略。

3.3 重建奖励消融

图2:不同重建奖励的归一化奖励分布

图 2:不同重建目标的归一化奖励分布。每个指标被归一化使成功奖励 = 1,失败奖励按比例缩放。误差棒 = 2× 标准差。Optical Flow 在成功/失败之间展现出最强的分离度,但其下游任务成功率反而不如 Pixel MSE。

方法 RLBench 成功率
Base 19%
$\pi_{\text{RL}}$ 18%
Pixel MSE 21%
Optical Flow MSE 19%
DINO MSE 16%
V-JEPA2 17%

核心洞察——奖励区分度 ≠ 优化有效性:

论文解释:

  1. Pixel MSE 与世界模型的训练目标(预测视觉观测)对齐更好,因此奖励信号与模型内部表征更一致
  2. Pixel MSE 对 OOD 动作施加更强惩罚:当 actor 偏离世界模型预测时,像素差异急剧增大,起到了正则化效果
  3. 这表明:奖励信号的"与模型训练目标的一致性"比"成功/失败区分度"更重要

四、视频 SFT 定性分析

图3:视频 SFT 对单次开环执行块行为的影响

图 3:有无视频 SFT 在一个 action chunk 中的行为对比。上图:无视频 SFT,模型在初始错误后无法恢复,快速漂移为 OOD 行为。下图:有视频 SFT,模型学会预见失败并生成恢复动作,最终成功完成任务。

无视频 SFT(上图):

有视频 SFT(下图):

这一定性对比直接说明了:视频 SFT 让世界模型不再假设理想执行,而是学会推理可能的偏差及纠正方式。


五、核心洞察与技术亮点

  1. 首次将 RL 引入 WA 范式:WAM-RL 是第一个在 World-Action 框架下实现强化学习的工作,提出联合优化世界模型和动作模型的训练策略。

  2. 关键发现:长时域任务必须联合优化:Actor-only RL 在短时域有效但在长时域失效——因为世界模型的累积预测误差是根本瓶颈。只有协同改进世界模型才能突破。

  3. KL 正则化稳定潜空间:通过将确定性 DiT 特征高斯化 + KL 约束,在允许世界模型适应新数据的同时防止潜空间 abrupt shift,解决了联合训练的不稳定性。

  4. 重建奖励 vs 任务奖励:用"想象 vs 执行的一致性"(重建奖励)而非传统任务奖励来优化 actor——使 actor 忠实于世界模型的规划,而非追逐可能误导的稀疏任务信号。

  5. 奖励区分度悖论:实验揭示了"奖励区分度 ≠ 优化有效性"的反直觉现象——Optical Flow 区分度最高但效果差,Pixel MSE 区分度低但效果好。奖励信号与模型训练目标的一致性才是关键。

  6. 恢复行为涌现:视频 SFT 后世界模型内化了失败-恢复动态,单个 action chunk 内即可预测纠错动作,无需显式规划或额外推理。


六、局限性

  1. KL 正则化的天花板:KL 约束在稳定训练的同时限制了世界模型的适应幅度——模型难以大幅超越预训练分布的能力边界,尤其在更大规模下这可能成为瓶颈。

  2. 重建奖励的判别力有限:当前奖励信号(Pixel MSE、Optical Flow 等)基于预训练表征或手工设计的相似度指标,成功/失败的对比度有限。论文建议未来探索更具表达力和任务感知的奖励形式。

  3. 任务覆盖面有限:仅在 LIBERO-Object 和 RLBench Water Plants 两个设定上评估,未测试更多样化的长时域操控任务和真实机器人场景。

  4. 计算成本高:8× A800 GPU 训练 8 小时——在线 RL + 视频 SFT 的混合训练对算力需求不低。

  5. 无代码开源:论文未提供代码实现,复现门槛较高。

  6. 未探索的维度:世界模型和 actor 的更新频率/比例、KL 系数的最优调度策略、更大规模下的 scaling 行为均未系统研究。


七、关键概念速查

概念 英文 简洁定义
WA 模型 World-Action Model 联合建模未来观测和动作的视频生成式策略模型
WAM-RL WAM Reinforcement Learning 首个将 RL 引入 WA 范式的框架,联合优化世界模型和 actor
Flow Matching — 学习向量场将噪声传输到数据的连续时间生成模型
Flow-SDE — 在 Flow ODE 中注入布朗运动,使其支持随机探索和似然估计
KL 正则化 KL Regularization 约束更新后世界模型的潜分布接近预训练模型,防止分布漂移
重建奖励 Reconstruction Reward 衡量世界模型预测与真实执行结果一致性的稠密奖励
在线视频 SFT Online Video SFT 用 RL 交互中收集的成功轨迹在线微调世界模型的视频预测能力
恢复行为 Recovery Behavior 世界模型在预测失败后自动生成纠错动作(如重新抓取)的能力
$\mathcal{L}_{\text{WM}}$ World Model Loss $\mathcal{L}_{\text{video}} + \lambda_{\text{KL}}\mathcal{L}_{\text{KL}}$,世界模型的联合训练目标
RL 算法 RL Algorithm REINFORCE 风格策略梯度,非 PPO/GRPO;Flow-SDE 提供 $\log \pi_\phi$,重建奖励提供 $A_t$
Actor-only RL — 仅优化动作模型而固定世界模型的 RL 方法,长时域下失效