# AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers

arXiv: https://arxiv.org/abs/2411.11188 · HTML: https://mlsanity.com/paper/2411.11188/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 93 of its 109 references that are in the map.
- First posted on arXiv 17 Nov 2024.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 93)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs](https://mlsanity.com/paper/1903.00161/) (2019)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [A Survey on Multi-Task Learning](https://mlsanity.com/paper/1707.08114/) (2017)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)

## Built on by (most cited first, 15 of 15)

- [Meta-World+: An Improved, Standardized, RL Benchmark](https://mlsanity.com/paper/2505.11289/) (2025)
- [Vintix: Action Model via In-Context Reinforcement Learning](https://mlsanity.com/paper/2501.19400/) (2025)
- [Mixture-of-Experts Meets In-Context Reinforcement Learning](https://mlsanity.com/paper/2506.05426/) (2025)
- [Zero-Shot Reinforcement Learning Under Partial Observability](https://mlsanity.com/paper/2506.15446/) (2025)
- [ELMUR: External Layer Memory with Update/Rewrite for Long-Horizon RL Problems](https://mlsanity.com/paper/2510.07151/) (2025)
- [Scalable In-Context Q-Learning](https://mlsanity.com/paper/2506.01299/) (2025)
- [Human-Level Competitive Pok'emon via Scalable Offline Reinforcement Learning with Transformers](https://mlsanity.com/paper/2504.04395/) (2025)
- [Towards Provable Emergence of In-Context Reinforcement Learning](https://mlsanity.com/paper/2509.18389/) (2025)
- [Safe In-Context Reinforcement Learning](https://mlsanity.com/paper/2509.25582/) (2025)
- [Vintix II: Decision Pre-Trained Transformer is a Scalable In-Context Reinforcement Learner](https://mlsanity.com/paper/2604.05112/) (2026)
- [Latent Q-Barrier Shielding for Safe In-Context Reinforcement Learning](https://mlsanity.com/paper/2605.25267/) (2026)
- [Context Value Informed In-Context Reinforcement Learning](https://mlsanity.com/paper/2509.23209/) (2025)
- [MARL-GPT: Foundation Model for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2604.05943/) (2026)
- [MATE: Solving Contextual Markov Decision Processes with Memory of Accumulated Transition Embeddings](https://mlsanity.com/paper/2605.17431/) (2026)
- [Towards Scalable Multi-Task Reinforcement Learning with Large Decision Models](https://mlsanity.com/paper/2606.24962/) (2026)

## Abstract (arXiv)

Language models trained on diverse datasets unlock generalization by in-context learning. Reinforcement Learning (RL) policies can achieve a similar effect by meta-learning within the memory of a sequence model. However, meta-RL research primarily focuses on adapting to minor variations of a single task. It is difficult to scale towards more general behavior without confronting challenges in multi-task optimization, and few solutions are compatible with meta-RL's goal of learning from large training sets of unlabeled tasks. To address this challenge, we revisit the idea that multi-task RL is bottlenecked by imbalanced training losses created by uneven return scales across different tasks. We build upon recent advancements in Transformer-based (in-context) meta-RL and evaluate a simple yet scalable solution where both an agent's actor and critic objectives are converted to classification terms that decouple optimization from the current scale of returns. Large-scale comparisons in Meta-World ML45, Multi-Game Procgen, Multi-Task POPGym, Multi-Game Atari, and BabyAI find that this design unlocks significant progress in online multi-task adaptation and memory problems without explicit task labels.