Source-linked AI summary
Multi-Game Decision Transformers
Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, Igor Mordatch
TL;DR
The paper asks whether large transformer-based models trained on diverse offline experience can produce capable generalist reinforcement-learning agents. It trains a single decision-transformer model across Atari games and finds close-to-human aggregate performance, with scaling and fine-tuning trends resembling those in vision and language. The results favor decision transformers with guided generation among the compared multi-environment approaches.
Problem
The paper investigates whether one model can learn to act across many environments from diverse offline experience, extending large-scale generalist modeling to reinforcement learning.
Method
The method formulates offline reinforcement learning as autoregressive sequence modeling over observations, returns, actions, and rewards, with inference-time guided generation of expert-level returns and actions.
Results
A single agent achieves 126% of human-level performance simultaneously across all games, while decision-transformer models provide the best multi-environment performance and scaling properties among compared approaches.
Takeaways & Limitations
Performance scaling with model size and rapid fine-tuning to new games mirror trends observed in large-scale vision and language models.
Takeaways & Limitations
The conclusions are based largely on Atari, whose games share aligned action and observation spaces, and their applicability to other settings remains unclear.
Abstract
from arXiv · showhide
A longstanding goal of the field of AI is a method for learning a highly capable, generalist agent from diverse experience. In the subfields of vision and language, this was largely achieved by scaling up transformer-based models and training them on large, diverse datasets. Motivated by this progress, we investigate whether the same strategy can be used to produce generalist reinforcement learning agents. Specifically, we show that a single transformer-based model - with a single set of weights - trained purely offline can play a suite of up to 46 Atari games simultaneously at close-to-human performance. When trained and evaluated appropriately, we find that the same trends observed in language and vision hold, including scaling of performance with model size and rapid adaptation to new games via fine-tuning. We compare several approaches in this multi-game setting, such as online and offline RL methods and behavioral cloning, and find that our Multi-Game Decision Transformer models offer the best scalability and performance. We release the pre-trained models and code to encourage further research in this direction.
1 Introduction
The paper asks whether scaling transformer-based models across diverse offline experience can produce capable generalist reinforcement-learning agents. On 41 Atari games, a single agent reaches human-level aggregate performance, while decision transformers provide the strongest multi-environment performance and scaling properties.
- Training across environments with different dynamics, rewards, visuals, and embodiments has received less attention in reinforcement learning than in vision and language.
- The study trains one parameterized model on 41 Atari games using previously collected expert and non-expert trajectories, then evaluates it interactively.
- 126% human-level performance is achieved simultaneously across all games after offline training on expert and non-expert datasets.
- The study reports rapid fine-tuning to unseen games, performance scaling with model size, and faster training progress for larger models.
- Decision transformer models offer the best performance and scaling properties among the compared multi-environment approaches.The comparison includes online and offline reinforcement learning, contrastive representations, and behavioral cloning.
- The authors release pretrained models and code to support further research in generalist agents.
2 Related Work
Prior work established general agents and multi-task learning, but this paper targets a single agent acting across many environments with distinct dynamics, rewards, and embodiments. It also differs from concurrent transformer work by learning from diverse offline data without requiring expert-trajectory prompting at inference.
- The Atari suite was introduced as a testbed for empirically assessing agents designed for general competency, although early single-algorithm agents still required separate training and hand-crafted components.
- Related robotics work studies multiple tasks within the same or similar environments, whereas this paper addresses collections of environments with distinct dynamics, rewards, and agent embodiments.
- Unlike concurrent transformer work trained on near-optimal data and prompted with expert trajectories at inference, this approach learns from diverse expert and non-expert data and predicts returns for optimality-conditioned actions.
3 Method
The method formulates offline reinforcement learning as autoregressive sequence modeling over observations, returns, actions, and rewards. It trains on diverse Atari trajectories and uses inference-time return guidance to generate expert-level actions without changing standard training.
- 3.1 Reinforcement Learning as Sequence Modeling: Offline reinforcement learning is formulated as predicting each next sequence token conditioned on all preceding tokens.The sequence-modeling formulation follows decoder-only autoregressive models.
- 3.1 Reinforcement Learning as Sequence Modeling: The sequence preserves causal order by interleaving observation patches, target returns, actions, and rewards across timesteps.The target return represents the agent’s return for the remainder of the sequence.
- 3.1 Reinforcement Learning as Sequence Modeling: Returns, actions, and rewards are tokenized, and the model predicts their next discrete tokens using standard cross-entropy loss.This design permits return-distribution prediction and sampling rather than manually fixing an expert-level return at inference.
- 3.1 Reinforcement Learning as Sequence Modeling: Future observations are not predicted because their non-discrete image outputs would require additional model capacity.The paper leaves image-based forward prediction for future investigation.
- 3.2 Tokenization: Observations are divided into image patches with trainable positional encodings and linear projections into the token-embedding space.The experiments use 6x6 patches, each corresponding to 14x14 pixels.
- 3.3 Dataset: Training uses an existing Atari trajectory dataset spanning 41 training games and 5 held-out games for out-of-distribution generalization experiments.The trajectories include behaviors collected across DQN training progress, from non-expert to expert performance.
- 3.3 Dataset: The model learns from all available expert and non-expert behaviors rather than filtering or balancing the dataset.The motivation is that suboptimal behavior can provide diverse information about environments and poor decisions, while optimality is graded rather than binary.
- 3.4 Expert Action Inference: At inference, return logits are biased toward high returns, a target return is sampled, and an action is sampled conditioned on that return.This guided generation affects inference only; training remains standard next-token prediction.
4 Experiments
The experiments compare online and offline methods across multi-game Atari settings, testing performance, scaling, transfer, improvement over training data, action inference, and mixed-quality training data. Multi-Game Decision Transformers show the strongest overall performance and scaling, with benefits from larger models, pretraining, optimality conditioning, and expert-plus-non-expert data.
- 4.3 How do different online and offline methods perform in the multi-game regime?: Multi-game Decision Transformers come closest to single-game specialist performance, outperforming online non-transformer methods and offline non-transformer methods in the multi-game regime.The reported multi-game median scores are 68% for C51 DQN and 70% for a comparable Impala model.
- 4.4 How do different methods scale with model size?: Decision Transformer performance reliably increases across more than an order of magnitude of parameter scaling, while other methods saturate or grow more slowly.Larger models also reach higher in-game performance after observing the same number of tokens.
- 4.5 How effective are different methods at transfer to novel games?: DT pretraining performs best on held-out-game fine-tuning, and every pretrained method outperforms CQL trained from scratch on 1% of the held-out data.DT fine-tuning improves with model size, whereas CQL fine-tuning is inconsistent with model size.
- 4.8 Does training on expert and non-expert data bring benefits over expert-only training?: Training on full expert and non-expert data improves Decision Transformer, while expert-only training improves behavioral cloning; full-data DT outperforms expert-data BC.These results distinguish the data benefits observed for the two training approaches.
5 Conclusion
The results support scalable generalist agents in multi-game settings, while showing that online RL remains difficult and single-task agents can still outperform the generalists. The authors caution that these conclusions may not extend beyond Atari or to other forms of generalization, and emphasize the current agents’ limited scope.
- Large transformer-based models improve performance in the multi-game domain, with larger models and rapid fine-tuning mirroring trends in vision and language.
- Decision transformers achieve the best reported results despite learning through supervised sequence modeling rather than temporal-difference learning, policy gradients, or contrastive representation learning.
- Online RL algorithms struggle with the complexity of multi-game Atari training.
- The results are based largely on Atari, where action and observation spaces are aligned, so their applicability to other settings remains unclear.
- The agents are limited to self-contained video-game domains and are not intended for human interaction or use outside those domains.
Checklist
The checklist records affirmative responses for contribution and scope, limitations, societal impacts, reproducibility materials, training details, error bars, compute reporting, and asset citation. Human-subject and personally identifiable information items are marked not applicable.
- The paper states that its contributions and scope are described and that its limitations are discussed in Section 5.
- The paper states that potential negative societal impacts are discussed in Section 5.
- The authors report that code, data, reproduction instructions, training details, error bars, and compute resources are included.
- Existing assets are cited, while asset licensing is marked not applicable.
- Consent, personally identifiable or offensive content, human-participant instructions, participant risks, and compensation are marked not applicable.
B.1 Transformer network architecture
The transformer architecture tokenizes image observations, returns, actions, and rewards into sequences processed by a decoder with modified within-timestep attention. Experiments evaluate model scaling, fine-tuning, action and return sampling, augmentation, and Atari performance.
- Input representation: Observations are 84 × 84 grayscale images split into 36 non-overlapping 14 × 14 patches and projected into dmodel-dimensional tokens.
- Input representation: Returns are discretized into 120 buckets from −20 to 100, while rewards use the ternary values {−1, 0, +1}.
- Attention architecture: The model uses a standard transformer decoder with learned positional embeddings, and allows observation tokens within one timestep to attend to each other.
- Transfer evaluation: Fine-tuning reserves five games and uses 1% of the original dataset, roughly 500,000 transitions, for low-data transfer evaluation.
- Evaluation protocol: Atari evaluation uses 18 shared discrete actions, averages each game score over 16 rollout trials, and omits sticky actions.
- Data augmentation: Random cropping and random rotations work best among the evaluated image augmentations.
C Baseline Implementation Details
The baselines include behavioral cloning, C51 DQN, CQL, contrastive and masked representation learning, and UDRL variants. Comparisons indicate that Decision Transformer benefits from both its transformer architecture and its UDRL sequence-modeling formulation.
- Behavioral cloning: Behavioral cloning removes return tokens and directly predicts actions without return conditioning, while retaining the remaining Decision Transformer choices.
- Offline RL baselines: C51 DQN and CQL use offline temporal-difference learning, with CQL adding a behavioral cloning loss to the C51 categorical loss.
- Representation-learning baselines: CPC learns representations by contrasting consecutive states against randomly sampled states, whereas BERT uses masked self-prediction and ACL adds action prediction.
- UDRL baseline: UDRL with Impala replaces the Decision Transformer’s observation encoder with an Impala network while preserving return, action, and reward tokenizers.
- Transformer and convolution comparisons: Figure 10 shows Decision Transformer outperforming UDRL with Impala, while UDRL with Impala also outperforms CQL at matched Impala model sizes.
- Transformer and convolution comparisons: The authors attribute the observed trend partly to attention over the flat sequence of image patches, actions, and return tokens, while identifying the broader architectural explanation as an open question.
E Comparisons between methods using median human normalized scores
The paper reports median human-normalized scores for multi-game models across comparisons with specialist models, expert-data configurations, Decision Transformer, CQL, and UDRL. Aggregate-metric rankings remain unchanged across the expert-filtering configurations.
- Aggregate metric: Median human-normalized scores are reported alongside IQM because median can have high variability and remain unaffected by zero performance on nearly half the tasks.IQM was used for Figure 1; median scores are provided for completeness.
- Method comparisons: Figure 11 compares median human-normalized scores across 41 Atari games for single-game specialist models and generalists.Grey bars represent specialists, blue bars represent generalists, and single-game BCQ results come from Gulcehre et al. [25].
- Expert-data comparison: Expert-filtering configurations retain the same ranking across aggregate metrics, including median human-normalized scores in Figure 12.The figure compares 40M transformer models trained on full data with models trained only on expert data.
- Model-size comparison: Figure 13 compares how UDRL median human-normalized scores scale with model size against Decision Transformer and CQL using the Impala architecture.The comparison covers games in the training set.
F Details of Expert Dataset Generation
The expert dataset is constructed by filtering trajectories separately for each game according to episodic return. Only the top 10% of trajectories are retained, using game-specific 90th-percentile cutoffs.
- Filtering procedure: The expert dataset retains the top 10% of training trajectories from each game by episodic return.The filtered trajectories are used to produce the expert dataset for Section 4.8.
- Filtering procedure: Figure 14 plots rollout score histograms with unnormalized score density vertically and clipped-reward game scores horizontally.A red vertical line marks the 90th-percentile performance cutoff for each game.
- Filtering procedure: Rollouts exceeding each game’s 90th-percentile score threshold are included in the expert dataset.The threshold is determined separately for each game from the rollout score distribution.
G Effect of Model Size on Training Speed
The paper finds that larger transformer models reach higher scores after fewer multi-game training steps than smaller models. This reproduces the faster-training trend reported for large transformer-based language models.
- Training speed: The observed trend matches the reported behavior of large transformer-based language models reaching higher performance after observing similar numbers of tokens.The paper explicitly states that this trend holds in its setting as well.
- Training speed: Larger models reach higher scores per number of multi-game training steps taken, and therefore per number of tokens observed.Figure 15 illustrates this pattern on two example games.
- Training speed: Figure 15 shows example game scores for different model sizes as multi-game training progresses.The examples illustrate performance trajectories during training.
H Qualitative Attention Analysis
The qualitative analysis finds that Decision Transformer attention consistently focuses on image patches containing meaningful game entities. The appendix also reports raw scores for 41 training Atari games and visualizes attended patches across games.
- Attention patterns: Decision Transformer attention consistently targets observation patches containing meaningful game entities.Figure 16 visualizes selected attention heads and layers for various games.
- Attention patterns: Attended entities include the player character, its free movement space, non-player objects, and environmental features.The paper reports this pattern across attention heads and layers.
- Raw-score reference: Table 2 reports raw scores for 41 training Atari games for the best-performing multi-game model sizes.The table is identified as the raw-score reference for the training games.