# AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents

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

## Key facts

- Cited by 41 papers in the map (data through 14 Oct 2026).
- Builds on 95 of its 115 references that are in the map.
- First posted on arXiv 15 Oct 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [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)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Deep Reinforcement Learning that Matters](https://mlsanity.com/paper/1709.06560/) (2017)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [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)
- [MineDojo: Building Open-Ended Embodied Agents with Internet-Scale Knowledge](https://mlsanity.com/paper/2206.08853/) (2022)
- [Deep Recurrent Q-Learning for Partially Observable MDPs](https://mlsanity.com/paper/1507.06527/) (2015)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Overcoming Exploration in Reinforcement Learning with Demonstrations](https://mlsanity.com/paper/1709.10089/) (2017)

## Built on by (most cited first, 30 of 41)

- [Survey on Large Language Model-Enhanced Reinforcement Learning: Concept, Taxonomy, and Methods](https://mlsanity.com/paper/2404.00282/) (2024)
- [Efficient Online Reinforcement Learning Fine-Tuning Need Not Retain Offline Data](https://mlsanity.com/paper/2412.07762/) (2024)
- [LLMs Are In-Context Bandit Reinforcement Learners](https://mlsanity.com/paper/2410.05362/) (2024)
- [A Comprehensive Survey on Multi-Agent Cooperative Decision-Making: Scenarios, Approaches, Challenges and Perspectives](https://mlsanity.com/paper/2503.13415/) (2025)
- [Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling](https://mlsanity.com/paper/2406.00079/) (2024)
- [From Words to Actions: Unveiling the Theoretical Underpinnings of LLM-Driven Autonomous Systems](https://mlsanity.com/paper/2405.19883/) (2024)
- [Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement](https://mlsanity.com/paper/2410.11448/) (2024)
- [AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers](https://mlsanity.com/paper/2411.11188/) (2024)
- [XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning](https://mlsanity.com/paper/2406.08973/) (2024)
- [Efficient Recurrent Off-Policy RL Requires a Context-Encoder-Specific Learning Rate](https://mlsanity.com/paper/2405.15384/) (2024)
- [Dynamics as Prompts: In-Context Learning for Sim-to-Real System Identifications](https://mlsanity.com/paper/2410.20357/) (2024)
- [Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learning](https://mlsanity.com/paper/2510.19732/) (2025)
- [Mixture-of-Experts Meets In-Context Reinforcement Learning](https://mlsanity.com/paper/2506.05426/) (2025)
- [ReLIC: A Recipe for 64k Steps of In-Context Reinforcement Learning for Embodied AI](https://mlsanity.com/paper/2410.02751/) (2024)
- [N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs](https://mlsanity.com/paper/2411.01958/) (2024)
- [Zero-Shot Reinforcement Learning Under Partial Observability](https://mlsanity.com/paper/2506.15446/) (2025)
- [Text-to-Decision Agent: Offline Meta-Reinforcement Learning from Natural Language Supervision](https://mlsanity.com/paper/2504.15046/) (2025)
- [Scalable In-Context Q-Learning](https://mlsanity.com/paper/2506.01299/) (2025)
- [In-Context Reinforcement Learning via Communicative World Models](https://mlsanity.com/paper/2508.06659/) (2025)
- [Learning Policy Committees for Effective Personalization in MDPs with Diverse Tasks](https://mlsanity.com/paper/2503.01885/) (2025)
- [NOLO: Navigate Only Look Once](https://mlsanity.com/paper/2408.01384/) (2024)
- [Enabling Adaptive Agent Training in Open-Ended Simulators by Targeting Diversity](https://mlsanity.com/paper/2411.04466/) (2024)
- [Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcement Learning for Out-of-Distribution Tasks](https://mlsanity.com/paper/2502.02834/) (2025)
- [How Far Can LLMs Improve from Experience? Measuring Test-Time Learning Ability in LLMs with Human Comparison](https://mlsanity.com/paper/2506.14448/) (2025)
- [Towards Provable Emergence of In-Context Reinforcement Learning](https://mlsanity.com/paper/2509.18389/) (2025)
- [Random Policy Enables In-Context Reinforcement Learning within Trust Horizons](https://mlsanity.com/paper/2410.19982/) (2024)
- [Transformers as Game Players: Provable In-context Game-playing Capabilities of Pre-trained Models](https://mlsanity.com/paper/2410.09701/) (2024)
- [Disentangling Exploration of Large Language Models by Optimal Exploitation](https://mlsanity.com/paper/2501.08925/) (2025)
- [Enabling Autonomic Microservice Management through Self-Learning Agents](https://mlsanity.com/paper/2501.19056/) (2025)
- [ICPRL: Acquiring Physical Intuition from Interactive Control](https://mlsanity.com/paper/2603.13295/) (2026)

## 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.