SimPLe (Sample-efficient Model-based Policy Learning with World Models)

一句话定位

SimPLe（Kaiser et al., Google Brain, 2019）是首个将 model-based RL 应用于 Atari 全系列游戏的工作，提出了将世界模型解耦为 deterministic transition model 和 stochastic latent model 的混合架构，显著提升了样本效率。

前置依赖

World Models (Ha & Schmidhuber 2018)：latent space planning 的基本思想
VAE (Variational Autoencoder)：Kingma & Welling 2013/2014
SS-TRM (Stochastic Stationary Transition Regularized Model)：相关变分建模技术
Dreamer 早期思想：latent space 中进行策略学习的理念
Cross Entropy Method (CEM)：用于 rollout 评估的进化方法

核心思想

SimPLe 的核心思想是在像素级环境中将 model-based RL 与 world models 结合：

变分观察模型：使用 VAE 将高维图像压缩到低维潜在空间
混合动态模型：结合 deterministic 和 stochastic 成分预测下一帧
** Imagination-driven policy learning**：在模型内部进行 rollout 来训练策略
Cross Entropy Method (CEM) 用于策略优化

SimPLe 与 World Models 的关键区别：

引入变分推断来处理 uncertainty
使用更强大的神经网络架构
应用 CEM 而非 CMA-ES
直接在 latent space 中优化策略

模型结构图

┌─────────────────────────────────────────────────────────────────┐
│                       SimPLe 架构                               │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  真实环境                                                        │
│  ┌─────────┐    ┌─────────┐    ┌─────────┐    ┌─────────┐       │
│  │   o_t   │───▶│  Env    │───▶│  o_{t+1} │───▶│  ...    │       │
│  │(原始像素)    │         │    │(原始像素)    │    │         │       │
│  └────┬────┘    └────━────┘    └────━────┘    └────┬───┘       │
│       │                                           │             │
│       ▼                                           ▼             │
│  ┌─────────┐                              ┌─────────┐           │
│  │   VAE   │                              │ Policy  │           │
│  │         │                              │ 优化    │           │
│  └────┬────┘                              │ (CEM)   │           │
│       │                                        │                │
│       ▼                                        │                │
│  ┌─────────┐    ┌─────────────────┐          │                │
│  │   z_t   │───▶│ Transition Model │◀── a_t   │                │
│  │(潜在向量)    │ (混合动态模型)    │          │                │
│  └────┬────┘    └──────┬──────────┘          │                │
│       │                │                      │                │
│       │                ▼                      │                │
│       │         ┌─────────────────┐          │                │
│       │         │ 预测 z_{t+1},    │          │                │
│       │         │ o_{t+1} 的分布   │          │                │
│       │         └─────────────────┘          │                │
│       │                                        │                │
│       └────────────────────────────────────────┘                │
│              (在 latent space 中进行 Imagination Rollout)         │
└─────────────────────────────────────────────────────────────────┘

观察模型（Observation Model / VAE）

将原始图像编码为潜在向量：

过渡模型（Transition Model）

混合架构：

其中 transition model 结合了：

Deterministic 成分：
Stochastic 成分：

重建模型（Reconstruction Model）

从潜在向量重建图像：

数学推导

整体训练目标

SimPLe 的训练目标是最大化 Evidence Lower Bound (ELBO)：

简化表示为三个组件：

变分损失（Reconstruction + KL）

观察模型的 VAE 损失：

动态损失（Transition Loss）

过渡模型的负对数似然：

其中

Imagination Rollout 目标

在 latent space 中进行步 rollout：

其中是奖励函数（通过潜在向量预测）。

Cross Entropy Method (CEM)

CEM 用于在 latent space 中优化动作序列：

初始化动作序列分布
对每个动作序列进行 rollout（在模型内部）
评估累积奖励
选择 top-N% 的序列
用这些序列更新分布参数
重复直到收敛

训练细节

训练流程

阶段一：数据收集

使用随机策略或部分训练的策略在环境中交互
收集初始数据集
通常需要约 100K-200K 环境步骤

阶段二：训练世界模型

训练 VAE 观察模型
训练 transition model（给定 latent 轨迹）
联合训练所有组件

阶段三：Imagination Rollout

使用训练好的世界模型进行 latent space rollout
使用 CEM 优化策略
在真实环境中评估策略
如果性能提升，将成功轨迹添加到数据集
重复迭代

Atari 实验配置

环境：Atari 全套 50+ 游戏
观察：84×84×4 堆叠帧（跳帧后保留）
动作空间：离散，4-18 个动作不等
VAE 潜在维度：或
Transition Model：MLP 或 LSTM
Rollout 长度：10-20 步（受 compounding error 限制）
训练数据：约 100K 环境步骤（vs Model-free 需要 10M+）

奖励函数

在 latent space 中预测奖励：

或直接从潜在向量回归：

推理 / Rollout / Planning 过程

Imagination Rollout

Algorithm: Imagination Rollout in SimPLe

Input: Initial observation o_0, trained world model (VAE + Transition), policy π, horizon H
Output: Optimized action sequence

1. z_0 ← VAE.encode(o_0)
2. for h in 1 to H:
3.     a_h ← sample from current policy (or CEM optimized)
4.     z_{h+1} ← Transition_model.predict(z_h, a_h)
5.     r_h ← reward_model.predict(z_h, a_h)
6.     accumulate rewards
7. end for
8. return action sequence with highest cumulative reward

CEM 策略优化

从先验分布采样 N 条候选动作序列
对每条序列在模型中进行 rollout
选择 top 10% 的序列
用这些序列拟合新的高斯分布
重复 5-10 次迭代

优点与局限

优点

样本效率极高：在 Atari 上仅需 ~100K 步骤（vs model-free 的 10M+）
变分推断：显式建模观察和动态的不确定性
通用架构：适用于多种 Atari 游戏
Latent space planning：证明了在压缩空间中进行规划的可行性

局限

Compounding Error：长 rollout 误差累积严重，限制了在长horizon任务上的表现
Pixel-level 重建：VAE 难以保留所有视觉细节
混合模型复杂度：deterministic + stochastic 的组合调参困难
CEM 效率：CEM 在高维动作空间效率较低
无梯度策略优化：CEM 无法利用梯度信息

与 World Models 的对比

特性	World Models	SimPLe
动态模型	MDN-RNN	混合（det + stochastic）
不确定性	隐式（MDN）	显式（variational）
策略优化	CMA-ES	CEM
环境	Car Racing, Doom	Atari 全套
样本效率	高	更高

与前后内容的衔接

前承：World Models 的 latent space planning 框架
后续：
- DreamerV1：继承并改进了 latent dynamics model + policy gradient
- MuZero：将 model-based 与 tree search 结合
- Video Prediction Models：改进重建质量
相关：SVG、Dreamer 等继承了 SimPLe 的变分框架

可复现实现要点

核心超参数

参数	典型值
VAE 潜在维度	256 或 512
Transition hidden dim	512
Rollout 长度	10-15
CEM 迭代次数	5-10
CEM 群体大小	100-1000
数据集大小	100K-200K steps

PyTorch 伪代码

class SimPLeWorldModel(nn.Module):
    def __init__(self, obs_dim, action_dim, latent_dim=256, hidden_dim=512):
        # VAE: encoder + decoder
        self.encoder = ConvEncoder(latent_dim * 2)  # mu, logvar
        self.decoder = DeconvDecoder(latent_dim)
 
        # Transition model: z_t, a_t -> z_{t+1}
        self.transition = nn.Sequential(
            nn.Linear(latent_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim * 2)  # mean, logvar
        )
 
        # Reward model
        self.reward_head = nn.Linear(latent_dim + action_dim, 1)
 
    def encode(self, obs):
        q = self.encoder(obs)
        mu, logvar = q.chunk(2, dim=-1)
        z = reparameterize(mu, logvar)
        return z, mu, logvar
 
    def decode(self, z):
        return self.decoder(z)
 
    def transition_step(self, z, a):
        """One step latent dynamics prediction"""
        x = torch.cat([z, a], dim=-1)
        pred = self.transition(x)
        mu, logvar = pred.chunk(2, dim=-1)
        z_next = reparameterize(mu, logvar)
        return z_next, mu, logvar
 
    def imagine_rollout(self, z0, action_seq, horizon):
        """Perform imagination rollout in latent space"""
        z_preds = [z0]
        for t in range(horizon):
            z_t = z_preds[-1]
            a_t = action_seq[t]
            z_next, _, _ = self.transition_step(z_t, a_t)
            r_t = self.reward_head(torch.cat([z_t, a_t], dim=-1))
            z_preds.append(z_next)
        return z_preds
 
def cem_policy_optimization(world_model, z0, horizon, n_iters=10, n_samples=100):
    """Cross Entropy Method for policy optimization in latent space"""
    action_dim = ...  # from env
    action_seq_dim = action_dim * horizon
 
    # Initial distribution
    mu = torch.zeros(action_seq_dim)
    sigma = torch.ones(action_seq_dim)
 
    for _ in range(n_iters):
        # Sample candidate action sequences
        action_seqs = torch.normal(mu, sigma).reshape(n_samples, horizon, action_dim)
 
        # Evaluate each sequence
        rewards = []
        for seq in action_seqs:
            z_preds = world_model.imagine_rollout(z0, seq, horizon)
            total_reward = sum(z_preds)  # simplified
            rewards.append(total_reward)
 
        # Select top 10%
        top_indices = torch.argsort(torch.stack(rewards))[-n_samples//10:]
        top_seqs = action_seqs[top_indices]
 
        # Update distribution
        mu = top_seqs.mean(dim=0).flatten()
        sigma = top_seqs.std(dim=0).flatten()
 
    return mu.reshape(horizon, action_dim)

训练循环

def train_simple(env, world_model, policy, n_iterations=10, n_steps_per_iter=1000):
    dataset = []
 
    for iteration in range(n_iterations):
        # Collect data with current policy
        obs = env.reset()
        for step in range(n_steps_per_iter):
            action = policy(obs)
            next_obs, reward, done, _ = env.step(action)
            dataset.append((obs, action, reward, next_obs))
            obs = next_obs
            if done:
                obs = env.reset()
 
        # Train world model
        world_model.train_on_batch(dataset)
 
        # Imagination rollout with policy optimization
        for _ in range(10):
            obs = env.reset()
            z0 = world_model.encode(obs)
            optimized_actions = cem_policy_optimization(world_model, z0, horizon=10)
            # Execute in real env and potentially add to dataset

章节摘要

SimPLe（Kaiser et al. 2019）是将 world model 与 model-based RL 结合的标志性工作，首次在 Atari 全套游戏中展示了 model-based 方法的潜力。其核心贡献是提出了混合动态模型架构（deterministic + stochastic），结合变分推断来处理不确定性，并使用 CEM 在 latent space 中进行策略优化。SimPLe 将样本效率提升至约 100K 步骤（相比 model-free 的 10M+），但仍受困于 compounding error 和缺乏梯度策略优化的问题，这些问题在后续 Dreamer 系列中得到解决。

关键词

SimPLe
World Models for Atari
Model-Based RL
Variational Inference
Imagination Rollout
Cross Entropy Method
Sample-Efficient RL
Latent Space Planning
Stochastic Transition Model
Hybrid Dynamics Model

Huajishe Notes

Explorer

3-SimPLe (World Models for Model-Based RL)