Source-linked AI summary
Transformer-based World Models Are Happy With 100k Interactions
Jan Robine, Marc Höftmann, Tobias Uelwer, Stefan Harmeling
TL;DR
Deep reinforcement learning often requires vastly more environment interaction than human learning, creating a need for more sample-efficient methods. The paper builds an autoregressive Transformer-XL world model over latent states, actions, and rewards, then uses imagined experience to train a policy. TWM outperforms previous model-free and model-based methods on the 26 games of Atari 100k while keeping inference costs low.
Problem
Deep reinforcement learning methods often require hundreds of millions of environment interactions, far more than humans need, limiting potential real-world applications.
Method
TWM uses a causally masked Transformer-XL world model that autoregressively processes latent states, actions, and rewards, and trains a policy in latent imagination.
Results
TWM outperforms previous model-free and model-based methods in human normalized score on the 26 games of the Atari 100k benchmark.
Takeaways & Limitations
Using the transformer only during training keeps computational costs low during inference while imagined trajectories support policy learning.
Takeaways & Limitations
The method is formulated for a partially observable Markov decision process with high-dimensional image observations and discrete actions.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Reinforcement learning remains far more interaction-hungry than human learning, motivating world models that generate training experience in imagination. This work proposes a Transformer-XL-based autoregressive world model that uses latent states, actions, and rewards to improve sample-efficient learning.
- Most deep reinforcement learning algorithms require hundreds of millions of environment interactions, while humans can achieve similar results with less than 1% of those interactions.
- Recent sample-efficient approaches include model-free auxiliary objectives and data augmentation, plus model-based planning and learning behaviors in imagination.
- Model-based world models learn environment dynamics and iteratively predict states and rewards to create new trajectories without further real-environment interaction.This can provide potentially indefinite training data for reinforcement learning.
- The proposed model uses a causally masked Transformer-XL to autoregressively process latent states, actions, and rewards while learning long-term dependencies efficiently.Transformer-XL adds recurrence and relative positional encodings, reducing dependence on absolute time steps.
- The work presents an autoregressive Transformer-XL world model and latent-imagination policy, with reward feedback, revised losses, entropy stabilization, and experience sampling as additional contributions.The policy does not require the transformer during inference, keeping execution computationally efficient.
2 METHOD
The method combines a Transformer-XL world model with latent imagination and policy learning for partially observable reinforcement-learning environments. It models observations, latent states, actions, rewards, and discounts autoregressively, while using balanced sampling and entropy control during training.
- Problem: The problem is formulated as learning a policy that maximizes expected discounted rewards in a partially observable environment.Observations, rewards, and episode ends are generated by unknown environment dynamics.
- World model: The observation model encodes image observations into compact stochastic latent states and reconstructs them with a decoder.The latent state is discrete, comprising 32 categorical variables with 32 categories.
- World model: The dynamics model uses a causally masked Transformer-XL to autoregressively process latent states, actions, and rewards while capturing long-term dependencies.Transformer-XL adds recurrence and relative positional encodings; its predictors estimate the next latent state, reward, and discount from the hidden state.
- Policy: The policy is trained on imagined trajectories with advantage actor-critic, using discounts predicted by the world model rather than a fixed discount.The policy uses an actor and critic, with advantages computed through Generalized Advantage Estimation.
- Policy: The entropy penalty acts only below a threshold, preserving a controlled fraction of policy entropy across games without epsilon-greedy action selection.The threshold is normalized by the maximum entropy of the categorical action distribution.
- Training: Balanced dataset sampling oversamples newer experience, whereas τ = ∞ recovers uniform sampling and smaller τ increases oversampling.Visitation counts are converted into sampling probabilities with a softmax temperature; uniform sampling can focus too heavily on early experience as the dataset grows slowly.
3 EXPERIMENTS
The experiments evaluate TWM on the Atari 100k benchmark using aggregate score statistics, performance profiles, imagined trajectories, and targeted ablations. TWM benefits from balanced sampling and reward conditioning, while attention analyses show use of multiple past modalities.
- Atari 100k limits interaction to 100K per game, corresponding to roughly 2 hours of gameplay and 500 times fewer frames than the usual 200 million.
- TWM is compared with five sample-efficient model-free and model-based competitors: SimPLe, DER, CURL, DrQ, and SPR.
- 3.1 RESULTS: TWM significantly improves all four aggregate metrics over previous approaches and brings the optimality gap closer to zero.The metrics are median, interquartile mean, mean, and optimality gap, computed on human normalized scores with confidence intervals.
- 3.1 RESULTS: The performance profile plots the fraction of runs above each human normalized score and displays pointwise 95% confidence bands.
- 3.2 ANALYSIS: Imagined trajectories show the actions and rewards produced by the world model, while attention analyses reveal strong focus on the current action, recent states, and high positive rewards.The transformer also attends to past states and rewards, whereas past actions are mostly ignored.
- 3.2 ANALYSIS: Balanced dataset sampling significantly improves performance over uniform sampling on evaluated games, while lowering the dynamics loss at the end of training.Uniform sampling can overemphasize early experience and contribute to overfitting later training data.
- 3.3 ABLATION STUDIES: Feeding predicted rewards back into the transformer can significantly increase performance, although some games show roughly unchanged scores.The ablation studies also examine thresholded entropy loss, history length, policy conditioning, and sample efficiency.
4 RELATED WORK
Related work spans world-model-based imagination, inference-time planning, and transformer-based sequence modeling for reinforcement learning. These approaches differ in how learned models generate experience, select actions, or condition policies on trajectories.
- World-model methods learn environment dynamics to generate experience for training, including VAE-RNN imagination in Ha and Schmidhuber and iterative model-policy training in SimPLe.
- Planning methods use learned models during inference to look ahead, with MuZero combining a sequence model of rewards and values with Monte-Carlo Tree Search.
- Decision Transformer and Trajectory Transformer apply transformers to offline trajectory datasets, with Decision Transformer conditioning on states, actions, and returns.
5 CONCLUSION
The paper presents TWM as a transformer-based world-model approach for reinforcement learning that improves Atari 100k performance while keeping inference costs low. It also identifies reward feedback, loss design, and balanced sampling as important components.
- TWM outperforms previous model-free and model-based methods on human normalized score across the 26 Atari 100k games.
- Using the transformer only during training keeps computational costs low when the learned policy runs in the real environment.
- The paper reports benefits from feeding predicted rewards into the transformer and introduces balanced cross-entropy, thresholded entropy, and balanced sampling procedures.
A.1 EXTENDED EXPERIMENTS
The extended experiments examine imagined trajectories, attention, sample efficiency, ablations, runtime, and policy throughput. They show that the world model supports diverse and long imagined behavior, while design choices such as thresholded entropy, longer histories, and Transformer-XL affect performance or efficiency.
- Imagined trajectories: The world model generates additional imagined trajectories for evaluation.
- Sample efficiency: After 50K interactions, TWM has a higher mean normalized score than previous sample-efficient methods, and it exceeds DER, CURL, and SimPLe after 25K interactions.
- World-model behavior: Stochastic next-state prediction produces multiple trajectories from the same starting state, while long-sequence imagination generalizes beyond the training sequence length ℓ=16.
- World-model behavior: Frame-stack visualizations show that the world model encodes and predicts object motion.
- Ablations: Without thresholded entropy, policy entropy is more likely to collapse or diverge; the threshold improves scores in the examined setting, while reducing history length to ℓ=4 lowers human normalized score.
- Runtime and throughput: TWM is more than 20 times faster than SimPLe, Transformer-XL imagination runs at approximately 39,000 samples/s versus 19,900 for vanilla imagination, and z-conditioned policy execution reaches 653 frames/s versus 213 with [z, h].
A.2 DERIVATION OF BALANCED CROSS-ENTROPY LOSS
This section derives the balanced cross-entropy formulation from a balanced KL-divergence objective. The derivation exposes separate control over cross-entropy terms and explicitly penalizes observation-encoder entropy.
- KL formulation: The balanced KL objective weights divergences between the stop-gradient encoder and predictor distributions in both directions.
- Derivation: Using D_KL(q ∥ p)=H(q,p)−H(q), the gradient becomes weighted cross-entropies between the observation encoder and state predictor plus an explicit encoder-entropy term.
- Derivation: The coefficients λ1, λ2, and λ3 provide control over the individual cross-entropy components and the entropy penalty.
- Implementation: Because the observation encoder and state predictor have separate parameters, the objective is split into two loss functions for the world model.
A.3 ADDITIONAL TRAINING DETAILS
The training details specify the algorithmic workflow, initialization, and model bookkeeping. World-model updates use sampled experience sequences, then latent states seed imagined trajectories for actor-critic training.
- Initialization: Pretraining the world model on pre-collected data provides a reasonable initialization for latent states while using the included 100K interactions.
- Training procedure: Algorithm 1 alternates world-model training from sampled observations, actions, rewards, and discounts with actor-critic training on imagined trajectories.
- World-model training: The world model encodes observations, reconstructs them, predicts dynamics, rewards, discounts, and latent states, and returns latent states for imagination.
- Actor-critic training: Actor-critic training starts from encoded latent states and iteratively imagines states, rewards, actions, and discounts for a fixed horizon.
- Experimental configuration: The experiments document their hyperparameters and model parameter counts in Tables 4 and 5.