Huajishe Notes

2-DreamerV2 Mastering Atari with Discrete World Models

8 min read

一句话定位

DreamerV2 是首个在 Atari 55 个游戏中达到人类级别表现的基于世界模型的强化学习算法，核心改进是将连续高斯隐变量替换为离散的 Categorical 分布，显著提升了表征能力。

前置依赖

  • DreamerV1 (Hafner et al. 2019): 基于 RSSM 和隐空间想象的策略学习框架
  • PlaNet (Hafner et al. 2019): RSSM 架构与 KL Balancing 机制
  • SPR (Spinup Priors Representation): 离散表征的重要性先验工作

核心思想

DreamerV2 保留了 DreamerV1 的整体算法框架，核心改进在于将连续高斯隐变量替换为离散的 Categorical 分布。这一改变解决了两个关键问题：

  1. 多模态表征问题：连续高斯分布难以捕捉环境的多模态特性（如游戏中的多种策略选择）
  2. 信息压缩问题：离散分布可以更紧凑地表示丰富的语义信息

同时改进了 Actor-Critic 算法的价值目标计算，并引入了更稳定的 KL 权重调度机制。

模型结构图

输入图像 o_t
    ↓
┌─────────────┐    ┌─────────────────────────┐
│  Encoder   │ →  │ 离散隐变量 z_t ~ Cat(K) │
└─────────────┘    └─────────────────────────┘
                           ↓
    ┌───────────────────────────────────────┐
    │           RSSM (离散版)                │
    │                                        │
    │  h_t = f_phi(h_{t-1}, a_{t-1}, z_t)   │
    │  p_phi(z_t | h_t) = Cat(cat_probs)    │
    │  q_phi(z_t | h_t, o_t) = Cat(post)    │
    └───────────────────────────────────────┘
                    ↓
    ┌─────────────────────────────────────────┐
    │         Actor-Critic 网络               │
    │                                          │
    │  π_theta(a_t | h_t)                      │
    │  V_psi(h_t)                              │
    └─────────────────────────────────────────┘

关键区别：隐变量从连续高斯变为离散 Categorical

数学推导

1. RSSM with Categorical Latent Variables

状态转移方程（不变）：

先验分布（离散 Categorical）：

其中 是 维 logits，向量经过 softmax。

后验分布（推断）：

观测重构分布：

2. KL Divergence for Categorical Distributions

对于两个 Categorical 分布 和 ，KL 散度为：

由于 DreamerV2 使用 KL balancing，需要分别计算先验-后验和后验-先验的 KL。

3. KL Balancing (PlaNet 继承)

原始 KL balancing（V1 风格）：

DreamerV2 改进版： DreamerV2 将 KL 权重从固定值改为动态调度：

其中 是目标 KL 值， 是后验 EMA。

4. Actor-Critic 改进

价值目标计算（改进）：

而不是 V1 中的：

优势估计仍使用 GAE-：

其中

Critic 损失保持不变：

Actor 梯度（Reparameterization）： 由于 不可导，DreamerV2 仍使用 policy network 输出的动作 进行 reparameterization：

5. ELBO 损失函数（离散版）

重构损失正则项

形式与 V1 相同，但 KL 计算针对 Categorical 分布。

训练细节

离散隐变量设计

  • 类别数 K = 32：即每个隐变量有 32 个可能的取值
  • 这提供了 （理论上）的组合表征空间
  • 实际使用：32 个独立的 Categorical 分布，每个独立采样

KL 权重调度

  • 目标 KL：（相对于后验与先验的 KL 比值）
  • 动态权重：根据实际 KL 比值调整
  • 稳定训练：避免 KL 过大或过小导致的后验崩塌/先验失效

训练流程

  1. 阶段一：世界模型训练

    • 使用收集的观测序列训练 RSSM
    • Encoder + Decoder：VAE 重构损失
    • Dynamics Prior + Posterior：KL balancing
    • 优化器：Adam，lr =

  2. 阶段二：策略训练

    • 在隐空间中展开 步想象轨迹
    • 计算折扣回报和价值目标
    • 更新 Actor 和 Critic

关键超参数

参数值
隐变量类别数 K32
GRU hidden dim512
imagination_horizon H20
KL target1e-3
actor_lr1e-4
critic_lr1e-4
discount γ0.99
gae_lambda0.95

推理 / Rollout / Planning 过程

与 DreamerV1 类似，推理时：

  1. 观测编码到隐空间
  2. 使用策略网络输出动作
  3. 通过 RSSM 转移隐状态
  4. 无需 Decoder 或 Planning

在线规划能力：DreamerV2 虽然使用离散表征，但推理时仍是前向传递，无需 MPC 式的在线规划。

优点与局限

优点

  1. 更强的表征能力：离散 Categorical 能捕捉多模态环境
  2. Atari 基准的人类级表现：首个达到此水平的基于世界模型的方法
  3. 更稳定的训练：动态 KL 调度减少了训练不稳定性
  4. 组合表征：32 个独立类别提供了丰富的组合空间

局限

  1. 仍然依赖 KL balancing：调参仍然敏感
  2. 离散化的信息损失：某些连续信息难以用离散表征精确表达
  3. 计算开销增加：Categorical 采样比高斯采样略复杂

与前后内容的衔接

前身：DreamerV1

  • 保留了 DreamerV1 的整体框架（RSSM + Actor-Critic + 隐空间想象）
  • 核心改变：将连续高斯隐变量替换为 32 类离散 Categorical
  • 改进了 KL 调度机制

后继：DreamerV3

  • DreamerV3 将统一处理连续和离散动作
  • 移除了 KL balancing（改用 Symmetric log loss）
  • 引入 LayerNorm 和 skip connections
  • 实现了真正的超参数统一（无需针对不同领域调参）

可复现实现要点

离散隐变量采样

# 使用 Gumbel-Softmax 进行可导采样
# during training: use Gumbel-softmax
logits = prior_network(h_t)  # (B, K)
q = F.gumbel_softmax(logits, tau=1, hard=True)
# during eval: argmax
z_t = one_hot(F.softmax(logits, dim=-1).argmax(-1), K)

关键超参数

参数值
K (类别数)32
hidden_dim512
imagination_horizon H20
KL target1e-3
actor_lr1e-4
critic_lr1e-4

章节摘要

DreamerV2 在 DreamerV1 基础上进行了两项关键改进：一是将连续高斯隐变量替换为离散的 Categorical 分布（K=32），大幅提升了模型对多模态环境的表征能力，使其成为首个在 Atari 55 个游戏中达到人类级表现的基于世界模型的 RL 算法；二是引入了动态 KL 权重调度机制，使训练更加稳定。DreamerV2 展示了离散表征在世界模型中的巨大潜力，为后续的统一化方法奠定了基础。

关键词

DreamerV2、离散隐变量、Categorical 分布、RSSM、KL Balancing、Gumbel-Softmax、Atari、Actor-Critic、GAE-、价值目标、世界模型、隐空间想象、策略梯度

Graph View