World Models (Ha & Schmidhuber, 2018)
一句话定位
World Models 首次提出将世界模型分解为 VAE（观察压缩）+ MDN-RNN（动态预测）的双组件架构，并支持在纯 latent space 中进行 long-horizon rollout 与策略优化，是世界模型领域的开创性工作。
前置依赖
- VAE (Variational Autoencoder)：Kingma & Welling 2013/2014，关于变分自编码器的基本原理、重参数化技巧、ELBO 目标函数
- MDN (Mixture Density Network)：Bishop 1994，混合密度网络用于建模多模态输出的概率分布
- RNN/LSTM：长短期记忆网络用于序列建模
- CMA-ES (Covariance Matrix Adaptation Evolution Strategy)：Hansen 1996，用于无梯度策略优化（Hansen & Ostermeier 2001）
核心思想
World Models 的核心思想是将”世界”解耦为两个独立学习的组件：
- 观察模型（Vision Model）：使用 VAE 将高维图像压缩为低维潜在向量
- 动态模型（Dynamics Model）：使用 MDN-RNN 预测下一个潜在向量的分布
，其中 是 RNN 的隐藏状态
这种设计使得 agent 可以在完全离线的 latent space 中进行rollout 和规划，大幅降低了对真实环境的交互成本。
模型结构图
┌─────────────────────────────────────────────────────────────────┐
│ World Model 整体架构 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 真实环境 │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ s_t │───▶│ Env │───▶│ s_{t+1} │───▶│ ... │ │
│ │(原始图像) │ │ │(原始图像) │ │ │ │
│ └────┬────┘ └────━────┘ └────━────┘ └─────┬───┘ │
│ │ │ │
│ ▼ ▼ │
│ ┌─────────┐ ┌─────────┐ │
│ │ VAE │ │ Policy │ │
│ │ (Encoder) │ Controller │ │
│ └────┬────┘ └────┬────┘ │
│ │ │ │
│ ▼ │ │
│ ┌─────────┐ ┌─────────────┐ │ │
│ │ z_t │───▶│ MDN-RNN │◀── action a_t │ │
│ │(潜在向量) │ (动态模型) │ │ │
│ └────┬────┘ └──────┬──────┘ │ │
│ │ │ │ │
│ │ ▼ │ │
│ │ ┌─────────────┐ │ │
│ │ │ 预测 z_{t+1} │ │ │
│ │ └──────━──────┘ │ │
│ │ │ │
│ └────────────────────────────────────────┘ │
│ (在 latent space 中进行 rollout) │
└─────────────────────────────────────────────────────────────────┘
Vision Model (VAE)
- Encoder：将原始图像
编码为潜在向量 的均值 和方差 - Decoder：从潜在向量重建图像
- 潜在空间维度：通常
，原论文使用 或
Dynamics Model (MDN-RNN)
- 输入：潜在向量
+ 动作 + 隐藏状态 - 输出：下一个潜在向量
的概率分布（混合高斯分布） - 使用 MDN 输出混合系数
、均值 、方差 ：
数学推导
VAE 目标函数
VAE 的目标函数为最大化观测数据的边际对数似然的下界（ELBO）：
其中：
：编码器输出的近似后验 ：解码器给出的生成分布 ：先验，通常为单位高斯
重参数化技巧：
MDN-RNN 目标函数
给定序列
其中
MDN 输出的参数通过神经网络计算：
训练流程
阶段一：收集随机轨迹数据
使用随机策略或简单策略在环境中收集大量轨迹数据
阶段二：训练 VAE
使用收集的数据训练 VAE，最小化重建误差和 KL 散度：
阶段三：训练 MDN-RNN
- 使用训练好的 VAE Encoder 将所有图像序列转换为潜在向量序列
- 使用 MDN-RNN 训练动态模型，最小性负对数似然
阶段四：训练控制器（Policy）
使用 CMA-ES 进化算法优化线性控制器参数
目标函数为最大化累积奖励：
推理 / Rollout / Planning 过程
Rollout 流程（在 latent space 中）：
- 初始观察
通过 VAE Encoder 编码为 - 初始化 RNN 隐藏状态
- 对于每个时间步
： a. 控制器基于当前 输出动作 b. MDN-RNN 基于 预测下一潜在分布 c. 从该分布采样得到 和更新的隐藏状态 d. 在 latent space 中”幻想”下一帧，而非真实执行（规划模式） - 也可以在真实环境中执行动作（控制模式）
Planning via Imagination：
控制器可以在部署前在模型内部进行”mental simulation”：通过在 latent space 中进行多步 rollout 评估不同动作序列的预期回报，选择最优动作执行。
训练细节
Car Racing 实验
- 环境：Car Racing V0（OpenAI Gym）
- 观察：96×96×3 RGB 图像
- 动作空间：连续三维
- VAE 潜在维度：
- MDN-RNN 结构：LSTM + MDN 层，隐藏状态维度
- MDN 混合组分数量：
- 训练数据：使用随机策略收集约 10,000 条成功轨迹
- 控制器：线性控制器
，偏置 - 策略优化：CMA-ES，群体大小
- 奖励：每步存活 +1，每帧得分由环境提供
Doom Takearound 实验
- 环境：VizDoom Doom Takearound
- VAE：类似 car racing 配置
- MDN-RNN：相同配置
- 策略优化：相同 CMA-ES 配置
控制器学习算法
CMA-ES 是一种无需梯度的进化策略：
- 维护协方差矩阵
和步长 - 从多元高斯分布采样候选解
- 根据奖励排序并更新
和 - 适合低维参数的连续优化
优点与局限
优点
- 开创性架构：首次将 VAE + RNN 的组合应用于世界建模，为后续研究提供了基础框架
- Latent space planning：证明了在压缩的 latent space 中进行规划和控制是完全可行的
- 样本效率：通过离线训练世界模型，减少了对真实环境的交互需求
- 通用性：架构适用于离散和连续动作空间，以及各种视觉观察环境
局限
-
MDN-RNN 的局限性：
- 混合高斯模型难以精确捕捉复杂的多模态动态
- 训练不稳定，特别是在长序列上
- 自回归采样会导致误差累积（compounding error）
-
Latent space 的信息损失：
- VAE 压缩可能丢失关键视觉细节
- 潜在空间的不确定性难以正确表征
-
随机策略优化：
- CMA-ES 难以扩展到高维动作空间
- 线性控制器表达能力有限
-
长 horizon 规划困难：
- 长序列 rollout 误差累积严重
- 缺乏对不确定性的显式建模
与前后内容的衔接
- 前承：dreamer 系列（后续工作）在 latent space 中使用更强大的动态模型（RSSM）替代 MDN-RNN
- 后续：DreamerV1/V2/V3 继承并改进了 World Models 的 latent planning 框架
- 相关：SimPLe 引入了 variational approach 来处理 uncertainty
可复现实现要点
核心超参数
| 组件 | 参数 | 典型值 |
|---|---|---|
| VAE | 潜在维度 | 32 |
| VAE | Encoder 结构 | Conv(32,64,128) + FC |
| VAE | Decoder 结构 | FC + Deconv(128,64,32) |
| MDN-RNN | 隐藏维度 | 256 |
| MDN-RNN | 混合数 | 5 |
| MDN-RNN | RNN 类型 | LSTM |
| 控制器 | 策略优化器 | CMA-ES |
| 控制器 | 群体大小 | 64 |
PyTorch 伪代码
class VAE(nn.Module):
def __init__(self, latent_dim=32):
# Encoder: Conv layers -> FC -> (mu, logvar)
# Decoder: FC -> Deconv layers
self.latent_dim = latent_dim
def encode(self, x):
# Return z ~ N(mu, sigma)
pass
def decode(self, z):
# Return reconstructed x
pass
def forward(self, x):
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
return self.decode(z), mu, logvar
class MDNRNN(nn.Module):
def __init__(self, latent_dim, action_dim, hidden_dim=256, n混合=5):
self.rnn = nn.LSTM(latent_dim + action_dim, hidden_dim)
self.mdn = nn.Linear(hidden_dim, n混合 * (2 * latent_dim + n混合 + 1))
def forward(self, z, a, h):
# z: (batch, latent_dim)
# a: (batch, action_dim)
# h: (batch, hidden_dim)
x = torch.cat([z, a], dim=-1)
out, h_next = self.rnn(x, h)
return self.mdn(out), h_next
def sample(self, z, a, h):
# Sample from MDN mixture
pi, mu, sigma = self.get_parameters(z, a, h)
# Mix up a mixture component
k = torch.multinomial(pi, 1)
# Sample from selected Gaussian
z_next = torch.normal(mu[k], sigma[k])
return z_next训练流程
- 收集随机轨迹数据
- 训练 VAE（使用重建 loss + KL loss）
- 生成 latent 轨迹数据
- 训练 MDN-RNN（最大化对数似然）
- 使用 CMA-ES 优化控制器参数
章节摘要
World Models（Ha & Schmidhuber 2018）提出了世界模型的基本框架：使用 VAE 压缩高维观察为低维潜在向量，使用 MDN-RNN 建模潜在空间中的动态变化。控制器通过 CMA-ES 在 latent space 中进行策略优化。该工作首次展示了在压缩 latent space 中进行 long-horizon rollout 和 planning 的可行性，但也暴露了 MDN-RNN 的局限性：难以处理多模态动态、不确定性建模不足、长序列误差累积严重。
关键词
- World Models
- VAE
- MDN-RNN
- Latent Space Planning
- Mixture Density Network
- CMA-ES
- Trajectory Optimization
- Model-Based Reinforcement Learning
- Dreamer