Source-linked AI summary
Model-Based Reinforcement Learning for Atari
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, Henryk Michalewski
TL;DR
Model-free Atari RL often needs far more interaction than humans, motivating learned models that predict game outcomes for more efficient learning. The paper introduces SimPLe, which trains policies inside stochastic video-prediction models and iteratively updates those models with real-environment data. In a 100K-interaction regime, SimPLe is more sample-efficient than model-free baselines on most tested games, while retaining lower final scores and substantial run-to-run variation.
Problem
Prior work had not clearly demonstrated competitive model-based control from predictive visual models in the Atari Learning Environment, despite the need for more interaction-efficient learning.
Method
SimPLe learns a stochastic video-based world model from real-environment data, trains a policy inside that model, and iteratively aggregates additional data from the original game.
Results
At 100k interactions, SimPLe is more sample-efficient than a highly tuned Rainbow baseline on almost all games, exceeding 10x efficiency on Freeway.
Takeaways & Limitations
SimPLe demonstrates that learned video models can support effective Atari policies with substantially fewer environment interactions than prior model-free methods in many games.
Takeaways & Limitations
Final scores were generally lower than the best model-free methods, and performance varied substantially across runs on the same game.
Abstract
from arXiv · showhide
Model-free reinforcement learning (RL) can be used to learn effective policies for complex tasks, such as Atari games, even from image observations. However, this typically requires very large amounts of interaction -- substantially more, in fact, than a human would need to learn the same games. How can people learn so quickly? Part of the answer may be that people can learn how the game works and predict which actions will lead to desirable outcomes. In this paper, we explore how video prediction models can similarly enable agents to solve Atari games with fewer interactions than model-free methods. We describe Simulated Policy Learning (SimPLe), a complete model-based deep RL algorithm based on video prediction models and present a comparison of several model architectures, including a novel architecture that yields the best results in our setting. Our experiments evaluate SimPLe on a range of Atari games in low data regime of 100k interactions between the agent and the environment, which corresponds to two hours of real-time play. In most games SimPLe outperforms state-of-the-art model-free algorithms, in some games by over an order of magnitude.
1 INTRODUCTION
SimPLe addresses the gap between humans’ rapid Atari learning and the much larger interaction budgets of model-free RL by using learned video models for model-based control. It introduces an iterative system that trains policies inside a learned model and reports strong sample-efficiency gains across Atari games.
- Motivation: Humans can learn Atari games in minutes, whereas leading model-free methods may require tens or hundreds of millions of time steps.The paper frames this difference as motivation for learning predictive models of game dynamics.
- Research gap: Prior predictive-model work had not clearly demonstrated competitive model-based control from images in the Atari Learning Environment.The paper identifies successful planning with learned models in ALE as an open challenge.
- Motivation: Model-based learning is motivated by its potential to improve deep-RL sample efficiency while supporting prediction, exploration, intrinsic motivation, and counterfactual reasoning.The paper emphasizes sample efficiency as a central benefit among several possible applications of environment models.
- Approach: SimPLe trains video prediction models, learns a policy inside the learned model, and repeatedly aggregates data collected from the original game.The system includes several stochastic video-prediction approaches, including a novel discrete-latent model.
- Results: At 100k samples, SimPLe is more sample-efficient than a highly tuned Rainbow baseline on almost all games and exceeds 10x efficiency on Freeway.On more than half of the games, SimPLe reaches a score that Rainbow requires at least twice as many samples to achieve.
2 RELATED WORK
Related work established that Atari video prediction is feasible, but offered limited evidence for model-based control from images. SimPLe builds on predictive modeling and alternating model-learning/policy-learning approaches while comparing against model-free and other model-based baselines.
- Atari reinforcement learning: Deep model-free methods such as DQN and actor-critic algorithms became the dominant successful approaches for learning Atari from images.The related-work discussion positions model-free Atari methods as the prevailing benchmark standard.
- Video prediction: Prior work showed that deep architectures can predict Atari frames over long horizons, motivating learned simulators for game-playing.Some reported predictions maintained low L2 error over hundreds of steps, but the cited work focused on prediction rather than effective control.
- Model-based control: Model-based reinforcement learning from images in video games remained comparatively rare, with some methods modeling rewards without predicting future frames.The paper distinguishes reward augmentation from full predictive modeling of game dynamics.
- Broader model-based RL: Image-based model-based RL had also been explored in robotics, including both real-world and simulated settings.These studies provide related precedents outside Atari games.
- Algorithmic structure: A common model-based RL structure alternates between learning an environment model and optimizing a policy with model-free reinforcement learning.The paper relates this structure to Dyna Q and later deep-RL methods.
- Evaluation context: Some comparison scores for Dyna-DQN and GATS were approximate because their authors did not provide tabular results.The comparison used random-normalized scores at 100K interactions.
3 SIMULATED POLICY LEARNING (SIMPLE)
SimPLe trains a policy primarily in a learned visual environment while using limited interaction with the original Atari environment for data collection and model updates. The procedure alternates real-environment data collection, supervised world-model training, and policy optimization inside the model.
- Problem formulation: The Atari task uses visual observations rather than emulator RAM, with four consecutive frames stacked to reduce partial observability.A single image does not determine the underlying game state.
- World model: The learned world model shares the original environment’s action and reward spaces and produces visual observations in the same format.It is trained to mimic the original environment using data initially collected from random rollouts.
- Objective: SimPLe aims to train a policy in the learned environment while using as few interactions with the original environment as possible.The policy ultimately needs to perform well in the original Atari environment.
- Algorithm: The algorithm initializes a policy and model, collects observations from the real environment, trains the model on accumulated data, and updates the policy inside the model.These steps repeat until the procedure terminates.
4 WORLD MODELS
The paper evaluates multiple world-model architectures and proposes a stochastic video-prediction model with discrete latent variables. Its design addresses stochasticity and inference-time latent mismatch while incorporating training techniques for stable long-horizon prediction.
- Architecture search: Architecture search produced a novel stochastic video-prediction model that outperformed the other evaluated architectures.The proposed architecture is described in detail and compared experimentally in a later section.
- Deterministic model: The deterministic model maps four consecutive frames and an action to the next frame and reward using convolutional processing and an action embedding.Actions condition the visual decoder through channel-wise multiplication with the embedded action.
- Architecture variants: The experiments varied network details including filter counts, output representations, reward prediction, dropout, and layer normalization.Visual outputs could use continuous pixel values or categorical 256-color distributions.
- Training: Clipped visual losses improved reward prediction and successful training by reducing the influence of gradients from large background regions.The authors conjecture that optimization can thereby focus on small but important areas.
- Training: Scheduled sampling replaces some training inputs with prior predictions to mitigate compounding-error drift during multi-step model use.The replacement probability increases linearly to 100% around the middle of the first training-loop iteration.
- Stochastic modeling: The stochastic model uses discrete latent variables because the variational alternative required game-dependent KL weighting and could produce unseen latent values at inference.The problematic KL weights were typically in the range [10^-3, 10^-5].
- Stochastic modeling: The proposed model discretizes latent values into bits and uses an auxiliary LSTM to predict those bits autoregressively at inference time.Uniform noise before discretization and dropout after discretization improve robustness to unseen latent bits.
5 POLICY TRAINING
SimPLe repeatedly trains a policy inside a learned world model, using short simulated rollouts and periodically refreshed real-environment data to limit compounding model errors.
- Policy optimization: PPO trains the policy π using rollouts collected from the simulated environment env′.The simulated environment shares the original environment’s action and reward spaces and produces visual observations.
- Policy optimization: Short rollouts restart every N = 50 steps from uniformly sampled states in the ground-truth buffer D.This design mitigates the effect of imperfections compounding over time in env′.
- Iterative training: The main training loop runs for 15 iterations, with 45K world-model training steps initially and 15K in each later iteration.Later iterations extend the learned dynamics to novel situations rather than relearning them from scratch.
- Policy optimization: Each iteration trains PPO with 16 parallel agents collecting 25, 50, or 100 simulated steps per epoch.The number of PPO epochs is scaled across selected loop passes, producing 800K·z simulated-environment interactions per pass.
6 EXPERIMENTS
SimPLe is evaluated on Atari games under a low-interaction budget and generally achieves higher sample efficiency than tuned model-free baselines. Additional experiments examine data scale, stochasticity, model design, and exploration.
- Experimental setup: 102,400 Atari-environment interactions are used during training, equivalent to 409,600 game frames or 114 minutes at 60 FPS.The evaluation loop collects 6,400 environment interactions across 16 collection stages.
- Experimental setup: SimPLe is evaluated on 26 Atari games selected for non-random performance with existing model-free methods, using Rainbow and PPO as comparisons.Rainbow uses a Dopamine implementation tuned for sample efficiency.
- Sample efficiency: SimPLe is more sample-efficient than tuned Rainbow on almost all games, uses less than half as many samples on more than half, and exceeds 10x efficiency on Freeway.It outperforms PPO by an even larger margin in the reported comparisons.
- Sample efficiency: SimPLe outperforms Rainbow and PPO even when those baselines receive twice as many interactions.Figure 4 normalizes SimPLe’s 100K score against Rainbow and PPO scores at 100K and 200K.
- Number of frames: Performance is poor at 20K interactions, nearly reaches the 100K level at 50K, and improves through 500K, where it matches model-free PPO.The low-data advantage disappears with larger amounts of data.
- Number of frames: A SimPLe policy trained with 100K interactions provides a useful initialization for model-free PPO, although its lower asymptotic performance is attributed to worse exploration.The pretrained policy’s low entropy at 100K hinders further PPO training.
- Environment stochasticity: Stochastic world models address uncertainty from limited frame histories, such as unpredictable opponent configurations in Kung Fu Master.The stochastic model samples potential opponents and renders them sharply instead of producing blurred averages.
- Environment stochasticity: With sticky actions, the stochastic world model accounts for action stickiness and usually achieves results similar to the deterministic setting without tuning.The full world-model and policy pipeline is trained in the sticky-action environment.
7 CONCLUSIONS AND FUTURE WORK
SimPLe learns Atari policies from raw pixels with far fewer interactions than prior methods, while remaining limited by lower final scores, run variability, and substantial training costs. The paper focuses on Atari and identifies broader applications as future work.
- 100K environment interactions, corresponding to 2 hours of play, suffice for SimPLe to learn many Atari games.In many cases, prior methods require several times more samples to reach the same reward.
- SimPLe uses stochastic latent variables, enabling application in highly stochastic environments.
- SimPLe’s final scores are generally lower than those of the best state-of-the-art model-free methods.The authors suggest better dynamics models as one route for improvement.
- Performance varies substantially across runs on the same game, likely because of complex interactions among the model, policy, and data collection.The paper suggests Bayesian parameter posteriors or ensembles could improve robustness.
- Training inside the world model has substantial computational and time requirements, motivating lighter models.
- The demonstrated scope is a suite of Atari games, while applying similar methods to other environments and tasks remains future work.
A ABLATIONS
The ablations examine model architecture, rollout resets, discounting, training iterations, model-training duration, and random starts. Results favor the stochastic discrete model, moderate reset intervals, more training loops, longer model training, and random-start rollouts.
- Models: The proposed stochastic discrete model performs best among the evaluated deterministic, deterministic recurrent, and stochastic discrete models.The experiments also examine the roles of stochasticity and recurrence.
- Steps: N = 100 ground-truth reinitialization steps performs worse than N = 25 or N = 50, likely because of compounding model errors.The effect is smaller than the effect of model architecture.
- Gamma: γ = 0.95 performs slightly better than other tested discount factors, although all three values perform comparably at the same number of steps.
- Model-based iterations: 15 training loops are beneficial when the total number of samples is fixed at 100K.
- Long model training: The best results use world-model training that is 5 times longer than the shorter training setting.Resource constraints limited other ablations to the short-training setting.
- Random starts: Randomly selected starting states from the real data buffer improve exploration and help mitigate compounding errors from short simulated rollouts.Removing random starts substantially deteriorates results in the reported Seaquest comparison.
B QUALITATIVE ANALYSIS
SimPLe achieved strong results on some Atari games while revealing distinct weaknesses in exploration, long-range prediction, and visually difficult settings.
- Solved games: SimPLe achieved the maximum score on Pong and Freeway within the 100K-interaction budget.
- Exploration: Freeway demonstrates that SimPLe can learn from rare successful events and incorporate them into a policy.
- Exploration: Performance was not reliable across runs when rapid entropy decay limited experience and produced a weak world model.
- Pixel-perfect games: Pixel-perfect predictions were observed in Pong, Freeway, and Breakout for episodes lasting up to 50 time-steps.
- Benign errors: Some model errors were harmless or mildly harmful, allowing policies trained on plausible but inaccurate predictions to play the original game.
- Failures on hard games: Models struggled with large global scene changes in Private Eye and with tiny relevant objects such as bullets in Atlantis and Battle Zone.
C ARCHITECTURE DETAILS
The world model uses a deep convolutional prediction architecture with latent-variable components, residual connections, normalization, and dropout.
- The model has approximately 74M parameters and supports visual prediction beyond Atari-specific designs.
- The frame prediction network combines convolutional downscaling, action and latent-predictor embeddings, recurrent latent prediction, and reward prediction.
- A separate latent inference network is used only during training.
- The architecture uses ReLU activations, softmax output layers, residual encoder-decoder connections, dropout, and layer normalization.
D NUMERICAL RESULTS
The experiments evaluate SimPLe across multiple runs and configurations, comparing model choices and tuned model-free baselines in the low-data setting.
- SimPLe experiments use five runs per configuration, with evaluation policies formed using a temperature of T = 0.5.
- Going down to T = 0 was detrimental in many cases, possibly because policies became more sensitive to model imperfections.
- Table 2 reports mean and standard deviation over five experiments, while median and maximum scores are also reported elsewhere.
- Rainbow was tuned specifically for the low-data regime through a hyperparameter search, with candidate settings including replay history, target updates, and replay scheme.
- The Rainbow hyperparameter selection used five agents trained on Pong for one million environment interactions.
F RESULTS AT DIFFERENT NUMBERS OF INTERACTIONS
The paper normalizes SimPLe’s performance against Rainbow, PPO, and human scores to compare learning progress across interaction budgets.
- Rainbow-relative performance is computed from SimPLe, random, and Rainbow scores at each sample count.
- PPO-relative performance uses the analogous normalization against PPO and random scores, with an adjustment when the denominator is negative.
- The cross-baseline comparison uses (SimPLe_score@100K − baseline_score)/human_score and normalizes points by average human score.