AI Research 技术调研

Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning

16分钟阅读

一句话定位

Diffusion World Model（DWM）用一个条件扩散模型一次前向预测未来多步的状态与奖励（而非逐步递归查询单步动力学模型），把它接入 Dyna 式 model-based value expansion，在离线强化学习（offline RL）里训练价值函数；在 D4RL 9 个 locomotion 数据集上比传统单步动力学模型高出 44% 的性能，达到与最强 model-free 方法（TD3+BC/IQL/Decision Diffuser）持平的 SOTA。ICML 2024。

背景与定位

Model-based RL 的核心瓶颈是单步动力学模型 f(s_{t+1}, r_t | s_t, a_t) 的复合误差（compounding error）：多步 rollout 需要递归调用模型，误差随 horizon 指数累积，mbpo 等方法因此把 rollout horizon 限制得很短。DWM 的问题设定与该系列一脉相承，但换了解法：借助近年”序列建模做决策”的思路（Decision Transformer、Trajectory Transformer、Diffuser、Decision Diffuser 等自回归/扩散序列模型），提出让扩散模型一次采样整段未来轨迹，从根源上消除递归查询——这是与 mbpo、dreamer-v3、td-mpc2 等依赖单步/RNN 状态转移模型的路线的关键分野。

论文明确将自己定位为对 Feinberg et al. (2018) 提出的 Model-Based Value Expansion（MVE）的扩散化替代：标准 MVE 用单步模型递归 rollout 做 TD 目标值估计；DWM 提出 Diffusion Model Value Expansion（Diffusion-MVE），用扩散世界模型一次采样的多步轨迹代替递归 rollout。论文强调这一算法在 offline RL 语境下可以有两种解读：(a) 因为扩散模型只在离线数据上训练，等价于对行为策略的生成式建模，Diffusion-MVE 构成一种通过生成模型实现的价值正则化；(b) 也可以视为用扩散模型生成的合成数据做离线 Q-learning。论文也与同样”预测未来轨迹”的 Diffuser（Janner et al. 2022，扩散模型身兼世界模型与策略、直接预测并执行动作）和 Decision Diffuser（Ajay et al. 2022，扩散模型只建模状态序列、动作由单独的逆动力学模型给出）做了对比——DWM 与 Decision Diffuser 一样不对动作序列建模，但用途不同：Decision Diffuser 的扩散轨迹直接驱动策略执行，DWM 的扩散轨迹只用于训练价值函数，得到的策略本身是 model-free 的，推理时不再依赖世界模型。

模型架构

世界模型 p_θ：条件扩散模型，backbone 为时序 U-Net（非 RSSM/DiT/AR-transformer）

  • 建模对象：长度 T 的子轨迹 x⁽⁰⁾ = (s_t, a_t, r_t, s_{t+1}, r_{t+1}, …, s_{t+T−1}, r_{t+T−1})，不建模未来动作（原因：Diffusion-MVE 用不到未来动作；且已有工作发现扩散建模连续动作精度较差）。
  • 条件变量 y = 当前 RTG（return-to-go）g_t；训练时以 classifier-free guidance（Ho & Salimans 2022）联合学习条件/无条件噪声预测器，null-conditioning 概率 p_uncond = 0.25。
  • 推理时额外固定 s_t、a_t 作为条件（每步去噪都把 x^(k) 的前 dim(s_t)+dim(a_t) 维强制替换为真实的 (s_t, a_t)，见 Algorithm 3），即”部分变量钳制式”条件采样，而非把 (s_t,a_t) 简单编码进 y。
  • 无 latent/tokenizer：直接在原始连续状态-奖励空间扩散，不经 VAE/离散码本。
  • 去噪网络 ε_θ：时序 U-net（沿用 Diffuser/Decision Diffuser 的架构），6 个重复残差块，每块 2 个时序卷积 + group norm + Mish 激活；扩散步 k 经正弦位置编码 + 2 层 MLP 投影，RTG 经 3 层 MLP 投影为 latent，动作 a_t 作为额外条件也经 3 层 MLP 投影，三者拼接注入。
  • 噪声调度：cosine schedule（Nichol & Dhariwal 2021）；训练扩散步数 K=5；采样用 stride sampling 加速，推理步数 N=3（r_infer=0.5，即 N≈K/2 是消融确定的关键分界点，见”训练方法”消融）；低温采样 α=0.5。
  • 序列长度 T=8（主实验）；另训练 T=32 变体验证长 horizon 鲁棒性，H 最长测到 31 步。
  • 折扣 γ=0.99；RTG 按任务归一化，reward scale：Hopper 400、Walker 550、Halfcheetah 1200。
  • 对比架构消融：把扩散模型替换为 4 层 / 4 head 的 Transformer（自回归预测状态-奖励序列，动作除首步外全部 mask 为 0，遵循 Online Decision Transformer 的优化设置），记为 T-TD3BC / T-IQL——用于验证扩散序列建模优于自回归序列建模。

下游策略：actor-critic（TD3+BC / IQL / PQL 三种，均为 model-free 策略，推理不依赖世界模型）

  • Actor π_ψ、Critic Q_ϕ（TD3+BC/PQL 为双 Q）均为 3 层 MLP，256 隐藏单元/层，ReLU；IQL 额外一个状态价值 V_ξ 网络（同规格）做 expectile 回归。
  • Action head：TD3+BC 用确定性策略 + Tanh 输出；IQL 用 Tanh-Normal 随机策略，AWR（Advantage-Weighted Regression）提取策略。
  • 单步动力学基线 f_θ(s_{t+1}, r_t|s_t, a_t)：两个 4 层 MLP（256 隐藏单元/层，ReLU）分别参数化高斯均值/方差（方差层用 SoftPlus 保正）。

数据

  • 来源：D4RL（Fu et al. 2020）9 个 MuJoCo locomotion 数据集——hopper / walker2d / halfcheetah × {medium, medium-replay, medium-expert}，均为 offline 静态数据集（非在线交互）。
  • 规模：论文未披露各 D4RL 子集的具体轨迹数/token 数（沿用 D4RL 官方发布规模，未在文中重述）。
  • 无 sim-to-real / 无跨任务 co-training：DWM 按环境单独训练，论文在 Conclusion 中明确指出”currently trained for each individual environment and is task-agnostic”，把多环境/多任务扩展列为未来方向。
  • 动作标注：数据集自带 (s,a,r,s′) 轨迹，无需额外动作标注流程；训练子轨迹为随机采样的长度-T 滑窗，条件变量 RTG 由窗口内真实 reward 折扣求和现算（g_t = Σ_{h=0}^{T−1} γ^h r_{t+h}）。
  • 数据配比/清洗：论文未做混合采样或过滤，直接在各任务的完整 D4RL 数据集上训练。

训练方法

  • 两阶段 Dyna 式流水线（Algorithm 1）：Stage 1 训练扩散世界模型 p_θ（式 3 的 classifier-free guidance 噪声预测损失）；Stage 2 用预训练好的 p_θ 做 Diffusion-MVE，训练 actor-critic 策略。
  • Diffusion Model Value Expansion（Definition 3.1）：给定 (s_t,a_t)，从 p_θ(·|s_t,a_t,g_eval) 采样整段 (r̂_t,ŝ_{t+1},…,ŝ_{t+T−1},r̂_{t+T−1})，H-step 目标值 Q̂ = Σ_{h=0}^{H−1} γ^h r̂_{t+h} + γ^H Q̂(ŝ_{t+H}, π(ŝ_{t+H}))。与标准 MVE 的本质区别：MVE 用 policy 预测的动作 â_t=π(ŝ_t) 递归查询单步模型（on-policy 式模拟）；Diffusion-MVE 是off-policy的——策略不参与采样过程，扩散模型只在 g_eval 上条件生成。
  • OOD RTG 条件是关键设计：为避免 offline RL 常见的过度悲观，训练/评测时刻意让 g_eval 取分布外（OOD）的较高值，让扩散模型输出”乐观”轨迹；但 g_eval 也不能过高（消融显示 halfcheetah-mr 上 g_eval>0.4 会让实际回报下降）。
  • 三个下游算法实例化：DWM-TD3BC（动作正则化，Eq.11 的 BC 项）、DWM-IQL（价值正则化，expectile τ=0.7 的 V_ξ 回归 + AWR 策略提取，β=3.0）、DWM-PQL（奖励正则化，仿 MOPO：MOPO 原用高斯动力学模型集成、以协方差矩阵的 Frobenius 范数衡量不确定性，DWM 无此参数化，改为从同一 DWM 采 m 条轨迹，取同一时间步跨样本奖励/下一状态差异的最大 L2 范数平方作为不确定性度量，对预测奖励做悲观惩罚，Eq.16，惩罚系数 κ∈{0.01,0.1,1.0} 网格搜索）。
  • λ-return 变体：仿 Dreamer 系列引入 λ-return（λ=0.95）递归计算目标值（Eq.5，λ=1 退化为 vanilla Diffusion-MVE）。消融显示 λ-return 对 DWM-IQL 有益（0.57→0.64）、对 DWM-TD3BC 有害（0.68→0.62）——因为 λ-return 依赖迭代调用 Q̂/V̂，对价值估计更准的方法（IQL 有显式正则化）更友好，TD3+BC 只做策略正则化、价值更容易过估计；据此 DWM-IQL 默认开启 λ-return，DWM-TD3BC 与 DWM-PQL 用 vanilla Diffusion-MVE。
  • 超参数：扩散模型 Adam，lr=1e-4，EMA β=0.995（每 10 iter 更新一次），训练 2×10⁶ iterations，batch size 64。单步动力学基线 Adam，lr=1e-4，训练 1×10⁶ iterations。下游 RL：TD3+BC/IQL 独立训练 1×10⁶ iterations（lr=3e-4，batch 128）；DWM 系算法（用扩散世界模型做 MVE）训练 5×10⁵ iterations；单步基线系算法（O-TD3BC/O-IQL/O-PQL）训练 2×10⁵ iterations（收敛更快）。TD3+BC 超参：policy noise 0.2、noise clip 0.5、policy/target update freq 2、α(BC 权重)=2.5；IQL 超参：expectile 0.7、β=3.0、max weight 100、policy update freq 1。simulation horizon H 在 {1,3,5,7} 上网格搜索，评测 RTG g_eval 按任务网格搜索（如 hopper-medium 搜 [0.6,0.7,0.8]）。
  • 关键消融——扩散步数 K 与推理步比 r_infer：训练 K∈{5,10,20,30,50,100}，推理步比 r_infer∈{0.2,0.3,0.5,1.0}（N=⌈r_infer·K⌉）。发现 r_infer=0.5 是分界岭：N<K/2 时预测精度显著下降；而 r_infer≥0.5 时小 K=5 与更大 K 表现相近。故主实验选 K=5、r_infer=0.5（N=3）。T=32 变体上该结论依然支持 r_infer=0.5，但更偏好 K=10。

Infra（训练 / 推理工程）

  • 论文未披露 GPU 型号、GPU 数量、GPU-小时或墙钟训练时间；也未报告推理 FPS/延迟（策略在推理时是纯 model-free 前向，不再调用扩散模型，故运行时开销主要来自策略网络本身，论文未单独测量）。
  • 已知的间接工程信息：训练与评测均基于自研 PyTorch 实现，架构参考 Decision Diffuser、Diffuser、SSORL 三个开源代码库；扩散推理用 stride sampling（N/K 的采样步比）把全 K 步反向过程压缩到 N=3 步，是文中唯一给出的显式”加速推理”机制。
  • 评测协议：每个算法在每个任务上训练 5 个随机种子实例，每个实例评测 10 个 episode，报告 5 个种子的均值±标准差。

评测 benchmark

D4RL 9 个 locomotion 任务（medium / medium-replay / medium-expert），归一化回报 0–1（1 为专家水平）：

EnvTD3+BCIQLDDO-TD3BCO-IQLO-PQLDWM-TD3BCDWM-IQLDWM-PQL
hopper-m0.58±0.110.48±0.080.49±0.070.39±0.040.45±0.050.63±0.120.65±0.100.54±0.110.50±0.09
walker2d-m0.77±0.090.75±0.150.67±0.160.39±0.150.52±0.240.74±0.140.70±0.150.76±0.050.79±0.08
halfcheetah-m0.47±0.010.46±0.070.49±0.010.44±0.050.44±0.030.45±0.010.46±0.010.44±0.010.44±0.01
hopper-mr0.53±0.190.25±0.020.66±0.150.26±0.050.25±0.030.32±0.030.53±0.090.61±0.130.39±0.03
walker2d-mr0.75±0.190.48±0.230.44±0.260.23±0.130.24±0.070.62±0.220.46±0.190.35±0.140.35±0.13
halfcheetah-mr0.43±0.010.44±0.010.38±0.060.43±0.010.42±0.020.42±0.010.43±0.010.41±0.010.43±0.01
hopper-me0.90±0.280.86±0.221.06±0.110.31±0.180.39±0.190.43±0.181.03±0.140.90±0.250.80±0.18
walker2d-me1.08±0.011.09±0.000.99±0.150.60±0.250.57±0.180.61±0.221.10±0.001.04±0.101.10±0.01
halfcheetah-me0.73±0.160.60±0.230.91±0.010.27±0.120.61±0.220.61±0.220.75±0.160.71±0.140.69±0.13
Average0.690.610.680.370.430.540.680.640.61

Model-free 组平均 0.660，单步动力学 model-based 组平均 0.447，DWM 组平均 0.643 —— DWM 相对单步模型提升 (0.643−0.447)/0.447 ≈ 44%（与摘要”44% performance gain”一致），且追平 model-free SOTA。

长 horizon 鲁棒性（Figure 4.1/Table 4.2）：O-IQL、O-TD3BC 随 simulation horizon 增大明显性能下降（大多数任务在 H=1~2 见顶后掉头），暴露单步模型复合误差；DWM-TD3BC/DWM-IQL 在 H 一路加到 31 仍不明显衰减。用 T=32 训练的 DWM 在 H∈{1,3,7,15,31} 上 9 任务平均最佳回报：DWM-TD3BC T=8 时 0.68±0.10 vs T=32 时 0.60±0.12；DWM-IQL(w/o λ) T=8 时 0.57±0.09 vs T=32 时 0.61±0.10——序列长度超过 T=8 并不进一步提升表现，故主实验选 T=8。

扩散 vs Transformer 序列模型消融（Table 4.3）：T-TD3BC 均值 0.44、T-IQL 均值 0.52，均明显低于 DWM-TD3BC 0.68 / DWM-IQL 0.64；T-IQL 与 O-IQL 相当，T-TD3BC 优于 O-TD3BC 但仍逊于 DWM，验证了扩散序列建模优于自回归 Transformer 序列建模的假设（Transformer 仍受自回归结构的误差累积影响）。

OOD 评测 RTG（Figure 4.3）：DWM-TD3BC/DWM-IQL 在分布外 g_eval 下表现鲁棒，但实际回报与指定 g_eval 不总是精确匹配（return-conditioned RL 的已知问题）；g_eval 过高（如 halfcheetah-mr 上 >0.4）会让性能下降，最优 RTG 因任务而异。

创新点与影响

  • 核心贡献：把”用扩散模型一次预测多步未来”引入 model-based value estimation，提出 Diffusion Model Value Expansion（Diffusion-MVE），从根本上避开单步动力学模型的递归复合误差，同时保持下游策略是纯 model-free（推理时不依赖世界模型，速度不受影响）。
  • 给出两种理论视角统一现象：Diffusion-MVE 既可读作通过生成式建模实现的价值正则化（扩散模型只见过离线数据，等价于行为策略的合成），也可读作用扩散模型生成的合成数据做离线 Q-learning；并论证这与传统”越悲观越安全”的 offline RL 正则化范式不同——DWM 依赖条件在 OOD 高 RTG 上生成”乐观”想象轨迹。
  • 系统性验证扩散世界模型对 long-horizon 模拟的鲁棒性（H up to 31 不掉分），且用架构消融（换成 Transformer）证明这一鲁棒性来自”整段并行生成”而非”序列建模”本身。
  • 作者自述局限（Conclusion）：① DWM 目前按单环境单独训练，是 task-agnostic 的，未验证多环境/多任务泛化；② 论文只在 offline 设定下验证，未探究 DWM 在 online RL 中的表现（避免探索副作用是刻意的取舍）；③ 尽管用了 stride sampling 加速，扩散模型的计算开销仍然偏高，作者认为进一步加速采样是把 DWM 推向更大规模问题的关键。

原始链接

一手源存档（sources/）

  • 未发现独立于 arXiv 的官方博客、GitHub 代码仓库、HF 模型卡或项目主页——已核实 arXiv abstract 页的 “Code, Data, Media” 关联区块（alphaXiv/CatalyzeX/DagsHub/HF/ScienceCast 均为 arXiv 自动生成的第三方索引，非作者提供的官方链接）。
  • 一手源即 arXiv 论文全文（HTML v1，2402.03570），已通读，未入 git；引用见上方 arXiv URL（arXiv 原文 PDF，不入 git）。

关系图谱

反向链接