[Paper Review] Learning and Querying Fast Generative Models for Reinforcement Learning

Lars Buesing, Théophane Weber+9 more|arXiv (Cornell University)|Feb 8, 2018

Reinforcement Learning in Robotics32 references70 citations

TL;DR

The paper proposes fast, state-space environment models (deterministic and stochastic) learned from pixels to enable efficient model-based RL, and demonstrates that imagination-augmented agents using these models outperform strong model-free baselines on MS_PACMAN, with further gains when learning to query the model.

ABSTRACT

A key challenge in model-based reinforcement learning (RL) is to synthesize computationally efficient and accurate environment models. We show that carefully designed generative models that learn and operate on compact state representations, so-called state-space models, substantially reduce the computational costs for predicting outcomes of sequences of actions. Extensive experiments establish that state-space models accurately capture the dynamics of Atari games from the Arcade Learning Environment from raw pixels. The computational speed-up of state-space models while maintaining high accuracy makes their application in RL feasible: We demonstrate that agents which query these models for decision making outperform strong model-free baselines on the game MSPACMAN, demonstrating the potential of using learned environment models for planning.

Motivation & Objective

• Promote computationally efficient environment models by learning compact state representations (state-space models).
• Investigate deterministic and stochastic state-space models and their pixel-based learning capacity.
• Demonstrate planning and decision-making benefits by querying state-space models in reinforcement learning agents.
• Showcase how uncertainty modeling and temporal abstraction (jumpy models) affect speed and accuracy in ALE tasks.

Proposed method

• Define and compare auto-regressive, recurrent auto-regressive, and state-space models (deterministic and stochastic) for environment modeling.
• Use pixel-based encoders and convolutional transition/decode modules with pool-and-inject layers to capture long-range dependencies.
• Train models with maximum likelihood or ELBO objectives; use variational inference for stochastic models (sSSMs).
• Employ jumpy training to approximate longer horizons with temporal abstraction, reducing computation.
• Integrate state-space models into Imagination-Augmented Agents (I2A) to perform Monte-Carlo rollouts and inform policy/value functions.
• Explore learning-to-query approaches where rollout policies are trained by distillation or backpropagation through the model to improve planning.

Experimental results

Research questions

• RQ1 Can state-space models (deterministic and stochastic) learned from pixel observations provide accurate and computationally efficient environment predictions?
• RQ2 Do state-space models enable effective model-based RL, and can Imagination-Augmented Agents leverage such models to outperform model-free baselines on challenging ALE tasks?
• RQ3 Does incorporating uncertainty and temporal abstraction (jumpy models) improve speed-accuracy trade-offs in rollouts and planning?
• RQ4 Is learning to query the model (via distillation or backprop) beneficial for decision making in model-based RL?

Key findings

• State-space models (both dSSM and sSSM) deliver substantial speed-ups over pixel-space autoregressive models while maintaining competitive accuracy.
• Stochastic state-space models with uncertainty (sSSM) achieve higher log-likelihoods than deterministic counterparts across several ALE games.
• Jumpy (temporally abstracted) SSMs provide large speed-ups (up to >10x) with comparable accuracy to non-jumpy variants.
• Imagination-Augmented Agents (I2A) using jumpy SSMs outperform model-free baselines on MS_PACMAN, and learning to query the model (distillation or backprop) yields best performance among I2A variants.
• Deterministic SSMs generally outperform stochastic SSMs in the I2A setting when rollouts are few (e.g., K=5) due to more stable signals, while sSSMs better capture environment uncertainty in isolation.