Transformers are Sample Efficient World Models

By Vincent Micheli, Eloi Alonso, Franccois Fleuret. International Conference on Learning Representations, 2022

339 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence

View Paper: https://www.semanticscholar.org/paper/235303a8bc1e4892efd525a38ead657422d8a519

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

Citations in the corpus

Learning Interactive Real-World Simulators (2023, background): “With large transformer architectures, learning image-based world models has become plausible [58, 59, 60, 61, 62, 63], but mostly in games or simulated domains with visually simplistic and abundant data.”
Mastering Diverse Domains through World Models (2023, methodology): “Without this complexity, Dreamer outperforms the best remaining methods, including the transformer-based IRIS and TWM agents, the model-free SPR, and SimPLe 45 .”
Vista: A Generalizable Driving World Model with High Fidelity and Versatile Controllability (2024, positive): “Such properties can be useful for representation learning [35, 45, 88, 105, 131], model-based reinforcement learning [39, 40, 42, 43, 94, 96, 98], as well as model-predictive control [27, 30, 35, 41, 92, 104, 142].”
A Survey on Vision--Language--Action Models for Embodied AI (2024, positive): “IRIS [74] employs a GPT-like autoregressive Transformer as the foundation of its world model, with a VQ-VAE serving as the vision encoder.”
Mastering diverse control tasks through world models (2025, neutral): “Besides that, we compare with a wide range of tuned expert algorithms reported in the literature 9,10,33,36,41,44,52–54 .”
Genie: Generative Interactive Environments (2024, background): “Note that a common practice for training world-models, including transformer-based models, is to concate-nate the action at time 𝑡 to the corresponding frame (Micheli et al., 2023; Robine et al., 2023).”
GAIA-1: A Generative World Model for Autonomous Driving (2023, methodology): “A recent line of work suggests casting world modeling as a single sequence model, treating states, actions and rewards as simply a stream of data [10, 19, 14, 88, 12, 89].”