mlsanityData to 14 Oct 2026

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

arXiv 2411.11188 ↗First posted 17 Nov 2024cs.LGcited by 15

On the map: Deep reinforcement learning methods → Deep RL algorithms and representations

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.

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.

Builds on

93 of 109 references.

Closest ancestors · 5

Same direction, newest first

  1. 2024MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning14 Mar
  2. 2024Stop Regressing: Training Value Functions via Classification for Scalable Deep RL6 Mar
  3. 2024SplAgger: Split Aggregation for Meta-Reinforcement Learning5 Mar
  4. 2024Hierarchical Transformers are Efficient Meta-Reinforcement Learners9 Feb
  5. 2024Sharing Knowledge in Multi-Task Deep Reinforcement Learning17 Jan

Other references · 88

  1. 2017Attention Is All You Needcited by 60,361
  2. 2020Language Models are Few-Shot Learnerscited by 32,419
  3. 2017Proximal Policy Optimization Algorithmscited by 15,215
  4. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 13,028
  5. 2020Measuring Massive Multitask Language Understandingcited by 7,963
  6. 2016Layer Normalizationcited by 5,669
  7. 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,371
  8. 2017Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networkscited by 4,897
  9. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,057
  10. 2015Continuous control with deep reinforcement learningcited by 3,407
  11. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problemscited by 1,920
  12. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,546
38 more
  1. 2018Group Normalizationcited by 1,524
  2. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  3. 2020Gradient Surgery for Multi-Task Learningcited by 1,256
  4. 2023Mastering Diverse Domains through World Modelscited by 1,253
  5. 2019Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learningcited by 1,204
  6. 2019DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphscited by 903
  7. 2018IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architecturescited by 883
  8. 2017A Survey on Multi-Task Learningcited by 840
  9. 2022A Generalist Agentcited by 808
  10. 2017Rainbow: Combining Improvements in Deep Reinforcement Learningcited by 786
  11. 2019Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learningcited by 782
  12. 2017A Distributional Perspective on Reinforcement Learningcited by 765
  13. 2020AWAC: Accelerating Online Reinforcement Learning with Offline Datasetscited by 671
  14. 2023TD-MPC2: Scalable, Robust World Models for Continuous Controlcited by 512
  15. 2016RL2: Fast Reinforcement Learning via Slow Reinforcement Learningcited by 503
  16. 2019Leveraging Procedural Generation to Benchmark Reinforcement Learningcited by 483
  17. 2021Conflict-Averse Gradient Descent for Multi-task Learningcited by 459
  18. 2019Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variablescited by 389
  19. 2016Learning to reinforcement learncited by 377
  20. 2021Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learningcited by 372
  21. 2017Overcoming Exploration in Reinforcement Learning with Demonstrationscited by 340
  22. 2018Maximum a Posteriori Policy Optimisationcited by 328
  23. 2023Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Taskscited by 327
  24. 2021Randomized Ensembled Double Q-Learning: Learning Fast Without a Modelcited by 288
  25. 2020Critic Regularized Regressioncited by 275
  26. 2015Contextual Markov Decision Processescited by 220
  27. 2020Keep Doing What Worked: Behavioral Modelling Priors for Offline Reinforcement Learningcited by 215
  28. 2022Multi-Game Decision Transformerscited by 214
  29. 2022In-context Reinforcement Learning with Algorithm Distillationcited by 190
  30. 2020Revisiting Fundamentals of Experience Replaycited by 175
  31. 2019Soft Actor-Critic for Discrete Action Settingscited by 174
  32. 2021Multi-Task Reinforcement Learning with Context-based Representationscited by 174
  33. 2020Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Modelscited by 169
  34. 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 161
  35. 2018Multi-task Deep Reinforcement Learning with PopArtcited by 141
  36. 2023Supervised Pretraining Can Learn In-Context Reinforcement Learningcited by 139
  37. 2023Human-Timescale Adaptation in an Open-Ended Task Spacecited by 138
  38. 2023Q-Transformer: Scalable Offline Reinforcement Learning via Autoregressive Q-Functionscited by 131

50 most cited of 88 All 88 (JSON)

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 10; 2026: 5.

15 papers.

Sort:

By area:

  1. 2025Meta-World+: An Improved, Standardized, RL Benchmarkcited by 43
  2. 2025Vintix: Action Model via In-Context Reinforcement Learningcited by 10
  3. 2025Mixture-of-Experts Meets In-Context Reinforcement Learningcited by 8
  4. 2025Zero-Shot Reinforcement Learning Under Partial Observabilitycited by 6
  5. 2025ELMUR: External Layer Memory with Update/Rewrite for Long-Horizon RL Problemscited by 6
  6. 2025Scalable In-Context Q-Learningcited by 5
  7. 2025Human-Level Competitive Pok'emon via Scalable Offline Reinforcement Learning with Transformerscited by 4
  8. 2025Towards Provable Emergence of In-Context Reinforcement Learningcited by 3
  9. 2025Safe In-Context Reinforcement Learningcited by 3
  10. 2026Vintix II: Decision Pre-Trained Transformer is a Scalable In-Context Reinforcement Learnercited by 1
  11. 2026Latent Q-Barrier Shielding for Safe In-Context Reinforcement Learningcited by 1
  12. 2025Context Value Informed In-Context Reinforcement Learningcited by 0
3 more
  1. 2026MARL-GPT: Foundation Model for Multi-Agent Reinforcement Learningcited by 0
  2. 2026MATE: Solving Contextual Markov Decision Processes with Memory of Accumulated Transition Embeddingscited by 0
  3. 2026Towards Scalable Multi-Task Reinforcement Learning with Large Decision Modelscited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2024Scrutinize What We Ignore: Reining In Task Representation Shift Of Context-Based Offline Meta Reinforcement Learning
  2. 2026MAGE: Meta-Reinforcement Learning for Language Agents toward Strategic Exploration and Exploitation
  3. 2024Theoretical Analysis of Meta Reinforcement Learning: Generalization Bounds and Convergence Guarantees
  4. 2026Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning
  5. 2026Reinforcement Learning from Meta-Evaluation: Aligning Language Models Without Ground-Truth Labels
  6. 2024Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement
  7. 2025Scalable Multi-Objective and Meta Reinforcement Learning via Gradient Estimation
  8. 2025Multi-Task Reinforcement Learning Enables Parameter Scaling