在 ICLR 2021 上发表的《MASTERING ATARI WITH DISCRETE WORLD MODELS》一文，为模型基强化学习领域带来了革命性突破。作为 Dreamer 系列的第二代作品，DreamerV2 首次实现了纯依赖世界模型的潜空间想象，在 55 个 Atari 游戏基准测试中达成人类级性能，且在单 GPU 训练条件下超越了 Rainbow、IQN 等顶尖无模型算法。

DreamerV1 虽在连续控制任务中展现出数据效率优势，但在复杂离散动作的 Atari 游戏中难以匹敌无模型方法的性能。DreamerV2 通过离散潜变量表征、KL 平衡优化等关键改进，解决了世界模型预测精度不足、训练不稳定等核心问题，不仅验证了模型基方法在高难度视觉控制任务中的潜力，更提供了一套兼顾数据效率与最终性能的通用强化学习框架。

原文链接：[2010.02193] Mastering Atari with Discrete World Models

代码链接：danijar/dreamerv2: Mastering Atari with Discrete World Models

沐小含持续分享前沿算法论文，欢迎关注...

1. 研究背景与核心挑战

1.1 强化学习的两大范式之争

强化学习方法长期分为两大阵营，各有显著优劣：

  • 无模型方法（Model-free）：如 DQN、Rainbow、IQN 等，直接从环境交互中学习策略，无需建模环境动态。虽在 Atari 等基准中表现出色，但存在数据效率极低、训练成本高昂的致命缺陷，往往需要上亿环境步才能收敛；

  • 模型基方法（Model-based）：通过学习世界模型模拟环境动态，在想象轨迹中优化策略，数据效率极高。但传统模型基方法（如 PlaNet、World Models）面临两大瓶颈：① 世界模型预测精度不足，难以支撑复杂任务的行为学习；② 多采用连续潜变量，对离散动作和非平滑环境的建模能力有限。

1.2 Atari 游戏的独特挑战

Atari 游戏作为强化学习的黄金基准，对智能体提出了多重考验：

  • 高维视觉输入（84×84 灰度图像），需有效提取关键特征；
  • 离散动作空间（通常 1-18 个动作），动作采样的不可微性增加优化难度；
  • 奖励稀疏且延迟，需长视域信用分配能力；
  • 环境动态存在非平滑转换（如场景切换、物体消失），对模型泛化性要求极高。

此前，仅有无模型方法能在 Atari 上达成人类级性能，而模型基方法因世界模型的预测误差累积和表征能力不足，始终难以突破性能瓶颈。

1.3 论文核心贡献

DreamerV2 通过四大关键创新，彻底改变了模型基方法的性能格局：

  1. 离散潜变量表征：用类别变量替代高斯潜变量，提升世界模型的预测精度和泛化能力；

  2. KL 平衡优化：分离先验与后验的 KL 损失权重，解决传统 ELBO 目标中先验训练滞后的问题；

  3. 混合梯度优化：结合 Reinforce 梯度与直通梯度（Straight-through Gradients），兼顾策略优化的无偏性与稳定性；

  4. 通用框架设计：同时适配离散动作（Atari）与连续动作（人形机器人控制）任务，验证了方法的通用性。

在 55 个 Atari 游戏中，DreamerV2 在 200M 环境步内达成人类级表现，训练效率较无模型方法提升一个数量级，成为首个在 Atari 基准中超越顶尖无模型算法的纯模型基智能体。

2. 离散世界模型与想象驱动学习

DreamerV2 延续了 "动力学学习 - 行为学习 - 环境交互" 的三阶段框架，但在每个环节都进行了突破性改进。其核心逻辑是：先学习一个精准的离散世界模型，再在该模型的潜空间中通过想象轨迹优化演员 - 评论家网络，最终将学到的策略部署到真实环境。

2.1 世界模型学习：离散表征与 KL 平衡

世界模型是 DreamerV2 的核心，负责将高维视觉输入抽象为紧凑的潜状态，并预测环境动态。与 DreamerV1 的连续潜变量不同，DreamerV2 采用离散潜表征，大幅提升了模型的预测精度和训练稳定性。

2.1.1 世界模型的五大组件

DreamerV2 的世界模型基于循环状态空间模型（RSSM）构建，工作架构如图：

通过几大核心组件，形成完整的预测闭环：

各组件的功能与实现细节如下：

  1. 循环模型（Recurrent Model）：采用门控循环单元（GRU），输入前一时刻的离散潜状态 z_{t-1}​、动作 a_{t-1}​ 和循环状态 h_{t-1}，输出当前确定性循环状态 h_t，负责捕捉时序依赖关系；

  2. 表征模型（Representation Model）：由 CNN 提取图像特征后，通过 MLP 将图像特征与循环状态 h_t 映射为离散潜变量 z_t​ 的分布参数。离散潜变量采用 32 个类别变量，每个变量包含 32 个类别，扁平化后形成 1024 维稀疏二进制向量（仅 32 个激活位）；

  3. 转移预测器（Transition Predictor）：基于循环状态 h_t 预测下一时刻的离散潜变量先验分布 \widehat{z}_t​，无需依赖图像输入，是潜空间想象的核心组件；

  4. 图像预测器（Image Predictor）：采用转置 CNN，从循环状态 h_t 和离散潜状态 z_t​ 重建输入图像，提供监督信号；

  5. 奖励 / 折扣预测器：均为 MLP 结构，分别预测当前潜状态对应的即时奖励 ​ 和 episode 终止概率（Bernoulli 分布），折扣因子默认设为 0.999，终止时置 0。

2.1.2 离散潜变量的优势与实现

DreamerV2 用类别分布替代 DreamerV1 的高斯分布作为潜变量，带来三大关键优势：

  1. 多模态建模能力：类别先验可完美拟合多模态后验分布（混合类别分布仍是类别分布），而高斯先验难以建模环境中的非平滑转换（如 Atari 游戏的场景切换）；

  2. 稀疏表征特性：32 个类别变量各选 1 个类别，形成天然的稀疏表征，利于提取任务关键特征，过滤冗余信息；

  3. 训练稳定性提升：离散潜变量的梯度通过直通梯度技巧计算，避免了高斯潜变量重参数化梯度可能出现的梯度消失 / 爆炸问题。

离散潜变量的梯度计算采用 Algorithm 1 的直通梯度技巧：

这种实现方式既保留了离散表征的建模优势，又保证了训练过程的可微性。

2.1.3 KL 平衡优化：解决先验训练滞后问题

传统变分自编码器（VAE）的 ELBO 目标中，KL 损失同时承担两个角色：训练先验逼近后验，以及正则化后验分布。但在时序建模中，先验（转移预测器）的训练难度远高于后验（表征模型），容易导致 "先验滞后" 问题 —— 后验过度依赖图像信息，先验无法学到有效的时序动态。

DreamerV2 提出 KL 平衡（KL Balancing）优化策略，通过分离 KL 损失的权重解决该问题，具体实现如 Algorithm 2：

其中 α=0.8，给予先验更强的更新权重。这一设计鼓励模型通过优化先验来降低 KL 散度，而非通过增加后验熵来逃避正则化，从而让先验学到更精准的时序动态模型，为后续的潜空间想象奠定基础。

2.1.4 世界模型的整体损失函数

世界模型通过最大化变分下界（ELBO）进行联合优化，损失函数定义为：

各损失项的作用：

  • 图像对数损失：确保潜状态能精准重建输入图像，保留关键视觉信息；

  • 奖励 / 折扣对数损失：让世界模型学习奖励机制和终止条件，为行为学习提供可靠信号；

  • KL 损失：通过 KL 平衡优化，协调先验与后验的训练进度。

超参数设置：Atari 任务中 ，连续控制任务中 ，采用 Adam 优化器，学习率为 2\times 10^{-4}。

2.2 行为学习：想象 MDP 中的演员 - 评论家优化

DreamerV2 的行为学习完全在世界模型的潜空间中进行，通过想象轨迹优化演员（Actor）和评论家（Critic）网络，无需真实环境交互，这是其实现超高数据效率的核心。工作架构如图：

2.2.1 想象 MDP 的构建

世界模型的离散潜状态满足马尔可夫性，因此可构建完全可观测的想象 MDP：

  • 初始状态分布：世界模型训练过程中收集的潜状态分布 ；

  • 状态转移：由转移预测器 定义；

  • 奖励信号：奖励预测器输出的均值 ；

  • 折扣因子：折扣预测器输出的 ，用于处理 episode 早期终止。

想象视界（Imagination Horizon）设为 H=15，即每个想象轨迹包含 15 步预测，通过累积折扣因子加权，软考虑远期奖励。

2.2.2 演员 - 评论家网络设计

  • 演员网络（Actor）：输入为离散潜状态 ，输出离散动作的类别分布（Atari 任务）或截断正态分布（连续控制任务）。采用 4 层 MLP 结构，每层 400 个单元，ELU 激活函数，共 1M 可训练参数；

  • 评论家网络（Critic）：输入为离散潜状态 ​，输出确定性的状态价值估计 ，结构与演员网络一致，同样为 1M 可训练参数。

2.2.3 \lambda- 目标价值估计

为平衡价值估计的偏差与方差，DreamerV2 沿用 DreamerV1 的 \lambda- 目标（\lambda-target），通过递归方式融合多步回报：

​其中 \lambda=0.95，更侧重长视域回报。\lambda- 目标本质是不同步长多步回报的指数加权平均，既利用了多步奖励的低偏差特性，又通过价值估计降低了方差。

2.2.4 混合梯度优化策略

DreamerV2 创新地结合了两种梯度估计方法，兼顾无偏性与稳定性：

  1. Reinforce 梯度：无偏但方差高，通过状态价值基线降低方差，适用于离散动作场景；

  2. 直通梯度（Dynamics Backprop）：有偏但方差低，通过世界模型的可微性反向传播梯度，适用于连续动作场景。

演员网络的损失函数定义为：

其中：

  • \rho 为梯度混合系数，Atari 任务中 \rho=1（纯 Reinforce 梯度），连续控制任务中 \rho=0（纯直通梯度）；

  • sg(⋅) 表示停止梯度操作，避免目标值与评论家网络参数联动；

  • 熵正则项 鼓励探索，Atari 任务中 \eta=10−3。

评论家网络的损失函数为均方误差损失：

为提升训练稳定性，采用目标网络技术，每 100 步更新一次目标评论家网络参数。

2.3 环境交互：数据收集与闭环学习

DreamerV2 的环境交互遵循 "少量真实交互 + 大量想象学习" 的模式：

  1. 经验数据集：采用先进先出（FIFO）策略，存储最近 2\times 10^{6} 条经验序列，每次采样 50 个长度为 50 的序列用于训练；

  2. 动作生成：基于当前观测的潜状态 z^t，通过演员网络采样动作，并添加探索噪声（Atari 任务中通过熵正则项实现探索）；

  3. 闭环更新：每收集一定量经验后，交替更新世界模型和行为网络，形成 "数据收集 - 模型更新 - 策略优化" 的闭环。

值得注意的是，由于世界模型已能有效整合时序信息，DreamerV2 无需采用传统无模型方法的帧堆叠技术，进一步简化了模型架构。

3. 关键改进与消融实验

DreamerV2 的成功源于多项关键改进的协同作用。论文通过详尽的消融实验，量化了各组件对最终性能的贡献，为后续研究提供了重要参考。

3.1 核心改进的性能贡献

关键结论：

  1. 图像梯度至关重要：移除图像重建的梯度信号后，性能几乎完全丧失（Clipped Record Mean 从 0.25 降至 0.01），证明基于图像重建的自监督信号是学习有效潜表征的核心；

  2. 离散潜变量不可或缺：改用高斯潜变量后，性能下降 24%，验证了离散表征对复杂环境建模的优势；

  3. KL 平衡优化关键：移除 KL 平衡后，性能下降 36%，说明解决先验滞后问题对世界模型精度的重要性；

  4. Reinforce 梯度主导：移除 Reinforce 梯度后，性能下降 40%，表明无偏梯度估计对离散动作策略优化的必要性；

  5. 奖励梯度非必需：移除奖励预测的梯度信号后，性能几乎无损失，甚至略有提升，说明纯图像重建已能学到足够的任务相关表征，避免了奖励过拟合。

3.2 离散 vs 连续潜变量对比

论文通过实验验证了离散潜变量的显著优势：在 55 个 Atari 游戏中，离散潜变量在 42 个游戏中表现更优，仅在 8 个游戏中性能略低，5 个游戏表现相当。离散潜变量的优势主要体现在：

  • 对非平滑环境转换的建模能力更强（如游戏场景切换、物体消失）；

  • 潜表征更稀疏，利于提取任务关键特征；

  • 训练过程更稳定，梯度消失 / 爆炸问题显著缓解。

3.3 不同聚合指标下的性能表现

为全面评估算法性能，论文采用四种指标聚合方式（表 1）：

  • Gamer Median：基于专业玩家得分归一化，取任务中位数，避免极端值影响；

  • Gamer Mean：基于专业玩家得分归一化，取任务均值；

  • Record Mean：基于人类世界纪录得分归一化，取任务均值；

  • Clipped Record Mean：基于人类世界纪录归一化后裁剪至 [0,1]，取任务均值，避免超人类表现的过度影响。

结果显示，DreamerV2 在所有指标中均大幅领先，其中 Clipped Record Mean 指标较 Rainbow 提升 33%，充分证明其在不同评估标准下的鲁棒性。论文提出推荐将 Clipped Record Mean 作为强化学习算法的标准评估指标，因其能更公平地反映算法在所有任务上的综合表现。

4. 实验验证：全方位性能评估

4.1 实验设置

  • 环境：55 个 Atari 游戏（Bellemare et al., 2013），采用 sticky actions 设置（动作有 25% 概率重复），更贴近真实游戏场景；

  • 观测：84×84 灰度图像，下采样至 64×64 输入模型；

  • 动作：离散动作空间（1-18 个动作，因游戏而异）；

  • 训练配置：单 NVIDIA V100 GPU + 10 CPU 核心，单环境实例，200M 环境步训练（约 10 天）；

  • 基线方法：Rainbow、IQN、C51、DQN（均采用 Dopamine 框架的官方实现与得分）。

4.2 核心性能对比

4.2.1 总体性能超越顶尖无模型方法

如图 4 所示，DreamerV2 在 200M 环境步内的性能全面超越所有基线方法：

  • Gamer Median 得分 2.15，远超 Rainbow 的 1.29 和 IQN 的 1.47，达到人类玩家水平的 215%；

  • 在 Record Mean 指标上，DreamerV2 的 0.44 是 Rainbow（0.21）的两倍多，表明其在接近人类世界纪录的表现上具有显著优势；

  • 训练效率方面，DreamerV2 仅需 200M 环境步即可收敛，而 Rainbow 等方法通常需要 1e8 环境步以上，数据效率提升一个数量级。

4.2.2 单游戏性能分析

DreamerV2 在大多数游戏中表现出色，尤其在以下场景中优势明显：

  • 长视域任务：如 James Bond（得分 40445 vs Rainbow 的 1097）、Up N Down（653662 vs Rainbow 的 82840），体现了世界模型的长视域规划能力；

  • 探索密集型任务：如 Gopher（92282 vs Rainbow 的 355040）、Demon Attack（82263 vs Rainbow 的 1556345），熵正则化与世界模型的结合提升了探索效率；

  • 非平滑环境：如 Ice Hockey（26 vs Rainbow 的 36）、Assault（23625 vs Rainbow 的 8647），离散潜变量的建模优势凸显。

仅在 Video Pinball 等少数游戏中表现稍逊，论文推测原因是该游戏的核心目标（球）仅占 1 个像素，图像重建损失难以捕捉这一关键特征。

4.2.3 连续控制任务拓展验证

DreamerV2 并非局限于离散动作任务，在连续控制的人形机器人任务中同样表现出色：

  • 任务：DeepMind Control Suite 的 Humanoid Walk 任务，21 维连续动作，仅输入像素观测；

  • 改进：演员网络输出截断正态分布，采用纯直通梯度优化（\rho=0），调整超参数 \eta=10^{-5}、\beta=2；

  • 结果：成功学会站立和行走动作，是首个仅依赖像素输入解决该任务的模型基方法，验证了框架的通用性。

4.3 鲁棒性与泛化性验证

4.3.1 稀疏奖励任务表现

在 Montezuma's Revenge 这一经典稀疏奖励游戏中，DreamerV2 通过调整折扣因子 \gamma=0.99，无需额外探索机制，即可达到与 Rainbow+ICM（显式探索算法）相当的性能（图 B.2），证明世界模型的泛化能力有助于缓解稀疏奖励问题。

4.3.2 超参数敏感性分析

DreamerV2 对关键超参数具有较强鲁棒性：

  • KL 损失系数 \beta：在 0.1-3.0 范围内性能稳定；

  • 演员熵正则化 \eta：在 3e-5-1e-3 范围内表现良好；

  • 折扣因子 \gamma：0.99-0.999 均能收敛，仅在极端值下性能下降。

4.4 计算效率优势

  • 训练时间：每 1e6 环境步训练时间约 3 小时，远低于 PlaNet 的 11 小时 / 1e6 步；

  • 想象效率：单 GPU 可并行生成 2500 条想象轨迹，每轮行为学习可利用 468B 个想象潜状态，是真实环境交互数据的 10000 倍；

  • 推理速度：因世界模型的潜空间想象已完成大部分计算，实际与环境交互时的推理速度远超无模型方法。

5. 方法对比与相关工作

5.1 与 DreamerV1 的核心差异

DreamerV2 作为 DreamerV1 的迭代改进，关键差异体现在：

特性 DreamerV1 DreamerV2
潜变量类型 连续高斯分布 离散类别分布（32×32）
KL 损失优化 标准 ELBO KL 平衡（α=0.8）
策略梯度 重参数化梯度 混合梯度（Reinforce + 直通）
主要任务 连续控制 离散 Atari 游戏 + 连续控制
模型参数 13M 22M（世界模型 20M + 演员评论家 2M）
探索机制 动作噪声 熵正则化

5.2 与其他模型基方法的对比

5.2.1 SimPLe

  • 核心差异：SimPle 在像素空间进行视频预测，而非潜空间；采用 PPO 优化策略，而非想象驱动的演员 - 评论家；

  • 性能差距：SimPle 仅在 36 个简单游戏中评估，400k 环境步的 Gamer Median 仅 0.28，远低于 DreamerV2 的 2.15；

  • 效率差距：SimPle 训练 4M 环境步需 40 天，计算效率远低于 DreamerV2。

5.2.2 MuZero

  • 核心差异：MuZero 无显式图像重建，仅通过价值梯度学习任务相关表征；依赖蒙特卡洛树搜索（MCTS）规划，计算成本极高；

  • 性能差距：MuZero 虽在部分游戏中表现出色，但需要 2 个月以上的多 GPU 训练，且代码未开源，难以复现；

  • 优势互补：MuZero 的 MCTS 规划可与 DreamerV2 的高精度世界模型结合，进一步提升性能。

5.3 与无模型方法的本质区别

对比维度 DreamerV2（模型基） Rainbow/IQN（无模型）
数据效率 极高（200M 环境步收敛） 极低（1e8 + 环境步）
计算成本 训练期高（想象轨迹计算），推理期低 训练期低，推理期高
泛化能力 强（世界模型捕捉环境通用动态） 弱（依赖任务特定交互）
长视域规划 天然支持（想象轨迹可长视域预测） 依赖多步回报，效果有限
离散动作适配 通过 Reinforce 梯度优化，效果优异 原生支持，但数据效率低

6. 结论与未来方向

6.1 主要结论

DreamerV2 通过离散世界模型与想象驱动学习，实现了模型基强化学习的重大突破，核心结论包括：

  1. 离散潜变量表征能显著提升世界模型的预测精度与泛化能力，尤其适合非平滑环境与离散动作任务；

  2. KL 平衡优化有效解决了时序世界模型的先验滞后问题，是提升模型预测精度的关键；

  3. 纯依赖世界模型的潜空间想象，可在复杂视觉控制任务中达成超越顶尖无模型方法的性能，且数据效率提升一个数量级；

  4. 图像重建提供的自监督信号足以学习任务相关表征，无需过度依赖奖励信号，降低了过拟合风险。

6.2 未来方向

论文指出了三个值得深入探索的研究方向：

  1. 多任务与迁移学习：世界模型捕捉的环境通用动态，有望实现跨任务知识迁移，降低新任务的训练成本；

  2. 更高效的离散表征学习：结合对比学习、自监督学习等最新进展，进一步提升离散潜变量的表征能力；

  3. 规划与想象的融合：将 MuZero 的 MCTS 规划与 DreamerV2 的高精度世界模型结合，有望进一步提升长视域任务的性能；

  4. 真实世界应用拓展：将框架应用于机器人视觉控制等真实场景，解决真实环境中的视觉输入、离散 / 连续动作混合等复杂问题。

DreamerV2 的发表标志着模型基强化学习正式迈入实用化阶段，其离散世界模型与想象驱动学习的核心思想，为后续研究奠定了重要基础。通过在 Atari 游戏中达成人类级性能，DreamerV2 不仅打破了无模型方法在复杂视觉控制任务中的垄断，更证明了模型基方法兼顾数据效率与最终性能的巨大潜力。随着表征学习与规划算法的进一步融合，模型基强化学习有望在机器人、自动驾驶等真实世界场景中实现更广泛的应用。

# 论文阅读
Logo
AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

cover

OpenHuman 智能体落地应用与价值实现指南

avatar AtomGit开源社区

基于遗传算法的配电网故障重构研究【IEEE33节点】（Matlab代码实现）

配电网故障重构是保障供电可靠性的关键技术，其核心是通过调整开关状态实现故障隔离与非故障区域供电恢复。针对传统方法易陷入局部最优、难以处理多约束等问题，本文以IEEE33节点系统为研究对象，提出基于遗传算法的故障重构模型。该模型通过二进制编码映射开关状态，采用多目标加权适应度函数（最小化停电负荷与网损），结合辐射状拓扑校验与潮流约束，实现全局优化求解。仿真结果表明，在单故障与多故障场景下，该算法较传

avatar AtomGit开源社区

【负荷预测】基于LSTM-KAN的负荷预测研究（Python代码实现）

LSTM-KAN模型结合了LSTM和注意力机制的特点。LSTM是一种特殊的循环神经网络（RNN），通过引入记忆单元和门控机制（遗忘门、输入门、输出门），有效解决了传统RNN在处理长序列数据时存在的梯度消失或梯度爆炸问题，能够捕捉时间序列数据中的长期依赖关系。而注意力机制则能够模拟人类在处理信息时的注意力分配过程，通过为不同时间步的输入数据分配不同的权重，突出重要信息对预测结果的影响。

avatar AtomGit开源社区