1. 论文要解决的核心问题

强化学习智能体传统上直接从像素学策略,存在两个大问题:

  1. 样本效率低——需要和真实环境交互海量次数
  2. 端到端训练难以拆解、难以复用

这篇论文提出:先让智能体学会”理解世界如何运作”,再基于这个内部理解去做决策,并验证了一个大胆的推论——如果这个内部理解足够准,智能体甚至可以完全在自己的”想象”里完成训练,不需要真实环境参与。

2. 三大模块:V / M / C

智能体被拆成三个独立训练的模块,分工明确:

V/M/C 三大模块架构

  • V(Vision):用 VAE 把每一帧高维图像压缩成一个低维潜变量 zz(论文里 D=32D=32 维)
  • M(Memory):用 RNN + MDN 学习环境的转移动态,预测下一时刻的 zz 可能长什么样
  • C(Controller):一个极简的线性模型,把 [zt,ht][z_t, h_t] 直接映射到动作 ata_t

动作 aa 的反馈回路:图中 C 输出的 ata_t 只画了向下的箭头(流向”环境”),但这个 ata_t 同时也会反馈进入下一步的 M——完整的因果链条是:

C(t):(zt,ht)→at  ⟶  M:(zt,at,ht)→ht+1  ⟶  C(t+1):(zt+1,ht+1)→at+1C(t): (z_t, h_t) \rightarrow a_t \;\longrightarrow\; M: (z_t, a_t, h_t) \rightarrow h_{t+1} \;\longrightarrow\; C(t+1): (z_{t+1}, h_{t+1}) \rightarrow a_{t+1}

ata_t 在时间 tt 是 C 的输出,在时间 t+1t+1 就变成了 M 需要的输入——同一个变量,在因果链条上的两个位置,这是标准的”智能体-环境交互循环”。

3. V 模块:VAE 的原理与训练

3.1 VAE 结构:Encoder + Decoder,联合训练

x  →Encoder  μ(x),σ(x)  →采样  z  →Decoder  x^x \;\xrightarrow{\text{Encoder}}\; \mu(x),\sigma(x) \;\xrightarrow{\text{采样}}\; z \;\xrightarrow{\text{Decoder}}\; \hat{x}

Encoder 不能单独训练——重建损失必须靠 Decoder 才能算出来,训练信号(梯度)从 Decoder 端产生,经反向传播一路传回 Encoder。如果没有 Decoder 参与,Encoder 收不到任何”该怎么编码才有意义”的反馈,唯一的驱动力(KL项)会让它退化成对所有输入都输出 μ=0,σ=1\mu=0,\sigma=1。

3.2 损失函数:重建项 + KL 正则,两者都源于同一个数学目标

LVAE=∥x−x^∥2⏟重建损失+DKL(q(z∣x) ∥ p(z))⏟KL正则\mathcal{L}_{\text{VAE}} = \underbrace{\|x-\hat{x}\|^2}_{\text{重建损失}} + \underbrace{D_{KL}\big(q(z|x)\,\|\,p(z)\big)}_{\text{KL正则}}

这两项不是拼凑出来的,而是从最大化 log⁡p(x)\log p(x) 的证据下界(ELBO)推导出来的:

log⁡p(x)≥Eq(z∣x)[log⁡p(x∣z)]−DKL(q(z∣x) ∥ p(z))\log p(x) \geq \mathbb{E}_{q(z|x)}[\log p(x|z)] - D_{KL}\big(q(z|x)\,\|\,p(z)\big)
  • 重建损失:强迫 zz 必须携带足够的关于 xx 的信息,否则 Decoder 无法还原
  • KL 正则:强迫每张图对应的分布 q(z∣x)q(z|x) 不能离先验 p(z)=N(0,I)p(z)=\mathcal{N}(0,I) 太远,防止退化、并让潜空间连续、稠密、无空洞

两项损失存在天然张力:重建损失想把不同输入的编码拉开,KL 项想把所有编码拉向公共的标准正态——训练收敛的结果是两者的折中。

3.3 先验 p(z)=N(0,I)p(z)=\mathcal{N}(0,I) 是什么

p(z)p(z) 不是一个向量,而是定义在整个 DD 维潜空间上的一个分布——DD 个相互独立的标准正态分布拼在一起:

p(z)=∏d=1DN(zd;0,1)p(z) = \prod_{d=1}^{D}\mathcal{N}(z_d;0,1)

下图直观展示了 D=2D=2 时,从这个分布采样1500个点是什么样子——散点云呈圆形对称分布,把每个点单独投影到某一个坐标轴上统计,得到的就是一维的标准正态曲线:

二维标准正态分布散点图

这个先验分布本身不携带任何预设的语义——它只是一个数学上方便处理的参照系(KL 散度有闭式解、支持重参数化技巧、采样简单、各向同性无偏好方向)。真正让潜空间”有意义”的,是 Encoder 和 Decoder 在训练中协调学出的映射关系。

3.4 “一张图对应哪个 z”——对应的是一个区域,不是一个点

Encoder 对每张具体的图 xx,输出的是这张图专属的 μ(x),σ(x)\mu(x),\sigma(x),再从 N(μ(x),σ(x)2)\mathcal{N}(\mu(x),\sigma(x)^2)(而不是原始的通用 N(0,1)\mathcal{N}(0,1))中采样,得到最终的 zz:

z=μ(x)+σ(x)⋅ϵ,ϵ∼N(0,1)z = \mu(x) + \sigma(x)\cdot\epsilon, \qquad \epsilon\sim\mathcal{N}(0,1)

这个采样公式,是标准化操作 z=(x−μ)/σz=(x-\mu)/\sigma 的逆运算:标准化是把任意正态分布”拉回”标准正态,这里则是把标准正态样本 ϵ\epsilon“变换”成指定均值方差的分布——都是在利用正态分布对线性变换封闭这条性质。这一写法叫重参数化技巧(reparameterization trick),它把不可导的”随机采样”操作,转化成对 μ,σ\mu,\sigma 可导的线性运算,让梯度能顺利反向传播。

训练前后的对比——训练开始时 Encoder 权重随机,同类图片的编码位置毫无规律;训练收敛后,相似的图片被推到潜空间中彼此靠近的区域,整体形状依然趋近先验 N(0,I)\mathcal{N}(0,I):

训练前后对比

这个”图 → 区域”的映射关系,不是人为设计或指定的,而是重建损失和 KL 正则两股力量共同”逼”出来的必然结构——具体每一维代表什么语义,人类通常无法预先解释,往往需要事后做探索性分析(如固定其他维度、遍历某一维观察解码结果)才能大致解读。不同的训练顺序、权重初始化、采样噪声,会让训练收敛到潜空间的不同”朝向”(比如整个潜空间可以任意旋转,先验分布保持不变),但只要重建质量高、潜空间连续,这些不同版本的 VAE 在功能上是同样合格的。

3.5 VAE vs 普通 Autoencoder

AutoencoderVAE
Encoder 输出一个确定的向量一个分布的参数 (μ,σ)(\mu,\sigma)
训练目标只有重建损失重建损失 + KL 散度
潜空间结构无约束,可能有空洞连续、稠密,适合采样
能否从中采样生成新内容不行(容易乱码)可以

4. M 模块:RNN + MDN

RNN 与 MDN 的详细结构

RNN 每一步吃进 (zt,at,ht−1)(z_t,a_t,h_{t-1}),更新出隐藏状态 hth_t:

ht=RNN(zt−1,at−1,ht−1)h_t = \text{RNN}(z_{t-1},a_{t-1},h_{t-1})

4.1 为什么用 MDN 而不是直接回归

环境的未来往往是多模态的(同一段历史可能对应好几种合理但截然不同的走向)。若直接回归输出一个 zz,MSE 训练会让网络学出”模糊的平均值”,谁都不像。MDN 让网络输出一个混合高斯分布的参数:

P(zt+1∣ht)=∑k=1Kπk⋅N(zt+1;μk,σk2)P(z_{t+1}\mid h_t) = \sum_{k=1}^{K}\pi_k\cdot\mathcal{N}(z_{t+1};\mu_k,\sigma_k^2)
参数数量含义约束
π\piKK 个选中每个分量的概率,所有维度共享同一套softmax,非负且和为1
μ\muK×DK\times D 个每个分量下,zz 每一维的期望值无约束
σ\sigmaK×DK\times D 个每个分量下,zz 每一维的离散程度需为正

下图展示了三个高斯分量按权重叠加成混合分布的直觉:

三个高斯分量叠加成混合分布

4.2 采样:两步走

k∼Categorical(π1,…,πK)⟹zd=μk,d+σk,d⋅ϵd,  ϵd∼N(0,1)k\sim\text{Categorical}(\pi_1,\ldots,\pi_K) \qquad\Longrightarrow\qquad z_d=\mu_{k,d}+\sigma_{k,d}\cdot\epsilon_d,\ \ \epsilon_d\sim\mathcal{N}(0,1)

先按 π\pi 抽签选一个分量(整个 DD 维向量只抽一次),再在选中的分量下,每一维独立采样——不是把各分量加权平均。π\pi 只负责”选择”这一步,选定之后就不再参与计算,这正是 MDN 能表达多模态、而不退化成模糊平均的关键。

4.3 训练:Teacher Forcing

训练时,RNN 每一步的输入始终是真实的 zt−1z_{t-1}(来自 V 编码真实观测),而不是模型自己采样出的 z′z',损失函数是负对数似然:

LM=−log⁡(∑k=1Kπk⋅N(zt真实;μk,σk2))\mathcal{L}_M = -\log\left(\sum_{k=1}^K \pi_k\cdot\mathcal{N}(z_t^{\text{真实}};\mu_k,\sigma_k^2)\right)

这一阶段 MDN 采样出的 z′z' 只用来算一次损失就被丢弃,不参与后续前向传播——训练全程锚定在真实数据上,不存在自回归误差累积的问题。

4.4 推理:MDN 输出分布,如何落地成一个具体向量

推理时不直接”处理”密度值,而是执行标准的采样流程,最终得到一个确定性的向量:

pi = softmax(pi_logits / tau)          # tau: 温度参数,控制随机程度
sigma = exp(log_sigma) * sqrt(tau)
k = multinomial(pi, num_samples=1)     # 第一步: 选分量
epsilon = randn_like(mu[k])
z_next = mu[k] + sigma[k] * epsilon    # 第二步: 采样得到具体向量

温度参数 τ\tau:τ→0\tau\to 0 趋向于”总是选权重最大的分量、几乎不加噪声”(类似 LLM 里的贪心解码);τ\tau 越大随机性越强。做梦阶段若 τ\tau 太低,生成的虚拟环境过于规整,Controller 容易”钻空子”而非学到鲁棒策略。

与 LLM 的类比:MDN 和 LLM 的下一词预测本质上是同一类问题——都是”网络输出一个分布的参数,推理时按概率采样得到具体输出,用温度控制随机程度”。区别只在于 LLM 面对有限词表上的离散选择(Categorical + softmax),MDN 面对连续实数空间(混合高斯);两者训练时都直接用真实标签算损失(交叉熵 / 负对数似然),不依赖模型自己的采样结果。

5. 完整训练流程:三个模块分开训练,顺序进行

阶段训练对象数据来源zz 的角色优化方法
1V(VAE)随机策略采集的真实帧编码目标重建损失 + KL(梯度下降)
2M(RNN+MDN)V 编码出的真实 zz 序列输入用真实 zz,teacher forcing负对数似然(梯度下降)
3C(线性控制器)M 生成的虚拟轨迹(“做梦”)自回归采样的 z′z' 驱动整条轨迹CMA-ES(进化策略)

做梦阶段的核心循环:

zt′∼MDN(ht)→ht+1=RNN(ht,zt′,at)→zt+1′∼MDN(ht+1)→⋯z_t' \sim \text{MDN}(h_t) \to h_{t+1}=\text{RNN}(h_t,z_t',a_t) \to z_{t+1}'\sim\text{MDN}(h_{t+1}) \to \cdots

C 参数量极小(几百个,线性模型),用 CMA-ES 而非梯度类强化学习方法训练:维护一个参数分布,每代采样候选控制器、评估其累积奖励、据表现更新分布均值和协方差,逐代进化。训好后把 C 放回真实环境测试,依然表现良好——证明智能体可以完全脱离真实环境,只靠内部想象完成学习。这里”评估累积奖励”具体怎么算,以及训练究竟是在真实环境还是梦境里进行,详见第六节。

6. Action 与 Reward 是从哪里来的

这是让整套训练流程能跑通的两个容易被忽略、但缺一不可的细节。

6.1 训练 V、M 阶段:action 由随机策略生成

采集数据训练 V 和 M 时,论文用完全随机的策略去探索环境——不需要任何智能,每一步直接从动作空间里随机采样:

  • CarRacing(连续动作):转向、加速、刹车三个连续值,每步随机采样
  • VizDoom(离散动作):从有限几个动作选项里随机选一个

这一阶段的目标只是让 V、M 见识到足够多样的画面和状态转移,不需要”表现聪明”。社区复现时发现纯粹逐帧独立的白噪声动作会让赛车很快冲出赛道、探索范围太窄,因此常改用带时间相关性(“布朗运动式”)的随机策略,让动作在时间上平滑漂移,以采集到更连贯、更有代表性的轨迹。

6.2 训练 C 阶段:reward 的来源,因实验而异

CarRacing:C 的训练直接在真实环境里进行,reward 就是 Gym 环境自带的奖励函数(访问的赛道格子数、用时),不涉及”梦境里怎么算奖励”的问题。

VizDoom(真正”完全在梦境里训练”的实验):该环境本身没有显式奖励,论文把奖励重新定义为存活的时间步数。为了让这个定义能在梦境里同样成立,M 模块被扩展为额外预测一个”done”(是否死亡)信号,与预测 zt+1z_{t+1} 的分布并列输出:

ht=RNN(zt−1,at−1,ht−1)  ⟶  (P(zt∣ht), P(donet∣ht))h_t=\text{RNN}(z_{t-1},a_{t-1},h_{t-1}) \;\longrightarrow\; \big(P(z_t\mid h_t),\ P(\text{done}_t\mid h_t)\big)

做梦时,一旦这个信号判定”死亡”,虚拟轨迹终止——轨迹走了多少步,就是这次评估的累积奖励,不需要额外训练一个”reward数值预测头”,巧妙地绕开了”如何在梦境里凭空生成奖励数值”这个问题。

对于奖励结构更复杂(不只是”死没死”)的环境,更通用的做法(后续 PlaNet、Dreamer 采用)是再给 M 加一个奖励预测头,和预测 zz、predict done 并列,用真实环境观测到的 reward 做监督训练,做梦时每一步直接输出一个预测的 reward 数值、累加得到总奖励。

实验C 在哪训练reward 来源
CarRacing真实环境环境自带奖励函数,直接读取
VizDoom(做梦实验)完全在 M 生成的虚拟环境里奖励=存活步数,由 M 额外预测的”done”信号间接决定

7. 为什么这篇论文重要

它把”感知压缩、动态预测、决策”三件事解耦成三个独立训练的小模块,大幅降低了训练难度和参数量;更重要的是,它第一次具体地证明了”在想象中训练”这件事是可行的。

这条思路直接启发了后续 DeepMind 的 PlaNet、Dreamer 系列(把”做梦训练”发展得更彻底,完全在潜空间里做多步展开,连解码回像素这一步都可以省略),也是这几年”世界模型”重新成为大模型/生成式建模热点话题(如 Sora、Genie、JEPA 等)的重要源头之一。