Reinforcement Learning with Inverse Rewards for World Model Post-training (RLIR)

Authors: Yang Ye, Tianyu He, Shuo Yang, Jiang Bian Affiliations: Microsoft Research arXiv: 2509.23958 Year: 2025

1. Motivation (研究动机)

• World model 的 action-following 能力不足：当前视频世界模型在视觉质量和时间一致性方面取得了显著进展，但准确执行人类指定动作（action-following）的能力仍被忽视。这对可控和交互式世界建模至关重要。
• RL post-training 难以直接迁移到 world model：在 NLP 领域，RLHF 已被广泛用于 LLM 对齐，但迁移到视频世界模型面临两大障碍：
  1. 人工偏好标注成本极高：为视频数据收集大规模人类偏好标注昂贵且存在偏差（如 VideoAlign 使用 180k 标注）
  2. 无法构建 rule-based verifier：与代码/数学不同，视频生成质量无法通过规则直接验证
• 现有 reward 方案的缺陷：
  • Human preference reward（如 VideoAlign）：粒度粗（video-level），受人类评估者的偏差和噪声影响
  • Pixel-level reward（如 RLVR-World 的 $L_1$ + LPIPS）：与 pre-training 目标高度相关导致 RL 探索受限；对所有像素均匀加权忽略 action-relevant 区域；易受 reward hacking（如整体变暗以降低 loss）

核心问题：如何为视频世界模型设计一种客观、可验证、无偏差的 reward signal，以通过 RL post-training 提升 action-following 能力？

---

2. Idea (核心思想)

核心洞察：与其在高维视频空间直接评估生成质量（不可验证），不如将视频映射回低维动作空间进行评估（可验证）。具体来说，用一个 Inverse Dynamics Model (IDM) 从生成的视频中反推动作序列，与输入的 ground-truth 动作对比，即可获得精确、客观、frame-level 的 reward signal。

与现有方法的本质区别：

• VideoAlign 等方法在高维视频空间评估 $\to$ 主观、粗粒度、有偏
• RLVR-World 用像素级 loss $\to$ 与 pre-training 目标冗余、忽略语义、易 reward hack
• RLIR 将评估空间从视频降维到动作 $\to$ 客观、可验证、语义对齐、frame-level 粒度

---

3. Method (方法)

3.1 整体框架

Figure 2 解读：RLIR 的整体流程分为三步：(1) World model 接收初始状态 $x_0$ 和动作序列 $a_1,\dots ,a_n$，生成视频帧序列 ${\hat{x}}_1,\dots ,{\hat{x}}_n$；(2) Inverse Dynamics Model (IDM) 从生成的视频中反推动作 ${\hat{a}}_1,\dots ,{\hat{a}}_n$；(3) 比较反推动作与 ground-truth 动作，计算 verifiable reward，通过 GRPO 优化 world model。右侧展示了两种 world model 架构：自回归模型（VQ-VAE + Transformer Decoder）和扩散模型（VAE + DiT Blocks）。

3.2 Inverse Dynamics Model (IDM) 作为 Reward Model

IDM 的定义：给定 $T$ 个观测 $o_t:t\in [\mathrm{1...}T]$，IDM 估计从 $o_t$ 到 $o_{t+1}$ 的转移动作：

$$p_{\text{IDM}}(a_t|o_t,o_{t+1})$$

IDM 通过最小化 ground-truth 动作的负对数似然进行训练。

Figure 1 解读：展示了 IDM 对细微视觉变化和动作差异的高灵敏度。左侧（Environmental Mismatch）：手动修改树干裂痕模拟局部生成失败，IDM 检测到不一致并输出错误动作 “attack,back”（而非 ground-truth “attack”）。右侧（Magnitude Deviation）：IDM 能可靠区分视觉差异极小的 “forward” 和 “sprint” 动作。VPT 的 IDM 在 Minecraft 中达到 90.6% 按键预测准确率，鼠标回归 $R^2=0.97$。

Reward 函数定义：

对于生成的轨迹 $T_j=[x_0,{\hat{x}}_1,\dots ,{\hat{x}}_n]$，IDM 预测动作 ${\hat{a}}_1,\dots ,{\hat{a}}_n$，reward 为：

$$R_{T_j}=\frac{1}{n}\sum _{i=1}^{n}r({\hat{a}}_i,a_i),r({\hat{a}}_i,a_i)\triangleq \{\begin{array}{ll}1,& \text{if }{\hat{a}}_i=a_i,\\ 0,& \text{otherwise.}\end{array}$$

关键优势：

• 可验证：动作匹配是 binary 判断，无需主观评价
• 无额外偏差：IDM 仅在视频+动作标注上训练，不引入人类偏好偏差
• Frame-level 粒度：每帧独立评估，语义对齐
• 自然关注 action-relevant 区域：IDM 聚焦于与动作相关的视觉变化

3.3 Group Relative Policy Optimization (GRPO)

GRPO 不需要 value network，通过组内归一化估计 advantage：

$${\hat{A}}_{i,t}=\frac{R_i-\text{mean}(\{R_i{\}}_{i=1}^G)}{\text{std}(\{R_i{\}}_{i=1}^G)}$$

优化目标（带 clipped objective 和 KL penalty）：

$${\mathcal{J}}_{\text{GRPO}}(\theta )={\mathbb{E}}_{q\sim \mathcal{D},\{o_i{\}}_{i=1}^G\sim p_{{\theta }_{\text{old}}}(\cdot |q)}\left[\frac{1}{G}\sum _{i=1}^G\frac{1}{|o_i|}\sum _{t=1}^{|o_i|}\left(\min \left(\frac{p_{\theta }^{i,t}}{p_{{\theta }_{\text{old}}}^{i,t}}{\hat{A}}_{i,t},\text{clip}\left(\frac{p_{\theta }^{i,t}}{p_{{\theta }_{\text{old}}}^{i,t}},1-\epsilon ,1+\epsilon \right){\hat{A}}_{i,t}\right)-\beta D_{\text{KL}}[p_{\theta }\Vert p_{\text{ref}}]\right)\right]$$

3.4 应用于自回归 World Model (MineWorld)

• Baseline: MineWorld（基于 LLaMA 架构的 visual-action autoregressive Transformer）
• Tokenization: VQ-VAE 编码视觉帧 $x_i$ 为 quantized codes $t$，action tokenizer 编码动作 $a_i$ 为离散 tokens
• 数据格式：$(t_0^{x_i},\cdots ,t_n^{x_i},[\text{aBOS}],t_0^{a_i},\cdots ,t_8^{a_i},[\text{aEOS}])$
• 关键设计 — Loss Masking：仅对 visual tokens 计算 policy gradient loss，对 action tokens 进行 loss masking。原因：action tokens 来自外部输入，不应被优化；若优化 action tokens 会导致不良训练动态

注：代码搜索未找到开源实现。以下伪代码基于论文描述重建，不是对真实源码的映射。

# Pseudocode: RLIR for Autoregressive World Model (MineWorld)
def rlir_autoregressive_train_step(world_model, idm, batch):
    """One RLIR training step for autoregressive world model."""
    for (x0, actions) in batch:
        # Step 1: Rollout — 用当前 policy 采样 G 条轨迹
        trajectories = []
        for g in range(G):  # G = rollout (24)
            frames = [x0]
            for i in range(n):  # n = 15 actions → 16 frames
                # Tokenize: VQ-VAE encode frame, action tokenizer encode action
                visual_tokens = vqvae_encode(frames[-1])
                action_tokens = action_tokenize(actions[i])
                # Autoregressive next-token prediction (top-p = 0.8)
                next_visual_tokens = world_model.generate(
                    context=[visual_tokens, action_tokens],
                    sample_method="top_p", p=0.8
                )
                next_frame = vqvae_decode(next_visual_tokens)
                frames.append(next_frame)
            trajectories.append(frames)
 
        # Step 2: IDM 计算 reward
        rewards = []
        for traj in trajectories:
            r_total = 0
            for i in range(n):
                a_hat = idm.predict(traj[i], traj[i+1])  # 预测动作
                r_total += 1.0 if a_hat == actions[i] else 0.0
            rewards.append(r_total / n)
 
        # Step 3: GRPO advantage 估计 (组内归一化)
        mean_r = mean(rewards)
        std_r = std(rewards)
        advantages = [(r - mean_r) / std_r for r in rewards]
 
        # Step 4: Policy gradient update (仅对 visual tokens)
        for g, traj in enumerate(trajectories):
            for t, token in enumerate(traj.tokens):
                if token.is_visual:  # Loss masking: 跳过 action tokens
                    ratio = world_model.prob(token) / old_model.prob(token)
                    clipped_ratio = clip(ratio, 1-eps, 1+eps)
                    loss += -min(ratio * advantages[g],
                                 clipped_ratio * advantages[g])
            loss += beta * kl_divergence(world_model, ref_model)
 
        optimizer.step(loss)

3.5 应用于扩散 World Model (NFD)

• Baseline: NFD (Next-Frame Diffusion)，使用 Diffusion Transformer Blocks + Block-wise Causal Attention
• 关键适配 — Flow-ODE 转 SDE：将 NFD 的确定性 Flow-ODE 转换为等价的 SDE（保持边际概率密度不变），使去噪过程可被建模为 Markov Decision Process (MDP)

MDP 形式化：

$${\mathbf{s}}_t\triangleq (\mathbf{c},t,{\mathbf{z}}_t),\pi ({\mathbf{a}}_t\mid {\mathbf{s}}_t)\triangleq p({\mathbf{z}}_{t-1}\mid {\mathbf{z}}_t,\mathbf{c})$$ $$P({\mathbf{s}}_{t+1}\mid {\mathbf{s}}_t,{\mathbf{a}}_t)\triangleq ({\delta }_{\mathbf{c}},{\delta }_{t-1},{\delta }_{{\mathbf{z}}_{t-1}})$$ $$R({\mathbf{s}}_t,{\mathbf{a}}_t)\triangleq \{\begin{array}{ll}r({\mathbf{z}}_0,\mathbf{c}),& \text{if }t=0,\\ 0,& \text{otherwise.}\end{array}$$

其中 $\mathbf{c}$ 为 action-conditioning 输入，reward 仅在最终去噪完成时（$t=0$）给出。

# Pseudocode: RLIR for Diffusion World Model (NFD)
def rlir_diffusion_train_step(nfd_model, idm, batch):
    """One RLIR training step for diffusion world model (NFD)."""
    for (x0, actions) in batch:
        # Step 1: Rollout — 逐帧生成，每帧通过去噪生成
        trajectories = []
        for g in range(G):  # G = rollout (24)
            frames = [x0]
            for i in range(n):  # n = 15 frames
                # Encode conditioning: action → linear layer → AdaLN-Zero
                c = action_embed(actions[i])
                # 采样初始噪声 z_T ~ N(0, I)
                z_T = sample_noise()
                # SDE 去噪 (T steps, 默认 10 steps)
                z_t = z_T
                for t in reversed(range(T)):  # T = sampling steps
                    # 添加 SDE 噪声 (noise level epsilon_t = 0.75)
                    noise = epsilon_t * randn_like(z_t) if t > 0 else 0
                    # 去噪一步: policy π(a_t | s_t)
                    z_t = nfd_model.denoise_step(z_t, t, c,
                                                  prev_frames=frames) + noise
                frame = vae_decode(z_t)  # z_0 → pixel space
                frames.append(frame)
            trajectories.append(frames)
 
        # Step 2: IDM 计算 reward (同自回归版本)
        rewards = compute_idm_rewards(trajectories, actions, idm)
 
        # Step 3: GRPO 优化
        advantages = group_normalize(rewards)
 
        # Step 4: 对去噪过程的每个 timestep 计算 policy gradient
        # Timestep selection: 从 [0, tau*T] 中均匀采样 (tau = 0.6)
        for g, traj in enumerate(trajectories):
            selected_t = uniform_sample(0, int(tau * T))
            # 重新注入噪声到 z_0 得到 z_selected_t
            z_t = add_noise(traj.z_0, selected_t)
            # 计算 log prob ratio
            ratio = nfd_model.log_prob(z_t, selected_t) - old_model.log_prob(z_t, selected_t)
            loss += -min(exp(ratio) * advantages[g],
                         clip(exp(ratio), 1-eps, 1+eps) * advantages[g])
 
        optimizer.step(loss)

3.6 Code-to-Paper 映射表

Paper 概念	对应组件	说明
Inverse Dynamics Model (IDM)	VPT pretrained model (Baker et al., 2022)	在 2000 小时 Minecraft gameplay 上训练，输入连续帧预测动作
Autoregressive World Model	MineWorld (Guo et al., 2025b)	基于 LLaMA 架构，VQ-VAE tokenization，300M/700M/1.2B
Diffusion World Model	NFD (Cheng et al., 2025)	Diffusion Transformer + Block-wise Causal Attention，310M/774M
Reward Function $R_{T_j}$	Eq. 3: frame-level action match accuracy	Binary reward，IDM 预测动作 == ground-truth 则 1，否则 0
GRPO Advantage	Eq. 1: 组内归一化	${\hat{A}}_i=(R_i-\mu )/\sigma$
GRPO Objective	Eq. 2: clipped + KL penalty	PPO-style clipping + optional KL divergence penalty
Loss Masking (AR)	仅对 visual tokens 计算 loss	避免优化外部输入的 action tokens
Flow-ODE → SDE (Diffusion)	Eq. 5: MDP formulation	加入噪声使确定性去噪变随机，支持 RL exploration
Timestep Selection (Diffusion)	$\tau =0.6$，从 $[0,\tau T]$ 采样	聚焦早期去噪步骤（对最终质量影响最大）

---

4. Experimental Setup (实验设置)

数据集

• VPT Dataset (Baker et al., 2022)：Minecraft gameplay 视频
• 预处理：移除 GUI 交互帧和静态场景帧
• 分辨率：$384\times 224$
• 约 1000 个训练样本即可收敛

Baseline 模型

模型	类型	参数量	架构
MineWorld	Autoregressive	300M / 700M / 1.2B	LLaMA + VQ-VAE
NFD	Diffusion	310M / 774M	DiT + Block-wise Causal Attention

评估指标

• Action-Following: F1, Precision, Recall（通过 IDM 推断动作后计算分类指标，9 类动作）
• 视频质量: FVD (Frechet Video Distance), PSNR
• 图像质量: VBench Image Quality (MUSIQ)
• 动态性: VBench Dynamic Score
• 人类评估: Action-following 能力 + Visual quality（两维度打分）

训练配置

超参数	MineWorld	NFD
Learning rate scheduler	cosine	cosine
Learning rate	$3\times 10^{-5}$	$1\times 10^{-5}$
Optimizer	AdamW	AdamW
Rollout	24	24
Clip Ratio ($\epsilon$)	0.2	0.2
Samples per iteration	32	16
Sampling steps	-	10
Noise level (${ϵ}_t$)	-	0.75
Timestep selection ($\tau$)	-	0.6
硬件	AMD MI300X GPUs	AMD MI300X GPUs
推理帧数（训练时）	16 frames	16 frames

---

5. Experimental Results (实验结果)

5.1 主实验结果 (Table 1)

Model	Param.	F1↑	Recall↑	Precision↑	FVD↓	PSNR↑	Img. Qual.↑	Dynamic
MineWorld	300M	0.70	0.71	0.72	246	15.13	0.675	0.97
+ RLIR	300M	0.77	0.76	0.79	231	15.58	0.672	0.97
MineWorld	700M	0.70	0.71	0.72	231	15.32	0.677	0.96
+ RLIR	700M	0.81	0.80	0.84	207	15.78	0.678	0.97
MineWorld	1200M	0.76	0.73	0.73	227	15.69	0.682	0.97
+ RLIR	1200M	0.81	0.81	0.83	205	15.99	0.684	0.96
NFD	310M	0.69	0.69	0.71	212	16.46	0.678	1.00
+ RLIR	310M	0.76	0.76	0.77	195	17.38	0.687	0.99
NFD	774M	0.77	0.78	0.78	184	16.95	0.692	0.99
+ RLIR	774M	0.83	0.83	0.85	180	17.48	0.688	1.00
GT	-	0.87	0.86	0.88	-	-	0.704	1.00

关键发现：

• Action-following F1 一致提升 5-11%（从 0.70→0.81, 0.69→0.76 等）
• 视觉质量同步提升：FVD 降低，PSNR 提升，接近 GT 上界
• Post-training 后模型几乎达到 IDM 准确率上界（GT 的 F1=0.87）

5.2 与其他 Reward 方法的对比 (Table 2)

Figure 5 解读：三种 reward 方法的训练曲线对比。RLIR 的 reward 稳定上升并收敛；VideoAlign 的 reward 上升但实际效果不佳（人类评估者引入偏差和噪声）；$L_1$ + LPIPS 的 reward 反而下降（reward hacking，模型通过变暗画面来降低 loss）。

Model	Method	F1↑	FVD↓	PSNR↑	IQ↑
MineWorld 700M	Baseline	0.70	231	15.32	0.677
	+ $L_1$ + LPIPS	0.71	228	15.47	0.673
	+ VideoAlign	0.73	219	15.50	0.669
	+ RLIR	0.81	207	15.78	0.678
NFD 774M	Baseline	0.77	184	16.95	0.692
	+ $L_1$ + LPIPS	0.77	193	17.09	0.645
	+ VideoAlign	0.76	181	17.45	0.689
	+ RLIR	0.83	180	17.48	0.688

RLIR 在所有关键指标上全面优于 pixel-level reward 和 human preference reward。

5.3 人类评估

Figure 4 解读：人类评估结果（Win/Tie/Lose）。在 action-following 维度，MineWorld 46.7% Win vs 11.1% Lose，NFD 52.2% Win vs 10.0% Lose。在 visual quality 维度同样显示明显优势（MineWorld 50.0% Win, NFD 73.3% Win）。

5.4 定性分析

Figure 3 解读：三个定性对比案例。Case 1：baseline 无法准确描绘挖掘动作；Case 2：baseline 距离感知不精确，角色位置偏移；Case 3：快速运动时 baseline 产生局部像素模糊。RLIR post-training 后均有效解决。

5.5 Ablation Studies

KL Penalty (MineWorld)：

• 小模型（300M）受益于 KL penalty ($\beta =1e\text{-}4$)：F1 从 0.69→0.77
• 大模型（700M, 1200M）不需要 KL penalty，去掉后效果更好

Denoising Steps (NFD)：

• 10 步最优，20 步接近，40 步性能提升缓慢且边际递减
• Noise level ${ϵ}_t$: 0.5-0.75 效果相似，过低（0.25）则增益减弱

5.6 Limitations

1. 性能上界受限于 IDM 准确率：IDM 并非完美（GT 的 F1 = 0.87 而非 1.0），这限制了 RLIR 能达到的理论最优
2. 模型规模受限：受算力限制，最大模型仅 1.2B 参数，可能无法充分展现 RLIR 在大模型上的 scaling 潜力
3. 仅在 Minecraft 环境验证：尚未在更多样化的环境（如真实世界视频、机器人操控）中验证泛化性