Model-Based Reinforcement Learning for Atari (SimPLe)

01 动机Motivation

人类玩家可以在几分钟内学会 Atari 游戏，而最优秀的 model-free RL 算法需要数千万乃至数亿次交互才能达到相近水平——相当于数周的实时训练。这种巨大的样本效率差距是 model-based RL 研究的核心动机。Human players can learn an Atari game within a few minutes, whereas the best model-free RL algorithms need tens of millions or even hundreds of millions of interactions to reach a comparable level — equivalent to weeks of real-time training. This enormous sample efficiency gap is the core motivation of model-based RL research.

"So far, there has been no clear demonstration of successful planning with a learned model in the ALE." — Machado et al. (2018) 对 Atari 基准上 model-based 控制的挑战性评述"So far, there has been no clear demonstration of successful planning with a learned model in the ALE." — Machado et al. (2018), a remark on how challenging model-based control is on the Atari benchmark

论文的核心假设：人类之所以能快速学会游戏，部分原因在于拥有对物理过程的直觉理解，能够预测动作的结果。SimPLe 通过学习视频预测模型来实现类似的能力，从而大幅降低与真实环境的交互次数。The paper's central hypothesis: humans can learn games quickly in part because they possess an intuitive understanding of physical processes and can predict the outcomes of their actions. SimPLe attains a similar ability by learning a video prediction model, thereby drastically reducing the number of interactions with the real environment.

SimPLe 主循环 — **Figure 1（论文原图）**：SimPLe 的主循环。① agent 按当前 policy 与真实环境交互，收集观测。② 用收集的数据更新 world model（对观测帧做自监督预测，对奖励做监督训练）。③ agent 在 world model 内部用 RL 更新 policy。新 policy 再次被送回真实环境评估和采样数据（回到①），循环迭代 15 次。**Figure 1 (from the paper)**: the main loop of SimPLe. ① The agent interacts with the real environment under the current policy and collects observations. ② The collected data are used to update the world model (self-supervised prediction of observed frames, supervised training for rewards). ③ The agent updates its policy with RL inside the world model. The new policy is sent back to the real environment for evaluation and data sampling (back to ①), iterating 15 times.

100K真实环境交互上限Cap on real-environment interactions

>10×Freeway 游戏上的样本效率提升（对比 Rainbow）Sample efficiency gain on Freeway (vs. Rainbow)

26评测的 Atari 游戏数Atari games evaluated

15SimPLe 迭代训练轮次SimPLe training iterations

02 方法Method

SimPLe 的核心是将 world model 训练与 policy 训练交替进行（类似 Dyna-Q 框架）。world model 是一个视频预测网络，接收 4 帧 stacked 输入和 action，预测下一帧画面与奖励。policy 则在 world model 内部通过 PPO 训练，避免大量真实环境交互。The core of SimPLe is to alternate world model training with policy training (similar to the Dyna-Q framework). The world model is a video prediction network that takes 4 stacked frames and an action as input and predicts the next frame and the reward. The policy is trained inside the world model with PPO, avoiding a large number of real-environment interactions.

世界模型架构 — **Figure 2（论文原图）**：提出的 stochastic model with discrete latent 的架构图。输入为 4 帧 stacked frames 及 agent 的 action，输出为下一帧预测及奖励预测。网络主体为 skip-connected 卷积 encoder-decoder；action 通过 embedding 与 decoder 各层输出相乘来条件化生成。推理网络（inference network）在训练时估计后验分布并将 latent 离散化为 bits；推理时由辅助 LSTM 自回归预测这些 bits（取代从先验采样）。整个模型约 74M 参数。**Figure 2 (from the paper)**: architecture of the proposed stochastic model with discrete latent. The input consists of 4 stacked frames and the agent's action; the outputs are the next-frame prediction and the reward prediction. The backbone is a skip-connected convolutional encoder-decoder; the action is embedded and multiplied with the outputs of each decoder layer to condition generation. The inference network estimates the posterior during training and discretizes the latent into bits; at inference time an auxiliary LSTM predicts these bits autoregressively (replacing sampling from the prior). The whole model has about 74M parameters.

World Model：随机离散 latent 模型World Model: stochastic discrete latent model

论文提出三种 world model 架构，最优为 stochastic discrete (SD) model：The paper proposes three world model architectures, the best of which is the stochastic discrete (SD) model:

Deterministic model：基础 skip-connected conv encoder-decoder，action embedding 通过乘法注入 decoder 各层。Deterministic model: a basic skip-connected conv encoder-decoder, with the action embedding injected multiplicatively into each decoder layer.
Stochastic VAE model：引入 variational autoencoder 建模环境随机性，但 KL 权重调参困难，推理时 latent 分布外问题严重。Stochastic VAE model: introduces a variational autoencoder to model environment stochasticity, but the KL weight is hard to tune and the out-of-distribution latent problem at inference time is severe.
Stochastic Discrete (SD) model（最优）：将 latent 值离散化为 bits（0/1），训练辅助 LSTM 自回归预测这些 bits，规避 VAE 的 OOD latent 问题。为提升鲁棒性，在离散化前加 uniform noise，离散化后加 dropout。Stochastic Discrete (SD) model (the best): discretizes the latent values into bits (0/1) and trains an auxiliary LSTM to predict these bits autoregressively, circumventing the OOD latent problem of the VAE. For robustness, uniform noise is added before discretization and dropout after it.

Loss 设计：视觉输出采用 per-pixel softmax（256 色空间）或 L2，并使用 clipped loss max(Loss, C)，其中对 L2 取 C=10，对 softmax 取 C=0.03（意味着置信度 >97% 时不产生梯度，避免大面积背景主导优化）。Loss design: the visual output uses a per-pixel softmax (256-color space) or L2, together with a clipped loss max(Loss, C), where C=10 for L2 and C=0.03 for softmax (meaning no gradient is produced once confidence is >97%, preventing large background regions from dominating the optimization).

Scheduled Sampling：为降低 compounding error，训练时按线性增长的概率将输入帧替换为模型自身的预测帧，到第一次迭代中期将混合率提升至 100%。Scheduled Sampling: to reduce compounding error, input frames are replaced with the model's own predicted frames with a linearly increasing probability during training, raising the mixing rate to 100% by the middle of the first iteration.

Policy Training：PPO in World ModelPolicy Training: PPO in World Model

采用 PPO（Proximal Policy Optimization，γ=0.95）在 world model 内训练 policy：The policy is trained inside the world model with PPO (Proximal Policy Optimization, γ=0.95):

每 N=50 步从真实数据 buffer D 中均匀采样起始状态重启模拟环境，避免长 rollout 带来的 compounding error。Every N=50 steps the simulated environment is restarted from a starting state sampled uniformly from the real-data buffer D, avoiding the compounding error caused by long rollouts.
rollout 末尾加入 value function 估值以弥补短 rollout 的远见不足。A value function estimate is added at the end of each rollout to compensate for the limited foresight of short rollouts.
每次 PPO epoch 使用 16 个并行 agent，收集 25/50/100 步（默认 50）。Each PPO epoch uses 16 parallel agents collecting 25/50/100 steps (50 by default).
SimPLe 在模拟环境内共执行约 15.2M 次交互，以 100K 真实步为代价换取大量模拟经验。SimPLe performs about 15.2M interactions in total inside the simulated environment, trading 100K real steps for a large amount of simulated experience.

Algorithm 1（SimPLe 伪代码）：Algorithm 1 (SimPLe pseudocode):

初始化 policy π 和 world model 参数 θ，初始化空数据集 D。Initialize the policy π and the world model parameters θ, and initialize an empty dataset D.
循环（15 轮）：①用 π 从真实环境 env 收集数据加入 D；②有监督地训练 world model：θ ← TRAIN_SUPERVISED(env', D)；③在 world model 内更新 policy：π ← TRAIN_RL(π, env')。Loop (15 iterations): ① collect data from the real environment env with π and add them to D; ② train the world model in a supervised manner: θ ← TRAIN_SUPERVISED(env', D); ③ update the policy inside the world model: π ← TRAIN_RL(π, env').

03 实验Experiments

在 Atari Learning Environment (ALE) 的 26 款游戏上评测，训练预算限定为 100K 次真实环境交互（= 400K 帧 = 约 114 分钟 @ 60 FPS）。基线方法为高度调优的 Rainbow（Q-learning SOTA）和 PPO（model-free policy gradient）。评测采用 5 次运行取平均，使用 softmax(logits(π)/T)（T=0.5）的确定性 policy 评测。Evaluation is carried out on 26 games of the Atari Learning Environment (ALE), with the training budget limited to 100K real-environment interactions (= 400K frames = about 114 minutes @ 60 FPS). The baselines are a heavily tuned Rainbow (Q-learning SOTA) and PPO (model-free policy gradient). Results are averaged over 5 runs, evaluated with a deterministic policy given by softmax(logits(π)/T) (T=0.5).

与 Rainbow 和 PPO 的比较 — **Figure 3（论文原图）**：与 Rainbow 和 PPO 的样本效率比较。每根柱子表示 Rainbow（左）或 PPO（右）需要多少次与环境的交互才能达到 SimPLe 在 100K 步时的得分。红线标注 SimPLe 使用的 100K 阈值——柱高超过红线意味着该 model-free 方法需要更多样本。SimPLe 在几乎所有游戏上都更高效，Freeway 上超过 10 倍。**Figure 3 (from the paper)**: sample efficiency comparison against Rainbow and PPO. Each bar shows how many environment interactions Rainbow (left) or PPO (right) needs to reach the score SimPLe attains at 100K steps. The red line marks the 100K threshold used by SimPLe — a bar above the line means that model-free method needs more samples. SimPLe is more efficient on almost every game, by more than 10 times on Freeway.

主要结果（选取部分游戏的 SD 模型均值得分，来自论文 Table 2）Main results (mean SD-model scores on a selection of games, from Table 2 of the paper)

游戏Game	Ours, SD（均值）Ours, SD (mean)	Ours, det. recurrent（均值）Ours, det. recurrent (mean)	Ours, deterministic（均值）Ours, deterministic (mean)
Freeway	20.3	23.7	5.9
Pong	12.8	-11.6	-17.4
Boxing	9.1	-3.1	-9.3
Breakout	12.7	10.2	6.1
CrazyClimber	39827.8	54700.3	19380.0
KungFuMaster	17257.2	4086.6	10340.9
RoadRunner	5169.4	1228.8	5724.4
Seaquest	370.9	289.6	419.5

注：数值来自论文 Table 2，为 5 次实验均值。SD = stochastic discrete（提出的最优模型）。Note: the numbers come from Table 2 of the paper and are means over 5 runs. SD = stochastic discrete (the best proposed model).

整体比较（论文摘要引述）Overall comparison (quoted from the paper's abstract)

"In most games SimPLe outperforms state-of-the-art model-free algorithms, in some games by over an order of magnitude."

更新至 v5 版本的结论：经 van Hasselt et al. (2019) 和 Kielak (2020) 改进后的 Rainbow 在低数据量下与 SimPLe 持平——两种 model-free 方法各在 13 款游戏上胜出，SimPLe 在另外 13 款游戏上胜出（共 26 款）。Conclusion as updated in version v5: after the improvements of van Hasselt et al. (2019) and Kielak (2020), Rainbow is on par with SimPLe in the low-data regime — each of the two model-free methods wins on 13 games, and SimPLe wins on the other 13 games (26 in total).

样本量 vs 得分分数图 — **Figure 4 & 5（论文原图）**：上方：以公式 (SimPLe_score@100K - random_score)/(baseline_score - random_score) 计算的得分比例，从左到右分别对比 Rainbow@100K、Rainbow@200K、PPO@100K、PPO@200K——SimPLe 在对手获得双倍交互时仍胜出。下方：随样本量增长，PPO 追上 SimPLe 所需帧数的变化曲线（a）及 SimPLe 预热后继续做 PPO 微调的效果（b）。**Figure 4 & 5 (from the paper)**. Top: score fractions computed with the formula (SimPLe_score@100K - random_score)/(baseline_score - random_score), compared from left to right against Rainbow@100K, Rainbow@200K, PPO@100K and PPO@200K — SimPLe still wins when the opponent is given twice the interactions. Bottom: as the sample budget grows, the curve of how many frames PPO needs to catch up with SimPLe (a), and the effect of continuing with PPO fine-tuning after a SimPLe warm-up (b).

消融实验Ablation studies

论文对 7 种配置各跑 5 次，汇总于 Table 1：The paper runs each of 7 configurations 5 times, summarized in Table 1:

模型配置Model configuration	最优（26 游戏中取最高的数量）Best (number of the 26 games with the highest score)	至少达中位数At least median
deterministic	0	7
det. recurrent	3	13
SD（默认）SD (default)	10	21
SD γ=0.9	1	14
SD 100 steps	0	14
SD 25 steps	4	19

模型架构：SD 模型全面优于 deterministic 和 det. recurrent 版本，是"the most important single factor"（inferred from ablation Table 1）。Model architecture: the SD model outperforms the deterministic and det. recurrent variants across the board, and is "the most important single factor" (inferred from ablation Table 1).
Rollout 长度：N=50（默认）与 N=25 性能相近，N=100 略差（compounding error 更严重）。Rollout length: N=50 (default) performs similarly to N=25, while N=100 is slightly worse (compounding error is more severe).
折扣因子 γ：γ=0.95 略好于 γ=0.9 和 γ=0.99，作者推测原因是对 model 不完美的更好容忍。Discount factor γ: γ=0.95 is slightly better than γ=0.9 and γ=0.99; the authors conjecture that this comes from better tolerance of an imperfect model.
迭代轮次：15 轮迭代优于 3 轮和 1 轮，且 random starts 对于避免 compounding error 至关重要。Number of iterations: 15 iterations beat 3 and 1, and random starts are essential for avoiding compounding error.
样本量扩展：SimPLe 在极低数据量（50K~100K）下优势明显，但随样本量增加至约 500K 时与 model-free PPO 持平，之后 PPO 反超——这与 model-based vs model-free 的经典规律吻合。Scaling with sample budget: SimPLe has a clear advantage in the very low data regime (50K~100K), but becomes on par with model-free PPO as the budget grows to about 500K, after which PPO overtakes it — consistent with the classic model-based vs model-free pattern.

04 局限性Limitations

Note：以下限制均为论文作者在 "Conclusions and Future Work" 及正文中明确陈述（stated），非推断。Note: all limitations below are explicitly stated by the authors in "Conclusions and Future Work" and in the main text, not inferred.

渐近性能低于 model-free SOTA（stated）Asymptotic performance below model-free SOTA (stated)

"The final scores are on the whole lower than the best state-of-the-art model-free methods." 论文作者指出这在 model-based RL 中普遍存在，需要更好的 dynamics model 来弥补。在极高样本量时，PPO 等方法的最终得分仍超过 SimPLe。"The final scores are on the whole lower than the best state-of-the-art model-free methods." The authors note that this is common in model-based RL and has to be made up for by a better dynamics model. With very large sample budgets, the final scores of methods such as PPO still exceed those of SimPLe.

训练方差大、稳定性差（stated）Large training variance and poor stability (stated)

"The performance of our method generally varied substantially between different runs on the same game." 多轮训练的互相影响（policy 训练、world model 训练和数据收集之间的耦合）以及 world model 与真实环境之间的 domain shift 是高方差的主因。作者建议未来使用 Bayesian 参数后验或 ensemble 方法提升鲁棒性。"The performance of our method generally varied substantially between different runs on the same game." The mutual influence across training rounds (the coupling between policy training, world model training and data collection), together with the domain shift between the world model and the real environment, is the main cause of the high variance. The authors suggest using a Bayesian posterior over parameters or ensemble methods in the future to improve robustness.

计算代价高昂（stated）High computational cost (stated)

"The computational and time requirement of training inside world model are substantial." World model（约 74M 参数）的推理约 32ms/帧（batch size=16，NVIDIA Tesla P100），而真实 ALE 模拟器仅约 0.4ms/步——相差约 80 倍。这使得开发更轻量的 world model 成为重要研究方向。"The computational and time requirement of training inside world model are substantial." Inference with the world model (about 74M parameters) takes roughly 32ms per frame (batch size=16, NVIDIA Tesla P100), whereas the real ALE simulator needs only about 0.4ms per step — a gap of roughly 80 times. This makes developing lighter world models an important research direction.

对小物体和全局变化场景预测失败（stated）Prediction fails for small objects and globally changing scenes (stated)

在 Atlantis 和 Battle Zone 等游戏中，子弹等极小但重要的物体容易从模型预测帧中消失。在 Private Eye 等需要场景切换（大范围全局变化）的游戏中，模型同样难以捕捉。这类游戏上的最终得分接近随机。In games such as Atlantis and Battle Zone, tiny but important objects such as bullets tend to disappear from the model's predicted frames. In games such as Private Eye that require scene transitions (large-scale global changes), the model likewise struggles to capture them. Final scores on such games are close to random.

低数据量时优势显著，高数据量时消失（stated）A clear advantage in the low data regime that vanishes with more data (stated)

"This demonstrates that SimPLe excels in a low data regime, but its advantage disappears with a bigger amount of data." 在约 500K 样本时 SimPLe 与 PPO 打平，之后 model-free 方法反超。作者将此归因于 SimPLe policy 熵过低导致的探索不足，制约了后续 PPO 微调的潜力。"This demonstrates that SimPLe excels in a low data regime, but its advantage disappears with a bigger amount of data." At about 500K samples SimPLe ties with PPO, after which the model-free method overtakes it. The authors attribute this to insufficient exploration caused by the overly low entropy of the SimPLe policy, which limits the potential of subsequent PPO fine-tuning.