WMPO：基于视频世界模型的 VLA On-policy 强化学习

论文：WMPO: World Model-based Policy Optimization for Vision-Language-Action Models

作者：Fangqi Zhu、Zhengyang Yan、Zicong Hong、Quanxin Shou、Xiao Ma、Song Guo

机构：香港科技大学（HKUST）、字节跳动 Seed

发布时间：2025 年 11 月

论文链接：arXiv:2511.09515 ｜ 项目主页

发表状态：arXiv 预印本

一句话总结

WMPO 用一个像素级视频生成世界模型当想象环境，完全在"脑内轨迹"上跑 on-policy GRPO 来微调 VLA，既省去真机交互又避免 latent WM 与 VLA 预训练表征的失配；额外通过 Policy Behavior Alignment、噪声帧条件、帧级动作控制三板斧让世界模型能稳定长时推演、并忠实模拟失败模式。

一、问题与动机

1.1 IL 的脆弱性 + 真机 RL 的样本成本

现有 VLA 基本都是 IL 路线，一旦 OOD 就会陷入 compounding error，没能力从失败中恢复。真机 RL 样本代价又过高，现有缓解方案（人类干预或仿真）要么耗人力要么难构造高保真仿真。

1.2 为什么不用 latent world model

Dreamer 系列这类 RSSM/latent WM 训练高效，但它学的隐空间与 VLA 基础模型在网络级图像上预训练得到的视觉表征对不上。论文的核心论点是：

要让 VLA 在"想象环境"里训练而不是丢弃其 web-scale 视觉先验，世界模型就必须输出像素空间的帧，让 VLA 看到的仍是它熟悉的分布。

1.3 WMPO 的三个设计关键

1. Pixel-space video-generative WM：视频扩散模型作环境，解码回像素供 VLA 使用。
2. Policy Behavior Alignment：WM 只在专家数据上训练时根本不会生成失败，要用 policy 自采轨迹微调才能覆盖失败分布。
3. Clip 级自回归 + noisy-frame conditioning + frame-level action control：保证长时展开稳定、动作-帧对齐精准。

这三点齐备后，就可以在"想象环境"里跑 on-policy GRPO（VLA 里少有的非 off-policy 方案），得到比 DPO、真机 GRPO 更强的性能和样本效率。

二、问题形式化

将 VLA 操作建模为 MDP $\mathcal{M}=(\mathcal{S},\mathcal{A},P,R)$：

• 状态空间 $\mathcal{S}=\mathcal{I}\times \mathcal{G}$：图像序列 $I_{0:K}$ + 语言指令。
• 动作空间 $\mathcal{A}$：长度 $K$ 的 action chunk $a_t\in {\mathbb{R}}^{K\times D}$，每维离散化为 256 个 bin（与 OpenVLA-OFT 保持一致）。
• 转移函数：由世界模型 $s_{t+1}\sim p_{\varphi }(s_{t+1}\mid s_t,a_t)$ 实现。
• 奖励函数：学习到的轨迹级二分类器 $R_{\psi }(\tau )\in \{0,1\}$。

目标：

$$\underset{\theta }{max}{\mathbb{E}}_{\tau \sim {\pi }_{\theta },p_{\varphi }}[R_{\psi }(\tau )].$$

这一形式化的关键在于：整个优化循环完全脱离真实环境，靠 $p_{\varphi }$ 展开 + $R_{\psi }$ 打分。

三、核心方法

3.1 Imagined Trajectory Generation

给定 $c$ 帧初始观测 $I_{0:c}$，策略从最近 $m$ 帧 + 指令 $g$ 预测一个动作 chunk：

$$a_{i:i+K}\sim {\pi }_{\theta }(I_{i-m:i},g),$$

世界模型再条件于最近 $c$ 帧 + 当前动作 chunk 生成未来 $K$ 帧：

$$\begin{array}{}\text{(2)}& I_{i:i+K}\sim p_{\varphi }(I_{i-c:i},a_{i:i+K}).\end{array}$$

反复执行至最大长度 $N$，得到一条想象轨迹 $\tau =\{I_{0:N},a_{0:N}\}$，再由 $R_{\psi }$ 给出二元成功/失败标签。实现中 $c=4$、$K=8$。

3.2 世界模型架构：OpenSora + 两处改动

Backbone：基于 OpenSora 的视频扩散模型。

改动 1：2D VAE 替换 3D VAE。OpenSora 原本用 3D VAE 做时空压缩，WMPO 换成 SDXL 的 2D VAE，保留更细的运动细节，避免过度时间压缩引入的伪影。扩散在 VAE 隐空间跑，但最终解码回像素空间供 VLA 消费，这是与传统 latent WM 的本质区别。

改动 2：Noisy-Frame Conditioning。自回归生成会累积误差。训练时对条件帧 $I_{i-m:i}$ 故意加 50/1000 步的扩散噪声，而非使用干净条件，使模型对不完美条件鲁棒。效果是能稳定生成几百帧而无明显质量衰减。

改动 3：Frame-level Action Control (扩展 AdaLN)。动作-帧对齐关键。借鉴 IRASim（Zhu et al. 2025），对每个动作 $a_i$ 用一个 MLP 生成三组调制系数：LayerNorm 前的 ${\gamma }_1^i$ 和 ${\beta }_1^i$、残差连接的 ${\alpha }_1^i$。每个 transformer block 的更新规则为：

$${\mathbf{x}}^i={\mathbf{x}}^i+(1+{\alpha }_1^i)\cdot \text{Block}({\gamma }_1^i\cdot \text{LayerNorm}({\mathbf{x}}^i)+{\beta }_1^i).$$

这样每一帧独立接收对应动作信号与扩散时间步嵌入。

3.3 Policy Behavior Alignment

两阶段训练：

1. 在 Open X-Embodiment (OXE) 数据上预训练，获得泛用机器人动力学先验；
2. 在当前 policy 实际 rollout 的数据（含失败）上微调，对齐下游任务的 (state, action) 分布、并让模型能生成失败模式。

为什么必要：OXE 和下游任务的专家轨迹都以成功为主，若不做这一步，WM 对失败的想象是"虚的"，GRPO 拿不到真实的负样本信号。

3.4 Reward Model

轻量级 VideoMAE 编码器 + 线性头，二分类头预测每个片段是否为成功结尾。

• 正样本：成功轨迹的末尾 clip $c_N=I_{N-L:N}$。
• 负样本：成功轨迹的中段 clip $\{c_i:L\le i\le N-L\}$ 或失败轨迹的任意 clip。
• 训练损失：batch 内正负样本均衡 + BCE。
• 推理：用步长 $s$ 的滑窗扫整条想象轨迹，任一 clip 成功概率超过阈值 ${\tau }_{\text{thr}}$ 就判成功。clip 长度 $L=8$，推理步长为 1。

实验中 F1 > 0.95，足以缓解 reward hacking。

3.5 On-Policy GRPO in Imagination

WMPO 强调on-policy 而非 DPO/离线 PPO 这类 off-policy 做法——理由：off-policy value estimation 天然有偏，物理交互瓶颈使 on-policy 在真机上难实现，但放到想象空间就没有瓶颈了。

采样：从真实初始帧 $I_{0:c}\sim \mathcal{D}$ 出发，当前 ${\pi }_{{\theta }_{\text{old}}}$ 在 WM 中采 $G$ 条想象轨迹 $\{{\tau }_1,\dots ,{\tau }_G\}$，$R_{\psi }$ 为每条打二元标签。

Dynamic Sampling（follow DAPO）：若某组 $G$ 条全部成功或全部失败，则整组丢弃，继续采直到 batch 填满——防止组内方差为零导致梯度消失。

log-prob 预计算：

$$\begin{array}{}\text{(3)}& \log {\pi }_{{\theta }_{\text{old}}}(a_t\mid s_t)=\sum _{i=1}^K\sum _{j=1}^D\log {\pi }_{{\theta }_{\text{old}}}(a_t^{i,j}\mid s_t),\end{array}$$

按 chunk 内每个动作、每个 DoF 逐项相加。

目标函数（去 KL，follow DAPO）：

$$\begin{array}{}\text{(4)}& \mathcal{J}(\theta )={\mathbb{E}}_{s_0\sim \mathcal{D},\{{\tau }_i{\}}_{i=1}^G\sim {\pi }_{{\theta }_{\text{old}}}}[\frac{1}{G}\sum _{i=1}^G\frac{1}{T}\sum _{t=0}^{T}min(r_{i,t}(\theta ){\hat{A}}_i,\text{clip}(r_{i,t}(\theta ),1-{ϵ}_{\text{low}},1+{ϵ}_{\text{high}}){\hat{A}}_i)],\end{array}$$$$\begin{array}{}\text{(5)}& r_{i,t}(\theta )=\frac{{\pi }_{\theta }(a_{i,t}\mid s_{i,t})}{{\pi }_{{\theta }_{\text{old}}}(a_{i,t}\mid s_{i,t})},{\hat{A}}_i=\frac{R_i-\text{mean}(\{R_i{\}}_{i=1}^G)}{\text{std}(\{R_i{\}}_{i=1}^G)}.\end{array}$$

注意几个选择：

• 去掉 KL 正则，不需要 reference model，节省显存并鼓励探索新行为。
• 双 clip 阈值 ${ϵ}_{\text{low}}=0.20$、${ϵ}_{\text{high}}=0.28$（正负偏移不对称，偏向鼓励上移）。
• 优势是轨迹级常量，在整条想象轨迹的每个时间步 $t$ 与每个动作 token 上共享同一个 ${\hat{A}}_i$。

3.6 训练循环（伪代码叙述化）

1. 初始化 ${\theta }_{\text{old}}\leftarrow \theta$。
2. 反复填充 batch $\mathcal{B}$：从数据集采初始态 $s_0$，用 WM+${\pi }_{{\theta }_{\text{old}}}$ 生成 $G$ 条想象轨迹，用 $R_{\psi }$ 打分；若该组全同标签，丢弃重采。
3. 计算每组的均值 $\mu$、标准差 $\sigma$，得到每条轨迹的归一化优势 ${\hat{A}}_i=(R_i-\mu )/\sigma$。
4. 预存 $\{\log {\pi }_{{\theta }_{\text{old}}}(a_t^i\mid s_t^i)\}$。
5. 在 $\mathcal{B}$ 上跑 $E$ 个 epoch，每个 epoch 切成 mini-batch $M$，按式 (4) 更新 $\theta$。
6. ${\theta }_{\text{old}}\leftarrow \theta$，回到步骤 2。

四、实验

4.1 设置

• Base policy：OpenVLA-OFT，在目标任务上做 IL 微调（每任务 300 条专家轨迹；去掉本体感受与腕部相机以简化）。
• 仿真：Mimicgen 四个精细操作任务 —— Coffee_D0、StackThree_D0、ThreePieceAssembly_D0、Square_D0；每任务 128 个初始态评估。
• 真机：Cobot Mobile ALOHA，任务 "Insert the square into the stick"（方块与柱之间仅 5 mm 间隙）。
• rollout 预算 $P$：policy 真实采集的轨迹数（用于 WM 微调 + 训 reward model），实验对比 $P=128$ 和 $P=1280$。
• 基线：
  • Online GRPO：在真实仿真器里跑 GRPO（对照 WMPO 的优势：不依赖仿真）。
  • Offline DPO：用相同的 $P$ 条 rollout 构造成功/失败偏好对。
• 算力：32× H100 训练 WM 与策略。

4.2 主结果（Mimicgen，成功率 %）

$P$	方法	Coffee	StackThree	ThreePieceAssembly	Square	平均
–	Base policy	43.8	46.9	19.5	24.2	33.6
128	GRPO	38.3	52.3	17.2	25.0	33.2
128	DPO	43.8	53.9	23.4	28.1	37.3
128	WMPO	61.7	56.3	37.5	32.8	47.1
1280	GRPO	47.7	54.7	20.3	25.8	37.1
1280	DPO	52.3	57.0	26.7	33.6	42.4
1280	WMPO	75.0	64.1	46.1	45.3	57.6

• $P=128$ 时 WMPO 比最强基线高 +9.8pp，$P=1280$ 时差距扩大到 +15.2pp。
• Online GRPO 在小预算下甚至不如 base（更新次数太少），DPO 则因静态数据复用迅速饱和，WMPO 随 $P$ 增长持续受益。

4.3 涌现行为

• 自纠正（Fig. 3，Square 任务）：base policy 一旦把方块推错碰到柱子，就会"认死理"继续往下压直到超时；WMPO 学到的策略会抬起、重新对齐、再插入。原因在于 WM 生成的失败想象轨迹让 policy 见过并学会纠正——这是 IL 无法提供的经验。
• 更短更顺（Fig. 5）：WMPO 成功轨迹相对长度 ~95–98%，显著短于 base/GRPO/DPO（均接近 100%），说明 WMPO 抑制了"卡住"行为。

4.4 泛化（扰动）

三种 OOD：位置扰动（固定柱 → 随机位置）、背景扰动（白桌面 → 灰背景）、纹理扰动（红底座 → 黑木底座）。成功率（%）：

方法	Pos. Dis.	Bg. Dis.	Tex. Dis.	Mean
Base policy	14.1	46.1	10.9	23.7
GRPO	15.6	47.7	10.9	24.7
DPO	16.4	34.4	7.8	19.5
WMPO	22.3	50.0	16.4	29.6

DPO 在背景/纹理扰动下甚至退化到低于 base（提示它学到了虚假视觉线索），WMPO 各项都最好。

4.5 Lifelong Learning

迭代采集：policy rollout 128 条 → WMPO 更新 → 新 policy 再 rollout 128 条 … 对照加更多专家数据（300 / 428 / 556 条）训练的 IL base。WMPO 在 StackThree 上随迭代稳步上升至 ~63%，DPO 不稳定甚至下降，说明 WM-RL 的可扩展性不靠人类供数。

4.6 真机实验

• 平台：Cobot Mobile ALOHA；200 条专家数据训 base，再用 128 条 policy rollout 同时微调 WM 与训策略。
• 30 trial 成功率：Base 53% ｜ DPO 60% ｜ WMPO 70%。
• Fig. 7/8 展示世界模型能忠实预测成功与失败轨迹；Fig. 9 给出失败案例——WM 预测到最后几帧还是没能捕捉方块被柱卡住的微扰动。

五、局限性

• 策略类限制：当前只处理离散化 action token；flow-matching 类连续策略（如 π₀）的集成留作未来工作，需结合 Flow-GRPO 等。
• 失败轨迹的稀有边缘模式：Fig. 9 中方块卡柱这类极细节物理现象仍会让 WM 失配。
• 对 OXE 预训练的依赖：未讨论若无 OXE 规模的数据，WM 是否仍能可靠泛化到下游任务分布。
• 本体感受与腕部相机被简化掉：真实部署往往需要这些模态。
• Reward model 的轨迹级稀疏信号：对长时多阶段任务，成功/失败二元信号可能不足以驱动精细优化。

六、个人思考

6.1 为什么它敢喊 "pixel-space matters"

这篇论文和近两年 Dreamer 谱系形成了明显对立。Dreamer 的核心赌注是：状态抽象得越狠、latent 越紧凑，越利于策略优化。但这个赌注只在从零训练策略时成立。一旦策略本身是带有 web-scale 视觉先验的 VLA，换成一个陌生的 latent 分布，等价于把 VLA 视觉表征的大部分价值丢掉了。WMPO 的观察是恰当的：world model 的表征空间必须服从 policy 的表征空间，而不是反过来。

6.2 Policy Behavior Alignment 的巧思

这一步很像 RLHF 里的 "reward hacking 防护"——若 WM 只在成功轨迹上训练，一旦 policy 进入失败分支就会被 WM 送回"幻觉成功"，GRPO 拿到的就是虚假正样本。作者让 policy 自己先去踩坑，把踩坑数据喂回 WM，让 WM 学会想象失败——这是 WMPO 的隐形关键（和 3.3 中提及的一致）。

6.3 去 KL 的代价

follow DAPO 去掉 KL 正则，在 LLM 里常见；但在 VLA 上，策略完全可以往 WM 的盲点漂移（类似 model-based RL 经典的 exploitation of WM error）。WMPO 用 Dynamic Sampling 与 noisy-frame conditioning 抵消这种风险，但没有显式约束策略留在支持集内，这比 LLM GRPO 更危险。可能的后续工作：加个 behavior-cloning regularizer 或小 KL，以及把优势从轨迹级细化到 step/chunk 级。

6.4 自回归视频 WM 与 action-chunk VLA 的天作之合

Action chunk 长度 $K=8$，WM 也一次生成 8 帧——两者节奏完全对齐。这不是偶然：VLA 的 chunk-level 推理让 WM 不需要一次展开上百帧，把 compound error 压到了一个可控窗口内。换言之，action chunking 本身是 WM-RL 能跑通的一个前置条件。

6.5 与同领域工作的关系

维度	WMPO	VLA-RL / SimpleVLA-RL	GRAPE
环境	学到的像素 WM	真仿真器	离线轨迹对
算法	GRPO on-policy	PPO/GRPO online	DPO offline
数据成本	中（需少量真轨迹 + OXE 预训 WM）	高（真机/仿真在线）	低（复用轨迹）
对 sim-to-real 依赖	无（无需建仿真器）	高	无
核心贡献	像素 WM + on-policy 想象 RL	真环境 scale-up	偏好对齐

WMPO 本质上是"用 WM 模拟把 online RL 的样本成本降到 offline RL 水平"。若 WM 本身的 scale laws 可被继续放大（Cosmos、Genie-3 这样的世界基础模型），这条路径的天花板很高。

参考

• Kim et al., Fine-tuning Vision-Language-Action Models: Optimizing Speed and Success, 2025 —— OpenVLA-OFT，WMPO 的 base policy。
• Zheng et al., Open-Sora: Democratizing Efficient Video Production for All, 2024 —— WMPO 世界模型的 backbone。
• Zhu et al., IRASim: A Fine-Grained World Model for Robot Manipulation, 2025 —— frame-level action control 的思路源头。
• Shao et al., DeepSeekMath, 2024 —— GRPO 算法来源。
• Yu et al., DAPO: An Open-Source LLM Reinforcement Learning System at Scale, 2025 —— 去 KL、动态采样、双 clip 的依据。
• Mandlekar et al., MimicGen, 2023 —— 仿真评测基准。
• Rafailov et al., Direct Preference Optimization, 2023 —— DPO 基线。
• Hafner et al., Mastering Diverse Domains through World Models, 2023 —— 被 WMPO 明确对立的 latent WM 代表。