[Paper Review] Model-Based Reinforcement Learning for Atari

Łukasz Kaiser, Mohammad Babaeizadeh+12 more|arXiv (Cornell University)|Mar 1, 2019

Reinforcement Learning in Robotics47 references420 citations

TL;DR

SimPLe uses a learned stochastic video-prediction world model to train policies inside a simulated Atari environment, achieving strong sample efficiency at 100K interactions and often outperforming model-free baselines like Rainbow and PPO.

ABSTRACT

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.

Motivation & Objective

• Reduce sample complexity for Atari play from pixel observations.
• Demonstrate a complete model-based RL loop (learn world model, train policy in the model, iterate with real data).
• Compare against state-of-the-art model-free baselines on a fixed 100K interaction budget.

Proposed method

• Introduce Simulated Policy Learning (SimPLe) that alternates between collecting real data, training a world model, and training a policy in the world model via PPO.
• Use a stochastic video-prediction world model with discrete latent variables to predict next frames and rewards conditioned on actions.
• Employ scheduled sampling to mitigate compounding errors and short rollouts to limit error accumulation during planning in the learned model.
• Train the world model with a mix of supervised learning (from real data) and predictive modeling, using clipping losses for stability.
• Aggregate data across iterations (dataset aggregation) to extend the world model’s coverage of the environment dynamics.
• Evaluate performance with 100K real-environment interactions across 26 Atari games, comparing to Rainbow and PPO.

Experimental results

Research questions

• RQ1 Can a learned video-prediction model enable effective policy learning in Atari with a fixed, low interaction budget (100K)?
• RQ2 How does model-based learning with stochastic discrete latent dynamics compare to advanced model-free baselines (Rainbow, PPO) in sample efficiency?
• RQ3 What is the impact of stochasticity, rollout length, and iterative data collection on policy performance in a learned world model?
• RQ4 Does pretraining or initializing with SimPLe data provide advantages for subsequent model-free fine-tuning?

Key findings

• SimPLe is more sample-efficient than a highly tuned Rainbow baseline on almost all tested games at 100K interactions.
• In the best case (Freeway), SimPLe is more than 10x more sample-efficient than Rainbow.
• Across 26 games, SimPLe outperforms model-free methods in over half the games, with some results exceeding human scores.
• For several games, Rainbow/PM PPO require more than double the samples to reach the same score as SimPLe at 100K.
• The method achieves performance competitive with model-free approaches in many settings, and can approach or surpass PPO in sample efficiency.
• Results are reported as averages over 5 runs per game, with notable variance across runs.