# Retrieval-Augmented Decision Transformer: External Memory for In-context RL arXiv: https://arxiv.org/abs/2410.07071 ยท HTML: https://mlsanity.com/paper/2410.07071/ ## Key facts - Cited by 21 papers in the map (data through 14 Oct 2026). - Placed in Agentic RAG, agents and memory, one of 15 directions in Retrieval-augmented generation. - Builds on 91 of its 110 references that are in the map. - First posted on arXiv 9 Oct 2024. ## Placed in - Retrieval-augmented generation / [Agentic RAG, agents and memory](https://mlsanity.com/rag/agentic-rag-agents-and-memory/) - Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/) ## Builds on (most cited first, 30 of 91) - [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017) - [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018) - [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020) - [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019) - [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017) - [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019) - [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022) - [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020) - [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017) - [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023) - [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019) - [Generative Agents: Interactive Simulacra of Human Behavior](https://mlsanity.com/paper/2304.03442/) (2023) - [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023) - [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023) - [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023) - [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017) - [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020) - [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022) - [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020) - [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020) - [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021) - [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021) - [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018) - [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020) - [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019) - [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017) - [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018) - [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022) - [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018) - [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022) ## Built on by (most cited first, 21 of 21) - [Towards Agentic RAG with Deep Reasoning: A Survey of RAG-Reasoning Systems in LLMs](https://mlsanity.com/paper/2507.09477/) (2025) - [LLMs are Greedy Agents: Effects of RL Fine-tuning on Decision-Making Abilities](https://mlsanity.com/paper/2504.16078/) (2025) - [Large language models for artificial general intelligence (AGI): A survey of foundational principles and approaches](https://mlsanity.com/paper/2501.03151/) (2025) - [A Survey of Continual Reinforcement Learning](https://mlsanity.com/paper/2506.21872/) (2025) - [REGENT: A Retrieval-Augmented Generalist Agent That Can Act In-Context in New Environments](https://mlsanity.com/paper/2412.04759/) (2024) - [A Large Recurrent Action Model: xLSTM enables Fast Inference for Robotics Tasks](https://mlsanity.com/paper/2410.22391/) (2024) - [Episodic memory in AI agents poses risks that should be studied and mitigated](https://mlsanity.com/paper/2501.11739/) (2025) - [Yes, Q-learning Helps Offline In-Context RL](https://mlsanity.com/paper/2502.17666/) (2025) - [Memory-Augmented Transformers: A Systematic Review from Neuroscience Principles to Enhanced Model Architectures](https://mlsanity.com/paper/2508.10824/) (2025) - [N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs](https://mlsanity.com/paper/2411.01958/) (2024) - [Behavioral Exploration: Learning to Explore via In-Context Adaptation](https://mlsanity.com/paper/2507.09041/) (2025) - [BiTrajDiff: Bidirectional Trajectory Generation with Diffusion Models for Offline Reinforcement Learning](https://mlsanity.com/paper/2506.05762/) (2025) - [IMWM: Intuition Models Complement World Models for Latent Planning](https://mlsanity.com/paper/2606.01626/) (2026) - [Disentangling Exploration of Large Language Models by Optimal Exploitation](https://mlsanity.com/paper/2501.08925/) (2025) - [In-Context Compositional Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2509.24067/) (2025) - [In-Context Reinforcement Learning under Non-Stationarity: A Survey](https://mlsanity.com/paper/2607.11906/) (2026) - [Learning in Context, Guided by Choice: A Reward-Free Paradigm for Reinforcement Learning with Transformers](https://mlsanity.com/paper/2602.08244/) (2026) - [DRIVE: Distributional and Retrieval-Augmented Bidding with Value Evaluation](https://mlsanity.com/paper/2606.14192/) (2026) - [Breaking Feedback-Blindness: Utility-Augmented Transformer for Sequential Decision Making](https://mlsanity.com/paper/2607.18910/) (2026) - [Downstream-Aware Context Selection for Online In-Context Reinforcement Learning](https://mlsanity.com/paper/2609.33166/) (2026) - [Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation](https://mlsanity.com/paper/2609.35333/) (2026) ## Abstract (arXiv) In-context learning (ICL) is the ability of a model to learn a new task by observing a few exemplars in its context. While prevalent in NLP, this capability has recently also been observed in Reinforcement Learning (RL) settings. Prior in-context RL methods, however, require entire episodes in the agent's context. Given that complex environments typically lead to long episodes with sparse rewards, these methods are constrained to simple environments with short episodes. To address these challenges, we introduce Retrieval-Augmented Decision Transformer (RA-DT). RA-DT employs an external memory mechanism to store past experiences from which it retrieves only sub-trajectories relevant for the current situation. The retrieval component in RA-DT does not require training and can be entirely domain-agnostic. We evaluate the capabilities of RA-DT on grid-world environments, robotics simulations, and procedurally-generated video games. On grid-worlds, RA-DT outperforms baselines, while using only a fraction of their context length. Furthermore, we illuminate the limitations of current in-context RL methods on complex environments and discuss future directions. To facilitate future research, we release datasets for four of the considered environments.