World Models：基于生成式世界模型的强化学习框架——从视觉编码到梦境训练的完整范式

论文信息

标题：World Models
会议：arXiv preprint (cs.LG) 2018
单位：Google Brain、NNAISENSE、Swiss AI Lab, IDSIA (USI & SUPSI)
代码：https://worldmodels.github.io（交互演示）
论文：https://arxiv.org/pdf/1803.10122.pdf

---

一、引言：我们都活在自己大脑构建的“虚拟世界”里

你有没有好奇过，为什么职业棒球手能击中时速160公里的快球？光靠眼睛看、脑子反应肯定来不及——等视觉信号完全传到大脑皮层，球都已经飞到脸上了。答案就是我们大脑里自带的世界模型：靠过往经验预判球的轨迹，肌肉形成条件反射直接挥棒，整个过程连意识都插不上手。

图1 人类心智模型的示意
出处：原文Figure 1，引自Scott McCloud《理解漫画》
解读：我们大脑里的世界从来不是真实世界的完整复刻，而是一个经过筛选、抽象的简化模型。我们记住的不是所有细节，而是关键概念和它们之间的关系，并靠这个模型来做所有决策。

不止棒球，我们日常的走路、开车、躲避障碍物，本质上都是靠内部的预测模型在“凭感觉行动”。那强化学习里的AI能不能也搞这么一套？

传统强化学习有个老大难问题：信用分配——奖励来得晚，根本不知道是哪一步动作该领奖、哪一步该背锅。模型越大，参数越多，这个问题就越棘手，所以大家往往只能用小网络训练策略。

这篇论文给出了一个极其优雅的解法：把智能体拆成两部分

1. 一个大的世界模型：用无监督学习单独训练，负责“看懂画面、预判未来”，把高维的原始像素压缩成低维的抽象表示；
2. 一个极小的控制器：只负责根据世界模型给出的抽象特征做决策，参数量只有几百到几千个，用进化算法就能轻松优化。

更绝的是：世界模型本身就能模拟整个环境，我们可以直接让AI在自己“脑补”的梦境里训练策略，练完直接搬到真实环境里就能用，大幅节省真实环境的训练成本。

图3 世界模型训练与应用流程
出处：原文Figure 3
解读：先用真实环境收集的数据训练RNN世界模型，训练完成后，世界模型就可以完全模拟环境，用来训练智能体的策略。

---

二、智能体的三大核心组件

整个智能体由三个模块协同工作，对应人类的视觉感知、记忆预测、决策执行三个能力：

• V（Vision）：视觉编码器，用VAE实现，负责把每帧画面压缩成低维隐向量；
• M（Memory）：记忆预测器，用MDN-RNN实现，负责根据历史信息预测未来的隐向量；
• C（Controller）：控制器，单层线性模型，负责根据视觉和记忆信息输出动作。
  

图4 智能体三大组件架构
出处：原文Figure 4
解读：三个模块各司其职：V处理空间信息，M处理时间信息，C做最终决策。绝大多数复杂度和参数都集中在V和M里，C保持极简。

2.1 V模型：用VAE给画面“写摘要”

环境给的输入是高维的RGB图像（比如64×64×3的像素），直接丢给控制器根本处理不动。V模型的任务就是做有损压缩：把一整张图压缩成一个短短的隐向量$z$，丢掉没用的细节，保留和任务相关的核心信息。

这里用的是变分自编码器（VAE），它和普通自编码器最大的区别是：隐向量不是一个固定的值，而是服从正态分布的随机变量，这样隐空间会更连续、更适合后续的预测模型采样。

核心公式：VAE的损失函数

$${\mathcal{L}}_{\text{VAE}}={\mathcal{L}}_{\text{recon}}+\beta \cdot {\mathcal{L}}_{\text{KL}}$$
符号解释：

• ${\mathcal{L}}_{\text{VAE}}$：VAE训练的总损失，我们的目标是最小化它
• ${\mathcal{L}}_{\text{recon}}$：重建损失，论文里用L2距离计算输入原图和解码器输出图像的像素差异，衡量“压缩后还原得像不像”
• $\beta$：KL散度的权重系数，用来平衡重建质量和隐空间的规整程度
• ${\mathcal{L}}_{\text{KL}}$：KL散度，衡量隐向量的分布和标准正态分布的差异，作用是把隐空间“掰”成规整的正态分布，方便后续模型采样

通俗解释：VAE就像一个“图片摘要生成器”。编码器负责把大图缩写成几十个数的摘要，解码器负责根据摘要还原回图片。训练的时候一边要求还原得尽量像，一边要求摘要的分布符合正态分布，这样后续生成新摘要的时候就不会乱。

图5 VAE的流程示意
出处：原文Figure 5
解读：输入图像经过编码器得到均值μ和标准差σ，采样得到隐向量z，再经过解码器重建出图像。整个过程是无监督的，不需要标注，只需要原始画面就能训练。

在赛车实验里，隐向量$z$的维度是32；在Doom实验里是64。也就是说，一张64×64×3=12288维的图片，被压缩成了32或64维，压缩比超过300倍。

2.2 M模型：用MDN-RNN预判未来

光有单帧的视觉压缩还不够，智能体需要知道“接下来会发生什么”。M模型的任务就是建模时序，预测未来的隐向量。

为什么不直接预测下一帧像素？因为像素太高维了，RNN很难学。有了VAE的压缩，我们只需要预测低维的$z$向量，难度直接降了几个数量级。

而且很多环境本身是随机的（比如怪物什么时候发火球），没法精准预测“下一个状态一定是什么”，所以我们不用确定性输出，而是输出下一状态的概率分布——这就是混合密度网络（MDN）的作用。

核心公式：下一时刻隐向量的概率分布

$$P(z_{t+1}|a_t,z_t,h_t)=\sum _{i=1}^K{\pi }_i\cdot \mathcal{N}({\mu }_i,{\sigma }_i^{2}I)$$
符号解释：

• $P(z_{t+1}|a_t,z_t,h_t)$：条件概率，给定当前动作$a_t$、当前隐向量$z_t$、RNN隐藏状态$h_t$时，下一时刻隐向量$z_{t+1}$的概率密度
• $K$：高斯混合分量的数量，论文中设置为5
• ${\pi }_i$：第i个高斯分量的权重，所有权重加和为1，代表“这种情况发生的概率”
• ${\mu }_i$：第i个高斯分量的均值向量
• ${\sigma }_i$：第i个高斯分量的标准差
• $\mathcal{N}$：正态分布
• $I$：单位矩阵，论文采用对角协方差假设，也就是隐向量各维度相互独立

通俗解释：MDN-RNN不直接说“下一秒画面一定是A”，而是说“下一秒有30%概率是A，50%概率是B，20%概率是C”。就像天气预报不说“明天一定下雨”，而是说“降水概率70%”，这样更符合真实环境的随机性。

图6 带MDN输出层的RNN结构
出处：原文Figure 6
解读：RNN的隐藏层处理时序信息，输出层不是直接预测z，而是输出混合高斯的三个参数：权重π、均值μ、标准差σ，采样后得到预测的z。

这个结构和当年用来生成手写、生成简笔画的SketchRNN是同款，只不过这里用来预测游戏画面的隐向量。我们还可以通过温度参数τ来控制随机性：τ越小，分布越尖锐，预测越确定；τ越大，分布越平缓，环境越随机。

2.3 C模型：极简的线性控制器

控制器C是整个系统里最“朴素”的部分——就是一个单层线性模型，没有隐藏层，没有激活函数（输出会加tanh限制范围）。

核心公式：控制器的动作输出

$$a_t=W_c\left[z_t{h}_t\right]+b_c$$
符号解释：

• $a_t$：t时刻智能体输出的动作向量。赛车任务里对应转向、油门、刹车三个连续值；Doom任务里对应左右移动的离散动作
• $W_c$：控制器的权重矩阵，是唯一需要通过强化学习优化的参数
• $z_t$：t时刻VAE编码得到的视觉隐向量，代表“现在看到了什么”
• $h_t$：t时刻MDN-RNN的隐藏状态，代表“过去的记忆和对未来的预判”
• $[z_t{h}_t]$：把两个向量拼接在一起作为输入
• $b_c$：控制器的偏置向量，也是可训练参数

通俗解释：控制器就是个“条件反射中枢”，拿到视觉信息和预测信息，直接线性组合输出动作。因为参数量太少了（赛车任务只有867个参数），根本不需要反向传播，用进化算法就能轻松优化。

为什么要把控制器做这么小？因为强化学习的信用分配问题在小参数空间里会简单很多。把复杂度都交给无监督训练的世界模型，控制器只需要在低维的抽象空间里学策略，效率极高。

2.4 三个模块的协同工作流程

完整的运行逻辑是：

1. 环境输出原始画面，VAE编码得到当前隐向量$z_t$；
2. 把$z_t$和RNN隐藏状态$h_t$拼起来，输入控制器得到动作$a_t$；
3. 动作作用于环境，环境返回新的画面和奖励；
4. 把$a_t$和$z_t$输入RNN，更新隐藏状态得到$h_{t+1}$，用于下一时刻。

图8 智能体完整运行流程
出处：原文Figure 8
解读：原始观测先经过VAE压缩成z，z和RNN的隐藏状态h一起输入控制器输出动作a；动作作用于环境的同时，也和z一起输入RNN，更新隐藏状态用于下一步。

核心代码：单轮rollout的实现

import torch
import numpy as np

def rollout(controller, vae, rnn, env, max_steps=1000):
    """
    执行一轮环境交互，返回累计奖励
    controller: 线性控制器
    vae: 视觉编码器
    rnn: MDN-RNN记忆模型
    env: 强化学习环境
    """
    obs = env.reset()
    h, c = rnn.init_hidden(batch_size=1)  # 初始化LSTM的隐藏状态和细胞状态
    done = False
    cumulative_reward = 0.0
    
    while not done:
        # 1. VAE编码当前画面得到隐向量z
        obs_tensor = torch.FloatTensor(obs).unsqueeze(0) / 255.0
        z, _, _ = vae.encode(obs_tensor)
        z = z.squeeze(0).numpy()
        
        # 2. 拼接z和隐藏状态h，输入控制器得到动作
        h_np = h.squeeze(0).numpy()
        input_vec = np.concatenate([z, h_np])
        a = controller.forward(input_vec)
        
        # 3. 动作作用于环境
        obs, reward, done, _ = env.step(a)
        cumulative_reward += reward
        
        # 4. 更新RNN的隐藏状态
        a_tensor = torch.FloatTensor(a).unsqueeze(0)
        z_tensor = torch.FloatTensor(z).unsqueeze(0)
        rnn_input = torch.cat([z_tensor, a_tensor], dim=-1).unsqueeze(0)
        _, (h, c) = rnn(rnn_input, (h, c))
    
    return cumulative_reward

---

三、实验一：CarRacing-v0 赛车任务

第一个实验是OpenAI Gym里的CarRacing-v0，一个俯视角赛车游戏。赛道每次都是随机生成的，要求100轮平均得分超过900才算“解决”，之前的方法一直差口气。

3.1 实验流程

整个训练分五步走，完全是“先学世界，再学控制”的思路：

1. 用随机策略跑10000局，收集所有的画面和动作数据；
2. 用收集的画面训练VAE，把每帧压缩成32维的隐向量；
3. 用压缩后的隐向量序列和动作序列训练MDN-RNN，学会预测下一帧的z；
4. 定义控制器为线性模型$a_t=W_c[z_t{h}_t]+b_c$；
5. 用CMA-ES进化算法优化控制器的参数，最大化累计奖励。

三个模块的参数量对比：

模块	参数数量
VAE	4,348,547
MDN-RNN	422,368
控制器	867

可以看到，控制器的参数量连世界模型的零头都不到。世界模型占了99%以上的参数，但它是无监督训练的，不需要奖励信号，训起来又快又稳。

3.2 实验结果与分析

我们做了三组对照实验，来验证每个模块的作用：

1）只用V模型（不带记忆）

控制器只能看到当前的z，看不到RNN的隐藏状态h。结果就是开车摇摇晃晃，急转弯经常冲出赛道，平均得分只有632±251。

图11 仅使用V模型的驾驶效果
出处：原文Figure 11
解读：没有记忆和预测能力的智能体，只能看到当前这一帧，就像新手开车只盯着车头前的地面，遇到弯道反应不过来，自然走不稳。

给控制器加一层隐藏层，得分提升到788±141，但还是达不到900的及格线。

2）完整世界模型（V+M）

给控制器加上RNN的隐藏状态h之后，驾驶稳定性直接上了一个台阶，过弯流畅，很少冲出赛道，最终平均得分达到906±21，首次正式解决了这个任务。

图12 完整世界模型的驾驶效果
出处：原文Figure 12
解读：有了记忆和预测能力的智能体，就像老司机，能预判赛道走向，提前打方向，自然开得又稳又快。

和其他方法的对比如下：

表1 CarRacing-v0各方法得分对比
出处：原文Table 1

方法	平均得分
DQN	343 ± 18
A3C（连续动作）	591 ± 45
A3C（离散动作）	652 ± 10
Gym排行榜最优	838 ± 11
仅V模型	632 ± 251
V模型+隐藏层控制器	788 ± 141
完整世界模型（本文）	906 ± 21

为什么加个h就能提升这么多？因为$h_t$里包含了对未来的预测信息。控制器不需要自己去想“接下来会怎样”，RNN已经帮它预判好了，它只需要根据预判做条件反射就行——就像棒球手靠肌肉记忆挥棒，不需要脑子里算抛物线。

3.3 在梦境里开车

既然MDN-RNN能预测下一帧的z，那我们完全可以不用真实环境，让RNN自己“脑补”出整个赛道：

• 初始状态给一个z0；
• 控制器输出动作a；
• RNN根据z和a预测下一个z的分布，采样得到$z_{t+1}$；
• 循环往复，整个过程完全在隐空间里进行。

如果想看画面，还可以用VAE的解码器把z还原成像素。这就是AI的“梦境”——完全由自己的世界模型生成的虚拟环境。

图13 智能体在梦境中驾驶
出处：原文Figure 13
解读：整个环境都是MDN-RNN生成、VAE渲染出来的，没有调用真实的游戏引擎。智能体在自己脑补的世界里照样能开车。

---

四、实验二：VizDoom 躲火球任务

如果说赛车实验只是证明了“世界模型能提取好用的特征”，那Doom实验就更硬核了：完全在梦境里训练控制器，然后直接迁移到真实环境。

4.1 任务设置

任务叫Take Cover：智能体在一个房间里，对面的怪物会发射火球，智能体需要左右移动躲避。活的时间越长，得分越高；活过750步就算通关。

和赛车实验不同，这里的MDN-RNN除了预测下一帧的z，还要多预测一个死亡信号d（done），这样才能完整模拟一个强化学习环境的所有输出：观测、终止信号。

4.2 完全在梦境中训练

训练流程和赛车几乎一样，但关键的区别是：控制器全程只在虚拟的梦境环境里训练，完全不碰真实环境。训练好之后，直接把控制器搬到真实的VizDoom里测试。

图15 智能体在梦境中学习躲避火球
出处：原文Figure 15
解读：火球、怪物、墙壁全都是RNN脑补出来的。智能体在这个虚拟世界里学会了左右移动躲火球。

结果非常惊人：

• 梦境环境里训练的控制器，在真实环境里平均存活1092±556步；
• 远高于750步的通关线，也超过了Gym排行榜的820±58步。

也就是说，AI在自己做的梦里练出来的本事，到真实世界里照样能用，甚至还更强。

4.3 一个有趣的问题：AI会“卡bug作弊”

你小时候玩游戏有没有卡过bug？比如穿墙、无限血、利用游戏引擎的漏洞逃课通关。我们的控制器也会干这事——它会专门找世界模型的漏洞来刷分。

最初实验的时候，研究者发现控制器学会了一种“超能力”：用特定的走位，能让梦里的怪物永远不发火球，甚至火球飞过来还能凭空消失。这样它就能轻松拿满分。

图18 控制器发现的“作弊”策略
出处：原文Figure 18
解读：控制器找到了世界模型的盲区，用特殊动作让火球直接消失。这种策略在梦里无敌，但放到真实环境里完全没用。

为什么会这样？因为世界模型是对真实环境的近似，不可能100%准确。控制器又极其聪明，会想尽一切办法最大化奖励，自然就会去钻模型的空子。这种“在模拟里很强，现实里拉胯”的问题，是模型基强化学习的经典坑。

4.4 解法：提高温度，让梦更“混乱”

怎么防止控制器作弊？答案是给梦境增加随机性——调高MDN采样的温度参数τ。

τ越低，模型越确定，越容易被找到漏洞；τ越高，环境越随机，漏洞就越难卡。控制器想在混乱的环境里活下去，就必须学真本事，而不是投机取巧。

不同温度下的迁移效果：

表2 不同温度τ下的训练与迁移得分
出处：原文Table 2

温度τ	虚拟环境得分	真实环境得分
0.10	2086 ± 140	193 ± 58
0.50	2060 ± 277	196 ± 50
1.00	918 ± 546	1145 ± 690
1.15	732 ± 269	1092 ± 556
1.30	1145 ± 690	753 ± 139
随机策略	-	210 ± 108
Gym排行榜最优	-	820 ± 58

可以看到：

• 低温的时候，虚拟环境得分极高，但真实环境连随机策略都不如——全是作弊出来的假成绩；
• 温度升到1.0以上，虚拟环境得分降下来了，但真实环境得分暴涨；
• τ=1.15的时候真实环境效果最好；
• 温度再高，环境太乱，也学不到有效策略。

通俗解释：这就像训练特种兵，不能在太平稳的训练场练，要故意增加各种随机干扰、恶劣条件。虽然训练场里成绩没那么好看，但上了真实战场反而更能打。

---

五、迭代训练：从“新手梦”到“高手梦”

上面两个实验都是简单任务，用随机策略收集的数据就能训出够用的世界模型。但如果是更复杂的环境呢？比如开放世界游戏，随机乱跑根本逛不到多少地方，世界模型就学不全。

这时候就需要迭代训练流程，和人类的学习过程一模一样：

1. 初始化世界模型和控制器；
2. 用当前控制器去真实环境探索，收集新的观测数据；
3. 用新数据更新世界模型，让它更准确、覆盖更多场景；
4. 在更新后的世界模型里训练更好的控制器；
5. 回到第2步循环，直到完成任务。
   

图19 信息转化为记忆的迭代过程
出处：原文Figure 19
解读：每一轮探索都给世界模型带来新的信息，模型升级后又能支撑更好的策略，形成正向循环。

这里还可以引入好奇心机制：把世界模型的预测误差当作内在奖励。模型预测不准的地方，就是智能体没见过的“新地方”，鼓励智能体多去探索，这样就能自动收集更多有价值的数据，让世界模型越来越完善。

---

六、相关研究脉络

世界模型不是凭空蹦出来的，这条研究线已经发展了几十年：

1. 早在1990年，Schmidhuber（本文作者之一，LSTM之父）就提出了用RNN做世界模型的思路，让控制器和RNN世界模型配合；
2. 后来的PILCO用高斯过程做动力学模型，在低维控制任务上效果很好，但没法处理高维图像；
3. 再后来大家开始用自编码器先压缩图像，再在隐空间里学动力学；
4. Graves在2015年演示过用RNN“幻觉”生成Atari游戏画面；
5. 本文的贡献是把VAE+MDN-RNN+进化控制器这套组合拳打顺了，并且第一次完整实现了“全梦境训练+零样本迁移到真实环境”。

进化策略用来训练小控制器也不是新鲜事，它的好处是不需要反向传播、不需要可微奖励、天生容易并行，特别适合这种几百几千个参数的小模型优化。

---

七、讨论与局限

7.1 这个框架的优势

1. 训练效率高：世界模型用无监督学习，训得快；控制器参数少，进化算法优化快；
2. 节省算力：可以在梦境里训练策略，不用跑昂贵的真实游戏引擎或物理仿真；
3. 可扩展性强：世界模型可以用GPU批量加速，未来换更大的模型、更好的架构，控制器不用改。

7.2 现存的局限

1. VAE无监督训练，可能抓不住重点：它会平等地压缩所有像素，可能把无关紧要的花纹细节编码进去，却漏掉了任务关键信息。如果让VAE和奖励预测一起训练，可能会更聚焦，但也会失去通用性；
2. 世界模型容量有限：LSTM能记住的东西有限，复杂的开放世界肯定装不下，还会有灾难性遗忘问题；
3. 没有分层规划：现在的模型是逐帧预测，就像只会走一步看一步，不会像人一样做长期、抽象的规划。

7.3 未来方向

• 换更大容量的模型（比如Transformer、混合专家模型），或者加外部记忆模块；
• 引入分层规划能力，让控制器能调用世界模型的子程序做高层推理；
• 把控制器和世界模型合并成一个大网络，用行为回放避免遗忘，也就是One Big Net的思路。

---

八、核心代码实现

8.1 VAE核心实现（PyTorch）

import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvVAE(nn.Module):
    def __init__(self, z_dim=32):
        super().__init__()
        self.z_dim = z_dim
        
        # 编码器：4层卷积
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 4, stride=2, padding=1),  # 64x64 -> 32x32
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2, padding=1), # 32x32 -> 16x16
            nn.ReLU(),
            nn.Conv2d(64, 128, 4, stride=2, padding=1),# 16x16 -> 8x8
            nn.ReLU(),
            nn.Conv2d(128, 256, 4, stride=2, padding=1),# 8x8 -> 4x4
            nn.ReLU(),
        )
        # 输出均值和对数方差
        self.fc_mu = nn.Linear(256*4*4, z_dim)
        self.fc_logvar = nn.Linear(256*4*4, z_dim)
        
        # 解码器：4层转置卷积
        self.fc_decode = nn.Linear(z_dim, 256*4*4)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1), # 4x4 -> 8x8
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),  # 8x8 -> 16x16
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),  # 16x16 -> 32x32
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 4, stride=2, padding=1),   # 32x32 -> 64x64
            nn.Sigmoid() # 输出0-1的像素值
        )
    
    def encode(self, x):
        h = self.encoder(x).view(x.size(0), -1)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        # 重参数化技巧
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mu + eps * std
        return z, mu, logvar
    
    def decode(self, z):
        h = self.fc_decode(z).view(z.size(0), 256, 4, 4)
        return self.decoder(h)
    
    def forward(self, x):
        z, mu, logvar = self.encode(x)
        recon = self.decode(z)
        return recon, mu, logvar

# VAE损失函数
def vae_loss(recon_x, x, mu, logvar, beta=1.0):
    recon_loss = F.mse_loss(recon_x, x, reduction='sum') / x.size(0)
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / x.size(0)
    return recon_loss + beta * kl_loss, recon_loss, kl_loss

8.2 MDN-RNN核心实现

class MDNRNN(nn.Module):
    def __init__(self, z_dim=32, action_dim=3, hidden_dim=256, n_mixtures=5):
        super().__init__()
        self.z_dim = z_dim
        self.action_dim = action_dim
        self.hidden_dim = hidden_dim
        self.n_mixtures = n_mixtures
        
        # LSTM层
        self.lstm = nn.LSTM(
            input_size=z_dim + action_dim,
            hidden_size=hidden_dim,
            num_layers=1,
            batch_first=True
        )
        
        # MDN输出层：输出每个高斯分量的权重、均值、对数标准差
        self.mdn_head = nn.Linear(hidden_dim, n_mixtures * (1 + z_dim + z_dim))
        # 死亡预测头（Doom任务用）
        self.done_head = nn.Linear(hidden_dim, 1)
    
    def forward(self, x, hidden=None):
        # x: [batch, seq_len, z_dim + action_dim]
        output, hidden = self.lstm(x, hidden)
        # output: [batch, seq_len, hidden_dim]
        
        mdn_out = self.mdn_head(output)
        # 拆分参数
        pi, mu, log_sigma = torch.split(
            mdn_out,
            [self.n_mixtures, self.n_mixtures*self.z_dim, self.n_mixtures*self.z_dim],
            dim=-1
        )
        # 权重用softmax归一化
        pi = F.softmax(pi, dim=-1)
        # reshape成 [batch, seq, n_mixtures, z_dim]
        mu = mu.view(*mu.shape[:-1], self.n_mixtures, self.z_dim)
        log_sigma = log_sigma.view(*log_sigma.shape[:-1], self.n_mixtures, self.z_dim)
        sigma = torch.exp(log_sigma)
        
        # 死亡概率
        done_logit = self.done_head(output)
        
        return (pi, mu, sigma), done_logit, hidden
    
    def init_hidden(self, batch_size):
        h0 = torch.zeros(1, batch_size, self.hidden_dim)
        c0 = torch.zeros(1, batch_size, self.hidden_dim)
        return h0, c0

# MDN损失函数：计算真实z在混合高斯分布下的负对数似然
def mdn_loss(pi, mu, sigma, target):
    # target: [batch, seq, z_dim] -> 扩展成 [batch, seq, 1, z_dim]
    target = target.unsqueeze(-2)
    # 计算每个高斯分量的概率密度
    norm = torch.distributions.Normal(mu, sigma)
    log_prob = norm.log_prob(target).sum(dim=-1)  # [batch, seq, n_mixtures]
    # 加权求和
    log_pi = torch.log(pi + 1e-8)
    loss = -torch.logsumexp(log_pi + log_prob, dim=-1)
    return loss.mean()

8.3 线性控制器实现

class LinearController:
    def __init__(self, z_dim=32, hidden_dim=256, action_dim=3):
        self.input_dim = z_dim + hidden_dim
        self.action_dim = action_dim
        # 参数扁平化，方便CMA-ES优化
        self.W = np.zeros((action_dim, self.input_dim))
        self.b = np.zeros(action_dim)
    
    def get_params(self):
        return np.concatenate([self.W.flatten(), self.b])
    
    def set_params(self, params):
        w_size = self.W.size
        self.W = params[:w_size].reshape(self.action_dim, self.input_dim)
        self.b = params[w_size:]
    
    def forward(self, x):
        # x: [input_dim]
        a = self.W @ x + self.b
        return np.tanh(a)  # 限制输出在-1到1之间

---

写在最后

World Models这篇论文之所以经典，不在于提出了什么惊世骇俗的新算法，而在于它用极其简洁优雅的架构，把“世界模型”这个认知科学里的老概念，在深度强化学习里落地得特别漂亮。

它告诉我们：智能不一定非要端到端一个大网络莽上去。先让AI无监督地认识世界、形成内部模型，再在这个内部模型上学决策，既符合人类的认知规律，工程上也高效好用。

从这篇论文之后，基于世界模型的强化学习一路发展，到后来的Dreamer系列、基于Transformer的世界模型，本质上都是在这条路上越走越远——让AI先学会“做梦”，再学会做事。

需要我补充CMA-ES训练控制器的完整代码示例吗？