PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformers are Sample-Efficient World Models

Vincent Micheli, Eloi Alonso, François Fleuret

arXiv:2209.00588v2cs.LGcs.AIcs.CV

TL;DR

Deep reinforcement learning is limited by poor sample efficiency, while imagination-based methods require accurate world models. IRIS combines a discrete autoencoder with an autoregressive Transformer to learn in imagined trajectories, achieving strong Atari 100k performance after two hours of gameplay. Its results establish a new state of the art without lookahead search, while performance remains constrained by world-model accuracy and exploration of rare game events.

  • Problem

    Deep reinforcement learning often requires extensive gameplay, motivating methods that improve sample efficiency for environments where interaction is costly or unsafe.

  • Method

    IRIS learns entirely in the imagination of a world model composed of a discrete autoencoder and an autoregressive Transformer.

  • Results

    IRIS sets a new state of the art without lookahead search on Atari 100k, achieving a mean human normalized score of 1.046 and outperforming humans on 10 out of 26 games.

  • Takeaways & Limitations

    IRIS demonstrates that Transformer-based world models can support sample-efficient reinforcement learning with only two hours of gameplay.

  • Takeaways & Limitations

    IRIS is constrained by world-model accuracy and struggles when unlikely events unlock new levels or game mechanics, creating a double exploration problem.

Abstract

from arXiv · show

Deep reinforcement learning agents are notoriously sample inefficient, which considerably limits their application to real-world problems. Recently, many model-based methods have been designed to address this issue, with learning in the imagination of a world model being one of the most prominent approaches. However, while virtually unlimited interaction with a simulated environment sounds appealing, the world model has to be accurate over extended periods of time. Motivated by the success of Transformers in sequence modeling tasks, we introduce IRIS, a data-efficient agent that learns in a world model composed of a discrete autoencoder and an autoregressive Transformer. With the equivalent of only two hours of gameplay in the Atari 100k benchmark, IRIS achieves a mean human normalized score of 1.046, and outperforms humans on 10 out of 26 games, setting a new state of the art for methods without lookahead search. To foster future research on Transformers and world models for sample-efficient reinforcement learning, we release our code and models at https://github.com/eloialonso/iris.

1 INTRODUCTION

Deep RL achieves strong performance but requires prohibitively large amounts of interaction, motivating sample-efficient world-model learning. IRIS addresses this challenge by combining discrete image tokens with an autoregressive Transformer for imagination-based Atari control.

  • Motivation: Deep RL agents may require months or thousands of years of gameplay, making sample efficiency necessary for real-world deployment.Environmental interactions can also involve speed, cost, and safety constraints.
  • Motivation: Learning in a world model removes direct interaction constraints during behavior training but depends heavily on accurate long-horizon predictions.The policy is trained purely in imagination, so model errors can affect learned behavior.
  • Motivation: Transformers model high-dimensional sequences with self-supervised objectives, while discrete autoencoders reduce images to manageable token sequences.This avoids the computational difficulty of applying quadratic-scaling Transformers directly to pixels.
  • Contribution: IRIS combines a discrete autoencoder and autoregressive Transformer to simulate trajectories and learn behaviors in imagination.The autoencoder creates an image-token language, while the Transformer composes it over time.
  • World-model capabilities: IRIS generates diverse KungFuMaster futures from the same initial frame while preserving the mechanic that hit enemies disappear.The variation reflects missing information about future enemies in the conditioning frame.

2 METHOD

IRIS repeatedly collects real experience, updates a discrete-token world model, and improves behavior inside imagined trajectories. Its autoencoder reconstructs observations, while a GPT-like Transformer predicts future tokens, rewards, and episode terminations.

  • Training loop: IRIS alternates real-environment experience collection, world-model updates, and policy and value-function updates in imagination.Real experience is used to learn environment dynamics, while the agent acts exclusively within the world model.
  • From image observations to tokens: The discrete autoencoder converts each image into K tokens from a vocabulary of size N and decodes those tokens back into an image.The encoder uses a CNN and nearest embedding vectors; the decoder is also a CNN.
  • From image observations to tokens: The autoencoder is trained on collected frames using reconstruction, commitment, and perceptual losses with a straight-through estimator.The three losses are equally weighted.
  • Modeling dynamics: The Transformer models interleaved frame and action-token sequences autoregressively to predict transitions, rewards, and episode terminations.Token predictions condition on previously predicted tokens within the next frame.
  • Learning in imagination: During imagination, the policy acts on reconstructed observations while the world model predicts rewards, episode ends, and subsequent reconstructed observations.Rollouts begin from a real observation, continue for H steps, and stop when termination is predicted.
  • Learning in imagination: A value network bootstraps returns beyond the fixed imagination horizon, while REINFORCE with a value baseline and entropy maximization trains the policy.The baseline reduces gradient variance and entropy encourages exploration.

3 EXPERIMENTS

The experiments evaluate IRIS on the Atari 100k benchmark using aggregate performance analyses and qualitative world-model examples. IRIS achieves leading learning-only results, but performance depends on whether the world model can represent rare transitions and fine visual details.

  • 3.2 RESULTS: IRIS outperforms learning-only methods in number of superhuman games, mean, interquartile mean, and optimality gap.Table 1 compares returns across all 26 games and human-normalized aggregate metrics.
  • 3.1 BENCHMARK AND BASELINES: Atari 100k evaluates agents on 26 games with only 100k actions per environment, roughly equivalent to two hours of human gameplay.Unconstrained Atari agents are usually trained for 50 million steps, a 500-fold increase in experience.
  • 3.2 RESULTS: 1.046 mean human-normalized score, 0.501 IQM, and 10 of 26 games above human performance are reported for IRIS after two hours of gameplay.The reported relative improvements are computed with respect to SPR.
  • 3.3 WORLD MODEL ANALYSIS: The world model struggles with Frostbite’s rare level transition but handles Krull’s more frequent transitions, while visually detailed games remain difficult to simulate.Increasing frame-token counts improves performance on visually challenging environments, at increased computational cost.
  • 3.2 RESULTS: IRIS matches the strongest baselines on its bottom 50% of games and then stochastically dominates the other methods; its improvement probability exceeds 0.5 for every baseline.These comparisons use performance profiles and probabilities of improvement.
  • 3.3 WORLD MODEL ANALYSIS: Because IRIS learns behaviors entirely in imagination, failures to reconstruct objects or capture reward and termination mechanics can produce suboptimal policies despite many imagined trajectories.Qualitative analyses examine future generation, pixel predictions, rewards, and episode terminations.

4 RELATED WORK

Prior work established learning in imagined world models for visual reinforcement learning, while recent research explored Transformers for reinforcement learning and video generation.

  • SimPLe trained a PPO policy in a video prediction model and outperformed humans in some Atari games.
  • Recent reinforcement-learning Transformer work includes architectural modifications for optimization stability and sequence-modeling approaches such as Trajectory Transformer.
  • Discrete autoencoders compress frames into token sequences that Transformers model autoregressively, extending to video generation with spatial and temporal attention.

5 CONCLUSION

IRIS learns entirely in the imagination of a discrete-autoencoder and autoregressive-Transformer world model, achieving a new Atari 100k state of the art without lookahead search. The authors identify future scaling, internal representations, and MCTS combination as open directions, alongside reproducibility and safety considerations.

  • IRIS learns purely in imagination using a world model composed of a discrete autoencoder and an autoregressive Transformer.
  • IRIS sets a new state of the art on Atari 100k among methods without lookahead search.
  • The world model achieves pixel-perfect predictions in some games and supports rich gameplay experiences during imagination-based training.
  • Future work includes scaling IRIS, using internal world-model representations, and combining imagination learning with MCTS.
  • The authors release open-source code with minimal dependencies, configuration-based customization, visualization scripts, and interactive world-model tools.

A.1 DISCRETE AUTOENCODER

IRIS uses a VQGAN-derived discrete autoencoder that reconstructs frames while learning discrete representations. Its objective combines reconstruction, commitment, and perceptual losses, and frames pass through the autoencoder during real-environment experience collection.

  • The autoencoder removes VQGAN’s discriminator, yielding a vanilla VQVAE with an additional perceptual loss.
  • The training objective combines reconstruction loss, commitment loss, and perceptual loss.
  • During real-environment experience collection, frames still pass through the autoencoder to keep the policy’s input distribution unchanged.

A.2 TRANSFORMER

IRIS’s autoregressive Transformer embeds interleaved action and frame-token sequences, processes them through GPT2-like blocks, and models environment dynamics over time.

  • The Transformer receives a sequence of L(K + 1) tokens and embeds it into an L(K + 1) × D tensor using action and frame-token embedding tables.
  • The model uses M GPT2-like Transformer blocks, each containing self-attention with input layer normalization and a residual connection.

A.3 ACTOR-CRITIC

The actor-critic processes 64 × 64 × 3 frames through a convolutional block and an LSTM, with shared actor and critic weights except for their final layers.

  • A.3 ACTOR-CRITIC: The actor-critic receives 64 × 64 × 3 frames and processes them with four repeated convolution, ReLU, and max-pooling layers followed by a 512-dimensional LSTM.The convolution uses 3 × 3 kernels with stride 1 and padding 1; pooling uses 2 × 2 windows with stride 2.
  • A.3 ACTOR-CRITIC: The actor and critic share weights except for the last layer.

B ACTOR-CRITIC LEARNING OBJECTIVES

IRIS trains its value and policy networks on imagined trajectories, using λ-returns for value regression and a variance-reduced, entropy-regularized REINFORCE objective for the actor.

  • B ACTOR-CRITIC LEARNING OBJECTIVES: The value network minimizes expected squared error between its predictions and λ-returns computed over imagined trajectories.The λ-return is used as a regression target to balance bias and variance.
  • B ACTOR-CRITIC LEARNING OBJECTIVES: The gradient-stopping operation treats optimization targets as constants during gradient-based updates.
  • B ACTOR-CRITIC LEARNING OBJECTIVES: The policy uses REINFORCE over imagined trajectories, with the value function as a baseline to reduce gradient variance and entropy maximization to support exploration.
  • B ACTOR-CRITIC LEARNING OBJECTIVES: The training procedure collects environment experience, updates the autoencoder and dynamics model, then updates policy and value functions from imagined rollouts.

E AUTOENCONDING FRAMES WITH VARYING AMOUNTS OF TOKENS

Increasing tokens per frame improves reconstruction quality but raises computational cost; in visually complex games, the resulting world-model accuracy can translate unevenly into performance.

  • E AUTOENCONDING FRAMES WITH VARYING AMOUNTS OF TOKENS: Increasing tokens per frame improves reconstructions while requiring more compute and memory.
  • E AUTOENCONDING FRAMES WITH VARYING AMOUNTS OF TOKENS: 64 tokens per frame correctly reconstruct the player, enemies, and rewards in Alien, whereas 16 tokens can misencode these elements.
  • E AUTOENCONDING FRAMES WITH VARYING AMOUNTS OF TOKENS: +36% is the marginal performance increase for Alien when IRIS uses 64 instead of 16 tokens per frame.
  • E AUTOENCONDING FRAMES WITH VARYING AMOUNTS OF TOKENS: +121% for Asterix and +432% for BankHeist show larger benefits from 64 tokens per frame.

F BEYOND THE SAMPLE-EFFICIENT SETTING

IRIS can be scaled beyond the sample-efficient regime by increasing environment data, although gains vary across games and training remains computationally demanding.

  • F BEYOND THE SAMPLE-EFFICIENT SETTING: Increasing environment steps from 100k to 10M drastically improves performance for most games.The results provide evidence that IRIS can be scaled beyond the sample-efficient setting.
  • F BEYOND THE SAMPLE-EFFICIENT SETTING: Some games show only marginal gains from additional data, likely because of hard exploration problems or visually challenging domains needing more frame tokens.
  • F BEYOND THE SAMPLE-EFFICIENT SETTING: Training uses five random seeds per Atari environment on eight Nvidia A100 40GB GPUs, averaging 3.5 days per environment.
  • F BEYOND THE SAMPLE-EFFICIENT SETTING: Freeway remains difficult under a 100k-frame budget because rewards are sparse and random policies almost surely obtain none.
Loading 2209.00588v2…