# Large-Scale Retrieval for Reinforcement Learning

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

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 27 of its 55 references that are in the map.
- First posted on arXiv 10 Jun 2022.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 27 of 27)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [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)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Neural Episodic Control](https://mlsanity.com/paper/1703.01988/) (2017)
- [Model-Free Episodic Control](https://mlsanity.com/paper/1606.04460/) (2016)
- [VLM: Task-agnostic Video-Language Model Pre-training for Video Understanding](https://mlsanity.com/paper/2105.09996/) (2021)
- [Playing hard exploration games by watching YouTube](https://mlsanity.com/paper/1805.11592/) (2018)
- [Online and Offline Reinforcement Learning by Planning with a Learned Model](https://mlsanity.com/paper/2104.06294/) (2021)
- [Control of Memory, Active Perception, and Action in Minecraft](https://mlsanity.com/paper/1605.09128/) (2016)
- [Unsupervised Predictive Memory in a Goal-Directed Agent](https://mlsanity.com/paper/1803.10760/) (2018)
- [Retrieval-Augmented Reinforcement Learning](https://mlsanity.com/paper/2202.08417/) (2022)
- [Episodic Memory Deep Q-Networks](https://mlsanity.com/paper/1805.07603/) (2018)
- [Towards mental time travel: a hierarchical memory for reinforcement learning agents](https://mlsanity.com/paper/2105.14039/) (2021)
- [Memory-based Parameter Adaptation](https://mlsanity.com/paper/1802.10542/) (2018)
- [Been There, Done That: Meta-Learning with Episodic Recall](https://mlsanity.com/paper/1805.09692/) (2018)
- [Rapid Task-Solving in Novel Environments](https://mlsanity.com/paper/2006.03662/) (2020)
- [DeepAveragers: Offline Reinforcement Learning by Solving Derived Non-Parametric MDPs](https://mlsanity.com/paper/2010.08891/) (2020)
- [Practical Kernel-Based Reinforcement Learning](https://mlsanity.com/paper/1407.5358/) (2014)
- [Fast deep reinforcement learning using online adjustments from the past](https://mlsanity.com/paper/1810.08163/) (2018)

## Built on by (most cited first, 24 of 24)

- [ExpeL: LLM Agents Are Experiential Learners](https://mlsanity.com/paper/2308.10144/) (2023)
- [Policy Regularization with Dataset Constraint for Offline Reinforcement Learning](https://mlsanity.com/paper/2306.06569/) (2023)
- [Retrieval-Augmented Decision Transformer: External Memory for In-context RL](https://mlsanity.com/paper/2410.07071/) (2024)
- [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)
- [Offline Reinforcement Learning with Imbalanced Datasets](https://mlsanity.com/paper/2307.02752/) (2023)
- [Memory-Enhanced Neural Solvers for Routing Problems](https://mlsanity.com/paper/2406.16424/) (2024)
- [Learning, Fast and Slow: A Goal-Directed Memory-Based Approach for Dynamic Environments](https://mlsanity.com/paper/2301.13758/) (2023)
- [Universal Retrieval for Multimodal Trajectory Modeling](https://mlsanity.com/paper/2506.22056/) (2025)
- [Using Memory-Based Learning to Solve Tasks with State-Action Constraints](https://mlsanity.com/paper/2303.04327/) (2023)
- [K-Nearest-Neighbor Resampling for Off-Policy Evaluation in Stochastic Control](https://mlsanity.com/paper/2306.04836/) (2023)
- [RANa: Retrieval-Augmented Navigation](https://mlsanity.com/paper/2504.03524/) (2025)
- [ALCAP: Alignment-Augmented Music Captioner](https://mlsanity.com/paper/2212.10901/) (2022)
- [Contrastive Distillation Is a Sample-Efficient Self-Supervised Loss Policy for Transfer Learning](https://mlsanity.com/paper/2212.11353/) (2022)
- [Heterogeneous Knowledge for Augmented Modular Reinforcement Learning](https://mlsanity.com/paper/2306.01158/) (2023)
- [Enhancing Sample Efficiency and Exploration in Reinforcement Learning through the Integration of Diffusion Models and Proximal Policy Optimization](https://mlsanity.com/paper/2409.01427/) (2024)
- [Linear Alignment of Vision-language Models for Image Captioning](https://mlsanity.com/paper/2307.05591/) (2023)
- [Heterogenous Memory Augmented Neural Networks](https://mlsanity.com/paper/2310.10909/) (2023)
- [Augmenting Unsupervised Reinforcement Learning with Self-Reference](https://mlsanity.com/paper/2311.09692/) (2023)
- [Generalizable Agent Modeling for Agent Collaboration-Competition Adaptation with Multi-Retrieval and Dynamic Generation](https://mlsanity.com/paper/2506.16718/) (2025)
- [Tactile Memory with Soft Robot: Robust Object Insertion via Masked Encoding and Soft Wrist](https://mlsanity.com/paper/2601.19275/) (2026)
- [Textual Planning with Explicit Latent Transitions](https://mlsanity.com/paper/2602.04557/) (2026)
- [ASH: Agents that Self-Hone in Long-Horizon Worlds](https://mlsanity.com/paper/2605.14211/) (2026)
- [Vector Search As Nearest Neighbor Matching: RAG-based Policy Learning in Causal Inference](https://mlsanity.com/paper/2607.18225/) (2026)

## Abstract (arXiv)

Effective decision making involves flexibly relating past experiences and relevant contextual information to a novel situation. In deep reinforcement learning (RL), the dominant paradigm is for an agent to amortise information that helps decision making into its network weights via gradient descent on training losses. Here, we pursue an alternative approach in which agents can utilise large-scale context sensitive database lookups to support their parametric computations. This allows agents to directly learn in an end-to-end manner to utilise relevant information to inform their outputs. In addition, new information can be attended to by the agent, without retraining, by simply augmenting the retrieval dataset. We study this approach for offline RL in 9x9 Go, a challenging game for which the vast combinatorial state space privileges generalisation over direct matching to past experiences. We leverage fast, approximate nearest neighbor techniques in order to retrieve relevant data from a set of tens of millions of expert demonstration states. Attending to this information provides a significant boost to prediction accuracy and game-play performance over simply using these demonstrations as training trajectories, providing a compelling demonstration of the value of large-scale retrieval in offline RL agents.