PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning and Querying Fast Generative Models for Reinforcement Learning

Lars Buesing, Theophane Weber, Sebastien Racaniere, S. M. Ali Eslami, Danilo Rezende, David P. Reichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, Daan Wierstra

arXiv:1802.03006v1cs.LG

TL;DR

Model-based RL seeks accurate, efficient environment models to support learning policies from less experience. The paper develops state-space models that predict through compact representations and agents that query them for decision making. The models capture Atari dynamics from pixels, and query-based agents outperform strong model-free baselines on MS PACMAN.

  • Problem

    Model-based RL needs environment models that predict future outcomes both efficiently and accurately to support sample-efficient policy learning.

    Source: §1 · p. 1

  • Method

    The paper develops deterministic and stochastic state-space models that predict through compact latent states, and agents that query those models to anticipate action outcomes.

    Source: §1 · p. 1, STATE-SPACE MODELS: ABSTRACTION IN SPACE · p. 2

  • Results

    State-space models match frame-autoregressive accuracy while enabling efficient rollouts, and I2A agents outperform strong model-free baselines; learning rollout policies adds modest, consistent gains.

    Source: §1 · p. 1, §6 · p. 8

  • Takeaways & Limitations

    State-space models learned from raw pixels are promising for model-based RL, providing efficient rollouts and dynamic state representations useful for evaluating and anticipating outcomes.

    Source: §6 · p. 8

    Model-based planning can be costly and brittle, and unavoidable model imperfections in complex environments can often lead to catastrophic outcomes.

    Source: §3 · p. 4

Abstract

from arXiv · show

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.

1. Introduction

The paper addresses the need for accurate, computationally efficient environment models in model-based RL, proposing state-space models and agents that query them for decisions. Its experiments report improved MS PACMAN performance over strong model-free baselines, with further gains from learning to query the model.

1. Introduction

  • Model-based RL aims to improve sample efficiency by using explicit environment models to learn task-specific policies rapidly.Such models must predict future environment states both efficiently and accurately.Source: §1 · p. 1
  • The paper proposes compact state-space models that reduce prediction computation, examines explicit transition uncertainty, and uses model queries to aid agent decisions.The agent anticipates action outcomes by querying a state-space model.Source: §1 · p. 1
  • State-space models yield MS PACMAN results that significantly improve on strong model-free baselines, and learning to query the model further increases policy performance.The paper also reports state-of-the-art environment-modeling accuracy with stochastic state-space models that produce diverse yet consistent rollouts.Source: §1 · p. 1
  • The paper compares deterministic and stochastic, pixel-space and state-space models for speed and accuracy in challenging Arcade Learning Environment environments.It reports state-of-the-art log-likelihood accuracy for stochastic state-space models.Source: §1 · p. 1

2. Environment models

The section defines environment models as learned predictors of observations and rewards under actions, then contrasts autoregressive and state-space approaches. State-space models predict through compact latent states, while jumpy models additionally reduce rollout computation by modeling at a coarser time scale.

2. Environment models

  • A model-based RL environment model approximates the environment’s action-conditioned observation and reward predictions, with the goal of making predictions both computationally cheap and accurate.The section focuses on Monte Carlo predictions over a chosen rollout horizon for later decision making.Source: §2 · p. 2, §2.1 · p. 2

AUTO-REGRESSIVE MODELS

  • Autoregressive models predict future observations sequentially; recurrent autoregressive models reuse recurrent statistics but still render observations during rollouts.Regular AR models condition on a FIFO buffer of the last K observations and actions.Source: AUTO-REGRESSIVE MODELS · p. 2

STATE-SPACE MODELS: ABSTRACTION IN SPACE

  • State-space models roll out compact latent states from the previous state and action, and use those states to predict observations without directly sampling pixels.The section considers deterministic transitions and stochastic transitions that explicitly model uncertainty; its observation models may also differ in whether they model uncertainty across time.Source: STATE-SPACE MODELS: ABSTRACTION IN SPACE · pp. 2–3

2.2. Jumpy models: abstraction in time

  • Jumpy models subsample observations and concatenate actions into coarser time steps, reducing Monte Carlo rollout computation roughly by a factor of c.The coarser time scale is set by the subsampling factor c.Source: §2.2 · p. 3

2.3. Model architectures, inference and training

  • The section introduces parametric architectures for the model families and describes adapting the stochastic state-space model into deterministic variants and a recurrent autoregressive model.For fair comparison, the deterministic variants use the latent distribution’s mean in place of latent samples.Source: §2.3 · pp. 3–4
  • The models use convolutional feature maps and pool-and-inject layers for long-range spatial dependencies; latent-variable models are trained by maximizing an evidence lower bound when likelihood is not generally available in closed form.The posterior uses an observation encoder, and training uses a single posterior sample for unbiased ELBO gradient estimates.Source: §2.3 · pp. 3–4

3. RL agents with state-space models

This section frames model-based planning as costly and brittle when models are imperfect, then describes I2A designs that combine model-free estimates with rollout information and approaches for training rollout policies or informative imagination distributions.

3. RL agents with state-space models

  • Brute-force model-based search is costly and brittle, while I2A combines model-free learning with rollout imaginations that RL learns to interpret, reducing susceptibility to model imperfections (Weber et al., 2017).Source: §3 · p. 4

3.1. Imagination-Augmented Agent

  • I2A actor-critic combines a CNN model-free pathway with a fixed-model pathway that samples K rollouts, summarizes imaginations, and joins them to estimate policy and value.Source: §3.1 · pp. 4–5
  • For SSM rollouts, the rollout policy depends only on abstract state, and predicted abstract-state imaginations are cheaper to compute than pixels.Source: §3.1 · pp. 4–5

3.2. Distillation

  • Distillation trains rollout policy πr to imitate model-based behavioral policy π by minimizing KL divergence, with λD trading off reward maximization and distillation loss (Weber et al., 2017).Source: §3.2 · p. 5

3.3. Learning to Query by Backpropagation

  • Relaxing discrete rollout actions to probability vectors lets policy gradients pass through the differentiable model into πr, while keeping model parameters fixed.The model is not guaranteed to generalize to continuous action probabilities, so an entropy penalty encourages rollout probabilities to remain close to one-hot actions.Source: §3.3 · p. 5

3.4. Modulation agent

  • The modulation agent replaces an unconditional sSSM's learned prior with trainable pimag, letting the I2A search latent trajectories for informative, potentially extreme outcomes.Source: §3.4 · pp. 5–6

4. Results

Across four ALE games, state-space models offer competitive or better likelihoods than autoregressive models while substantially speeding rollouts. In MS PACMAN, I2A agents outperform the model-free baseline, with abstract-state rollouts outperforming pixel-level reasoning.

4.1. Comparison of environment models

  • State-space models are competitive with or outperform autoregressive models, while stochastic models outperform deterministic variants and SSMs substantially speed up rollouts.The table caption reports that uncertainty-bearing RAR and sSSM models outperform models without uncertainty, and that dSSM and sSSM are faster than autoregressive models.Source: §4.1 · p. 6, Table 1 · p. 7
  • SSM rollouts are over 5 times faster than the standard autoregressive model, and jumpy sSSM rollouts are over an order of magnitude faster with comparable performance.The autoregressive implementation already reuses vision preprocessing where possible.Source: §4.1 · p. 6

4.2. RL with state-space models on MS PACMAN

  • The MS PACMAN comparison uses a simplified five-action variant, jumpy SSMs, five rollouts per step, and horizons of 2, 3, or 4.The tested horizons correspond to 24, 36, and 48 environment steps; autoregressive-model I2As were too expensive to run.Source: §4.2 · p. 6
  • I2A agents outperform the model-free baseline, and abstract-state variants outperform pixel-level reasoning on MS PACMAN.The pixel-level agent eventually beats the copy and model-free baselines, but all abstract-level I2A variants outperform it.Source: §4.2 · pp. 6–7
  • I2As with deterministic state-space models outperform stochastic-model I2As by a large margin, although the stochastic models better capture environment dynamics.The authors hypothesize that only five stochastic samples yield highly variable features for the I2As.Source: §4.2 · p. 7
  • For deterministic models, learning to query narrowly beats distillation, both beat uniform-random rollouts, and depth τ = 4 beats τ = 2 or 3.The authors report that imagining behaviors different from the agent’s policy can benefit planning.Source: §4.2 · p. 7

5. Related Work

The paper positions its generative sequence models as pixel-based stochastic state-space models and applies them to model-based RL. It contrasts this explicit, unsupervised environment model with prior implicit-planning approaches and notes limits of prior model-based RL results on pixels.

Generative sequence models

  • Prior stochastic sequence-model work largely focused on music, speech, or low-dimensional data; this paper studies high-dimensional pixel observations and stochastic state-space models as action-conditional environment models.The authors describe the action-conditional application as the first of its kind to their knowledge.Source: §5 · p. 7

Model-based reinforcement learning

  • The paper reports that I2As outperform the model-free baseline on MS PACMAN, with learned rollout policies providing the best results.Source: Figure 3 · p. 8

Generative sequence models

Model-based reinforcement learning

  • Unlike implicit-planning approaches, this agent uses an explicit environment model learned unsupervised from sensory observations.The comparison includes value-iteration architectures, the Predictron, and auxiliary prediction losses.Source: §5 · p. 8
  • Prior pixel-based model-based RL work was under-explored because of high accuracy and computational demands; cited robot-arm prediction required explicit pixel rendering and did not use abstract space.The paper also distinguishes its approach from model-free MS PACMAN results that rely on privileged information.Source: §5 · p. 8

6. Discussion

The discussion concludes that pixel-trained state-space models capture complex dynamics at accuracy similar to frame-autoregressive models, support efficient Monte Carlo rollouts, and provide useful dynamic features for I2As. It also reports modest gains from learning a rollout policy and identifies pretrained-model reliance and further speedups as open directions.

6. Discussion

  • State-space models capture complex dynamics with accuracy similar to frame-autoregressive models, support efficient Monte Carlo rollouts, and provide better features for anticipating outcomes than raw pixels.These properties supported Imagination-Augmented Agents outperforming strong model-free baselines.Source: §6 · p. 8
  • Learning a rollout policy by backpropagating policy gradients yields consistent, if modest, improvements.Source: §6 · p. 8
  • The agents assume a pre-trained environment model; jointly learning the model and agent and further speeding models remain future directions, including adaptive temporal abstractions for flexible reaction.Source: §6 · pp. 8–9

A.1. Architectures

The appendix identifies the inference-model architectures for the dSSM-VAE and sSSM, then specifies how pixel channels and rewards are modeled.

A.1. Architectures

A.2. Detail in the observation model

  • Pixel color channels are treated as pseudo-probabilities and scored by KL divergence against model predictions.Source: §A.2 · p. 11
  • Rewards are modeled separately using independent Bernoulli coefficients for the binary reward representation, plus sign and zero-indicator variables.The coefficients are conditioned on s_t and z_t.Source: §A.2 · p. 11

A.3. Details of neural network implementations

The sSSM implementation is built from three higher-level blocks: convolutional and residual convolutional stacks, and a Pool & Inject layer; these blocks define the listed modules.

A.3. Details of neural network implementations

  • The sSSM uses a three-layer convolutional stack with specified kernel sizes and channel counts.Source: §A.3 · p. 11
  • A three-layer residual convolutional stack with fixed sizes is another building block.Source: §A.3 · p. 11
  • The Pool & Inject layer is the third higher-level building block.Source: §A.3 · p. 11
  • The implementation defines its modules from these building blocks, as shown in Figures 9–14.Source: §A.3 · p. 11

A.4. Collection of training data

The authors collect ALE transitions from trained DQN policies on four games, preprocess observations and actions, and model every fourth frame.

A.4. Collection of training data

  • For each of BOWLING, CENTIPEDE, MS PACMAN, and SURROUND, trained DQN policies supplied 108 training and 107 test transitions.The agents used the original 18-action space.Source: §A.4 · p. 11
  • Actions were one-hot encoded, and pixel observations were cropped to 200 × 160 and normalized by 255 into [0, 1]³.Source: §A.4 · p. 11
  • Because DQN used an action repeat of four, the models included only every fourth frame.Source: §A.4 · p. 11

A.5. Training details

Stochastic state-space models produce plausible rollouts where deterministic models blur under uncertain dynamics; the experiments also specify training and a simplified MS PACMAN setting.

A.5. Training details

  • All models were optimized with Adam using mini-batches of 16.Source: §A.5 · p. 11

A.6. Comparison of deterministic and stochastic state-space models

  • On bouncing-ball sequences with diffusion, sSSM rollouts depict plausible trajectories, while dSSM rollouts are blurry because the ball’s state is not fully predictable.The passage notes a similar rollout result for models trained on ALE games.Source: §A.6 · p. 11

B.1. MS PACMAN environment variant

  • The MS PACMAN variant uses five actions, action-repeat four, and observations from every fourth environment frame.Source: §B.1 · p. 12

B.2. Architecture

The architecture pairs a model-free CNN with an I2A path that summarizes imagined rollouts, alongside convolutional modules for encoding observations and modeling state transitions.

B.2. Architecture

  • The model-free baseline uses a four-layer CNN with ReLU activations, followed by a 512-unit fully connected layer that outputs a value estimate and policy logits.Source: §B.2 · pp. 12–13
  • The model modules encode observations without reward, compute prior and posterior latent statistics, initialize state from three preceding observations, and decode pixel and reward distributions.The posterior also receives prior statistics; the prior is conditioned on the previous state and action.Source: Figure 11 · p. 12, Figure 10 · p. 12, Figure 12 · p. 13
  • The I2A model-based path encodes each rollout step with a two-layer CNN, summarizes rollout statistics with an LSTM, and concatenates them with model-free path outputs.The LSTM has 256 hidden units; the rollout encoder’s fully connected layer has 128 outputs.Source: §B.2 · p. 13
  • State-based and pixel-based rollout policies use CNNs followed by fully connected layers, with the pixel-based policy matching the baseline except for its last two layers’ feature-map counts.Source: §B.2 · p. 13

C. Results I2A with stochastic state-space models

For I2A agents using stochastic state-space models on MS PACMAN, learning to query eventually outperforms the other tested agents, while both learning-to-query and distillation beat uniform random rollout policies.

C. Results I2A with stochastic state-space models

  • Both learning-to-query and distillation I2As outperform uniform-random-rollout I2As, while learning-to-query starts weakly but eventually outperforms the other agents.The passage attributes the performance benefit to learning to sample informative outcomes.Source: C. Results I2A with stochastic… · p. 13, Figure 18 · p. 15
Loading 1802.03006v1…