3-DreamerV3 Mastering Diverse Domains through World Models

一句话定位

DreamerV3 是首个使用单一超参数配置在超过 150 个不同领域任务（涵盖连续控制、离散动作、图像观测、向量观测等）中超越专用方法的通用世界模型强化学习算法。

前置依赖

DreamerV2 (Hafner et al. 2020): 离散 Categorical 隐变量架构，Atari 上的人类级表现
DreamerV1 (Hafner et al. 2019): RSSM + Actor-Critic + 隐空间想象框架
PlaNet (Hafner et al. 2019): KL Balancing 机制（V3 移除了此机制）

核心思想

DreamerV3 的核心贡献是统一化：使用同一组超参数处理截然不同的任务类型。为此引入三项关键改进：

对称对数损失（Symmetric Log Loss）：替换 KL balancing，无需调参即可平衡重构与正则化
LayerNorm + 跳跃连接：使 dynamics model 更易训练，支持深度堆叠
简化的动作空间处理：统一框架同时支持连续动作（Gaussian）和离散动作（Categorical）

模型结构图

输入观测 o_t
    ↓
┌──────────────────────────────────────────┐
│            Encoder (共享权重)             │
│                                          │
│  图像: Conv Encoder                      │
│  向量: FC Encoder                        │
└──────────────────────────────────────────┘
                    ↓
    ┌──────────────────────────────────────┐
    │         RSSM (改进版)                   │
    │                                        │
    │  h_t = LayerNorm(                     │
    │      f_phi(h_{t-1}, a_{t-1}, z_t)     │
    │      + Skip(h_{t-1}) )               │
    │                                        │
    │  p_phi(z_t | h_t) = Cat(K)           │
    │  q_phi(z_t | h_t, o_t) = Cat(K)     │
    └──────────────────────────────────────┘
                    ↓
    ┌──────────────────────────────────────────────┐
    │               Decoder (共享权重)               │
    │                                               │
    │  图像: Deconv Decoder                         │
    │  向量: FC Decoder                             │
    └──────────────────────────────────────────────┘
                    ↓
    ┌──────────────────────────────────────────────┐
    │              Actor-Critic                     │
    │                                               │
    │  连续动作: π(a|h) = N(mu(h), sigma(h))       │
    │  离散动作: π(a|h) = Cat(softmax(h))          │
    │  V_psi(h)                                     │
    └──────────────────────────────────────────────┘

数学推导

1. 对称对数损失（Symmetric Log Loss）— 替换 KL Balancing

KL balancing 是 DreamerV1/V2 中的核心技巧，但需要仔细调参。DreamerV3 用对称对数损失替换：

观测重构损失：

对称正则项（替换 KL balancing）：

其中对称对数为：

展开可简化为：

直观理解：

当时，
当与偏离时，损失增大
相比 KL，散度对极端概率值不敏感（对数惩罚）

完整世界模型损失：

其中（固定，无需调参）。

2. RSSM with LayerNorm and Skip Connections

状态转移方程：

其中：

：标准 GRU/MLP dynamics 转换
：skip connection 路径，直接将加到输出
LayerNorm：稳定训练，加速收敛

先验/后验分布（保持 Categorical）：

3. 统一动作空间处理

连续动作（Gaussian）：

使用 reparameterization trick：

离散动作（Categorical）：

使用 Gumbel-Softmax 或 Straight-Through estimator 处理不可导采样。

动作网络统一架构：

# 动作编码
if discrete:
    action_logits = MLP(h)  # (B, num_actions)
    action_dist = Categorical(logits=action_logits)
else:
    action_mean = MLP(h)  # (B, action_dim)
    action_std = softplus(MLP(h))  # (B, action_dim)
    action_dist = Normal(action_mean, action_std)

4. Actor-Critic（与 V2 类似但统一）

价值目标计算：

Critic 损失：

Actor 梯度（统一形式）：

5. 优势估计（GAE-λ）

与 V2 相同：

训练细节

统一超参数设计

DreamerV3 的核心创新是单一超参数配置适用于所有任务：

参数	值	说明
hidden_dim	512	GRU 隐状态维度
K (类别数)	32	离散隐变量类别数
imagination_horizon H	15-20	想象步数
actor_lr	3e-4	Actor 学习率
critic_lr	3e-4	Critic 学习率
world_model_lr	3e-4	世界模型学习率
discount γ	0.99	折扣因子
gae_lambda	0.95	GAE λ
reg_weight λ_reg	1.0	正则损失权重
grad_clip	100.0	梯度裁剪

关键改进：无需 KL Balancing

V1/V2 的 KL balancing 需要针对不同任务调参。V3 的 Symmetric Log Loss 消除了这一需求：

KL balancing: 需要调 α（后验权重）和 β（KL 权重）
Symmetric Log: 单一固定权重 λ_reg = 1.0

训练流程

数据收集：使用随机或先前策略收集观测序列
世界模型训练：
- Encoder：对观测编码
- RSSM：学习隐动态
- Decoder：重构观测
- 使用 Symmetric Log Loss 作为正则项
策略训练：
- 在隐空间中进行想象 Rollout
- 计算折扣回报和价值目标
- 更新 Actor 和 Critic

推理 / Rollout / Planning 过程

与 V2 完全一致：

观测 → Encoder →
执行并进入下一循环

无在线规划需求，这是 Dreamer 系列的特点。

优点与局限

优点

真正的统一性：同一超参数适用于所有领域（连续/离散、图像/向量）
无需 KL balancing：Symmetric Log Loss 自动平衡重构与正则
更强表征能力：LayerNorm + Skip 使更深的 dynamics 模型可训练
SOTA 表现：在 150+ 任务上超越专用方法

局限

离散隐变量的局限性：仍使用离散表征，某些信息可能损失
计算开销：LayerNorm 和 Skip 增加计算量
超参数敏感性：虽然比 V1/V2 稳定，但 lr 等仍有敏感区间

与前后内容的衔接

前身：DreamerV2

保留了离散 Categorical 隐变量
继承了 Actor-Critic 框架
关键改变：移除 KL balancing，引入 Symmetric Log Loss
新增：LayerNorm + Skip Connections 改进 dynamics

后续发展：

DreamerV3 确立了三代 Dreamer 的演进路径：
- V1: 连续隐变量 + KL balancing
- V2: 离散隐变量 + 改进 KL balancing
- V3: 统一超参数 + Symmetric Log Loss

可复现实现要点

Symmetric Log Loss 实现

def symmetric_log_loss(q_probs, p_probs):
    """
    q_probs: posterior distribution (B, K)
    p_probs: prior distribution (B, K)
    """
    log_q = torch.log(q_probs + 1e-8)
    log_p = torch.log(p_probs + 1e-8)
    
    cross_entropy = -torch.sum(log_q * log_p, dim=-1)
    entropy_q = 0.5 * torch.sum(log_q**2, dim=-1)
    entropy_p = 0.5 * torch.sum(log_p**2, dim=-1)
    
    return cross_entropy - entropy_q - entropy_p

LayerNorm + Skip 在 GRU 中的应用

class RSSMCell(nn.Module):
    def __init__(self, hidden_dim, action_dim, stoch_dim, K):
        super().__init__()
        self.gru = nn.GRUCell(hidden_dim + action_dim + stoch_dim, hidden_dim)
        self.prior_net = nn.Sequential(nn.Linear(hidden_dim, hidden_dim),
                                        nn.LayerNorm(hidden_dim),
                                        nn.ReLU(),
                                        nn.Linear(hidden_dim, K))
        
    def forward(self, h_prev, a_prev, z_prev):
        gru_input = torch.cat([h_prev, a_prev, z_prev], dim=-1)
        h_new = self.gru(gru_input, h_prev)
        
        # Skip connection
        h_new = h_new + h_prev  # 直接加 skip
        h_new = LayerNorm(h_new)  # 再 LayerNorm
        
        prior_logits = self.prior_net(h_new)
        prior_probs = F.softmax(prior_logits, dim=-1)
        return h_new, prior_probs

关键超参数（统一配置）

参数	值
hidden_dim	512
K (类别数)	32
imagination_horizon H	15
actor_lr	3e-4
critic_lr	3e-4
world_model_lr	3e-4
reg_weight	1.0
grad_clip	100.0

章节摘要

DreamerV3 实现了世界模型强化学习算法的重大统一化突破。通过引入对称对数损失（Symmetric Log Loss）替换 KL balancing，配合 LayerNorm 和跳跃连接改进 dynamics 模型，以及统一的动作空间处理框架，DreamerV3 仅需一组超参数即可在超过 150 个不同领域的任务中达到或超越专用方法的表现。这些任务涵盖连续控制（DeepMind Control Suite）、离散动作（Atari 游戏）、图像观测和向量状态空间等。DreamerV3 证明了世界模型方法具有良好的通用性和可扩展性。

关键词

DreamerV3、对称对数损失、Symmetric Log Loss、统一超参数、LayerNorm、跳跃连接、RSSM、离散 Categorical 隐变量、连续/离散动作统一、GAE-、Actor-Critic、世界模型、隐空间想象、KL Balancing 移除

Huajishe Notes

Explorer