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

01 动机

VLA 模型（Vision-Language-Action models）在通用机器人操作上展现出强大潜力，但对专家 demonstration 的依赖使其无法从失败中学习、也难以进行自我修正。强化学习（RL）虽然能解决这一问题，却在真实机器人上存在极高的 sample complexity。

"RL addresses these through self-improving interactions with the physical environment, but suffers from high sample complexity on real robots."

现有三类 VLA 训练范式各有局限：Imitation Learning（IL）依赖人类 demonstration，策略脆弱易出错；Off-policy RL（如 DPO）利用历史数据但性能受限；On-policy RL（如 GRPO）需要大量真实机器人交互，成本高昂。WMPO 通过像素级世界模型在"想象"空间中生成 on-policy 轨迹，从根本上解决了这一矛盾。

图 1：三种 VLA 训练范式对比。(a) Imitation learning 从人类 demonstration 学习；(b) Off-policy RL 利用历史收集的轨迹；(c) WMPO 利用世界模型生成 on-policy 的"想象"轨迹，无需真实机器人交互。

57.6%WMPO 仿真平均成功率（P=1280）

+15.2pp相对 DPO 的提升（P=1280）

70%真实机器人成功率（30次试验）

0.95+Reward model F1 score（全任务）

02 方法

WMPO 框架由三个核心模块组成：像素级视频世界模型、轻量级 reward model 和基于 GRPO 的 on-policy 策略优化。整体思路是让 policy 在世界模型生成的"想象轨迹"中进行 on-policy 更新，从而规避真实机器人交互的高成本。

WMPO 训练流程 — 图 2：WMPO 整体训练流程。从初始状态 s₀ 出发，世界模型生成 G 条想象轨迹；reward model 对每条轨迹评分；GRPO 计算 normalized advantage 并更新 VLA policy。三个组件交替迭代，实现 on-policy 的自我提升。

像素级视频世界模型（Pixel-Space World Model）

与常见的 latent world model 不同，WMPO 采用像素级预测，使"想象"轨迹与 VLA 预训练特征（基于海量网络图像）保持对齐。世界模型基于 video diffusion 架构：将 OpenSora 的 3D VAE 替换为来自 SDXL 的 2D VAE，以更好地保留运动细节。模型首先在 Open X-Embodiment（OXE）数据集上预训练，再通过 Policy Behavior Alignment 在策略自身采集的真实轨迹上微调，以解决 state-distribution mismatch 问题——确保模型能准确模拟当前 policy 产生的失败模式。两项技术增强保证了质量：(1) Noisy-frame conditioning：条件帧加入 50/1000 步扩散噪声，增强鲁棒性；(2) Frame-level action control：通过扩展的 AdaLN block 实现精确的动作-帧对齐，支持 clip-level 自回归视频生成。

轻量级 Reward Model

Reward model 在真实轨迹上训练，以二元信号（成功/失败）预测任务完成情况，为 GRPO 提供 sparse reward 信号。在所有仿真任务上的 F1 score 均超过 0.95，确保了训练信号的可靠性。

On-Policy GRPO 策略优化

WMPO 采用 Group Relative Policy Optimization（GRPO）进行 on-policy 学习。对每个初始状态采样 G 条想象轨迹，reward model 评分后计算 normalized advantage，再通过 clipped policy gradient 更新 policy：

𝒥(θ) = 𝔼[1/G ∑ min(r_i,t(θ) Â_i, clip(r_i,t(θ), 1−ε_low, 1+ε_high) Â_i)]

为保证有效学习，采用 dynamic sampling 策略：丢弃结果完全一致（全成功或全失败）的 group，确保每个 batch 包含成功与失败轨迹的混合，从而产生有意义的梯度信号。VLA backbone 为 OpenVLA-OFT，动作空间离散化为每维度 256 个 bin。

自我纠错行为分析 — 图 3：Square 任务中 WMPO 的自我纠错能力分析。WMPO 在遭遇失败后能主动调整策略、恢复到正确轨迹，而 base policy 和 DPO 则无法做到。

03 实验

实验在 Mimicgen benchmark 的四个精细操作任务（Coffee_D0、StackThree_D0、ThreePieceAssembly_D0、Square_D0）上进行仿真评估，并在真实机器人上验证 square 插入任务。Baselines 包括：base policy（OpenVLA-OFT imitation learning）、GRPO（on-policy，真实机器人交互）和 DPO（off-policy）。

Table 1：仿真基准成功率（%）

Rollout Budget	方法	Coffee	StackThree	ThreePieceAssembly	Square	Mean (%)
—	Base policy	43.8	46.9	19.5	24.2	33.6
P=128	GRPO	38.3	52.3	17.2	25.0	33.2
	DPO	43.8	53.9	23.4	28.1	37.3
	WMPO	61.7	56.3	37.5	32.8	47.1
P=1280	GRPO	47.7	54.7	20.3	25.8	37.1
	DPO	52.3	57.0	26.7	33.6	42.4
	WMPO	75.0	64.1	46.1	45.3	57.6

WMPO 在小 rollout 预算（P=128）下即超越基线 +9.8pp（vs. DPO 37.3%），在大预算（P=1280）下领先 +15.2pp（vs. DPO 42.4%），充分体现出 on-policy 更新的 sample efficiency 优势。

Table 2：泛化能力（%）

方法	Position Disruption	Background Disruption	Texture Disruption	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

图 4a：Position Disruption — 物体位置扰动下的泛化测试。WMPO 22.3% vs. DPO 16.4%。

图 4b：Background Disruption — 背景替换下的泛化测试。WMPO 50.0% vs. DPO 34.4%。

真实机器人实验

持续学习（Lifelong Learning）

轨迹长度分析 — 图 5：不同策略在成功试验中的相对平均轨迹长度。WMPO 的自我纠错能力使其路径更高效。

持续学习实验中，迭代收集 P=128 条真实轨迹后执行 WMPO 更新，再用更新后的策略继续采集。结果表明 WMPO 实现了"稳定且显著的提升"，而 DPO 由于训练不稳定无法持续改进（StackThree 任务上验证）。

局限性消融：世界模型失败案例

世界模型预测失败案例 — 图 9：世界模型预测失败示例。在某些复杂场景下，世界模型生成的视频与真实物理动态存在偏差，导致 reward 信号失真。

04 局限性

说明：以下局限性部分来自论文明确陈述，部分为从方法设计推断（已标注）。

仅支持离散化动作表示（论文明确陈述）

"While the WMPO framework can in principle support flow-based policies, this work focuses on discretized action representations."论文将 flow-matching based policies 的扩展列为未来工作，当前版本每维度离散为 256 个 bin，限制了连续动作空间场景的适用性。

世界模型预测质量受限（论文明确陈述）

论文在 Figure 9 中展示了世界模型预测失败的案例：在某些复杂操作场景下，生成的视频与真实物理动态存在偏差，导致 reward model 收到失真信号，影响策略优化质量。

Policy Behavior Alignment 带来额外训练成本（从设计推断）

为缓解 state-distribution mismatch，世界模型需要在每轮迭代中用当前策略的真实轨迹进行微调（Policy Behavior Alignment），这意味着每次策略更新后世界模型也需要同步更新，增加了整体训练复杂度和计算开销。

仅在有限任务集上验证（从设计推断）

实验仅在 Mimicgen 的四个操作任务和单一真实机器人场景上验证，尚未扩展到更多元的任务类型（如灵巧手操作、移动操作等），泛化能力有待进一步验证。