研一刚进组那阵我啃的第一批深度强化学习论文里就有这篇《Mastering Atari with Discrete World Models》。说实话当时看到标题里Discrete World Models几个词我连个大致印象都没有只记得读完摘要就愣住了一个模型从头学玩几十个Atari游戏分数还能压过人类玩家后来花了将近两周细读、画图、推伪代码才慢慢把世界模型到底怎么帮智能体做决策这层窗户纸捅破。这篇文章我就用研一阶段最容易踩坑的视角把这篇论文从问题背景、核心改动、训练流程到实验结果一点点拆开最后给你一条可以直接照做的精读和复现路径。这篇论文值得精读不只是因为它当时的分数好看而是它代表了一条和主流的无模型强化学习完全不同的路子先让智能体学会世界长什么样、会怎么变再让它在这个内部模型里反复预演决策。对刚入门的同学来说它几乎是一份完美的论文学习样本——问题定义清楚方法有清晰的演进逻辑实验有消融对照代码还是开放的。读透它你对表示学习和决策学习是怎么结合的会有质的理解。1. 这篇论文解决的是强化学习里的学得太慢问题1.1 绕不开的样本效率困境先从一个研一新生最常见的问题说起为什么大家都在用几十万个样本训一个游戏智能体却没人想着让智能体少玩几局因为主流强化学习算法大多是无模型路线——智能体不建立任何关于环境的内部规律纯粹靠试错积累奖励信号来调整策略。以Atari游戏为例传统算法动辄需要几千万帧画面才能达到像样的水平一个游戏训练半天到几天是家常便饭。这种硬碰硬的学习方式有两个痛点一是真实环境交互成本高尤其换到机器人、自动驾驶这类场景几千万次试错根本不现实二是智能体学到的知识高度绑定在具体任务上换一个游戏往往从头再来。那让智能体先学习一个世界模型呢意思是智能体在与环境交互一段时间后能不能把我做了什么动作画面会变成什么样、能拿多少奖励这些规律压缩进一个神经网络里然后用这个模型代替真实环境在内部预演未来的轨迹。这样真实环境只需要提供少量样本大量的策略试错在模型内部完成就够了。这个思路就是论文世界模型路线的起点。1.2 世界模型不是新鲜概念但拿来决策一直很难世界模型这个词在强化学习出现很早但绝大多数早期工作只把它当预测任务来做让模型能准确预测下一帧像素、预测奖励至于怎么用这个模型来改进策略始终没有特别成功的方案。原因有两个坎。第一像素级预测误差会随推演步数快速累积——模型在第一步预测偏一点第二步就偏更多几步之后想象出来的画面完全失真策略在这上面学就白学了。第二世界模型和策略之间怎么衔接是个工程难题早期的模型预测的是像素而策略需要的是对当前状态的理解两者之间的接口很模糊。作者团队的前作已经把循环状态空间模型简称RSSM这条路铺开了不直接预测像素而是让模型学一个压缩的隐状态再用这个隐状态去重建画面、预测奖励。这样预测误差不会完全压在像素层面隐状态本身相当于一个对世界的结构化摘要。但前作在离散控制任务上表现尚可一到Atari这类画面复杂、奖励稀疏的游戏就明显吃力。这篇论文就是在这个基础上做关键升级其中最重要的一刀就是把隐变量从连续分布换成了离散分布。1.3 为什么研一精读这篇论文的回报特别高我个人觉得研一阶段选论文精读挑的不应该是最新最热的而应该是结构完整、演进脉络清晰、还能动手复现的。这篇恰好满足三个条件。第一它的论文骨架非常规整开头谈样本效率问题中间讲离散世界模型的设计后半部分用大量消融实验说明每个设计选择各贡献了多少。你可以照它的结构学习怎么写一篇好的系统型论文。第二它的代码开放环境用标准的Atari接口就能跑你不需要自备特殊硬件一张普通显卡就能起来一个小规模的实验。第三它是后续一系列基于模型工作的起点读懂了这篇再去读后续的扩展版本、其他团队的类似工作基本就是一条顺畅的坡道不需要再重新爬一次。2. 核心改动把隐变量从连续池子换成离散词条2.1 先搞清楚论文里的state到底是个什么东西很多研一同学看到world model的图就晕不知道圆框里的state是真实画面还是别的表示。这里必须说清楚state不是图像本身而是编码器从图像里提取出的一个压缩向量。打个比方看一幅台球的桌面画面人类不会记住每一个像素的位置而是记住白球在左侧、红球靠近袋口、目前轮到我了这几个关键要素。latent state就是模型版本的这些关键要素——一段固定长度的向量里面编码了当前局面最重要的信息。在RSSM结构里latent state还不是一个静态向量它由两部分组成一部分是来自当前观测的表征也就是我看到什么另一部分是来自上一时刻状态和动作的预测也就是我预期会发生什么。模型把这两种信息融合起来既尊重事实又保留了对未来的预测能力。这个融合是RSSM能缓解预测误差累积的核心——哪怕预测偏了当前观测还能把状态拉回正轨。2.2 连续高斯的毛病表达力有余训练稳定性不足前作DreamerV1的latent state是一个多元高斯分布用均值和方差来描述状态可能的分布范围。高斯分布用起来方便梯度也容易回传但它有两个实际毛病。第一个毛病是点多面广但不够准。高斯假设状态服从一个单峰分布可真实世界经常是多模态的同一个画面可能对应球正在向左滑和球正在向右滑两种完全不同的走势单峰高斯只能模棱两可地覆盖中间地带。第二个毛病是训练稳定性。高斯的方差项在梯度反传里很容易出问题方差拉大一点、缩小一点整个损失函数的曲面都会剧烈变化模型往往要把大量容量花在维持分布稳定上而不是花在学准世界规律上。离散分布就换了一种思路把隐状态定义成一小串离散词条的组合。比如用一个长度为32的向量每位从32个离散类别里选一个那么总共能表达的组合空间非常大而且每个类别都是一个明确的状态桶。模型不用费劲去拟合一个光滑的钟形分布而是直接表达当前属于哪几个离散状态。论文作者的核心判断是对世界模型来说这种离散表达更适合长时间稳定地推演未来。2.3 代码层面的离散化实现直通梯度的巧劲这里有个工程上的难点离散采样本身不可导而神经网络训练需要梯度。你没法直接从按类别采样这个操作回传梯度。作者的解决办法是直通梯度straight-through前向传播时真的去采样一个one-hot向量反向传播时则假装这个向量就是它自己的softmax logits。这样既保证了前向时状态真的是离散的又让梯度能无障碍地流到生成logits的网络里。# 示意离散隐变量采样伪代码 logits transition_model(prev_state, action) # 预测下一状态的类别得分 sample one_hot(categorical_sample(logits)) # 前向真实采样 latent sample logits - stop_gradient(logits) # 反向使用logits的梯度第三行是整个技巧的核心前向传播时stop_gradient不产生任何影响结果就是sample反向传播时sample被视为常量可训练的梯度全部通过logits传导。用这个技巧离散世界模型可以端到端训练。2.4 离散化带来的彩蛋KL平衡训练更省心除了表达力更强离散化的另一个隐性收益是世界模型的训练更容易稳定。连续高斯版本的训练里KL散度项的尺度往往随训练波动稍不留神就出现后验塌缩——编码器不再向隐状态放信息解码器也摆烂整个模型退化成无意义的预测器。作者针对这个问题设计了一套KL平衡策略把KL散度拆成先验→后验和后验→先验两个方向一个方向用梯度更新另一个方向用停止梯度保持稳定再配合一个自由比特下限允许在训练初期容忍一定程度的KL偏差避免模型过早收敛到平庸解。这一套组合拳放在离散分类分布上格外顺手因为分类分布的KL分解非常干净可操作性比高斯分布强得多。这里如果你跳过不看后面复现时损失函数写不对模型就会训练得非常神经质。3. 训练三阶段预测世界、想象世界、在想象中做决策3.1 第一阶段世界模型学习让模型先学会读懂画面世界模型不是一个独立网络而是一组彼此配合的模块。先看数据流经验回放池里存了一批真实的交互序列——观测、动作、奖励、终止标志。编码器把观测x_t转换成后验隐状态q(s_t | s_{t-1}, a_{t-1}, x_t)这就是结合当前画面修正预测的环节。RSSM同时维护一个先验转移模型p(s_t | s_{t-1}, a_{t-1})作用是只看前一步状态和动作预测下一步可能的状态分布。在这之上还挂了三个输出头解码器负责从隐状态重建观测画面奖励预测器负责估计当前状态下能拿多少奖励继续预测器负责判断episode是否结束。三个头各管一件事对应损失函数的三项。如下表所示损失项作用通俗理解观测重建损失解码器还原画面隐状态有没有保留足够的视觉信息奖励预测损失奖励预测器逼近真实奖励隐状态能不能推算分数变化继续预测损失终止预测器判断是否结束隐状态能不能感知是否game overKL散度损失让后验靠近先验预测的未来别和实际观测差太远世界模型的训练目标就是把这四项损失加在一起做梯度下降。这里有一个关键细节所有训练都在回放池的真实经验上进行不需要在训练世界模型时额外消耗环境交互次数。3.2 第二阶段在梦境里想象轨迹世界模型训练到一定程度后训练循环进入第二阶段想象。想象的意思是从某个真实经验片段中取出一个隐状态作为起点然后不再读取真实环境只靠转移模型在潜在空间里向后滚动。每一步转移模型预测下一个隐状态的分布从中采样得到虚拟的下一帧状态同时用奖励预测器给出虚拟的奖励。这一步要特别注意整个想象过程完全发生在latent空间不经过解码器重建图像也不经过真实环境。解码器只在可视化时用相当于给你看智能体做过的梦长什么样。为什么能这么做因为世界模型学到的隐状态已经包含了决策需要的全部信息状态演变规律、奖励信号、终止条件。策略关心的是我该选哪个动作它不需要看到真实像素只需要看到足够好的状态表示。想象轨迹的长度也叫想象视野论文里一般取十几步。这个长度要权衡太短策略只看得到眼前几步学不到长远规划太长预测误差累积变严重想象轨迹失真。作者后面对比了不同想象长度对分数的影响这也是一个值得你在复现时亲自调一调的超参数。3.3 第三阶段Actor-Critic在梦境里学策略有了想象轨迹强化学习算法就有用武之地了。在梦境轨迹上每一步都有一份隐状态和对应的虚拟奖励这样你可以直接计算折扣回报再用一个Actor网络输出动作分布一个Critic网络估计状态价值。Actor的更新目标是让策略在想象轨迹上的期望回报最大化Critic的更新目标则是让它的价值估计逼近λ-return。整个流程和常规的Actor-Critic没有本质区别唯一的区别是它训练用的环境是学出来的世界模型而不是真实环境。这正是论文方法能大幅提升样本效率的根本原因真实环境只用于收集训练世界模型的数据策略的大量试错转移到便宜得多的梦境推演里。3.4 三阶段合成的训练循环一轮又一轮地转把上面三个阶段串起来就是论文训练算法的核心循环智能体按当前策略与环境交互收集新的观测、动作、奖励存入经验回放池。从回放池采mini-batch训练世界模型更新编码器、转移模型、解码器和各个预测头。利用更新后的世界模型在潜在空间展开多条想象轨迹更新Actor和Critic。回到第1步用更新后的策略继续与环境交互。整个循环的每一步都是在真实数据→世界模型→想象轨迹→策略更新这条流水线上完成的。研一同学看这段时最容易觉得抽象我建议你亲手画一遍数据流向图从真实观测进入编码器开始到策略输出动作结束把每一步数据的形状标在旁边基本就通透了。4. 关于Atari实验分数会骗人人类水平也要问怎么定4.1 看懂了实验设置才算看懂论文的野心论文选用的基准是Atari 2600平台的57款游戏这是强化学习领域沿用多年的经典测试场涵盖格斗、射击、平台跳跃、策略等不同类型。每个游戏给智能体100万帧的训练预算——这个数字相对无模型算法动辄数千万帧动辄低了很多但相对极低样本强化学习比如100K帧又宽松不少。换句话说论文的目标不是做最极限的样本效率而是在合理预算下达到甚至超过人类玩家的水平。理解这个设置很重要你在复现时不需要照搬100万帧预算用小游戏先跑通流程即可。4.2 消融实验告诉你改进到底堆在哪里论文最值得精读的部分其实是消融实验。作者做了几个版本的对比连续潜变量的DreamerV1、加大规模的连续版本、以及离散潜变量的最终版本。论文报告的结果显示仅仅把潜变量从连续换成离散平均分数就有非常明显的提升在此基础上再增大模型规模又能进一步拉升成绩。也就是说离散表示和模型容量是两个各自独立起作用的因素合在一起才达到人类水平。对比项相对表现说明连续隐变量小模型基线在部分游戏上能玩但整体远低于人类连续隐变量大模型中等规模提升有帮助但离散化收益更关键离散隐变量大模型最高论文最终报告达到人类水平这张表是作者用实验画出来的“贡献分解图”。研一读者请记住这个习惯不看消融实验你就只知道这个方法有效不知道这个方法为什么有效、哪个设计最值钱。4.3 人类归一化得分只报均值是典型的幸存者偏差论文用人类归一化得分来统一衡量57个游戏的表现把人类玩家的平均水平当作100%智能体的得分按比例换算。最终论文报告的均值落在人类水平附近这个成绩确实漂亮。但你必须看清楚这是一个平均值。逐游戏细看智能体在一些游戏上远超人类在另一些游戏上则远低于人类。比如动作简单、规则明确的游戏模型能靠预演推算出极优策略而需要长时记忆和抽象推理的游戏世界模型也会露怯。所以论文标题说的Mastering Atari要理解为整体达到人类级别的掌控力而不是每个游戏都超越人类。研一同学读任何论文的实验结果都要先看单个任务的表现分布再回来看均值否则很容易被一个漂亮的平均值带偏。4.4 与无模型算法对比位置决定价值把DreamerV2放到当时的算法版图里看最合理的参照是那些同样用100万帧预算的深度强化学习算法。在同样的交互预算下DreamerV2的平均人类归一化得分和最强的无模型算法不相上下甚至更优。这点很关键一个带世界模型的算法在同等样本预算下追平了那些专注试错的无模型算法这本身就是世界模型路线值得做的强烈信号。更难得的是因为世界模型的存在你还多了可解释这层收益——你可以把模型想象的画面解码出来看观察智能体脑子里在预测什么。这种可读性对调试和理解学习过程非常有帮助无模型算法是很难提供这种视角的。5. 研一精读实操从伪代码到复现的完整路径5.1 我的精读顺序先图表、后伪代码、再公式研一的同学最容易犯的错是从第一页开始逐字读到最后一页结果到了方法章节已经晕掉。我建议你按这个顺序走先读摘要和引言用三句话写下作者要解决什么问题、他提出什么方案、实验结论是什么。直接跳到方法部分的结构图不看公式只看图中模块之间的数据流向。把训练伪代码手抄一遍想象你要用代码实现它哪些地方你会卡住。再回头看公式这时已经知道每个符号对应哪个模块公式会好理解得多。最后读实验和消融重点关注对比表和消融结论。这套顺序的核心逻辑是先用图像和伪代码建立整体认知再让公式为整体认知填空。直接啃公式对新手极其不友好还容易半途放弃。5.2 伪代码逐行过一遍训练主循环其实就三个块论文给出的训练伪代码信息量很大但拆解后就三块。我用示意伪代码帮你梳理# DreamerV2 训练主循环示意结构 for iteration in range(train_steps): # 阶段1与环境交互收集经验 for env_step in range(steps_per_iteration): latent world_model.encode(obs) action actor(latent) # 加探索噪声后执行 obs, reward, done env.step(action) replay.add(obs, action, reward, done) # 阶段2在回放数据上训练世界模型 world_model.train(replay.sample_batch()) # 阶段3在潜在空间想象并训练行为 for dream_step in range(imagination_steps): imag_latent world_model.rollout(actor, start_state) actor.critic.update(imag_latent)第一块负责产生数据第二块负责让世界模型更准确第三块负责让策略更聪明。这三块各训练各的损失互不冲突。研一阶段理解到这个粒度就够了再往下的实现细节就交给代码吧。5.3 复现时最容易翻车的四个坑我后来照着论文跑过一个小规模的复现踩了不少坑给你列出来提前避雷。第一个坑是奖励没做缩放。Atari游戏的奖励范围很大Actor的梯度容易被极端奖励值带崩。论文在实验中对回报做了归一化你复现时也要处理这一步否则训练曲线会剧烈震荡。第二个坑是KL平衡缺了一个方向。有些初学者实现世界模型时只写了后验逼近先验那项忽略了先验向真实经验那项的梯度传导结果模型早期预测能力很差表现为训练曲线长期不动。这个方向很容易在复现时漏掉你要对照论文公式一个个核对。第三个坑是discrete采样梯度断流。用了one-hot采样但忘了直通梯度的stop_gradient技巧模型整体不更新或者报错说梯度无法通过。前馈数据看了一眼是离散向量但你要确保它的配套梯度计算是连续的。第四个坑是对训练时间预估不足。完整跑完57个Atari游戏需要很多卡时个人复现非常不现实。我建议你先选一个简单游戏比如经典的Pong或Breakout跑通全流程验证管道无误之后再考虑扩展。复现的目的是理解不是为了秒掉完整论文。5.4 读完之后往哪走顺着这篇论文的引用树继续爬精读完这篇我建议你顺手做两件事。第一去读论文引用了哪些关键工作尤其是世界模型的历史脉络和RSSM的原始出处补上前情提要。第二看看后续版本是怎么扩展它的。后续工作把离散世界模型推向了更通用的方向还引入了无需任务奖励就能学习等更前沿的设定。你会发现这篇论文里的很多设计思路被保留了下来但训练稳定性和通用性又上了一个台阶——这本身就是一段完整的科学演进脉络。我现在回想研一读这篇论文的过程最受益的不是记住了几个公式而是学会带着作者到底在赌什么去读论文。这篇论文赌的是与其让智能体在真实世界里硬碰硬地学策略不如先在内部模型里把世界推演熟。这个判断影响了我后来做项目时对样本利用率的看法——很多问题不是数据不够而是我们有没有把已有数据的价值榨干。所以如果你也是研一新手我最后的建议是不要只停留在读懂了试着亲手实现一次它的最小版本。哪怕只是拿一个小游戏把三阶段循环跑通你对世界模型、想象轨迹、离散表示这些概念的理解都会比刷十篇综述都来得扎实。这大概就是论文精读的真正意义。 
长河君 · 2026/10/22 5:25:29 · 阅读 1.2k 离散世界模型如何提升强化学习样本效率？DreamerV2精读
研一刚进组那阵我啃的第一批深度强化学习论文里就有这篇《Mastering Atari with Discrete World Models》。说实话当时看到标题里Discrete World Models几个词我连个大致印象都没有只记得读完摘要就愣住了一个模型从头学玩几十个Atari游戏分数还能压过人类玩家后来花了将近两周细读、画图、推伪代码才慢慢把世界模型到底怎么帮智能体做决策这层窗户纸捅破。这篇文章我就用研一阶段最容易踩坑的视角把这篇论文从问题背景、核心改动、训练流程到实验结果一点点拆开最后给你一条可以直接照做的精读和复现路径。这篇论文值得精读不只是因为它当时的分数好看而是它代表了一条和主流的无模型强化学习完全不同的路子先让智能体学会世界长什么样、会怎么变再让它在这个内部模型里反复预演决策。对刚入门的同学来说它几乎是一份完美的论文学习样本——问题定义清楚方法有清晰的演进逻辑实验有消融对照代码还是开放的。读透它你对表示学习和决策学习是怎么结合的会有质的理解。1. 这篇论文解决的是强化学习里的学得太慢问题1.1 绕不开的样本效率困境先从一个研一新生最常见的问题说起为什么大家都在用几十万个样本训一个游戏智能体却没人想着让智能体少玩几局因为主流强化学习算法大多是无模型路线——智能体不建立任何关于环境的内部规律纯粹靠试错积累奖励信号来调整策略。以Atari游戏为例传统算法动辄需要几千万帧画面才能达到像样的水平一个游戏训练半天到几天是家常便饭。这种硬碰硬的学习方式有两个痛点一是真实环境交互成本高尤其换到机器人、自动驾驶这类场景几千万次试错根本不现实二是智能体学到的知识高度绑定在具体任务上换一个游戏往往从头再来。那让智能体先学习一个世界模型呢意思是智能体在与环境交互一段时间后能不能把我做了什么动作画面会变成什么样、能拿多少奖励这些规律压缩进一个神经网络里然后用这个模型代替真实环境在内部预演未来的轨迹。这样真实环境只需要提供少量样本大量的策略试错在模型内部完成就够了。这个思路就是论文世界模型路线的起点。1.2 世界模型不是新鲜概念但拿来决策一直很难世界模型这个词在强化学习出现很早但绝大多数早期工作只把它当预测任务来做让模型能准确预测下一帧像素、预测奖励至于怎么用这个模型来改进策略始终没有特别成功的方案。原因有两个坎。第一像素级预测误差会随推演步数快速累积——模型在第一步预测偏一点第二步就偏更多几步之后想象出来的画面完全失真策略在这上面学就白学了。第二世界模型和策略之间怎么衔接是个工程难题早期的模型预测的是像素而策略需要的是对当前状态的理解两者之间的接口很模糊。作者团队的前作已经把循环状态空间模型简称RSSM这条路铺开了不直接预测像素而是让模型学一个压缩的隐状态再用这个隐状态去重建画面、预测奖励。这样预测误差不会完全压在像素层面隐状态本身相当于一个对世界的结构化摘要。但前作在离散控制任务上表现尚可一到Atari这类画面复杂、奖励稀疏的游戏就明显吃力。这篇论文就是在这个基础上做关键升级其中最重要的一刀就是把隐变量从连续分布换成了离散分布。1.3 为什么研一精读这篇论文的回报特别高我个人觉得研一阶段选论文精读挑的不应该是最新最热的而应该是结构完整、演进脉络清晰、还能动手复现的。这篇恰好满足三个条件。第一它的论文骨架非常规整开头谈样本效率问题中间讲离散世界模型的设计后半部分用大量消融实验说明每个设计选择各贡献了多少。你可以照它的结构学习怎么写一篇好的系统型论文。第二它的代码开放环境用标准的Atari接口就能跑你不需要自备特殊硬件一张普通显卡就能起来一个小规模的实验。第三它是后续一系列基于模型工作的起点读懂了这篇再去读后续的扩展版本、其他团队的类似工作基本就是一条顺畅的坡道不需要再重新爬一次。2. 核心改动把隐变量从连续池子换成离散词条2.1 先搞清楚论文里的state到底是个什么东西很多研一同学看到world model的图就晕不知道圆框里的state是真实画面还是别的表示。这里必须说清楚state不是图像本身而是编码器从图像里提取出的一个压缩向量。打个比方看一幅台球的桌面画面人类不会记住每一个像素的位置而是记住白球在左侧、红球靠近袋口、目前轮到我了这几个关键要素。latent state就是模型版本的这些关键要素——一段固定长度的向量里面编码了当前局面最重要的信息。在RSSM结构里latent state还不是一个静态向量它由两部分组成一部分是来自当前观测的表征也就是我看到什么另一部分是来自上一时刻状态和动作的预测也就是我预期会发生什么。模型把这两种信息融合起来既尊重事实又保留了对未来的预测能力。这个融合是RSSM能缓解预测误差累积的核心——哪怕预测偏了当前观测还能把状态拉回正轨。2.2 连续高斯的毛病表达力有余训练稳定性不足前作DreamerV1的latent state是一个多元高斯分布用均值和方差来描述状态可能的分布范围。高斯分布用起来方便梯度也容易回传但它有两个实际毛病。第一个毛病是点多面广但不够准。高斯假设状态服从一个单峰分布可真实世界经常是多模态的同一个画面可能对应球正在向左滑和球正在向右滑两种完全不同的走势单峰高斯只能模棱两可地覆盖中间地带。第二个毛病是训练稳定性。高斯的方差项在梯度反传里很容易出问题方差拉大一点、缩小一点整个损失函数的曲面都会剧烈变化模型往往要把大量容量花在维持分布稳定上而不是花在学准世界规律上。离散分布就换了一种思路把隐状态定义成一小串离散词条的组合。比如用一个长度为32的向量每位从32个离散类别里选一个那么总共能表达的组合空间非常大而且每个类别都是一个明确的状态桶。模型不用费劲去拟合一个光滑的钟形分布而是直接表达当前属于哪几个离散状态。论文作者的核心判断是对世界模型来说这种离散表达更适合长时间稳定地推演未来。2.3 代码层面的离散化实现直通梯度的巧劲这里有个工程上的难点离散采样本身不可导而神经网络训练需要梯度。你没法直接从按类别采样这个操作回传梯度。作者的解决办法是直通梯度straight-through前向传播时真的去采样一个one-hot向量反向传播时则假装这个向量就是它自己的softmax logits。这样既保证了前向时状态真的是离散的又让梯度能无障碍地流到生成logits的网络里。# 示意离散隐变量采样伪代码 logits transition_model(prev_state, action) # 预测下一状态的类别得分 sample one_hot(categorical_sample(logits)) # 前向真实采样 latent sample logits - stop_gradient(logits) # 反向使用logits的梯度第三行是整个技巧的核心前向传播时stop_gradient不产生任何影响结果就是sample反向传播时sample被视为常量可训练的梯度全部通过logits传导。用这个技巧离散世界模型可以端到端训练。2.4 离散化带来的彩蛋KL平衡训练更省心除了表达力更强离散化的另一个隐性收益是世界模型的训练更容易稳定。连续高斯版本的训练里KL散度项的尺度往往随训练波动稍不留神就出现后验塌缩——编码器不再向隐状态放信息解码器也摆烂整个模型退化成无意义的预测器。作者针对这个问题设计了一套KL平衡策略把KL散度拆成先验→后验和后验→先验两个方向一个方向用梯度更新另一个方向用停止梯度保持稳定再配合一个自由比特下限允许在训练初期容忍一定程度的KL偏差避免模型过早收敛到平庸解。这一套组合拳放在离散分类分布上格外顺手因为分类分布的KL分解非常干净可操作性比高斯分布强得多。这里如果你跳过不看后面复现时损失函数写不对模型就会训练得非常神经质。3. 训练三阶段预测世界、想象世界、在想象中做决策3.1 第一阶段世界模型学习让模型先学会读懂画面世界模型不是一个独立网络而是一组彼此配合的模块。先看数据流经验回放池里存了一批真实的交互序列——观测、动作、奖励、终止标志。编码器把观测x_t转换成后验隐状态q(s_t | s_{t-1}, a_{t-1}, x_t)这就是结合当前画面修正预测的环节。RSSM同时维护一个先验转移模型p(s_t | s_{t-1}, a_{t-1})作用是只看前一步状态和动作预测下一步可能的状态分布。在这之上还挂了三个输出头解码器负责从隐状态重建观测画面奖励预测器负责估计当前状态下能拿多少奖励继续预测器负责判断episode是否结束。三个头各管一件事对应损失函数的三项。如下表所示损失项作用通俗理解观测重建损失解码器还原画面隐状态有没有保留足够的视觉信息奖励预测损失奖励预测器逼近真实奖励隐状态能不能推算分数变化继续预测损失终止预测器判断是否结束隐状态能不能感知是否game overKL散度损失让后验靠近先验预测的未来别和实际观测差太远世界模型的训练目标就是把这四项损失加在一起做梯度下降。这里有一个关键细节所有训练都在回放池的真实经验上进行不需要在训练世界模型时额外消耗环境交互次数。3.2 第二阶段在梦境里想象轨迹世界模型训练到一定程度后训练循环进入第二阶段想象。想象的意思是从某个真实经验片段中取出一个隐状态作为起点然后不再读取真实环境只靠转移模型在潜在空间里向后滚动。每一步转移模型预测下一个隐状态的分布从中采样得到虚拟的下一帧状态同时用奖励预测器给出虚拟的奖励。这一步要特别注意整个想象过程完全发生在latent空间不经过解码器重建图像也不经过真实环境。解码器只在可视化时用相当于给你看智能体做过的梦长什么样。为什么能这么做因为世界模型学到的隐状态已经包含了决策需要的全部信息状态演变规律、奖励信号、终止条件。策略关心的是我该选哪个动作它不需要看到真实像素只需要看到足够好的状态表示。想象轨迹的长度也叫想象视野论文里一般取十几步。这个长度要权衡太短策略只看得到眼前几步学不到长远规划太长预测误差累积变严重想象轨迹失真。作者后面对比了不同想象长度对分数的影响这也是一个值得你在复现时亲自调一调的超参数。3.3 第三阶段Actor-Critic在梦境里学策略有了想象轨迹强化学习算法就有用武之地了。在梦境轨迹上每一步都有一份隐状态和对应的虚拟奖励这样你可以直接计算折扣回报再用一个Actor网络输出动作分布一个Critic网络估计状态价值。Actor的更新目标是让策略在想象轨迹上的期望回报最大化Critic的更新目标则是让它的价值估计逼近λ-return。整个流程和常规的Actor-Critic没有本质区别唯一的区别是它训练用的环境是学出来的世界模型而不是真实环境。这正是论文方法能大幅提升样本效率的根本原因真实环境只用于收集训练世界模型的数据策略的大量试错转移到便宜得多的梦境推演里。3.4 三阶段合成的训练循环一轮又一轮地转把上面三个阶段串起来就是论文训练算法的核心循环智能体按当前策略与环境交互收集新的观测、动作、奖励存入经验回放池。从回放池采mini-batch训练世界模型更新编码器、转移模型、解码器和各个预测头。利用更新后的世界模型在潜在空间展开多条想象轨迹更新Actor和Critic。回到第1步用更新后的策略继续与环境交互。整个循环的每一步都是在真实数据→世界模型→想象轨迹→策略更新这条流水线上完成的。研一同学看这段时最容易觉得抽象我建议你亲手画一遍数据流向图从真实观测进入编码器开始到策略输出动作结束把每一步数据的形状标在旁边基本就通透了。4. 关于Atari实验分数会骗人人类水平也要问怎么定4.1 看懂了实验设置才算看懂论文的野心论文选用的基准是Atari 2600平台的57款游戏这是强化学习领域沿用多年的经典测试场涵盖格斗、射击、平台跳跃、策略等不同类型。每个游戏给智能体100万帧的训练预算——这个数字相对无模型算法动辄数千万帧动辄低了很多但相对极低样本强化学习比如100K帧又宽松不少。换句话说论文的目标不是做最极限的样本效率而是在合理预算下达到甚至超过人类玩家的水平。理解这个设置很重要你在复现时不需要照搬100万帧预算用小游戏先跑通流程即可。4.2 消融实验告诉你改进到底堆在哪里论文最值得精读的部分其实是消融实验。作者做了几个版本的对比连续潜变量的DreamerV1、加大规模的连续版本、以及离散潜变量的最终版本。论文报告的结果显示仅仅把潜变量从连续换成离散平均分数就有非常明显的提升在此基础上再增大模型规模又能进一步拉升成绩。也就是说离散表示和模型容量是两个各自独立起作用的因素合在一起才达到人类水平。对比项相对表现说明连续隐变量小模型基线在部分游戏上能玩但整体远低于人类连续隐变量大模型中等规模提升有帮助但离散化收益更关键离散隐变量大模型最高论文最终报告达到人类水平这张表是作者用实验画出来的“贡献分解图”。研一读者请记住这个习惯不看消融实验你就只知道这个方法有效不知道这个方法为什么有效、哪个设计最值钱。4.3 人类归一化得分只报均值是典型的幸存者偏差论文用人类归一化得分来统一衡量57个游戏的表现把人类玩家的平均水平当作100%智能体的得分按比例换算。最终论文报告的均值落在人类水平附近这个成绩确实漂亮。但你必须看清楚这是一个平均值。逐游戏细看智能体在一些游戏上远超人类在另一些游戏上则远低于人类。比如动作简单、规则明确的游戏模型能靠预演推算出极优策略而需要长时记忆和抽象推理的游戏世界模型也会露怯。所以论文标题说的Mastering Atari要理解为整体达到人类级别的掌控力而不是每个游戏都超越人类。研一同学读任何论文的实验结果都要先看单个任务的表现分布再回来看均值否则很容易被一个漂亮的平均值带偏。4.4 与无模型算法对比位置决定价值把DreamerV2放到当时的算法版图里看最合理的参照是那些同样用100万帧预算的深度强化学习算法。在同样的交互预算下DreamerV2的平均人类归一化得分和最强的无模型算法不相上下甚至更优。这点很关键一个带世界模型的算法在同等样本预算下追平了那些专注试错的无模型算法这本身就是世界模型路线值得做的强烈信号。更难得的是因为世界模型的存在你还多了可解释这层收益——你可以把模型想象的画面解码出来看观察智能体脑子里在预测什么。这种可读性对调试和理解学习过程非常有帮助无模型算法是很难提供这种视角的。5. 研一精读实操从伪代码到复现的完整路径5.1 我的精读顺序先图表、后伪代码、再公式研一的同学最容易犯的错是从第一页开始逐字读到最后一页结果到了方法章节已经晕掉。我建议你按这个顺序走先读摘要和引言用三句话写下作者要解决什么问题、他提出什么方案、实验结论是什么。直接跳到方法部分的结构图不看公式只看图中模块之间的数据流向。把训练伪代码手抄一遍想象你要用代码实现它哪些地方你会卡住。再回头看公式这时已经知道每个符号对应哪个模块公式会好理解得多。最后读实验和消融重点关注对比表和消融结论。这套顺序的核心逻辑是先用图像和伪代码建立整体认知再让公式为整体认知填空。直接啃公式对新手极其不友好还容易半途放弃。5.2 伪代码逐行过一遍训练主循环其实就三个块论文给出的训练伪代码信息量很大但拆解后就三块。我用示意伪代码帮你梳理# DreamerV2 训练主循环示意结构 for iteration in range(train_steps): # 阶段1与环境交互收集经验 for env_step in range(steps_per_iteration): latent world_model.encode(obs) action actor(latent) # 加探索噪声后执行 obs, reward, done env.step(action) replay.add(obs, action, reward, done) # 阶段2在回放数据上训练世界模型 world_model.train(replay.sample_batch()) # 阶段3在潜在空间想象并训练行为 for dream_step in range(imagination_steps): imag_latent world_model.rollout(actor, start_state) actor.critic.update(imag_latent)第一块负责产生数据第二块负责让世界模型更准确第三块负责让策略更聪明。这三块各训练各的损失互不冲突。研一阶段理解到这个粒度就够了再往下的实现细节就交给代码吧。5.3 复现时最容易翻车的四个坑我后来照着论文跑过一个小规模的复现踩了不少坑给你列出来提前避雷。第一个坑是奖励没做缩放。Atari游戏的奖励范围很大Actor的梯度容易被极端奖励值带崩。论文在实验中对回报做了归一化你复现时也要处理这一步否则训练曲线会剧烈震荡。第二个坑是KL平衡缺了一个方向。有些初学者实现世界模型时只写了后验逼近先验那项忽略了先验向真实经验那项的梯度传导结果模型早期预测能力很差表现为训练曲线长期不动。这个方向很容易在复现时漏掉你要对照论文公式一个个核对。第三个坑是discrete采样梯度断流。用了one-hot采样但忘了直通梯度的stop_gradient技巧模型整体不更新或者报错说梯度无法通过。前馈数据看了一眼是离散向量但你要确保它的配套梯度计算是连续的。第四个坑是对训练时间预估不足。完整跑完57个Atari游戏需要很多卡时个人复现非常不现实。我建议你先选一个简单游戏比如经典的Pong或Breakout跑通全流程验证管道无误之后再考虑扩展。复现的目的是理解不是为了秒掉完整论文。5.4 读完之后往哪走顺着这篇论文的引用树继续爬精读完这篇我建议你顺手做两件事。第一去读论文引用了哪些关键工作尤其是世界模型的历史脉络和RSSM的原始出处补上前情提要。第二看看后续版本是怎么扩展它的。后续工作把离散世界模型推向了更通用的方向还引入了无需任务奖励就能学习等更前沿的设定。你会发现这篇论文里的很多设计思路被保留了下来但训练稳定性和通用性又上了一个台阶——这本身就是一段完整的科学演进脉络。我现在回想研一读这篇论文的过程最受益的不是记住了几个公式而是学会带着作者到底在赌什么去读论文。这篇论文赌的是与其让智能体在真实世界里硬碰硬地学策略不如先在内部模型里把世界推演熟。这个判断影响了我后来做项目时对样本利用率的看法——很多问题不是数据不够而是我们有没有把已有数据的价值榨干。所以如果你也是研一新手我最后的建议是不要只停留在读懂了试着亲手实现一次它的最小版本。哪怕只是拿一个小游戏把三阶段循环跑通你对世界模型、想象轨迹、离散表示这些概念的理解都会比刷十篇综述都来得扎实。这大概就是论文精读的真正意义。