[Paper Review] Transformer-based World Models Are Happy With 100k Interactions

Jan Robine, Marc Höftmann+2 more|arXiv (Cornell University)|Mar 13, 2023

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes a Transformer-based World Model (TWM) that achieves state-of-the-art performance on the Atari 100k benchmark using only 100k environment interactions. By autoregressively modeling observations, actions, and predicted rewards through a Transformer-XL architecture, and feeding back predicted rewards into the model, TWM enables sample-efficient policy learning that outperforms prior model-free and model-based methods while keeping inference computationally efficient.

ABSTRACT

Deep neural networks have been successful in many reinforcement learning settings. However, compared to human learners they are overly data hungry. To build a sample-efficient world model, we apply a transformer to real-world episodes in an autoregressive manner: not only the compact latent states and the taken actions but also the experienced or predicted rewards are fed into the transformer, so that it can attend flexibly to all three modalities at different time steps. The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state. By utilizing the Transformer-XL architecture, it is able to learn long-term dependencies while staying computationally efficient. Our transformer-based world model (TWM) generates meaningful, new experience, which is used to train a policy that outperforms previous model-free and model-based reinforcement learning algorithms on the Atari 100k benchmark.

Motivation & Objective

• Address the sample inefficiency of deep reinforcement learning, which typically requires hundreds of millions of environment interactions—far more than human learners.
• Overcome limitations of recurrent world models by enabling direct, long-range attention to past states through a self-attention mechanism.
• Improve sample efficiency by training a policy in latent imagination using a generative world model that synthesizes new experience.
• Stabilize training through novel loss functions and sampling strategies to enhance generalization and hyperparameter robustness.
• Demonstrate that a transformer-based world model can achieve strong performance with minimal real environment interaction, specifically 100k interactions on Atari.

Proposed method

• Use a CNN to encode high-dimensional image observations into latent states, which are then embedded into the Transformer input sequence.
• Feed into the Transformer: latent states, actions, and predicted rewards as input tokens, enabling the model to attend flexibly across all three modalities.
• Employ the Transformer-XL architecture with relative positional encodings to model long-term dependencies efficiently and reduce computational cost at inference.
• Predict next latent state, reward, and discount factor using multi-layer perceptrons (MLPs) conditioned on the Transformer’s hidden state.
• Introduce a balanced KL divergence loss to fine-tune the trade-off between entropy and cross-entropy terms in the world model’s objective.
• Apply a thresholded entropy loss to stabilize policy entropy during training, reducing sensitivity to hyperparameter choices.
• Implement a balanced sampling procedure that prioritizes recent experience to correct data distribution shifts in the training dataset.

Experimental results

Research questions

• RQ1 Can a Transformer-based world model trained on only 100k environment interactions outperform existing model-free and model-based RL methods on the Atari 100k benchmark?
• RQ2 How does feeding back predicted rewards into the Transformer input affect the quality of the world model and downstream policy performance?
• RQ3 To what extent does the proposed balanced sampling strategy improve training stability and performance compared to uniform sampling?
• RQ4 Can the thresholded entropy loss stabilize policy training and reduce hyperparameter sensitivity across diverse Atari games?
• RQ5 Does the use of a Transformer-XL architecture enable better long-term dependency modeling than RNN-based world models while maintaining low inference cost?

Key findings

• The proposed TWM achieves state-of-the-art human normalized score on the Atari 100k benchmark, outperforming previous model-free and model-based methods.
• Conditioning the Transformer on predicted rewards leads to a significant performance boost in most games, as shown in ablation studies.
• The balanced sampling procedure improves policy performance by shifting focus toward more recent experience, reducing data distribution shift.
• The thresholded entropy loss stabilizes policy entropy during training, simplifying hyperparameter tuning and improving generalization across games.
• The balanced KL divergence loss enables fine-grained control over the entropy and cross-entropy trade-off, improving world model learning.
• The model achieves high performance with only 100k real environment interactions, demonstrating strong sample efficiency and computational efficiency at inference due to the absence of the Transformer during policy rollouts.