Huajishe Notes

2-Recurrent World Models

6 min read

Recurrent World Models

一句话定位

Recurrent World Models 是 World Models 的扩展版本，通过引入额外的循环架构改进动态建模质量，并优化了训练流程以支持更长的序列预测。

前置依赖

  • World Models (Ha & Schmidhuber 2018)：MDN-RNN 架构的基础
  • RNN/LSTM 基础：循环神经网络的基本原理
  • Sequence-to-Sequence 模型：编码器-解码器架构

核心思想

Recurrent World Models 的核心改进在于使用更强大的循环架构来处理世界动态建模。具体而言：

  1. 增强的循环架构：在原有 MDN-RNN 的基础上增加层归一化、残差连接等技术，改善梯度流动
  2. 改进的训练流程：引入 teacher forcing 和自由运行（free-running）相结合的策略
  3. 更精确的潜在变量建模：改进了 MDN 层的参数化方式

数学推导

增强的 RNN 动态模型

给定历史信息 ，RNN 维护隐藏状态 ：

预测分布：

训练目标

整体训练目标结合了重建损失和动态预测损失：

其中 为负对数似然：

训练细节

训练策略

Teacher Forcing + Free-Running 混合训练：

  1. Teacher Forcing 阶段：训练时使用真实的上一步输出作为下一步输入，加速收敛
  2. Free-Running 阶段：使用模型自己的预测作为下一步输入，训练模型从错误中恢复

模型架构

增强的 MDN-RNN 配置：

  • 层归一化（Layer Normalization）：应用于 RNN 的输入和输出
  • 残差连接（Residual Connections）：帮助深层网络的梯度流动
  • Dropout：防止过拟合

架构图

┌─────────────────────────────────────────────────────────────────┐
│                  Recurrent World Model 架构                     │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  输入层                                                          │
│  ┌─────────────────────────────────────┐                        │
│  │ [z_t, a_t] ──▶ LayerNorm ──▶ RNN    │                        │
│  └─────────────────────────────────────┘                        │
│         │                                    │                  │
│         ▼                                    ▼                  │
│  ┌─────────────┐                      ┌─────────────┐           │
│  │    h_t      │                      │   h_{t+1}   │           │
│  └──────┬──────┘                      └──────┬──────┘           │
│         │                                    │                  │
│         ▼                                    ▼                  │
│  ┌─────────────┐                      ┌─────────────┐           │
│  │    MDN      │                      │ 预测 z_{t+1} │           │
│  └─────────────┘                      └─────────────┘           │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

推理 / Rollout 过程

Latent Space Rollout

  1. 初始观察 通过 VAE 编码为
  2. 初始化 或学习到的初始状态
  3. 每个时间步：
    • 控制器输出
    • MDN-RNN 采样预测
    • 更新隐藏状态

长序列稳定性

Recurrent World Models 改进了长序列预测的稳定性：

  • 层归一化减少了内部协变量偏移
  • 残差连接缓解了梯度消失问题
  • 混合训练策略让模型学会从错误中恢复

优点与局限

优点

  1. 更好的梯度流动：层归一化和残差连接改善了深层 RNN 的训练
  2. 更长的记忆：改进了对长期依赖的建模能力
  3. 更稳定的 rollout：模型在长序列预测中更难出现坍缩

局限

  1. 仍然依赖 MDN 进行多模态建模，表达能力有限
  2. 缺乏对不确定性的显式处理（没有 variational inference）
  3. 控制器仍然使用 CMA-ES，难以扩展到高维动作空间

与前后内容的衔接

  • 前承：World Models 的基本框架
  • 后续：SimPLe 引入 variational approach 来处理 uncertainty
  • 相关：Dreamer 系列的 RSSM 模型继承了循环架构的思想并加入 variational

可复现实现要点

PyTorch 伪代码

class RecurrentWorldModel(nn.Module):
    def __init__(self, latent_dim, action_dim, hidden_dim=512, n_layers=2):
        self.encoder = VAE_Encoder(latent_dim)
        self.decoder = VAE_Decoder(latent_dim)
 
        # 增强的 RNN
        self.rnn = nn.LSTM(
            latent_dim + action_dim,
            hidden_dim,
            num_layers=n_layers,
            batch_first=True
        )
        self.layer_norm = nn.LayerNorm(hidden_dim)
        self.mdn = nn.Linear(hidden_dim, n_mix * (2 * latent_dim + n_mix + 1))
 
    def forward(self, z, a, h=None):
        # Layer norm before RNN
        x = torch.cat([z, a], dim=-1)
        x = self.layer_norm(x)
 
        if h is None:
            out, h_next = self.rnn(x)
        else:
            out, h_next = self.rnn(x, h)
 
        out = self.layer_norm(out)
        return self.mdn(out), h_next

训练流程

  1. 预训练 VAE（重建 + KL）
  2. 生成 latent 轨迹
  3. 使用 teacher forcing 预训练 RNN
  4. 切换到 free-running 微调
  5. CMA-ES 优化控制器

章节摘要

Recurrent World Models 是对原始 World Models 的改进版本，通过增强的循环架构（层归一化、残差连接）改善了动态建模质量和长序列预测的稳定性。虽然仍使用 MDN 而非 variational 方法，但其架构改进为后续 Dreamer 系列等更先进的工作奠定了基础。

关键词

  • Recurrent World Models
  • Layer Normalization
  • Residual Connections
  • Teacher Forcing
  • Free-Running
  • Mixture Density Network
  • Latent Space Dynamics

Graph View