# Retrieval-Augmented Reinforcement Learning

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

## Key facts

- Cited by 53 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 62 of its 96 references that are in the map.
- First posted on arXiv 17 Feb 2022.

## Placed in

- Retrieval-augmented generation / [Agentic RAG, agents and memory](https://mlsanity.com/rag/agentic-rag-agents-and-memory/)

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

- [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)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Big Self-Supervised Models are Strong Semi-Supervised Learners](https://mlsanity.com/paper/2006.10029/) (2020)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [Latent Retrieval for Weakly Supervised Open Domain Question Answering](https://mlsanity.com/paper/1906.00300/) (2019)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Interaction Networks for Learning about Objects, Relations and Physics](https://mlsanity.com/paper/1612.00222/) (2016)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [FeUdal Networks for Hierarchical Reinforcement Learning](https://mlsanity.com/paper/1703.01161/) (2017)
- [Maximum a Posteriori Policy Optimisation](https://mlsanity.com/paper/1806.06920/) (2018)
- [Leveraging Demonstrations for Deep Reinforcement Learning on Robotics Problems with Sparse Rewards](https://mlsanity.com/paper/1707.08817/) (2017)
- [ERNIE 3.0: Large-scale Knowledge Enhanced Pre-training for Language Understanding and Generation](https://mlsanity.com/paper/2107.02137/) (2021)
- [BabyAI: A Platform to Study the Sample Efficiency of Grounded Language Learning](https://mlsanity.com/paper/1810.08272/) (2018)
- [Distral: Robust Multitask Reinforcement Learning](https://mlsanity.com/paper/1707.04175/) (2017)
- [Accelerating Reinforcement Learning with Learned Skill Priors](https://mlsanity.com/paper/2010.11944/) (2020)
- [Recurrent Independent Mechanisms](https://mlsanity.com/paper/1909.10893/) (2019)
- [Keep Doing What Worked: Behavioral Modelling Priors for Offline Reinforcement Learning](https://mlsanity.com/paper/2002.08396/) (2020)
- [Search on the Replay Buffer: Bridging Planning and Reinforcement Learning](https://mlsanity.com/paper/1906.05253/) (2019)
- [RUDDER: Return Decomposition for Delayed Rewards](https://mlsanity.com/paper/1806.07857/) (2018)
- [Learning Continuous Control Policies by Stochastic Value Gradients](https://mlsanity.com/paper/1510.09142/) (2015)

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

- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (2023)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [From Play to Policy: Conditional Behavior Generation from Uncurated Robot Data](https://mlsanity.com/paper/2210.10047/) (2022)
- [Learning to Induce Causal Structure](https://mlsanity.com/paper/2204.04875/) (2022)
- [Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets](https://mlsanity.com/paper/2304.08742/) (2023)
- [PlayFusion: Skill Acquisition via Diffusion from Language-Annotated Play](https://mlsanity.com/paper/2312.04549/) (2023)
- [ASIF: Coupled Data Turns Unimodal Models to Multimodal Without Training](https://mlsanity.com/paper/2210.01738/) (2022)
- [Knowledge-Enhanced Dual-stream Zero-shot Composed Image Retrieval](https://mlsanity.com/paper/2403.16005/) (2024)
- [Memory, Benchmark & Robots: A Benchmark for Solving Complex Tasks with Reinforcement Learning](https://mlsanity.com/paper/2502.10550/) (2025)
- [Policy Regularization with Dataset Constraint for Offline Reinforcement Learning](https://mlsanity.com/paper/2306.06569/) (2023)
- [Large-Scale Retrieval for Reinforcement Learning](https://mlsanity.com/paper/2206.05314/) (2022)
- [Retrieval-Augmented Decision Transformer: External Memory for In-context RL](https://mlsanity.com/paper/2410.07071/) (2024)
- [Discrete Key-Value Bottleneck](https://mlsanity.com/paper/2207.11240/) (2022)
- [Natural Language Reinforcement Learning](https://mlsanity.com/paper/2402.07157/) (2024)
- [RetroAgent: From Solving to Evolving via Retrospective Dual Intrinsic Feedback](https://mlsanity.com/paper/2603.08561/) (2026)
- [SymbolicAI: A framework for logic-based approaches combining generative models and solvers](https://mlsanity.com/paper/2402.00854/) (2024)
- [Object-Centric Multiple Object Tracking](https://mlsanity.com/paper/2309.00233/) (2023)
- [Latent learning: episodic memory complements parametric learning by enabling flexible reuse of experiences](https://mlsanity.com/paper/2509.16189/) (2025)
- [Retrieval-Guided Reinforcement Learning for Boolean Circuit Minimization](https://mlsanity.com/paper/2401.12205/) (2024)
- [A Fine-tuning Enhanced RAG System with Quantized Influence Measure as AI Judge](https://mlsanity.com/paper/2402.17081/) (2024)
- [Hierarchical Planning for Complex Tasks with Knowledge Graph-RAG and Symbolic Verification](https://mlsanity.com/paper/2504.04578/) (2025)
- [Retrieved Sequence Augmentation for Protein Representation Learning](https://mlsanity.com/paper/2302.12563/) (2023)
- [ExpReS-VLA: Specializing Vision-Language-Action Models Through Experience Replay and Retrieval](https://mlsanity.com/paper/2511.06202/) (2025)
- [Semi-Parametric Inducing Point Networks and Neural Processes](https://mlsanity.com/paper/2205.11718/) (2022)
- [Offline Reinforcement Learning with Imbalanced Datasets](https://mlsanity.com/paper/2307.02752/) (2023)
- [From Standalone LLMs to Integrated Intelligence: A Survey of Compound Al Systems](https://mlsanity.com/paper/2506.04565/) (2025)
- [Using External Off-Policy Speech-To-Text Mappings in Contextual End-To-End Automated Speech Recognition](https://mlsanity.com/paper/2301.02736/) (2023)
- [Natural Language Reinforcement Learning](https://mlsanity.com/paper/2411.14251/) (2024)
- [IMWM: Intuition Models Complement World Models for Latent Planning](https://mlsanity.com/paper/2606.01626/) (2026)
- [RealDrive: Retrieval-Augmented Driving with Diffusion Models](https://mlsanity.com/paper/2505.24808/) (2025)

## Abstract (arXiv)

Most deep reinforcement learning (RL) algorithms distill experience into parametric behavior policies or value functions via gradient updates. While effective, this approach has several disadvantages: (1) it is computationally expensive, (2) it can take many updates to integrate experiences into the parametric model, (3) experiences that are not fully integrated do not appropriately influence the agent's behavior, and (4) behavior is limited by the capacity of the model. In this paper we explore an alternative paradigm in which we train a network to map a dataset of past experiences to optimal behavior. Specifically, we augment an RL agent with a retrieval process (parameterized as a neural network) that has direct access to a dataset of experiences. This dataset can come from the agent's past experiences, expert demonstrations, or any other relevant source. The retrieval process is trained to retrieve information from the dataset that may be useful in the current context, to help the agent achieve its goal faster and more efficiently. he proposed method facilitates learning agents that at test-time can condition their behavior on the entire dataset and not only the current state, or current trajectory. We integrate our method into two different RL agents: an offline DQN agent and an online R2D2 agent. In offline multi-task problems, we show that the retrieval-augmented DQN agent avoids task interference and learns faster than the baseline DQN agent. On Atari, we show that retrieval-augmented R2D2 learns significantly faster than the baseline R2D2 agent and achieves higher scores. We run extensive ablations to measure the contributions of the components of our proposed method.