研一刚接触深度强化学习的时候我拿到的第一份任务就是精读《MASTERING ATARI WITH DISCRETE WORLD MODELS》也就是 DreamerV2。说实话第一次翻论文的时候整个人是懵的世界模型、RSSM、categorical latent、想象轨迹……这些词连在一起和我之前看的 DQN、PPO 完全是另一套体系。但这篇论文的价值恰恰在于它把“从像素直接学环境动力学再在想象中规划行为”这条路线做到了 Atari 57 游戏上的 SOTA让基于模型的强化学习第一次在复杂视觉任务上达到人类水平。对于研一学生来说它不是最好读的论文但绝对是最值得啃完的一篇补全视野的论文。这篇文章我按自己当时精读的顺序拆开讲把核心方法、每个设计背后的动机、经典消融结果以及“小白如何不迷失在符号里”的实操方法一次讲透。1. 为什么世界模型必须改成离散变量写在论文前面1.1 先搞清楚世界模型到底在学什么传统强化学习里智能体直接看状态、选动作、拿奖励依靠 Q 值或策略梯度来改进策略这类方法统称 model-free。而基于模型的强化学习多了一步先学一个“环境模拟器”让智能体在脑内预演未来。DreamerV2 的世界模型做的事情可以概括为给定历史观测和动作序列预测下一时刻的观测和奖励。在 Atari 游戏里观测就是一帧游戏画面动作是摇杆或按键奖励是当前帧的游戏得分增量。关键难点在于原始像素空间维度高、冗余多直接在里面做预测既昂贵又脆弱。所以世界模型先把像素压缩成一个低维潜状态在潜状态里学习动力学。这个思路最早可以追溯到 latent space 里的各种 RNN 架构而 Dreamer 系列的特色在于它把“随机潜变量”和“确定性递归单元”揉在一起构成 RSSM。这就相当于我们把每一帧画面的有效信息提炼成一段话再用这段话去推演下一步会发生什么而不是把整张截图的内存全记下来。1.2 DreamerV1 到 DreamerV2 改了什么DreamerV1 已经把“世界模型训练 actor-critic”的框架跑通了但那时 RSSM 里的随机潜变量用的是高斯分布。高斯分布表达力有限而且在实际训练里很容易出现均值被推着走、方差趋近于零的“塌缩”问题。DreamerV2 的核心改动非常直接把潜变量从连续高斯改成离散的 categorical 分布并且把这个分布做成高维卡片组合。论文里默认实现是 32 个维度乘以 32 个类别也就是每个维度的离散取值有 32 种。你可以把每个维度的类别理解成某种“特征槽位”32 个槽位组合起来就是当前帧的一种编码。这个改动不是简单换个分布它对整个模型的表达能力影响非常深远。连续潜变量的每一个维度只能表示一个连续数值离散潜变量的每个维度却能让网络学习一种“离散语义”就像语音识别里把音频帧映射到有限的音素集合。不同的离散组合天然具备多模态能力能描述多种互斥的未来状态这对游戏环境这种高度非线性的动力学特别重要。1.3 离散潜变量为什么比连续潜变量好用我在读论文和后续复现时对离散潜变量的优势有两个特别明显的体感。第一是表达力的差异。环境中的未来往往是多分支的敌人可能向左跳也可能向右跳。高斯分布本质上只能表达一个以均值为中心的连续区域不适合表达“要么左、要么右”这种互斥假设。而 categorical 分布天然是多模态的它的每一个离散组合就是一个独立的假设分支这让模型能更精准地建模复杂动力学。第二是优化稳定性的差异。高斯潜变量要把均值方差都预测准训练时稍不注意方差就崩。离散潜变量则像分类任务网络只需要在一组有限的牌里选一张损失路径更平滑。论文里的消融实验也明确显示在同等条件下离散潜变量带来的最终回报远高于高斯潜变量这几乎是 DreamerV2 最重要的一个工程判断。注意很多第一次读这篇论文的人会误以为离散化只是为了简化计算。实际上离散化牺牲了部分连续表达能力换来的是更好的多模态拟合和更稳定的训练。在强化学习的复杂环境里这种取舍相当划算。2. 论文核心结构拆解RSSM 怎么把像素变成状态2.1 先备齐符号再谈模型结构RSSM 这个缩写经常把新手吓住拆开就是 Recurrent State-Space Model循环状态空间模型。它由两条路径组成。一条是确定性路径用 GRU 循环网络维护一个时刻变化的确定状态 h_t另一条是随机路径在确定状态之上再采样一个随机潜变量 z_t。状态 s_t 由 h_t 和 z_t 共同组成一个承载历史记忆一个承载当前帧的随机抽象。论文里定义了几个核心子模型表示模型也叫后验分布负责从当前观测 x_t 和历史信息 h_t 中推出 z_t 的真实分布。它的输入是图像经 CNN 编码后的特征。转移模型也叫先验分布只依赖 h_t 去预测 z_t不直接看观测。这样模型才能在没有真实画面的情况下推演未来。奖励模型从当前状态和动作中预测奖励值。解码器把潜状态还原成像素观测用来监督表示模型保留足够信息。你如果做过 VAE会发现这套结构和 VAE 非常像后验从观测里编码先验负责生成KL 散度约束两者靠近。区别在于这里多了一个 GRU 递归让模型能处理时序。2.2 表示模型从一帧游戏画面到离散编码Atari 画面输入一般会被缩放到 64×64 分辨率还是三通道 RGB。表示模型用卷积神经网络提取视觉特征然后把这些特征和 GRU 输出的确定状态拼在一起经过全连接层输出 32×32 的 logits。最后用 softmax 得到 32 个类别上的概率。到这里一帧画面就被编码成 32 个离散 token 的组合可以理解成把一帧画面翻译成了 32 个符号组成的一句话。这里有个容易忽略的细节表示模型要建模的是后验分布 p(z_t | h_t, x_t)。它在训练时会用到真实的当前帧所以推理幻觉时不能用它必须换成先验转移模型。这也是两个模型分工的意义后验负责校准先验负责预测。2.3 转移模型在潜空间里做时序预测转移模型是 RSSM 里的动力核心。它的输入是过去的确定状态 h_{t-1}、过去的随机潜变量 z_{t-1} 和过去动作 a_{t-1}GRU 更新出 h_t然后先验分布从这个 h_t 中采样出 z_t 的预测。注意这里有两个转移过程的区别训练世界模型时我们输入真实观测得到后验 z_t再用后验 z_t 去计算重建损失和转移损失而真正让智能体做“梦想”时没有真实观测只能从先验分布里采样 z_t并用采出的 z_t 作为下一步的输入不断向前滚动。这两个过程的差异是论文中反复出现的读代码时一定要分清哪条线用的是真实数据哪条线用的是预测数据。2.4 奖励预测多步直接回归更靠谱世界模型除了预测下一帧还要预测奖励。DreamerV2 里的奖励模型可以有两种做法。一种是单步预测每步单独预测 r_t然后在想象轨迹里把这些单步预测乘起来或加起来。另一种是多步直接预测也就是让奖励模型一次性输出未来 K 步的奖励序列。论文的实验结果显示直接预测未来 K 步奖励比逐步预测更高效最终性能也更好。这个结论反直觉但很实用因为单步预测的误差会随长度累积而直接多步让模型一次性建模整体回报模式有效缓解了累积误差。我在读代码时看到实现里奖励模型实际输入当前潜状态 s_t输出的是未来 K 步奖励的一个向量比如 K10就一次给 10 个标量。这个细节影响很大建议你读到这里时停下来想一想如果你来设计会让奖励模型输出什么形式。3. 训练流程拆开看世界模型和行为学习到底怎么同步3.1 第一阶段从真实轨迹里学世界模型世界模型的训练目标非常直白输入真实的观测、动作、奖励序列让模型不仅能重建观测还能预测潜状态和奖励。总的损失函数包含四部分。表示损失衡量解码器从潜状态重建像素的能力本质上就是 VAE 的重建误差。对像素输出DreamerV2 把每个像素看作 256 类别的分类任务用交叉熵计算这样比 MSE 更适合 RGB 离散值。转移损失约束先验分布 p(z_t | h_t) 尽量匹配后验分布 p(z_t | h_t, x_t)。优化时只把梯度传到转移模型后验用 stop-gradient避免两个分布互相追逐导致训练不稳。奖励损失回归真实奖励一般用交叉熵或 MSE。KL 项衡量先验和后验的差距为了防止 KL 被压得过小而让先验失去预测力论文引入了 free bits给 KL 项设定一个底线只优化超出底线的部分。这样一套损失组合下来世界模型就能做到三件事看懂画面、推演未来、估计回报。3.2 第二阶段让智能体在自己脑子里做实验世界模型学好后行为学习不再和真实环境直接交互。每一轮先在真实环境中收集一段数据存进回放缓冲区然后用这些数据训练世界模型。训练完世界模型就从当前真实的潜状态出发用先验转移模型往后展开多条想象轨迹。每条轨迹不长一般在 15 到 50 步之间但足够 actor-critic 学到有价值的策略信号。Critic 负责估计想象轨迹上每个潜状态的长期回报Actor 负责输出动作概率通过最大化 critic 给出的价值来更新。因为整个过程都在潜空间里完成不需要真的渲染图像也不依赖真实环境响应所以训练速度非常快。这就是“在梦里学习”的核心真实环境只用来采集数据真正的策略训练发生在世界模型生成的想象里。3.3 真实训练循环交替更新才稳定整个训练流程是循环的用当前策略在真实环境里跑若干步把观测、动作、奖励存入回放池。从回放池采样序列训练世界模型更新表示、转移、奖励三个模块。用更新后的世界模型生成想象轨迹在想象轨迹上更新 actor 和 critic。回到第一步继续收集新数据。论文里数据生成和模型更新是交错进行的batch size 与序列长度我记得大约是 50具体数值以原文附录为准。这个流程和标准 actor-critic 的最大不同在于critic 训练的是“想象中”的价值而不是真实环境里的价值因此世界模型的准确度直接决定了行为学习的上限。世界模型错了actor 再好也是白搭。提示如果你要复现训练过程最需要注意的就是两个阶段的学习率设置。世界模型学习率通常和 actor-critic 不同代码里常常分开配置。我见过很多尝试复现的同学把两组参数混用一个学习率最后模型怎么调都不收敛。4. 实验结果与消融对照哪些结论值得抄进自己的项目4.1 Atari 基准与训练协议论文选择 Atari 2600 的 57 个游戏作为测试场。这个基准之所以难是因为画面像素级输入、奖励稀疏且可能延迟加上游戏类型跨度大对算法的泛化能力要求很高。DreamerV2 在这一基准上的训练量非常大大约在 2 亿帧左右和人类持续游玩数十小时的体量相当这也是它能刷出惊艳结果的前提条件之一。对比对象既有 model-free 界的老牌选手如 DQN、Rainbow也有基于模型的先辈 SimPLe。在这样的大规模对比下DreamerV2 的平均人类标准化得分明显超过人类基线且远超 SimPLe。这里的“超过人类”不是说每个游戏都赢而是平均意义上的跨越在不少游戏上它已经能打出人类选手难以达到的分数。4.2 消融结论一离散潜变量完胜高斯潜变量论文中最有价值的消融实验就是把离散 categorical 潜变量换回高斯潜变量。结果显示离散潜变量带来的最终回报远高于高斯潜变量。这个结论几乎可以当作一条设计准则如果你在做一个基于视觉观测的世界模型优先考虑分类潜变量而不是直接套用 VAE 里的高斯假设。这背后有一个朴素逻辑游戏画面中的状态经常是离散的语义比如“站在台上”“跳在空中”“吃到金币”。离散编码天然贴合这种语义。高斯连续空间反而需要做额外工作才能表达这些边界清楚的状态。4.3 消融结论二多步奖励预测与参数敏感度另一个让我印象深刻的消融是把多步奖励预测改成单步预测。论文显示多步预测不仅收敛更快最终分数也更高。原因是 Atari 奖励延迟较长单步预测的误差会沿着时序滚雪球而多步预测让模型一次性看到未来多步的奖励模式相当于给了模型一个“远视眼”。此外categorical 分布的维度和类别数也不是越多越好。论文围绕 32×32 的默认设置做了敏感性分析少数类别不足以表达复杂状态类别太多又会让分布过于稀疏、训练不稳定。这个结论对做工程落地很有参考价值一开始直接用 32×32 往往比盲目调大更稳。4.4 这些结论对整个领域的启示DreamerV2 的意义不只是把 model-based RL 刷到人类水平它证明了“世界模型想象力训练”这条路能走通。很多原本属于 model-free 的场景现在都可以试着用世界模型去做。特别是在样本效率敏感的领域比如机器人控制、自动驾驶仿真、游戏 AI世界模型意味着智能体可以在自己脑子里模拟一万次碰撞而真实机器人只需要承担少数几次试错。同时也要冷静看待局限。世界模型的预测误差会随时间放大长视野规划仍是短板。如果你以后想做多步决策或者真实世界交互类任务需要时刻关注世界模型的校准质量这也是后续很多改进工作的突破口。5. 研一小白读这篇论文的顺序、笔记与代码配套方法5.1 不要第一章开始读先搭整体框架研一学生最大的误区是从摘要读到结论逐字推进结果在中途被公式劝退。我的建议是三轮读法。第一轮只读摘要、引言和结论目标是回答三个问题这篇论文解决了什么问题它的核心方法是什么方向实验结果证明了什么这一轮控制在半小时内。第二轮读图。重点看 Figure 1 的框架示意图、Figure 2 的训练流程、实验结果柱状图。把图看懂了论文的骨架就立起来了。然后带着图片的印象返回去读方法章节你会突然发现那些符号不是孤立的而是图中的一个个块。第三轮才是连公式带实验通读这时候你的目标是验证自己的理解而不是重新学一遍。如果哪里卡住去代码里看具体实现比死磕论文原文效率高得多。5.2 制作属于自己的符号表和流程图论文里符号很多我的笔记习惯是建一张符号对照表比如x_t 是观测帧a_t 是动作h_t 是确定状态z_t 是随机潜变量p_θ 是模型分布。遇到每个公式先把公式里的每个符号在表格里圈出来再按因果顺序在笔记本上画数据流箭头。画图的时候不会用太复杂的工具一支笔一张纸就够了目的只是把“真实输入 → 后验编码 → 先验预测 → 解码重建 → 奖励回归”这条路走通。这一步看似笨拙但对研一阶段建立底层模型直觉非常有帮助。我后来读很多其他论文时发现很多公式只是同一个框架的不同变体符号表让我能快速定位它们之间改了什么。5.3 配合开源代码读注意版本差异论文方法的很多细节只有在代码里才说得清楚。看代码时重点看四个文件或模块环境交互与回放缓存、世界模型的 loss 计算、actor-critic 的想象 rollout、训练循环的调度顺序。要注意网上流传的 DreamerV2 实现版本不少不同版本的实现细节可能不同。比如奖励预测的 K 步长度、free bits 的具体数值、策略是否使用 KL 正则不同代码仓库处理各异。如果你跑出来的结果和论文不一致先确认这些细节是否对齐而不要怀疑算法本身有问题。我当时就是从论文里一个“simplify”的表述跳到了代码的某个 loss 计算函数才真正理解了 KL stop-gradient 的设计动机。5.4 给零基础同学的前置知识补课清单如果你之前只接触过 DQN 或 PPO建议先补三块内容再回来啃论文一是 VAE 的基本原理搞清楚编码器和解码器的关系二是 RNN/GRU 的时序机制弄明白隐藏状态怎么沿时间传递三是 actor-critic 的基本框架知道策略网络和价值网络各自负责什么。这三块补完RSSM 就不会再让你头大。6. 常见阅读误区与避坑经验6.1 误区一分不清先验和后验的使用场景这是我见过最多人混淆的地方。后验模型只在训练世界模型时出现它的输入包含真实观测先验模型在决策和想象时出现它不依赖真实观测。如果你在测试阶段误用了后验等于偷看了标准答案数据效率的优势就不存在了。读代码时看到 encoder 只在训练分支被调用inference 分支只走 GRU 和先验网络就说明对了。6.2 误区二以为重建得越像越好世界模型的像素重建误差低不代表动力学预测准确。重建目标只是表示模型的监督信号之一真正影响下游表现的是潜状态里有没有保留决策所需的关键信息以及转移模型能不能准确预测状态变化。我在复现时遇到过重建 loss 很漂亮但策略分数极低的情况后来排查发现是潜变量维度太少把关键的动作依赖信息丢掉了。所以调参时别只盯重建曲线要同时看转移模型的预测准确率和最终回报。6.3 误区三直接套默认参数不理解指标含义论文里的超参数背后都有动机。比如 free bits 设成 1 nat是为了让先验分布有足够自由度来建模不确定性不至于因为 KL 惩罚太狠而退化成功底预测。如果你随手改成很大先验就会失去预测多样性改成 0又容易过拟合后验。调参前先去读自由比特相关的原理解释比盲目搜索高效得多。6.4 复现过程中容易踩的工程坑工程实现上我发现最常踩的坑有三个。第一个是 GRU 的初始状态处理训练和推理时隐状态初始化的方式不一致会导致转移模型在长序列上崩掉。第二个是并行环境数量作者在多环境并行采集时对回放数据的采样有特殊处理减少环境间不均衡的影响新手容易忽略。第三个是日志和评估频率世界模型训练很慢建议用较小的 batch 先跑通整个循环确认维度匹配无误后再放大规模。如果你打算在论文基础上做改进建议从最简单的消融开始比如把 categorical 维度从 32 改成 16 或者 64观察得分变化。这样能快速建立对模型行为的体感也比一上来就改 loss 更稳妥。常见困惑原因排查建议想象轨迹和真实轨迹差异大先验模型误差累积缩短想象长度或检查 KL free bits 是否过小重建 loss 降低但得分不涨表示模型丢失决策关键信息增加潜变量维度观察转移预测准确率训练初期 loss 波动剧烈学习率或 KL 权重设置不当分别调节两个阶段学习率降低 KL 权重复现得分始终不如论文超参数、并行环境数或奖励归一化细节不一致逐项对照代码仓库的配置文件而不是只调主学习率研一阶段啃这篇论文最大的收获其实不是记住 DreamerV2 的结构而是理解一种全新的训练范式先压缩世界再在想象中学习。读完全文再回头看RSSM 的离散潜变量、多步奖励预测、次元训练循环都是围绕着“让世界模型更准、让想象更可信”这两个目标服务。建议你效仿三轮读法配合开源代码把图画到自己手上遇到符号卡壳就先查笔记再查原理。这篇论文之后再看任何基于模型的工作视野都会开阔很多。