Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization

Paper: arXiv:2605.26282 Code: Edmond1Cheng/MBDPO Code reference: main @ 97c290bb (2026-05-27)

1. Motivation (研究动机)

现有 World Model RL 的可扩展性瓶颈不只是常见的 model bias / error compounding。论文指出，更结构性的矛盾是：许多基于 TD-MPC / TD-MPC2 的方法用 replay buffer 中的非搜索策略 $\beta$ 来训练 value function $\hat{Q}$，但真正和环境交互时又用 MPPI/CEM 这类 search policy $\pi$ 在 learned world model 中最大化 $\hat{Q}$。也就是说，value learning 的数据分布和 policy improvement 的行动分布并不一致。

这个错位会产生两个连锁问题。第一，$\hat{Q}$ 在 $\beta$ 分布附近被训练，但 search 会主动寻找 $\hat{Q}$ 的高值区域；如果这些高值来自 out-of-distribution 过估计，Bellman backup 会反复 bootstrap 错误峰值。论文把这个 gap 写成： $\Vert T^{\pi }\hat{Q}-T^{\beta }\hat{Q}{\Vert }_{\infty }\le \gamma \Vert \hat{Q}{\Vert }_{\infty }\sqrt{2D_{\mathrm{KL}}^{\max }(\pi \Vert \beta )}.$ 第二，MPPI/CEM 的 search 本质上是一次性 planner，不等价于一个被持续训练、受信赖域约束的 policy。随着 search policy $\pi$ 偏离 replay-buffer policy $\beta$，action drift 变大，value approximation error 又被进一步放大。对 offline pretraining 尤其危险：offline 数据固定，策略不能通过新环境交互快速修正偏移。

论文要解决的具体目标是：在 learned world model 内，把 trajectory search、value learning 和 policy optimization 统一成一个可扩展的 diffusion policy optimization 过程，让 policy improvement 不再是外置 heuristic planner，而是一个和 world model / value function 联合训练、并被 implicit energy / KL trust region 约束的生成式策略。

这个问题值得研究，因为如果 world model RL 想像 LLM / diffusion model 一样靠更大模型和更多任务持续 scaling，policy improvement 必须能随模型容量一起稳定变好。论文的实验证据覆盖 online-from-scratch、multi-task offline pretraining 和 offline-to-online adaptation：MBDPO 在小模型上尤其能缓解 world model 不准时的 policy drift，在大模型上还能保持随参数量提升的 monotonic scaling。

Figure 1 解读：左侧展示 multi-task offline pretraining 中，MBDPO 随参数量从 $1.7$M 到 $340$M 增大而呈现更清晰的 scaling 行为；右侧展示 online-from-scratch 中，MBDPO 在 DMControl、MetaWorld、ManiSkill2、MyoSuite 等控制任务上相对 TD-MPC2/SAC/DreamerV3 更稳定。它不是单个 benchmark 的小改进，而是在不同训练范式下都试图修正同一个“search-value mismatch”。

2. Idea (核心思想)

核心 insight 是：把 world model 中的 policy search 改写为 diffusion denoising 过程，并用 value-induced energy 与 behavior-policy energy 共同塑形 score field。这样，action sequence 从 Gaussian prior 逐步 denoise 成高价值轨迹时，不只是贪心追逐 $\hat{Q}$，还会被 replay buffer / behavior distribution 约束在可信区域附近。

MBDPO 的关键创新有三点。第一，它把 MPC 式的 action-sequence search 写成 DDPM reverse process，优化对象是 horizon $H$ 内的 action sequence $a_{t:t+H}$。第二，它引入 implicit energy function $E_{\theta }$，用 buffer action 作为正样本、策略/随机动作作为 hard negative，让 diffusion 的目标分布近似“高 return 但不偏离行为分布”的 Gibbs policy。第三，它把 world model update、value iteration、energy learning 和 diffusion score fitting 放在同一个交替训练循环里。

和 TD-MPC2 的根本差别是：TD-MPC2 仍依赖 MPPI planner 在每一步从候选轨迹里挑高值 action，planner 和 actor/value 的训练目标并没有严格对齐；MBDPO 则把“搜索”本身变成一个可训练的 diffusion policy distribution，并用 KL-style constraint 显式约束 ${\pi }_{\varphi }(\cdot |z_t)$ 到 $\beta (\cdot |z_t)$ 的偏移。它更接近把 offline RL 的 conservatism、TRPO/PPO 的 trust region、world-model rollout 的 trajectory evaluation 合在一个 score-based policy refinement 框架里。

3. Method (方法)

3.1 Overall framework：从 latent world model 到 diffusion action sequence

MBDPO 仍然使用 TD-MPC2 风格的 latent world model。观测 $s_t$ 先被 encoder $\mathcal{E}$ 映射到 latent state $z_t$；latent dynamics $\mathcal{F}$ 预测 $z_{t+1}$；reward head $\mathcal{R}$ 预测即时奖励；$\hat{Q}$ 提供 bootstrap value；implicit energy $E_{\theta }$ 约束 action distribution；score network $\hat{\varphi }$ 或 analytic model-based score 负责 denoising action sequence。

组件映射可以写成： $\begin{array}{rlrl}& \text{Encoder}:& & z_t=\mathcal{E}(s_t,e),\\ & \text{Latent Dynamics}:& & z_{t+1}=\mathcal{F}(z_t,a_t,e),\\ & \text{Latent Reward}:& & {\hat{r}}_t=\mathcal{R}(z_t,a_t,e).\end{array}$

Figure 2 解读：这张图是 MBDPO 的核心机制。上半部分表示在 learned world model 中 rollout 多条 action sequence；每条 sequence 由 reward/value 评估得到 return $G$。下半部分把这些 trajectory return 转成 diffusion score 的权重，让 action sequence 从 noise 逐步 denoise 到高价值、低 drift 的策略分布。implicit energy 相当于把 replay buffer 的行为分布注入 score field，避免策略只沿着 $\hat{Q}$ 的过估计方向漂移。

直觉上，MBDPO 有效的原因不是“diffusion 比 Gaussian policy 更强”这么简单，而是 diffusion 的多步 refinement 给了 policy improvement 一个可控的路径：每个 denoising step 都可以用 imagined rollout 的 return 重新加权，同时被 behavior energy 拉回可信区域。这样 policy 不会像 MPPI 一样每次临时搜索出一个和 actor/value training distribution 脱节的 action，而是把搜索结果沉淀进一个持续优化的 policy distribution。

3.2 Misalignment analysis：为什么 search 和 value learning 会错位

论文把 replay-buffer policy 记作 $\beta$，search-induced policy 记作 $\pi$。如果 value function 在 $\beta$ 生成的数据上训练，却在 Bellman update 或 action selection 中被 $\pi$ 使用，就会出现 operator gap： $\Vert T^{\pi }\hat{Q}-T^{\beta }\hat{Q}{\Vert }_{\infty }\le \gamma \Vert \hat{Q}{\Vert }_{\infty }\sqrt{2D_{\mathrm{KL}}^{\max }(\pi \Vert \beta )}.$ 这个式子的含义是：只要 search policy 和 behavior/replay policy 的最大 KL divergence 增大，value update 的 mismatch upper bound 就会变大。MPPI 会选择 $\hat{Q}$ 高的候选动作，但这些动作未必在 buffer 分布附近；如果 $\hat{Q}$ 在这些区域过估计，planner 会更频繁访问错误区域，Bellman target 又会把错误扩散回 value network。

Figure 3 解读：左侧 cross TD error 显示 MBDPO 比 TD-MPC2 更低、更稳定，说明它减少了 $T^{\pi }\hat{Q}$ 与 $T^{\beta }\hat{Q}$ 的错位。右侧 action drift 用 $\frac{1}{H}{\sum }_h\log \frac{\pi (a_h|z_h)}{\beta (a_h|z_h)}$ 统计策略偏移；MPPI/TD-MPC2 的 drift 持续较大，而带 contrastive energy 的 MBDPO（论文强调 $\eta =0.1$）更接近当前 policy distribution，训练到约 $1$M steps 后 drift 接近稳定。

3.3 Diffusion policy：把 action refinement 写成 reverse diffusion

MBDPO 不是直接输出单步 action，而是优化一个 horizon 为 $H$ 的 action sequence： $a_{t:t+H}=\{a_t,\dots ,a_{t+H}\}.$ DDPM reverse process 从 Gaussian noise ${\pi }^N(a_{t:t+H}^N|z_t)\sim \mathcal{N}(0,I)$ 开始，经过 $N$ 个 denoising step 得到 clean action sequence： ${\pi }_{\varphi }^0(a_{t:t+H}^0|z_t)={\int }_{a_{t:t+H}^{1:N}}{\pi }_{\varphi }^N(a_{t:t+H}^N|z_t){\prod }_{\tau =N}^1{\pi }_{\varphi }^{\tau -1}(a_{t:t+H}^{\tau -1}|a_{t:t+H}^{\tau },z_t)d{a}_{t:t+H}^{1:N}.$ 每个 reverse step 的 transition kernel 由 score function $\varphi (z_t,a_{t:t+H}^{\tau },\tau )$ 控制： ${\pi }_{\varphi }^{\tau -1}(a_{t:t+H}^{\tau -1}|a_{t:t+H}^{\tau },z_t)=\mathcal{N}\left(a_{t:t+H}^{\tau -1};\frac{1}{\sqrt{{\alpha }^{\tau }}}\left(a_{t:t+H}^{\tau }+(1-{\alpha }^{\tau })\varphi (z_t,a_{t:t+H}^{\tau },\tau )\right),{\sigma }^2(\tau )I\right).$ 实际采样形式是： $a_{t:t+H}^{\tau -1}=\frac{1}{\sqrt{{\alpha }^{\tau }}}\left(a_{t:t+H}^{\tau }+(1-{\alpha }^{\tau })\varphi (z_t,a_{t:t+H}^{\tau },\tau )\right)+\sigma (\tau )ϵ,ϵ\sim \mathcal{N}(0,I).$

3.4 Implicit energy：把 KL trust region 写进 score field

论文用 implicit energy function $E_{\theta }$ 来近似行为分布约束。训练目标是一个二分类/contrastive 形式：buffer action $(z,a)$ 是正样本，负样本 $a_j^{-}$ 来自策略或随机动作： ${\min }_{E_{\theta }}\mathcal{L}(E_{\theta })=-\mathbb{E}\left[\log \frac{\exp (-E_{\theta }(z,a))}{\exp (-E_{\theta }(z,a))+{\sum }_j\exp (-E_{\theta }(z,a_j^{-}))}\right].$ 在 KL-constrained policy optimization 中，MBDPO 形式化为： \begin{split} \max_{\pi_\phi}\quad &\mathbb E_{a_{t:t+H}\sim\pi_\phi(\cdot|z_t)}[G(z_t,a_{t:t+H})]\\ \text{s.t.}\quad &D_{\mathrm{KL}}(\pi_\phi(\cdot|z_t)\|\beta(\cdot|z_t))\le\varepsilon,\\ &z_{h+1}=\mathcal F(z_h,a_h),\quad h=t,\dots,t+H-1. \end{split} Theorem 1 的结论是：上述 KL-constrained optimum 可以写成一个 Gibbs distribution： ${\pi }_{\varphi }^{\ast }(a_{t:t+H}|z_t)\propto \left({\prod }_{h=t}^{t+H}\beta (a_h|z_h)\right)\exp \left(\frac{G(z_t,a_{t:t+H})}{\kappa }\right){\prod }_{h=t}^{t+H-1}{1}_{\{z_{h+1}=\mathcal{F}(z_h,a_h)\}}.$ 这说明目标策略既偏好高 return trajectory，又乘上行为策略密度作为 trust-region anchor。$\kappa$ 是 temperature：越小越强调高 return，越大越平滑。

3.5 Score estimation：用 imagined trajectories 做 Monte Carlo score target

MBDPO 不必先做 forward diffusion data corruption；它可以在 learned world model 中直接从 noisy action sequence 的 posterior 采样 clean candidates，并用 return-weighted average 估计 score： $\varphi (z_t,a_{t:t+H}^{\tau },\tau )=-\frac{a_{t:t+H}^{\tau }}{1-{\bar{\alpha }}^{\tau }}+\frac{\sqrt{{\bar{\alpha }}^{\tau }}}{1-{\bar{\alpha }}^{\tau }}{\mathbb{E}}_{a_{t:t+H}^{0|\tau }}\left[w(a_{t:t+H}^{0|\tau })\cdot a_{t:t+H}^{0|\tau }\right].$ 其中： $w(a_{t:t+H}^{0|\tau })=\frac{\exp (\tilde{G}(z_t,a_{t:t+H}^{0|\tau })/\kappa )}{\mathbb{E}[\exp (\tilde{G}(z_t,a_{t:t+H}^{0|\tau })/\kappa )]}.$ score network 的 supervised fitting loss 是： $\min \mathbb{E}\left[\Vert \hat{\varphi }(z_t,a_{t:t+H}^{\tau },\tau )-\varphi (z_t,a_{t:t+H}^{\tau },\tau ){\Vert }_2^2\right].$ 论文还给出 Monte Carlo 近似，并说明相对于目标 Gibbs distribution 的 distributional error 以 $\mathrm{Poly}(N,T,L_{\varphi })$ 速率下降，其中 $N$ 是 diffusion steps，$T$ 是 Monte Carlo samples，$L_{\varphi }$ 是 score function 的 Lipschitz constant。

3.6 Practical training loop：world model、value、energy、score 联合训练

MBDPO 的 joint loss 同时更新 world model、value 和 energy： $\min {\mathbb{E}}_{\tau \sim \mathcal{B}}\left[{\sum }_{h=t}^{t+H}{\gamma }^{h-t}\left(\Vert \mathcal{F}(z_h,a_h)-\mathrm{sg}(\mathcal{E}(s_{h+1})){\Vert }_2^2+\mathrm{CE}({\hat{r}}_h,r(s_h,a_h))+\mathrm{TD}(\hat{Q}(z_h,a_h))+\mathcal{L}(E_{\theta })\right)\right].$ 训练范式分三种：online from scratch 先随机 warmup 收集 replay buffer，再循环执行 diffusion action、环境交互、world model/value/energy update、score fitting；multi-task offline pretraining 将 replay buffer 初始化为静态 offline dataset 并省略环境交互；offline-to-online fine-tuning 则加载 offline pretrained world model / score network，跳过 warmup 后用少量新交互微调。

3.7 Released code implementation notes：论文公式与 released code 实现差异

论文公式与 released code 实现差异：论文 Algorithm 1 把 score network $\hat{\varphi }$ 作为核心输入，并在 Step 3 中更新 score network；released code 的 cfgs/config.yaml 默认 use_score_network: false。因此 MBDPO/diffusion.py::Diffusion.plan 默认走 analytic model-based score score_mb，只有显式打开 use_score_network=true 时才调用 WorldModel.score 并启用 MBDPO._mc_score_target / score_loss 路径。

论文公式与 released code 实现差异：论文 Appendix Table 4 写 energy regularized factor $\eta =0.1$，但 released cfgs/config.yaml 中对应的 contrastive_eta 默认是 0.01，cfgs/online_parallel_config.yaml 没有覆盖它。复现实验时应检查 checkpoint/config 中是否显式 override 到论文值；否则代码默认行为和论文表格存在差异。

另外，released implementation 中 WorldModel.F 是一个 contrastive behavior score head；MBDPO._estimate_value 将它归一化后作为 shaped reward 的一部分：shaped_reward = reward + contrastive_eta * f_norm。这和论文中 “implicit energy anchors policy updates” 的思想一致，但代码层面不是单独命名为 EnergyNetwork，而是放在 world_model.py 的 _F head 与 mbdpo.py 的 contrastive loss / reward shaping 中。

3.8 Pseudocode from released implementation

Diffusion planning / inference (MBDPO/diffusion.py)

import torch
import torch.nn.functional as F
 
def diffusion_plan(agent, obs, cfg, task=None):
    z0 = agent.model.encode(obs, task)
    x_tau = torch.zeros(cfg.horizon, cfg.action_dim, device=obs.device)
    betas = torch.linspace(cfg.diffusion_beta0, cfg.diffusion_betaT, cfg.diffusion_steps, device=obs.device)
    alphas = 1.0 - betas
    alpha_bar = torch.cumprod(alphas, dim=0)
 
    for tau in reversed(range(1, cfg.diffusion_steps)):
        alpha_bar_tau = alpha_bar[tau]
        a0_samples = (x_tau / torch.sqrt(alpha_bar_tau))[None] + torch.sqrt((1.0 - alpha_bar_tau) / alpha_bar_tau) * torch.randn(
            cfg.diffusion_num_samples, cfg.horizon, cfg.action_dim, device=obs.device
        )
        a0_samples = a0_samples.clamp(-1, 1)
        values = agent._estimate_value(z0.repeat(cfg.diffusion_num_samples, 1), a0_samples.permute(1, 0, 2), task).nan_to_num(0.0).squeeze(-1)
        logits = (values - values.mean()) / (values.std() + 1e-6)
        a_bar = (torch.softmax(logits / cfg.diffusion_temperature, dim=0)[:, None, None] * a0_samples).sum(dim=0)
        score_mb = (-x_tau + torch.sqrt(alpha_bar_tau) * a_bar) / (1.0 - alpha_bar_tau + 1e-8)
        score = agent.model.score(z0, x_tau[None], torch.tensor([tau], device=obs.device), task)[0] if cfg.use_score_network else score_mb
        x_tau = (x_tau + (1.0 - alpha_bar_tau) * score) / torch.sqrt(alphas[tau])
 
    return x_tau.clamp(-1, 1)[0]

Monte Carlo score target (MBDPO/mbdpo.py::_mc_score_target)

def monte_carlo_score_target(agent, z0, task):
    betas = torch.linspace(agent.cfg.diffusion_beta0, agent.cfg.diffusion_betaT, agent.cfg.diffusion_steps, device=z0.device)
    alphas = 1.0 - betas
    alpha_bar = torch.cumprod(alphas, dim=0)
    tau = torch.randint(1, agent.cfg.diffusion_steps, (1,), device=z0.device)
    alpha_bar_tau = alpha_bar[tau].squeeze(0)
    x_tau = torch.randn(agent.cfg.horizon, agent.cfg.action_dim, device=z0.device)
    a0_samples = (x_tau / torch.sqrt(alpha_bar_tau))[None] + torch.sqrt((1.0 - alpha_bar_tau) / alpha_bar_tau) * torch.randn(
        agent.cfg.diffusion_num_samples_mf, agent.cfg.horizon, agent.cfg.action_dim, device=z0.device
    )
    values = agent._estimate_value(z0.repeat(len(a0_samples), 1), a0_samples.permute(1, 0, 2), task).nan_to_num(0.0)
    logits = (values.reshape(len(a0_samples), -1).mean(dim=-1) - values.mean()) / (values.std() + 1e-6)
    weights = torch.softmax(logits / agent.cfg.diffusion_temperature, dim=0)
    a_bar = (weights[:, None, None] * a0_samples).sum(dim=0)
    return x_tau, tau, (-x_tau + torch.sqrt(alpha_bar_tau) * a_bar) / (1.0 - alpha_bar_tau + 1e-8)

Contrastive implicit energy / behavior score (MBDPO/mbdpo.py::_contrastive_loss)

def contrastive_energy_loss(model, z_seq, buffer_actions, task, noise_scale=0.5):
    z_flat = z_seq.reshape(-1, z_seq.shape[-1]).detach()
    a_pos = buffer_actions.permute(1, 0, 2).reshape(-1, buffer_actions.shape[-1])
    pos_logits = model.F(z_flat, a_pos, task)
    loss_pos = F.binary_cross_entropy_with_logits(pos_logits, torch.ones_like(pos_logits))
 
    with torch.no_grad():
        pi_action, _ = model.pi(z_flat, task)
        noisy_action = (pi_action + noise_scale * torch.randn_like(pi_action)).clamp(-1, 1)
        random_action = (2.0 * torch.rand_like(pi_action) - 1.0).clamp(-1, 1)
        candidates = torch.stack([pi_action, noisy_action, random_action], dim=1)
        candidate_logits = model.F(z_flat[:, None].expand(-1, 3, -1).reshape(-1, z_flat.shape[-1]), candidates.reshape(-1, candidates.shape[-1]), task).reshape(-1, 3)
        hard_neg = candidates[torch.arange(candidates.shape[0]), candidate_logits.argmax(dim=1)]
 
    neg_logits = model.F(z_flat, hard_neg, task)
    loss_neg = F.binary_cross_entropy_with_logits(neg_logits, torch.zeros_like(neg_logits))
    return 0.5 * (loss_pos + loss_neg)

Joint update (MBDPO/mbdpo.py::_update_core + _update)

def train_update(agent, replay_batch):
    obs, action, reward, terminated, task = replay_batch
    next_z = agent.model.encode(obs[1:], task).detach()
    td_targets = agent._td_target(next_z, reward, terminated, task)
    z = agent.model.encode(obs[0], task)
    zs, consistency_loss = [z], 0.0
    for t, (a_t, next_z_t) in enumerate(zip(action.unbind(0), next_z.unbind(0))):
        z = agent.model.next(z, a_t, task)
        consistency_loss = consistency_loss + F.mse_loss(z, next_z_t) * (agent.cfg.rho ** t)
        zs.append(z)
    zs = torch.stack(zs)
    reward_loss = reward_prediction_loss(agent.model.reward(zs[:-1], action, task), reward)
    value_loss = td_loss(agent.model.Q(zs[:-1], action, task, return_type="all"), td_targets)
    energy_loss = contrastive_energy_loss(agent.model, zs[:-1], action, task)
    total_loss = agent.cfg.consistency_coef * consistency_loss + agent.cfg.reward_coef * reward_loss + agent.cfg.value_coef * value_loss + agent._contrastive_coef * energy_loss
    if agent.cfg.use_score_network:
        x_tau, tau, target_score = agent._mc_score_target(zs[0, 0:1], task)
        pred_score = agent.model.score(zs[0, 0:1], x_tau[None], tau, task)[0]
        total_loss = total_loss + agent.cfg.score_loss_coef * F.mse_loss(pred_score, target_score.detach())
    total_loss.backward(); agent.optim.step()
    pi_info = agent.update_pi(zs.detach(), task)
    agent.model.soft_update_target_Q()
    return pi_info

3.9 Code-to-paper mapping

Code reference: main @ 97c290bb (2026-05-27) — pseudocode and mapping based on this commit

Paper Concept	Source File	Key Class/Function
Latent world model $\mathcal{E},\mathcal{F},\mathcal{R},\hat{Q}$	MBDPO/common/world_model.py	WorldModel.encode, next, reward, Q
Gaussian policy prior / actor	MBDPO/common/world_model.py	WorldModel.pi
Implicit energy / behavior score	MBDPO/common/world_model.py, MBDPO/mbdpo.py	WorldModel.F, MBDPO._contrastive_loss, MBDPO._estimate_value
Diffusion planner / reverse denoising	MBDPO/diffusion.py	Diffusion.plan
Monte Carlo score target	MBDPO/mbdpo.py	MBDPO._mc_score_target
Joint model/value/energy/score update	MBDPO/mbdpo.py	MBDPO._update_core, MBDPO._update, MBDPO.update
Online and offline launch paths	scripts/train.py, scripts/online_parallel_train.py, scripts/offline_to_online.py	train, parallel launcher, O2O fine-tuning entry
Experiment config values	cfgs/config.yaml, cfgs/online_parallel_config.yaml	horizon, diffusion steps, samples, LR, batch size, steps/seeds

4. Experimental Setup (实验设置)

4.1 Benchmarks and task scale

论文评测 MBDPO 在四类主要 control benchmarks 上的表现，覆盖 paper text 所述的 $121$ 个 diverse control tasks：DMControl、MetaWorld、ManiSkill2、MyoSuite。profiling table 另列 MetaWorld $50$、DMControl $39$、ManiSkill2 $5$、MyoSuite $10$、Visual RL $10$ 个 tasks；论文还包含 Locomotion、PickYCB 等高维变体。

Figure 12 解读：DMControl demo 展示典型低维状态控制任务，包括 locomotion、reacher、cartpole、finger 等，主要考察样本效率和连续控制稳定性。

Figure 13 解读：MetaWorld demo 覆盖按钮、抽屉、门、咖啡机、peg insertion 等机器人操作任务；这些任务对 goal-conditioned latent representation 和 sparse success signal 更敏感。

Figure 14 解读：ManiSkill2 demo 包含立方体堆叠、抓取、插接等复杂接触操作，是检验 learned world model 是否能处理 object-centric interaction 的重要补充。

Figure 15 解读：MyoSuite demo 使用肌骨手部控制任务，action/state 维度高且动力学更复杂；它说明论文不是只在常规 MuJoCo locomotion 上验证。

4.2 Baselines and metrics

Baselines 包括 SAC、DreamerV3、TD-MPC、TD-MPC2。Metrics 包括 episode return、success rate、cross TD error、relative action drift、training wall-clock profiling。对 MetaWorld / manipulation 任务，success rate 是关键指标；对 DMControl / MyoSuite / locomotion，episode return 更常用。Action drift 用 log-likelihood ratio 统计 $\pi$ 与 $\beta$ 的偏移，cross TD error 用来检验 value iteration gap 是否被放大。

4.3 Training config and hardware

论文 Appendix Table 4 给出的默认 6M online model 配置如下；released code 中对应主配置文件是 cfgs/config.yaml，parallel online dispatch 由 cfgs/online_parallel_config.yaml 生成每个任务/seed 的 scripts/train.py 调用。

• Diffusion search：$H=3$；diffusion steps 为 $5$–$20$（released default diffusion_steps: 20）；Monte Carlo samples $512$；temperature $\kappa =0.5$；论文表格 $\eta =0.1$，released default contrastive_eta: 0.01。
• Replay / optimization：buffer capacity $1,000,000$；uniform sampling；batch size online $256$ / offline $1024$；Adam；learning rate $3\times 10^{-4}$；encoder LR 约 $1\times 10^{-4}$；gradient clip norm $20$。
• Default 6M architecture：encoder dim $256$，MLP dim $512$，latent dim $512$，task embedding dim $96$，$5$ 个 $\hat{Q}$ functions，LayerNorm + Mish，$\hat{Q}$ dropout $1\%$。
• Scaling configs / hardware：multi-task 模型为 $1.7$M / $6$M / $21$M / $54$M / $340$M；online profiling 在单张 NVIDIA A800-80GB 上测量，offline 10M-step profiling 在单张 NVIDIA H200 上测量。

5. Experimental Results (实验结果)

5.1 Online-from-scratch：总体 benchmark 曲线

Figure 4 解读：该图汇总 online setting 下 DMControl、MetaWorld、ManiSkill2、MyoSuite 的 aggregate performance。论文未把每条曲线的最终数值整理成表格，但图中趋势显示 MBDPO 在多个 benchmark 上相对 TD-MPC2/SAC/DreamerV3 更快达到高 reward / success rate，并且早期训练方差更低。作者把这归因于 diffusion policy refinement 与 value/world model 同步更新，而不是每步临时调用外置 planner。

Figure 5 解读：Visual RL 的 $10$ 个 image-based control tasks 每个用 $5$ 个 seeds 平均。MBDPO 在高维视觉输入下提升明显，说明 diffusion policy optimization 不只是低维 state-based control 的技巧；当 observation encoder 更难学时，policy drift 约束和 trajectory-level denoising 仍能提供稳定收益。

5.2 Multi-task offline pretraining：scaling behavior

Figure 6 解读：offline pretraining 中，MBDPO 展示从 $1.7$M 到 $340$M 的 monotonic scaling。小模型上 MBDPO 相对 TD-MPC2 优势更明显，因为 world model 不够准时，MPPI 更容易利用 value overestimation；MBDPO 的 implicit KL / energy constraint 能防止 policy 过度偏离 offline data。随着模型变大，TD-MPC2 也受益于更准确的 world model，因此 gap 缩小；论文还观察到 $340$M 在 $30$ tasks 设置下可能受 dataset capacity saturation 影响。

5.3 Offline-to-online：70 tasks pretrain + 10 unseen tasks fine-tune

Figure 7 解读：O2O 设置先在 $70$ 个 tasks 上预训练 generalist agent，再在 $10$ 个 unseen tasks 上每个只用 $40k$ interactions fine-tune。右侧曲线表明 fine-tuning 明显优于 from-scratch，说明 MBDPO 学到的 latent representation 和 diffusion policy prior 能快速适配新任务，而不只是 memorizing offline tasks。

5.4 Bottleneck verification：TD error and action drift

Figure 21 解读：per-task cross TD-error 曲线补充了 Figure 3 左侧的 aggregate 结论。MBDPO 在多个任务上保持较低 TD error，支持论文对 Problem 1 的判断：search/value mismatch 会放大 Bellman operator gap，而 diffusion policy + energy anchor 能缓解这个问题。

Figure 22 解读：per-task action drift 展示不同 search / policy variants 的漂移分布。MPPI 风格 search drift 更大；带 contrastive energy 的 MBDPO 更接近 behavior / policy distribution。这个实验对应 Problem 2：外置 heuristic search 不能保证稳定 policy improvement。

5.5 Ablations：$\eta$、Monte Carlo samples、diffusion steps

Figure 8 解读：$\eta$ 控制 implicit energy / behavior constraint 的强度。论文报告在 $\eta \in [0.3,1]$ 内性能较稳定；$\eta =5$ 时 performance 低于 $70$，因为 behavior cloning 约束太强，限制了 policy 探索 high-value regions。一个关键数值是：即使 $\eta =0$，MBDPO 仍优于 TD-MPC2（$75.29$ vs. $57.1$），说明 diffusion policy optimization 本身也比 heuristic search 更稳定；但 energy term 对最优表现仍重要。

Figure 9 解读：Monte Carlo samples $T$ 增大时，return-weighted score estimate 更稳定，性能随之提高。由于 $T$ 维度可以并行，论文认为增加 $T$ 对实际 runtime 的边际影响比增加 diffusion steps 更小。

Figure 10 解读：diffusion denoise timesteps $N$ 增大时，reverse process 有更多 refinement 机会，因此性能提升；代价是 inference runtime 随 $N$ 近似线性增长。论文把这个现象和理论中的 $\mathrm{Poly}(N,T,L_{\varphi })$ error decay 对应起来。

5.6 Latent representation：物理结构是否被学出来

Figure 11 解读：latent trajectory visualization 用 locally linear embedding 展示 controlled latent rollouts。循环任务如 Cheetah Run Front、Reacher Hard、Cup Spin 中，MBDPO 的 latent trajectory 呈现更清晰的 closed-loop manifold；MetaWorld 顺序操作任务中，trajectory 更像从 initial state 到 target goal 的 directed path。TD-MPC2 对比更 noisy，说明 MBDPO 的 joint optimization 不只提高 reward，还让 representation 更符合物理/因果结构。

Figure 23 解读：Embedding 可视化说明 representation quality 是 MBDPO 的关键前提。论文特别提到 Hopper Hop 的失败与 latent embedding distribution deviation 有关：如果 representation 本身错位，diffusion policy 即使有 search / score mechanism，也很难在错误 latent space 中采样出有效 action。

5.7 Training efficiency and profiling

论文 Table 1 的 exact profiling numbers 如下。

Online-from-scratch training time（ms / env step，单张 NVIDIA A800-80GB）：

Domain	# tasks	TD-MPC2	MBDPO
MetaWorld	50	$1.08\pm 0.16$	$1.08\pm 0.15$
DMControl	39	$0.78\pm 0.71$	$0.76\pm 0.71$
ManiSkill2	5	$6.50\pm 1.83$	$6.91\pm 2.51$
MyoSuite	10	$1.58\pm 0.22$	$1.54\pm 0.19$
Visual RL	10	$3.55\pm 0.51$	$3.04\pm 0.37$

Multi-task offline training time（10M steps，hours，单张 NVIDIA H200）：

Parameters	30 tasks	80 tasks	80/30
1.7M	$34.8\pm 0.6$	$39.9\pm 0.4$	$1.15\times$
6M	$35.0\pm 0.5$	$40.5\pm 0.6$	$1.16\times$
21M	$39.8\pm 0.3$	$42.6\pm 0.4$	$1.07\times$
54M	$52.6\pm 0.1$	$69.6\pm 0.7$	$1.32\times$
340M	$202.9\pm 0.9$	$206.8\pm 0.8$	$1.02\times$

这些 profiling 结果说明 MBDPO 的主要额外成本没有吞掉收益：online 单步训练时间和 TD-MPC2 基本同量级，Visual RL 还从 $3.55$ ms 降到 $3.04$ ms；offline 从 $30$ tasks 扩到 $80$ tasks 的时间倍率为 $1.02\times$–$1.32\times$，$340$M 时几乎不增加（$202.9$ vs. $206.8$ hours）。

5.8 Limitations and conclusions

作者明确列出三类未来方向。第一，受计算资源限制，MBDPO 尚未扩展到 foundation-model-scale 的参数量，真正的 scaling limit 还未验证。第二，实验集中在 simulated environments；真实机器人部署是检验 transferability 和 robustness 的关键下一步。第三，作者认为更强的 pretrained representation（例如 visual / multimodal foundation encoders）可能进一步改善 latent space，从而推动 large-scale RL within world models。

总体结论是：MBDPO 把 world model、value learning、policy optimization 和 behavior-distribution constraint 合成一个 diffusion policy optimization 框架，缓解了 TD-MPC 系列中 search/value mismatch 导致的 value error amplification 与 action drift。它的优势不仅体现在 online reward 曲线，也体现在 offline scaling、O2O transfer、action drift / TD error diagnostics 和 latent trajectory structure 上。