Learning and Querying Fast Generative Models for Reinforcement Learning

TL;DR

This paper introduces state-space generative models that efficiently learn environment dynamics from raw pixels, enabling faster and accurate predictions for reinforcement learning agents, demonstrated on Atari games.

Contribution

The paper presents a novel approach to learning compact, accurate environment models that significantly reduce computational costs in model-based RL from raw pixel data.

Findings

State-space models accurately capture Atari game dynamics.

Models enable faster decision-making in RL agents.

Agents using these models outperform strong baselines on MSPACMAN.

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.