mlsanityData to 7 Oct 2026

AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents

arXiv 2310.09971 ↗First posted 15 Oct 2023cs.LGcited by 36

Not on the map. Why?

Abstract (arXiv)

We introduce AMAGO, an in-context Reinforcement Learning (RL) agent that uses sequence models to tackle the challenges of generalization, long-term memory, and meta-learning. Recent works have shown that off-policy learning can make in-context RL with recurrent policies viable. Nonetheless, these approaches require extensive tuning and limit scalability by creating key bottlenecks in agents' memory capacity, planning horizon, and model size. AMAGO revisits and redesigns the off-policy in-context approach to successfully train long-sequence Transformers over entire rollouts in parallel with end-to-end RL. Our agent is scalable and applicable to a wide range of problems, and we demonstrate its strong performance empirically in meta-RL and long-term memory domains. AMAGO's focus on sparse rewards and off-policy data also allows in-context learning to extend to goal-conditioned problems with challenging exploration. When combined with a multi-goal hindsight relabeling scheme, AMAGO can solve a previously difficult category of open-world domains, where agents complete many possible instructions in procedurally generated environments.

Key facts. Cited by 36 papers in the map (data through 7 Oct 2026). Builds on 95 of its 115 references that are in the map. First posted on arXiv 15 Oct 2023.

Builds on

95 of 115 references.

References · 95

  1. 2017Attention Is All You Needcited by 59,932
  2. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 12,865
  3. 2016Layer Normalizationcited by 5,579
  4. 2017Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networkscited by 4,861
  5. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,016
  6. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  7. 2015Continuous control with deep reinforcement learningcited by 3,279
  8. 2018Addressing Function Approximation Error in Actor-Critic Methodscited by 2,302
  9. 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
  10. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problemscited by 1,873
  11. 2016OpenAI Gymcited by 1,684
  12. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
38 more
  1. 2023Mastering Diverse Domains through World Modelscited by 1,193
  2. 2018Soft Actor-Critic Algorithms and Applicationscited by 1,179
  3. 2019Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learningcited by 1,168
  4. 2017Hindsight Experience Replaycited by 1,096
  5. 2019Dota 2 with Large Scale Deep Reinforcement Learningcited by 1,084
  6. 2020Mastering Atari with Discrete World Modelscited by 990
  7. 2020On Layer Normalization in the Transformer Architecturecited by 906
  8. 2022A Generalist Agentcited by 802
  9. 2017Rainbow: Combining Improvements in Deep Reinforcement Learningcited by 774
  10. 2017Deep Reinforcement Learning that Matterscited by 727
  11. 2020AWAC: Accelerating Online Reinforcement Learning with Offline Datasetscited by 647
  12. 2016RL2: Fast Reinforcement Learning via Slow Reinforcement Learningcited by 490
  13. 2019Leveraging Procedural Generation to Benchmark Reinforcement Learningcited by 478
  14. 2022MineDojo: Building Open-Ended Embodied Agents with Internet-Scale Knowledgecited by 469
  15. 2015Deep Recurrent Q-Learning for Partially Observable MDPscited by 457
  16. 2019Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variablescited by 372
  17. 2016Learning to reinforcement learncited by 369
  18. 2017Overcoming Exploration in Reinforcement Learning with Demonstrationscited by 322
  19. 2019Learning Latent Plans from Playcited by 310
  20. 2020Agent57: Outperforming the Atari Human Benchmarkcited by 300
  21. 2021Randomized Ensembled Double Q-Learning: Learning Fast Without a Modelcited by 278
  22. 2020Critic Regularized Regressioncited by 269
  23. 2019Improving Sample Efficiency in Model-Free Reinforcement Learning from Imagescited by 268
  24. 2019Stabilizing Transformers for Reinforcement Learningcited by 262
  25. 2018Visual Reinforcement Learning with Imagined Goalscited by 248
  26. 2022The Primacy Bias in Deep Reinforcement Learningcited by 242
  27. 2015Contextual Markov Decision Processescited by 213
  28. 2022Multi-Game Decision Transformerscited by 212
  29. 2020The NetHack Learning Environmentcited by 196
  30. 2022Online Decision Transformercited by 194
  31. 2021Benchmarking the Spectrum of Agent Capabilitiescited by 193
  32. 2022In-context Reinforcement Learning with Algorithm Distillationcited by 186
  33. 2020Revisiting Fundamentals of Experience Replaycited by 171
  34. 2019Soft Actor-Critic for Discrete Action Settingscited by 161
  35. 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 154
  36. 2019Transformers without Tears: Improving the Normalization of Self-Attentioncited by 153
  37. 2023Bigger, Better, Faster: Human-level Atari with human-level efficiencycited by 145
  38. 2022Prompting Decision Transformer for Few-Shot Policy Generalizationcited by 143

50 most cited of 95 All 95 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 16; 2025: 15; 2026: 5.

36 papers.

Sort:

By area:

  1. 2024Survey on Large Language Model-Enhanced Reinforcement Learning: Concept, Taxonomy, and Methodscited by 114
  2. 2024Efficient Online Reinforcement Learning Fine-Tuning Need Not Retain Offline Datacited by 58
  3. 2024LLMs Are In-Context Bandit Reinforcement Learnerscited by 27
  4. 2025A Comprehensive Survey on Multi-Agent Cooperative Decision-Making: Scenarios, Approaches, Challenges and Perspectivescited by 27
  5. 2024Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modelingcited by 18
  6. 2024From Words to Actions: Unveiling the Theoretical Underpinnings of LLM-Driven Autonomous Systemscited by 17
  7. 2024Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglementcited by 16
  8. 2024AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformerscited by 14
  9. 2024XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learningcited by 13
  10. 2024Efficient Recurrent Off-Policy RL Requires a Context-Encoder-Specific Learning Ratecited by 10
  11. 2024Dynamics as Prompts: In-Context Learning for Sim-to-Real System Identificationscited by 10
  12. 2025Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learningcited by 9
24 more
  1. 2025Mixture-of-Experts Meets In-Context Reinforcement Learningcited by 7
  2. 2024N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needscited by 6
  3. 2025Zero-Shot Reinforcement Learning Under Partial Observabilitycited by 6
  4. 2025Text-to-Decision Agent: Offline Meta-Reinforcement Learning from Natural Language Supervisioncited by 5
  5. 2025Scalable In-Context Q-Learningcited by 5
  6. 2025Learning Policy Committees for Effective Personalization in MDPs with Diverse Taskscited by 4
  7. 2024NOLO: Navigate Only Look Oncecited by 3
  8. 2024ReLIC: A Recipe for 64k Steps of In-Context Reinforcement Learning for Embodied AIcited by 3
  9. 2024Enabling Adaptive Agent Training in Open-Ended Simulators by Targeting Diversitycited by 3
  10. 2025Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcement Learning for Out-of-Distribution Taskscited by 3
  11. 2025How Far Can LLMs Improve from Experience? Measuring Test-Time Learning Ability in LLMs with Human Comparisoncited by 3
  12. 2025In-Context Reinforcement Learning via Communicative World Modelscited by 3
  13. 2025Towards Provable Emergence of In-Context Reinforcement Learningcited by 2
  14. 2024Transformers as Game Players: Provable In-context Game-playing Capabilities of Pre-trained Modelscited by 1
  15. 2024Random Policy Enables In-Context Reinforcement Learning within Trust Horizonscited by 1
  16. 2025Disentangling Exploration of Large Language Models by Optimal Exploitationcited by 1
  17. 2025Enabling Autonomic Microservice Management through Self-Learning Agentscited by 1
  18. 2025Towards Monotonic Improvement in In-Context Reinforcement Learningcited by 0
  19. 2025Meta-RL with Bayesian Linear Task Modelscited by 0
  20. 2026ICPRL: Acquiring Physical Intuition from Interactive Controlcited by 0
  21. 2026Towards Scalable Multi-Task Reinforcement Learning with Large Decision Modelscited by 0
  22. 2026Dynamic Context Scheduling: Learning Beyond the Static Universecited by 0
  23. 2026Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillationcited by 0
  24. 2026From Static Policies to Adaptive Priors in Offline Reinforcement Learningcited by 0