Diffusion for World Modeling: Visual Details Matter in Atari

01 动机（Motivation）Motivation

现有世界模型普遍依赖离散隐变量进行压缩，这种有损压缩会丢弃对策略决策至关重要的视觉信息。作者认为，小的视觉细节——例如远处的行人或交通灯——可能直接改变智能体的行为，而离散编码无法忠实还原这些细节。Existing world models generally rely on discrete latents for compression, and such lossy compression discards visual information that is crucial for policy decisions. The authors argue that small visual details, for example a distant pedestrian or a traffic light, can directly change an agent's behavior, and that discrete codes cannot faithfully reproduce them.

"Discrete latent representations involve lossy compression that may discard visual information crucial for learning. Small details in the visual input, such as a traffic light or a pedestrian in the distance, may change the policy of an agent."

DIAMOND imagination procedure — **Figure 1.** DIAMOND 的想象展开过程。顶行展示策略 π_ϕ 在学习到的扩散世界模型 D_θ 的"想象"中执行动作序列。横轴为环境时间 t，纵轴为扩散的去噪时间 τ（从 T 到 0 逆向流动）。每一帧的生成都经历一个完整的去噪过程，保留丰富的像素级细节。**Figure 1.** DIAMOND's imagination rollout. The top row shows the policy π_ϕ executing a sequence of actions inside the "imagination" of the learned diffusion world model D_θ. The horizontal axis is environment time t and the vertical axis is diffusion denoising time τ, flowing backwards from T to 0. Every frame is produced by a complete denoising process, which preserves rich pixel-level detail.

1.46Mean HNS（Atari 100k）
世界模型训练智能体最优Mean HNS (Atari 100k)
Best among agents trained in a world model

0.64IQM 得分
超越 STORM (0.636)IQM score
Ahead of STORM (0.636)

11/26超越人类水平的游戏数量Number of games above human level

3 NFE每帧去噪步数
vs. IRIS 的 16 NFEDenoising steps per frame
vs. 16 NFE for IRIS

扩散模型天然具备在连续空间建模、灵活捕捉多模态分布、无模式崩塌等优势。核心问题在于：能否将扩散模型高效、稳定地用于在线强化学习的世界模型，从而让智能体在高质量的想象轨迹中学习？Diffusion models are naturally suited to modeling in continuous space, capture multimodal distributions flexibly and do not suffer mode collapse. The core question is whether a diffusion model can serve efficiently and stably as the world model for online reinforcement learning, so that the agent learns inside high-quality imagined trajectories.

02 方法（Method）Method

DIAMOND 将扩散模型作为环境动态模型，在像素空间直接预测下一帧图像。智能体在世界模型生成的"梦境轨迹"中通过 REINFORCE 算法学习策略，无需访问真实环境（除数据收集阶段外）。DIAMOND uses a diffusion model as the environment dynamics model, predicting the next frame directly in pixel space. The agent learns its policy with the REINFORCE algorithm inside the "dream trajectories" generated by the world model, without accessing the real environment (except during data collection).

EDM 扩散框架选择Choosing the EDM diffusion framework

关键设计决策是采用 EDM（Elucidating the Design Space of Diffusion-Based Generative Models）而非 DDPM 框架。两者的核心差异在于预训练目标：The key design decision is to adopt EDM (Elucidating the Design Space of Diffusion-Based Generative Models) rather than the DDPM framework. Their core difference lies in the training objective:

DDPM：预测噪声，当噪声主导时模型趋向于学习恒等映射，导致长时序误差累积（compounding error），世界模型快速漂移出分布。DDPM: predicts the noise; when noise dominates, the model tends to learn the identity mapping, which leads to compounding error over long horizons and makes the world model drift out of distribution quickly.
EDM：预测干净图像，"adaptive mixing of signal and noise...means that the model is trained to predict the clean image when the noise is dominant."——即使在单步去噪下也能保持长时序稳定性。EDM: predicts the clean image, and its "adaptive mixing of signal and noise...means that the model is trained to predict the clean image when the noise is dominant." — long-horizon stability is preserved even under single-step denoising.

DDPM vs EDM trajectory stability — **Figure 3.** 基于 DDPM（左）与 EDM（右）的扩散世界模型在 Breakout 游戏上的想象轨迹对比，从同一初始帧出发，每行对应不同的去噪步数 n。DDPM 出现严重的误差累积，世界模型迅速漂移出分布；EDM 即使在单步去噪下也表现出长时序稳定性。**Figure 3.** Imagined trajectories on Breakout from diffusion world models based on DDPM (left) and EDM (right), all starting from the same initial frame, with each row corresponding to a different number of denoising steps n. DDPM exhibits severe compounding error and the world model drifts out of distribution quickly; EDM remains stable over long horizons even under single-step denoising.

网络架构与条件化Network architecture and conditioning

动态模型采用标准 U-Net 2D 架构，将历史 L 帧图像观测与动作拼接为条件输入：The dynamics model uses a standard U-Net 2D architecture, concatenating the past L frame observations with the actions as conditioning input:

图像条件化：将过去 L 帧观测按通道维度拼接后输入 U-Net。Image conditioning: the past L observations are concatenated along the channel dimension and fed into the U-Net.
动作条件化：通过 adaptive group normalization 层将动作注入残差块，实现动作对世界模型的精细控制。Action conditioning: actions are injected into the residual blocks through adaptive group normalization layers, giving fine-grained control of the world model by actions.
预条件化去噪器（Eq. 6）：自适应地混合有噪观测与网络预测输出，在高噪声时以干净信号为主要训练目标，低噪声时以差值（即高斯噪声）为目标。Preconditioned denoiser (Eq. 6): it adaptively mixes the noisy observation with the network prediction, so that at high noise the training target is mainly the clean signal and at low noise it is the residual (i.e. the Gaussian noise).

多步去噪与多模态分布Multi-step denoising and multimodal distributions

采用 n=3 步 Euler 方法去噪。单步采样在环境状态多模态（如 Boxing 中黑色选手移动方向不可预测）时产生模糊预测；多步采样将生成"驱向特定模式"，产生清晰图像。Denoising uses n=3 Euler steps. Single-step sampling yields blurry predictions when the environment state is multimodal (for instance the unpredictable movement direction of the black player in Boxing); multi-step sampling instead drives the generation towards a particular mode and produces sharp images.

Single vs multi-step sampling in Boxing — **Figure 4.** Boxing 游戏中单步（顶行）与多步（底行）采样对比。黑色选手的移动方向不可预测，单步去噪在多种可能结果间插值，导致模糊预测；多步采样则生成清晰图像，"driving the generation towards a particular mode"。**Figure 4.** Single-step (top row) versus multi-step (bottom row) sampling in Boxing. The movement direction of the black player is unpredictable, so single-step denoising interpolates between several possible outcomes and gives a blurry prediction; multi-step sampling produces a sharp image, "driving the generation towards a particular mode".

RL 智能体架构RL agent architecture

奖励与终止模型采用独立的 CNN-LSTM 网络；Actor-Critic 共用一个 CNN-LSTM 主干，分别接策略头与价值头。策略学习使用带价值基线的 REINFORCE 算法，价值更新使用 Bellman 误差与 λ-returns。The reward and termination models are separate CNN-LSTM networks; the actor-critic shares one CNN-LSTM trunk with a policy head and a value head on top. Policy learning uses REINFORCE with a value baseline, and the value function is updated with the Bellman error and λ-returns.

03 实验（Experiments）Experiments

在 Atari 100k 基准（26 款游戏，每款 5 个随机种子，每种子 100k 帧真实交互）上与 SimPLe、IRIS、DreamerV3、TWM、STORM 等世界模型基线对比。计算资源：每款游戏约 2.9 天（单张 Nvidia RTX 4090），共约 1.03 GPU 年。Evaluated on the Atari 100k benchmark (26 games, 5 random seeds per game, 100k real interaction frames per seed) against world-model baselines including SimPLe, IRIS, DreamerV3, TWM and STORM. Compute: about 2.9 days per game on a single Nvidia RTX 4090, roughly 1.03 GPU years in total.

Mean and IQM human normalized scores — **Figure 2.** 各智能体在 Atari 100k 上的 Mean HNS 与 IQM 得分（含置信区间）。DIAMOND（蓝色）获得 Mean HNS 1.46、IQM 0.64，超越所有世界模型内训练的智能体。**Figure 2.** Mean HNS and IQM scores (with confidence intervals) of each agent on Atari 100k. DIAMOND (blue) obtains Mean HNS 1.46 and IQM 0.64, surpassing all agents trained inside a world model.

智能体Agent	Mean HNS (↑)	IQM (↑)	超人类游戏数 (↑)Superhuman games (↑)
SimPLe	0.332	0.130	1
TWM	0.956	0.459	8
IRIS	1.046	0.501	10
DreamerV3	1.097	0.497	9
STORM	1.266	0.636	10
DIAMOND（本文）DIAMOND (ours)	1.459	0.641	11

IRIS vs. DIAMOND 视觉一致性对比IRIS vs. DIAMOND: visual consistency

定性分析展示了 IRIS 在连续帧中出现的视觉不一致现象，而 DIAMOND 无此问题：Qualitative analysis reveals visual inconsistencies across consecutive frames for IRIS, which DIAMOND does not exhibit:

IRIS inconsistencies — IRIS 生成的连续帧：白框标注处可见敌人与奖励显示混淆、砖块/分数渲染不一致等错误。Consecutive frames generated by IRIS: the white boxes mark errors such as an enemy confused with the reward display and inconsistent rendering of bricks and scores.

DIAMOND consistency — DIAMOND 生成的连续帧：无上述不一致现象，视觉细节保持忠实连贯。Consecutive frames generated by DIAMOND: none of these inconsistencies appear, and the visual details stay faithful and coherent.

CS:GO 扩展实验CS:GO scaling experiment

为验证方法的可扩展性，作者在 Counter-Strike: Global Offensive 的 Dust II 地图上训练了大规模扩散世界模型：87 小时（5M 帧）训练数据，动态模型参数量从 4M 扩展至 381M（含 51M 上采样器），训练耗时 12 天（RTX 4090），推理运行于 10Hz（RTX 3090）。To verify the scalability of the method, the authors trained a large-scale diffusion world model on the Dust II map of Counter-Strike: Global Offensive: 87 hours (5M frames) of training data, a dynamics model scaled from 4M to 381M parameters (including a 51M upsampler), 12 days of training on an RTX 4090, and inference running at 10Hz on an RTX 3090.

CS:GO world model gameplay — **Figure 6.** 玩家用键鼠在 DIAMOND 扩散世界模型内部游玩 CS:GO Dust II 地图的截图。该模型仅基于静态游戏录像训练，可生成"stable trajectories over hundreds of timesteps"，展现了扩散世界模型在复杂真实场景中的可扩展性。**Figure 6.** Screenshots of a player using keyboard and mouse to play the CS:GO Dust II map inside the DIAMOND diffusion world model. Trained only on static gameplay recordings, the model can generate "stable trajectories over hundreds of timesteps", showing the scalability of diffusion world models to complex realistic settings.

消融实验Ablations

核心消融结论：Key ablation findings:

EDM vs. DDPM：DDPM 在 ≤10 去噪步时出现"severe compounding error"，轨迹迅速崩溃；EDM 即使单步采样也能保持稳定，这是 DIAMOND 能够高效（3 NFE/帧）运行的基础。EDM vs. DDPM: with ≤10 denoising steps DDPM exhibits "severe compounding error" and the trajectory collapses quickly; EDM stays stable even under single-step sampling, which is the basis for DIAMOND running efficiently at 3 NFE per frame.
去噪步数 n 的影响：n=1 时多模态环境生成模糊帧，n≥3 时生成清晰图像且性能收敛，n=3 在质量与计算成本间取得最佳平衡。Effect of the number of denoising steps n: at n=1 multimodal environments produce blurry frames, at n≥3 the frames are sharp and performance converges, and n=3 strikes the best balance between quality and compute cost.
视觉质量与智能体表现的关联：IRIS 的离散压缩导致帧间视觉不一致，DIAMOND 的连续建模消除这些伪影，直接带来性能提升（如 Asterix：IRIS 853.6 → DIAMOND 3698.5；Breakout：IRIS 83.7 → DIAMOND 132.5）。Link between visual quality and agent performance: the discrete compression in IRIS causes visual inconsistency across frames, while DIAMOND's continuous modeling removes these artifacts and directly improves performance (e.g. Asterix: IRIS 853.6 → DIAMOND 3698.5; Breakout: IRIS 83.7 → DIAMOND 132.5).

04 局限性（Limitations）Limitations

说明：以下局限性均为论文 Section 8 中作者明确陈述（stated），非推断。Note: all limitations below are explicitly stated by the authors in Section 8 of the paper, not inferred.

仅评估离散动作控制，未扩展至连续控制域Only discrete-action control is evaluated, with no extension to continuous control domains

本文实验集中在 Atari 等离散动作游戏，尚未在连续控制任务（如 MuJoCo、机器人操控）上验证扩散世界模型的有效性。The experiments focus on discrete-action games such as Atari; the effectiveness of the diffusion world model has not yet been validated on continuous control tasks (e.g. MuJoCo or robot manipulation).

帧堆叠提供的记忆机制十分有限Frame stacking provides only a very limited memory mechanism

当前架构通过拼接过去 L 帧（frame stacking）提供历史信息，作者承认这"provides a minimal mechanism for memory"。基于 Transformer 的方法（如 STORM）在长时依赖建模上具有潜在优势，未来引入 Transformer 历史编码器可进一步提升性能。The current architecture provides history by concatenating the past L frames (frame stacking), which the authors acknowledge "provides a minimal mechanism for memory". Transformer-based methods (such as STORM) have a potential advantage in modeling long-range dependencies, and introducing a Transformer history encoder in future work could further improve performance.

奖励与终止预测独立于扩散模型，架构未完全统一Reward and termination prediction is separate from the diffusion model, so the architecture is not fully unified

奖励和终止信号由独立的 CNN-LSTM 网络预测，而非由扩散动态模型统一建模。作者认为将奖励/终止集成进扩散模型"would make our world model unnecessarily complex"，因此保持了模块化设计，但这也意味着整体架构并非端到端统一。Reward and termination signals are predicted by separate CNN-LSTM networks rather than modeled jointly by the diffusion dynamics model. The authors argue that integrating reward/termination into the diffusion model "would make our world model unnecessarily complex", so they keep a modular design, which also means the overall architecture is not unified end to end.

计算成本较高，不适合资源受限场景Compute cost is high, which makes it unsuitable for resource-constrained settings

每款 Atari 游戏的完整训练需要约 2.9 天（RTX 4090），26 款游戏合计约 1.03 GPU 年。扩散模型的逐步去噪推理比离散编码方案计算密集，尽管 3 NFE/帧已是较大优化，CS:GO 大规模版本（381M 参数）更需 12 天训练。Fully training each Atari game takes about 2.9 days (RTX 4090), roughly 1.03 GPU years for the 26 games in total. The step-by-step denoising inference of a diffusion model is more compute-intensive than discrete-code schemes; although 3 NFE per frame is already a substantial optimization, the large-scale CS:GO version (381M parameters) needs 12 days of training.