02.2 Dreamer 系列内部机制¶

flowchart LR
    obs["观测 o_t"]
    enc["Encoder<br/>CNN"]
    rssm["RSSM<br/>h_t (det) + z_t (stoch)"]
    rew["Reward Head"]
    cont["Continue Head"]
    dec["Decoder<br/>重建 o_t"]
    actor["Actor π(a|s)"]
    critic["Critic V(s)"]

    obs --> enc --> rssm
    rssm --> rew
    rssm --> cont
    rssm --> dec
    rssm --> actor --> critic

    classDef stage fill:#fff,stroke:#cc785c,color:#1a1a1a;
    class obs,enc,rssm,rew,cont,dec,actor,critic stage

最后更新: 2026-04-14 | 深度调研

Dreamer V3 RSSM 架构图（原论文 arXiv:2301.04104 Figure 2）

飞书图片缺失，待手动补

原飞书 wiki 此处嵌了 RSSM 架构图，import 时 token 空、未能下载。可参考论文 arXiv:2301.04104 Figure 2 自行查阅，或手画 mermaid 替换。

RSSM 数学结构¶

确定性路径: h_t = GRU(h_{t-1}, s_{t-1}, a_{t-1}) — 长程记忆
随机先验: p(z_t|h_t) = Categorical(32x32) — 纯预测(用于想象)
随机后验: q(z_t|h_t,o_t) — 贝叶斯修正(观测+历史)
训练: ELBO = 重建 + 奖励预测 - β·KL(后验||先验)

Dreamer V3 关键创新¶

Symlog: sign(x)·ln(|x|+1), 对称处理正负值, 跨环境统一
Two-Hot编码: 255个桶的分类分布输出, 非标量回归
32x32离散潜变量: 160 bits信息量, 比连续高斯更稳定
Free Bits: max(1 nat, KL), 防后验坍缩
KL Balancing: 0.5·L_dyn + 0.1·L_rep, 非对称权重
单一配置: γ=0.997, lr=1e-4, horizon=15步, 150+任务零调参

Dreamer V4 (2025) — 革命性突破¶

架构从GRU→Transformer, 2B参数
分辨率: 360x640真实游戏分辨率, 192帧上下文
训练: 2500小时人类Minecraft视频(仅4%有动作标注)
Flow Matching + Shortcut Forcing: 推理速度13x提升, 21 FPS实时
首个纯离线数据解决Minecraft Diamond(20000+步)
三阶段: 世界模型预训练→行为克隆→PMPO想象训练

MuZero vs Dreamer¶

MuZero: MCTS树搜索, 无解码器(价值等价), 适合离散动作
Dreamer: 想象展开, 有解码器(像素重建), 适合连续控制
EfficientZero V2: Gumbel搜索(仅8次模拟), Atari-100k +45%

TD-MPC2¶

无解码器, MPPI潜空间规划
SimNorm稳定化, 317M参数跨80任务
高维连续控制(38自由度Dog)超越DreamerV3

IRIS¶

VQ-VAE tokenize + GPT自回归, Atari-100k超DreamerV3
原因: Transformer并行训练, 离散token归纳偏置

Scaling Laws¶

世界模型遵循LLM式幂律: N∝C^0.49, D∝C^0.51
V4的2B参数是首次"大规模预训练"世界模型

参考¶

工程实现细节¶

RSSM 完整张量流¶

下面展示从观测 o_t 到潜在状态 z_t、隐状态 h_t 的完整数据流，所有张量形状基于 batch_size B, sequence length T。

# ============ 输入 ============
o_t: (B, T, H=64, W=64, C=3)              # 像素观测
a_t: (B, T, A)                             # 动作
r_t: (B, T)                                # 奖励
c_t: (B, T)                                # continue flag

# ============ 1. 编码器 ============
# image encoder: 4 层 3×3 Conv, stride 2, 通道 [48,96,192,384]
e_t = encoder(o_t)                          # (B, T, E=1024)

# ============ 2. RSSM 序列展开 ============
h_0 = zeros(B, H_dim=4096)
z_0 = zeros(B, 32, 32)

for t in range(T):
    # (A) Deterministic path: Block-Diagonal GRU (官方实现)
    # 注意: danijar/dreamerv3 用的是分组块线性 GRU, 非全局 Linear
    #   h_t 分 8 个块, 每块独立门控/更新、并行计算 (效率更高)
    x = concat([flatten(z_{t-1}), a_{t-1}])     # (B, 32*32 + A)
    h_t = BlockDiagonalGRUCell(x, h_{t-1})      # (B, H_dim)
    # 简化伪代码写法等价于: Linear 投影 + GRU, 实际实现用分组

    # (B) Prior p_θ(z_t | h_t)
    prior_logits = MLP(h_t).reshape(B, 32, 32)
    # Unimix: 99% 神经网络 + 1% 均匀分布 (防零概率)
    prior_probs = 0.99 * softmax(prior_logits, -1) + 0.01 / 32
    prior = Categorical(prior_probs)

    # (C) Posterior q_φ(z_t | h_t, o_t)
    post_input = concat([h_t, e_t[:, t]])
    post_logits = MLP(post_input).reshape(B, 32, 32)
    post_probs = 0.99 * softmax(post_logits, -1) + 0.01 / 32
    post = Categorical(post_probs)

    # (D) 采样 z_t 与直通估计器 (STE)
    # 官方写法: one_hot + (probs - sg(probs))
    z_one_hot = one_hot(argmax(post_logits, -1), 32)   # forward
    z_t = z_one_hot + (post_probs - post_probs.detach())  # backward via probs

# ============ 3. Decoder / Heads ============
state_t = concat([h_t, flatten(z_t)])          # (B, H_dim + 1024)
o_hat = ImageDecoder(state_t)                  # (B, T, 3, 64, 64)
# 重建 loss: MSE(ô_t, symlog(o_t))
r_logits = MLP(state_t)                        # (B, T, 255) two-hot
c_logits = MLP(state_t)                        # (B, T, 1) Bernoulli

Symlog Two-Hot 编码与解码¶

这是跨任务统一训练的关键，让同一网络处理 Atari (0-100) 到 Minecraft (0-1 稀疏) 等跨数量级的奖励。

# ============ Symlog 变换 ============
def symlog(x):
    return torch.sign(x) * torch.log(torch.abs(x) + 1)

def symexp(x):
    return torch.sign(x) * (torch.exp(torch.abs(x)) - 1)

# ============ Two-Hot 编码 (encoder) ============
# v: 目标标量值，形状 (B,)
# buckets: 255 个桶，在 symlog 空间均匀分布 [-20, +20]
def two_hot_encode(v, num_buckets=255, low=-20, high=20):
    v = symlog(v)                              # (B,) 先变到 symlog 空间
    v = torch.clamp(v, low, high)
    # 找到 v 在 buckets 中的位置
    bucket_width = (high - low) / (num_buckets - 1)
    idx_float = (v - low) / bucket_width        # (B,) 浮点位置
    idx_low = torch.floor(idx_float).long()     # 下界桶
    idx_high = idx_low + 1
    w_high = idx_float - idx_low.float()        # (B,) 权重
    w_low = 1.0 - w_high

    target = torch.zeros(B, num_buckets)
    target.scatter_(1, idx_low.unsqueeze(1), w_low.unsqueeze(1))
    target.scatter_(1, idx_high.unsqueeze(1), w_high.unsqueeze(1))
    return target                               # (B, 255)

# ============ Two-Hot 解码 (predict → value) ============
def two_hot_decode(logits, num_buckets=255, low=-20, high=20):
    probs = softmax(logits, dim=-1)             # (B, 255)
    bucket_centers = torch.linspace(low, high, num_buckets)
    # 期望值
    v_symlog = (probs * bucket_centers).sum(-1)  # (B,)
    return symexp(v_symlog)                      # (B,) 返回原始空间

# ============ Loss ============
target = two_hot_encode(r_true)                  # (B, 255)
logits = head(state)                             # (B, 255)
loss = -(target * log_softmax(logits, -1)).sum(-1).mean()  # 软 CE

Free Bits + KL Balancing¶

防止后验坍缩，同时让先验追赶后验。非对称权重是关键：让 dynamics 学习比 representation 正则化更重要。

# ============ KL Loss 计算（V3 关键设计）============
# 核实对照 danijar/dreamerv3 agent.py
# 注意: Dreamer 的 stop-gradient 位置和权重与我们直觉想的相反——
#   dyn_loss: KL(sg(post) || prior)  — 让 prior 学习追赶 posterior (sg 在 post 上)
#   rep_loss: KL(post || sg(prior))  — 让 post 正则化靠近 prior (sg 在 prior 上)
def compute_kl_loss(post_logits, prior_logits, free_nats=1.0):
    # Dynamics loss: 让 prior 学习预测 posterior
    # stop gradient on posterior
    kl_dyn = kl_divergence(
        Categorical(logits=post_logits.detach()),   # sg(post)
        Categorical(logits=prior_logits)            # prior (有梯度)
    )  # shape: (B, 32)

    # Representation loss: 让 post 正则化接近 prior
    # stop gradient on prior
    kl_rep = kl_divergence(
        Categorical(logits=post_logits),            # post (有梯度)
        Categorical(logits=prior_logits.detach())   # sg(prior)
    )  # shape: (B, 32)

    # Free bits: per-category，在求和前 apply max
    kl_dyn = torch.maximum(kl_dyn, torch.tensor(free_nats))
    kl_rep = torch.maximum(kl_rep, torch.tensor(free_nats))

    # 官方权重：dyn=1.0, rep=0.1 (从 configs.yaml: loss_scales.dyn=1.0, loss_scales.rep=0.1)
    # 之前版本写的 0.5/0.1 是错误的, 实际是 1.0/0.1
    loss_kl = 1.0  kl_dyn.sum(-1) + 0.1  kl_rep.sum(-1)  # (B,)
    return loss_kl.mean()

Imagination Rollout 与 λ-return¶

Actor-Critic 完全在世界模型的想象中训练，不与真实环境交互。\(\lambda\)-return 平衡偏差与方差。

# ============ Imagination Rollout (H=15 步) ============
def imagine(state_0, actor, world_model, horizon=15):
    """
    state_0: (B, H_dim + 1024) 从真实轨迹采样的起点
    返回想象轨迹的所有中间量
    """
    h = state_0.h  # (B, H_dim)
    z = state_0.z  # (B, 32, 32)
    states, actions, rewards, conts = [], [], [], []

    for t in range(horizon):
        state = concat([h, flatten(z)])
        a = actor.sample(state)                 # (B, A) 动作
        r = reward_head_decode(state)           # (B,) 奖励预测
        c = sigmoid(continue_head(state))       # (B,) 持续概率

        states.append(state); actions.append(a); rewards.append(r); conts.append(c)

        # 用 prior 前进一步 (纯想象，无观测)
        x = concat([flatten(z), a])
        h = GRUCell(Linear(x), h)               # 确定性转移
        prior_probs = softmax(prior_mlp(h))
        z = Categorical(prior_probs).sample()   # 想象的 z

    return stack(states), stack(actions), stack(rewards), stack(conts)

# ============ λ-return 计算 (GAE-style) ============
def lambda_return(rewards, values, conts, gamma=0.997, lam=0.95):
    """
    rewards, values, conts: (H, B)
    返回每个时刻的 λ-return target
    """
    returns = []
    R = values[-1]                              # bootstrap
    for t in reversed(range(H - 1)):
        R = rewards[t] + gamma  conts[t]  (
            (1 - lam)  values[t+1] + lam  R
        )
        returns.append(R)
    returns = stack(reversed(returns))          # (H-1, B)

    # Return Normalization (V3 新增)
    # 使用 5-95 百分位范围，避免离群点影响
    scale = ema_update(quantile(returns, 0.95) - quantile(returns, 0.05))
    R_norm = returns / max(1.0, scale)          # 归一化后的回报
    return returns, R_norm

# ============ Actor / Critic 更新 ============
def actor_loss(states, actions, R_norm):
    # Reinforce with baseline
    values = critic(states).detach()
    advantages = R_norm - (values / max(1.0, scale))
    log_probs = actor.log_prob(actions, states)
    L_actor = -(log_probs  advantages).mean() - 3e-4  actor.entropy()
    return L_actor

def critic_loss(states, R_normalized_target):
    # Two-hot regression
    logits = critic(states)
    target = two_hot_encode(R_normalized_target)
    L_critic = -(target * log_softmax(logits, -1)).sum(-1).mean()
    return L_critic

实际训练超参 (danijar/dreamerv3)¶

LaProp 优化器 (为什么不用 Adam?)¶

# LaProp 与 Adam 的关键区别
# Adam: m_t = β1m_{t-1} + (1-β1)g_t       # 一阶矩
#       v_t = β2v_{t-1} + (1-β2)g_t²      # 二阶矩
#       update = m_t / (sqrt(v_t) + eps)

# LaProp: 先对 g 做 RMSProp-style 归一化，再做动量
#         g_norm_t = g_t / (sqrt(v_t) + eps)
#         m_t = β1m_{t-1} + (1-β1)g_norm_t
#         update = m_t

# 优势:
# 1. 在稀疏奖励下比 Adam 更稳定
# 2. 对不同 loss 组件 (重建 + KL + 奖励) 量级差异更鲁棒
# 3. Dreamer 论文实验显示比 Adam 收敛更快

Dreamer V4 关键工程差异¶

Block-Causal Transformer 结构¶

# V4 摒弃 GRU，使用 4 层循环模式的 Block-Causal Transformer
# 192 帧上下文, 每帧 patch 化后形成 token 序列

class DreamerV4Block(nn.Module):
    def forward(self, x, layer_idx):
        """
        x: (B, T=192, HW_patches, D=dim)
        4 层为一循环单元:
          Layer 1-3: 空间注意力 (仅在同一帧内 HW patches 间)
          Layer 4:   时空联合注意力 (跨 T 帧 + 帧内 HW)
        """
        if layer_idx % 4 in [0, 1, 2]:
            # 空间注意力: 重排为 (B*T, HW, D)
            x_reshaped = x.reshape(B*T, HW, D)
            x_reshaped = self.spatial_attn(x_reshaped)  # attention over HW
            x = x_reshaped.reshape(B, T, HW, D)
        else:
            # 时空联合: 重排为 (B, T*HW, D) with causal mask
            x_flat = x.reshape(B, T*HW, D)
            x_flat = self.spatiotemporal_attn(x_flat, causal_mask=True)
            x = x_flat.reshape(B, T, HW, D)
        return x

# Grouped Query Attention (GQA): 减少 KV cache 4×
# Q: 16 heads (每个 head_dim=128)
# K/V: 4 heads (广播到 Q 的 4 倍)
class GQA(nn.Module):
    def forward(self, x):
        Q = self.q_proj(x).reshape(B, T, 16, 128)
        K = self.k_proj(x).reshape(B, T,  4, 128)  # 少 4×
        V = self.v_proj(x).reshape(B, T,  4, 128)
        # Repeat K, V 4 次以匹配 Q 的 head 数
        K = K.repeat_interleave(4, dim=2)  # (B, T, 16, 128)
        V = V.repeat_interleave(4, dim=2)
        # Flash Attention 实际实现不做真正的 repeat, 只在 kernel 内广播
        return flash_attention(Q, K, V, causal=True)

Shortcut Forcing (13× 加速)¶

# 标准 Flow Matching: 数百步去噪
# x_0 ~ N(0, I), iterate: x_{t+dt} = x_t + dt * v_θ(x_t, t)

# Shortcut Forcing: 训练时让模型学会不同步长的跳跃
def train_step_with_shortcut(x1, x0, v_model):
    # 采样随机时间 t 和随机跳跃步长 k
    t = torch.rand(B)                          # (B,)
    k = torch.randint(1, K_max, (B,))          # 跳跃长度

    # 目标: 一步跳跃 k 个时间步
    x_t = (1 - t)  x0 + t  x1
    t_next = t + k / K_max                     # 跳跃后的时间
    x_t_next = (1 - t_next)  x0 + t_next  x1

    # 训练: 从 x_t 预测 x_t_next (跳过中间步)
    v_pred = v_model(x_t, t, shortcut_k=k)
    loss = (x_t_next - (x_t + (k/K_max)  v_pred)) * 2
    return loss.mean()

# 推理: 用更大的 step
x = sample_noise()
for k in [K_big, K_big, K_med, K_small]:       # 非均匀步长调度
    dt = k / K_max
    v = v_model(x, t, shortcut_k=k)
    x = x + dt * v
    t += dt
# 推理从 ~50 步降到 13 步, 21 FPS 实时 on H100

关键数学公式¶

Symlog / Symexp 变换¶

\(\operatorname{symlog}(x) = \operatorname{sign}(x)\cdot\ln(|x|+1)\)

\(\operatorname{symexp}(x) = \operatorname{sign}(x)\cdot(e^{|x|}-1)\)

RSSM 联合分布¶

\(p_\theta(o_{1:T},z_{1:T}\mid a_{1:T}) = \prod_{t=1}^{T} p_\theta(o_t\mid h_t,z_t)\,p_\theta(r_t\mid h_t,z_t)\,p_\theta(z_t\mid h_t)\)

ELBO 训练目标¶

\(\mathcal{L}_{\text{ELBO}} = \mathbb{E}_{q_\phi}\!\left[\sum_{t=1}^{T}\ln p_\theta(o_t\mid h_t,z_t) + \ln p_\theta(r_t\mid h_t,z_t) - \beta\,D_{\text{KL}}\!\left[q_\phi(z_t\mid h_t,o_t)\,\|\,p_\theta(z_t\mid h_t)\right]\right]\)

KL Balancing（Free Bits 版本）¶

动态损失：\(\mathcal{L}_{\text{dyn}} = \max\!\left(1,\ D_{\text{KL}}\!\left[\operatorname{sg}(q_\phi)\,\|\,p_\theta\right]\right)\)

表示损失：\(\mathcal{L}_{\text{rep}} = \max\!\left(1,\ D_{\text{KL}}\!\left[q_\phi\,\|\,\operatorname{sg}(p_\theta)\right]\right)\)

总 KL 损失：\(\mathcal{L}_{\text{KL}} = 1.0\,\mathcal{L}_{\text{dyn}} + 0.1\,\mathcal{L}_{\text{rep}}\)

λ-Return（GAE-style）¶

\(G_t^\lambda = r_t + \gamma_t c_t\left[(1-\lambda) V(s_{t+1}) + \lambda\,G_{t+1}^\lambda\right]\)

其中 \(\gamma_t = \text{continue} \times \gamma = 0.997\)，\(\lambda = 0.95\)。

Return Normalization（V3 关键创新）¶

\(S = \operatorname{EMA}\!\left(\operatorname{Perc}_{95}(R^\lambda) - \operatorname{Perc}_{5}(R^\lambda)\right)\)

\(\tilde R = R^\lambda / \max(1, S)\)

使用 5-95 百分位数范围避免离群值污染。

Actor 损失（带熵正则）¶

\(\mathcal{L}_{\text{actor}} = -\mathbb{E}\!\left[\log\pi(a_t\mid s_t)\cdot\operatorname{sg}(\tilde R - \tilde V) + \eta\,\mathcal{H}[\pi(\cdot\mid s_t)]\right]\)

\(\eta = 3 \times 10^{-4}\)，Advantage 用归一化后的 return 和 value 差值。

Two-Hot 编码（解决跨任务奖励量级）¶

目标值 v 编码到相邻两桶：

\(w_k = \frac{b_{k+1} - v}{b_{k+1} - b_k},\quad w_{k+1} = 1 - w_k\)

解码用概率加权期望：

\(\hat v = \sum_{i=1}^{255} p_i\cdot\operatorname{symexp}(b_i)\)

LaProp 优化器（V3 首选）¶

先 RMSProp 归一化梯度，再做动量（与 Adam 顺序相反）：

二阶矩：\(v_t = \beta_2 v_{t-1} + (1-\beta_2)\,g_t^2\)

归一化梯度：\(\tilde g_t = g_t / (\sqrt{v_t} + \epsilon)\)

一阶矩：\(m_t = \beta_1 m_{t-1} + (1-\beta_1)\,\tilde g_t\)

参数更新：\(\theta_t = \theta_{t-1} - \alpha\,m_t\)

V4 Shortcut Forcing（13× 加速来源）¶

训练时让网络学习跨步长 k 的跳跃：

\(x_t = (1-t)x_0 + t x_1,\quad t\sim\mathcal{U}(0,1),\ k\sim\mathcal{U}\{1,\dots,K_{\max}\}\)

\(\mathcal{L}_{\text{SF}} = \left\|x_{t+k/K} - \left(x_t + \frac{k}{K}v_\theta(x_t,t,k)\right)\right\|_2^2\)

推理时用非均匀 step schedule，50 步降至 13 步。

Code 引用索引与可信度¶

✓=官方代码直接移植, △=基于论文描述重构, ⚠=推测填充 ——

注: 所有 "✓" 级代码的完整版本请参考对应 GitHub repo；本文的 PyTorch 伪代码经过简化以突出核心逻辑。

↑ 上级 · 02 技术架构与核心方法