Retrieval-Augmented Reinforcement Learning

Anirudh GoyalAbram L. FriesenAndrea BaninoTheophane WeberNan Rosemary KeAdrià Puigdomènech BadiaArthur GuezMehdi MirzaPeter Conway HumphreysKsenia Konyushkova

article2022ICML78 citations

Proposes an architecture that augments reinforcement learning agents with an attention-based neural retrieval process, allowing policies to dynamically query past trajectories to reduce task interference in multi-task offline settings and significantly accelerate learning in online Atari benchmarks.

Listen

Modern artificial intelligence systems using reinforcement learning typically compress all previous experience into fixed network parameters via repetitive gradient updates. This standard paradigm faces critical limitations: it demands excessive computation, requires numerous updates to integrate new experiences, and restricts agent behavior based on fixed model capacity. Consequently, agents frequently struggle to recall specific, relevant past situations when facing complex or multi-task scenarios.

The article demonstrates an alternative approach called Retrieval-Augmented Reinforcement Learning, which pairs an agent with a dedicated retrieval mechanism that dynamically fetches relevant contextual information from external datasets of past trajectories during decision-making.

The authors evaluate this framework across standard video game benchmarks (Atari) and complex multi-task offline environments (Gridroboman, BabyAI, and CausalWorld continuous control). The retrieval process uses attention mechanisms over summarized past experiences and regulates information exchange using an information bottleneck to avoid overload. Crucially, the retrieval process and the decision-making agent are parameterized as separate interacting components.

The findings establish that retrieval augmentation significantly enhances agent efficiency and capabilities. On Atari games, the retrieval-augmented agent achieved an 11.32% improvement in mean human-normalized score compared to a leading baseline, demonstrating particular strength in tasks requiring extended planning. In offline multi-task environments, retrieval-augmented models prevented task interference as the number of tasks scaled from 10 to 30, where standard architectures degraded. On compositional language-directed benchmarks, retrieval increased task success rates from 45% to 74% with multi-task data, as the architecture successfully reused sub-task experiences.

These results demonstrate that externalizing memory retrieval relieves network capacity constraints, reduces data collection burdens, and improves performance in multi-task operations. This shift suggests organizations can deploy more agile and data-efficient decision models by decoupling memory access from core parametric networks.

Organizations evaluating reinforcement learning for complex multi-task or robotic control should consider retrieval-augmented architectures over scaling model parameters alone. Future efforts should test retrieval across decentralized multi-agent deployments and validate adaptations in zero-shot or few-shot environments.

Confidence in these findings is high for simulated control and standard benchmarks across multiple tested random seeds. However, decision-makers should exercise caution because computational overhead increases when querying large trajectory batches, and the methodology has not yet been demonstrated in real-world physical systems or massive multi-agent competitive environments.

deep reinforcement learning demonstrations learning agents Retrieval-augmented reinforcement learning offline reinforcement learning episodic memory

Goyal et al (2022).pdf

Paper: Prioritized Experience Replay, Tom Schaul et al. (2016). Prioritized Experience Replay shows how selecting informative past transitions can improve RL learning, a useful precursor to retrieving relevant experience for decisions.
Paper: Experience Replay for Continual Learning, David Rolnick et al. (2018). Experience Replay for Continual Learning establishes how replaying stored experience can protect performance across sequential tasks, grounding the source’s concern with multi-task interference.
Paper: Deep Recurrent Q-Learning for Partially Observable MDPs, Matthew Hausknecht et al. (2015). Deep Recurrent Q-Learning explains how RL agents use memory to exploit past observations, clarifying the alternative to the source’s externalized retrieval.
Paper: Meta-Learning with Memory-Augmented Neural Networks, Adam Santoro et al. (2016). Meta-Learning with Memory-Augmented Neural Networks develops content-based access to external memory, a key conceptual foundation for retrieving information separately from model parameters.

Paper: Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning, Bowen Jin et al. (2025). Search-R1 extends retrieval-augmented decision-making into language-model RL, training agents to make and use multi-step external searches through reinforcement learning.

Cover for Retrieval-Augmented Reinforcement Learning

Abstract

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

Topics

deep reinforcement learning demonstrations learning agents retrieval performance software agents value functions reward function imitation learning

Keywords

Retrieval-augmented reinforcement learning offline reinforcement learning episodic memory episodic control experience replay recurrent DQN multi-task replay

1. Introduction
2. Retrieval-Augmented Agents
2.1. Retrieval-augmented agent
2.2. Retrieval batch sampling and pre-processing
2.3. Retrieving contextual information
3. Experimental Results
3.1. Atari: Single-task off-policy RL
3.1.1. Ablations and Analysis
3.2. Gridroboman: Multi-task offline RL with a task-specific retrieval dataset
3.3. BabyAI: Multi-task offline RL with a multi-task retrieval dataset
3.4. CausalWorld: Multi-task offline continuous control
4. Related Work
5. Conclusion
6. Acknowledgements
References
A. Appendix
A.1. Extended Relevant Work
A.2. Information Theoretic Formulation
A.3. Atari: Implementation details and raw scores for R2D2 and RA-R2D2
A.3.2. Details on Ablations
A.3.3. Additional Atari Ablations
A.3.4. Episodic Control Baseline
A.3.5. Retrieval Augmented RL Helping Frostbite
A.3.6. Auxiliary Losses Implementation
A.4. Gridroboman environment
A.4.1. Gridroboman Experiment Setup
A.4.2. Gridroboman Hyperparameters
A.4.3. Gridroboman Evaluation Curves When Training on All 30 Tasks
A.5. BabyAI environment
A.6. CausalWorld environment

Knowls

Knowl 1 — Retrieval-augmented agents keep experience directly available to behavior
model/method
A retrieval-augmented agent (R2A) has two interacting neural processes with separate internal states. At time $t$ , an encoder maps the agent input $x_t$ to an agent representation $s_t$ . A retrieval process takes $s_t$ , its previous memory state $m_{t-1}$ , and a dataset $B$ of experience trajectories, then returns updated memory $m_t$ and retrieved context $u_t$ . The agent process uses $s_t$ and $u_t$ to form its action or value predictions. The dataset may contain the agent’s own experience, demonstrations, or experience from other tasks or agents. Unlike a conventional parametric policy or value function that must absorb experience into its weights, R2A can condition behavior on information retrieved directly from this dataset. The retrieval process is learned to provide useful information without explicit knowledge of the agent’s policy, and the agent and retrieval processes can shape one another’s representations.
Knowl 2 — Bidirectional trajectory summaries make past and future information retrievable
model/method
R2A preprocesses sampled trajectories by encoding their observations with the agent’s encoder and then applying learned forward and backward sequence summarizers. For trajectory $i$ and position $j$ , the forward summary $h_{i,j}$ represents information from the trajectory’s beginning through position $j$ , while the backward summary $b_{i,j}$ represents information from position $j$ through the trajectory’s end. The two summaries let retrieval match a current query to contextual information in a trajectory while returning information that includes its future as well as its past. In the experiments, the summarizers were GRUs. Their representations were trained with auxiliary action or policy, value, and reward prediction losses; a BERT-style masked-representation loss was also tested. These auxiliary objectives encourage summaries to preserve information useful for downstream decisions.
Knowl 3 — Slot-based attention retrieves selected trajectory contexts and updates the agent
algorithm
At each timestep, R2A uses $n_f$ memory slots. Each slot combines the current agent representation $s_t$ with its previous slot state using a GRU, then produces a query. For every summarized state in a sampled retrieval batch, the method projects its forward summary $h_{i,j}$ into a key and scores the key against each slot’s query with scaled dot-product attention. It aggregates state scores by trajectory, keeps the $k_{\text{traj}}$ highest-scoring trajectories, then selects the $k_{\text{states}}$ highest-scoring states within those trajectories. The selected attention weights are renormalized, and the slot retrieves a weighted average of projections of the corresponding backward summaries $b_{i,j}$ , yielding a representation $g_t^k$ for slot $k$ . In the experiments, $k_{\text{traj}}=10$ and $k_{\text{states}}=10$ .

Each slot samples a bottlenecked representation $z_t^k$ from the retrieved information, adds it to the slot’s context-updated state, and then participates in a self-attention update with the other slots. Finally, the agent attends to the resulting slot representations: its current state supplies the query, and the retrieved representations supply keys and values. The attended result is added to the agent representation and used in the agent’s predictions; the value function is also conditioned on retrieved information. The slots maintain their own state and can retrieve different material from the same batch. In the reported experiments, R2A used eight slots for Atari and four for Gridroboman and BabyAI.
Knowl 4 — An information bottleneck limits dependence on retrieved context
theoretical result
R2A regularizes the information its policy takes from retrieved context. Let $S$ be the agent state, $G$ the retrieval information conditioned on $S$ , $A$ the action, and $\pi_\theta(A\mid S,G)$ the policy. The paper’s objective trades off expected reward against the conditional mutual information between actions and retrieved information given the agent state:
$J(\theta)=\mathbb{E}_{\pi_\theta}[r]-\beta I(A;G\mid S) =\mathbb{E}_{\pi_\theta}\!\left[r-\beta D_{\mathrm{KL}}\!\left(\pi_\theta(A\mid S,G)\,\|\,\pi_0(A\mid S)\right)\right],$
where $r$ is reward, $\beta>0$ controls the regularization strength, $D_{\mathrm{KL}}$ is Kullback–Leibler divergence, and $\pi_0(A\mid S)$ is the default policy with retrieval information marginalized out. The objective encourages the agent to rely on retrieved information when it improves reward, rather than letting retrieval arbitrarily alter its behavior. The paper also gives an upper bound on the conditional information in a latent retrieval representation $Z$ : $I(Z;G\mid S)\leq \sum_s p(s)\sum_g p(g\mid s)D_{\mathrm{KL}}(p(Z\mid s,g)\|r(Z))$ , for state and retrieval distributions $p(s)$ and $p(g\mid s)$ and a reference distribution $r(Z)$ . Operationally, the implementation samples $z_t^k$ from a distribution conditioned on the retrieved vector and penalizes its divergence from a reference distribution conditioned on the slot memory.
Knowl 5 — Retrieval improves Atari R2D2 performance and sample efficiency
empirical result
In single-task Atari 2600 experiments, retrieval-augmented R2D2 (RA-R2D2) retrieved from the agent’s current replay buffer and was compared with vanilla R2D2 over 2 billion environment steps. Across the reported games, RA-R2D2 increased mean human-normalized score by $11.32\pm1.2\%$ relative to the baseline. The experiment used a recurrent R2D2 agent and a separate slot-based retrieval process with eight slots; the forward and backward trajectory summarizers were 512-unit GRUs. The improvement was especially large on Frostbite, which the authors hypothesize benefits from retrieving information about temporally distant states. The result supports using an agent’s own replay buffer as a useful source of retrieval context, including at large data scales.
Knowl 6 — Atari ablations identify the contribution of retrieval design choices
empirical result
Ablations of RA-R2D2 showed that its gains were associated with retrieving useful experience through a separate retrieval process, rather than merely adding parameters or computation. Replacing the retrieval process’s separate memory state with a query computed directly from the agent state performed no better than baseline R2D2; the authors report that this non-parametric direct-access variant could hurt performance. Removing access to the retrieval batch while retaining an updated slot process also failed to reproduce the retrieval benefit. Shortening summarized trajectories from 80 to 5 timesteps reduced performance, with a notable decrease on Frostbite. Adding a BERT-style masked auxiliary loss to the action, reward, and value prediction losses could improve performance further, while using only the BERT-style loss performed worse than the combined-loss version but remained better than baseline R2D2. On five games where RA-R2D2 initially performed worst relative to baseline, independently varying the numbers of selected trajectories and states over $\{5,10,20\}$ could improve performance.
Knowl 7 — Task-specific retrieval helps offline multi-task Gridroboman learning
empirical result
Gridroboman is a 7-by-7 grid-world manipulation environment with a robot and three colored objects, designed with 30 tasks such as touching, lifting, moving, and stacking objects. The offline dataset was generated by training a DQN separately on each task and recording 100,000 episodes per task. A single DQN or retrieval-augmented DQN (RA-DQN) was trained on data from 10, 20, or 30 tasks; at evaluation, the retrieval dataset contained experience only from the task being evaluated. With 10 training tasks, DQN and RA-DQN performed similarly. As the number of training tasks increased to 20 and 30, RA-DQN learned more effectively than DQN. The authors interpret this pattern as evidence that directly retrieving task-relevant experience can help a shared agent learn across tasks without requiring it to encode all task-specific detail in its parametric model.

Knowl 8 — Multi-task retrieval raises BabyAI success, especially for compositional tasks

data/table

On the 40 partially observable BabyAI levels, a recurrent DQN (RDQN) and retrieval-augmented recurrent DQN (RA-RDQN) were evaluated with 50K or 200K training trajectories per task. Retrieval buffers contained either only the evaluated task’s data or data from multiple tasks. The values below are mean success rates, averaged over three random seeds and reported with standard errors. Retrieval improved on the RDQN baseline, and the multi-task retrieval buffer with the information bottleneck (IB) achieved the highest success rate at both data sizes. Removing IB from the multi-task retrieval version reduced its success rate.

Method	Success Rate (50K)	Success Rate (200K)
RDQN	$32\%\pm4\%$	$45\%\pm6\%$
RA-RDQN (single-task retrieval buffer)	$48\%\pm4\%$	$64\%\pm5\%$
RA-RDQN (multi-task retrieval buffer, without IB)	$47\%\pm3\%$	$59\%\pm6\%$
RA-RDQN (multi-task retrieval buffer)	$55\%\pm5\%$	$74\%\pm3\%$

In a separate analysis, the agent retrieved information from other tasks in 54% of retrievals while solving compositional tasks, compared with 21% for non-compositional tasks. This association is consistent with the possibility that experience from related subtasks supports compositional tasks; the analysis does not establish that this retrieval pattern caused the performance difference.

Knowl 9 — Retrieval-augmented behavior cloning improves CausalWorld success

empirical result

In offline continuous-control experiments on five CausalWorld object-manipulation tasks—Pushing, Picking, Pick and Place, Stacking2, and Stacked Blocks—the underlying learner was behavior cloning (BC). At the reported 50K-data condition, retrieval-augmented BC (RA-BC) outperformed BC with either a task-specific or multi-task retrieval buffer. Success rates are means over three random seeds with standard errors:

The highest reported success rate used a multi-task retrieval buffer, showing that retrieval augmentation benefited behavior cloning in this benchmark even when the buffer included data from other tasks.

Knowl 10 — Evaluated scope leaves several extensions untested
limitation
The paper identifies several directions that its experiments do not resolve. R2A queries only a subset of the retrieval dataset, which limits how broadly the method can search. The study does not evaluate few-shot adaptation to a new task using a pre-trained trajectory encoder, nor does it establish how the approach scales to more complex multi-agent settings. The authors also leave open whether separately training the retrieval and agent processes with different objectives would be preferable to end-to-end training. These are proposed future investigations, not demonstrated shortcomings in the reported benchmarks.

Coverage note — The appendix’s per-game Atari score list, full optimizer settings, and exhaustive Gridroboman task definitions are omitted because they provide supporting experimental detail rather than distinct main contributions.

References

1.Abdolmaleki, A., Springenberg, J. T., Tassa, Y., Munos, R., Heess, N., and Riedmiller, M. Maximum a posteriori policy optimisation. In International Conference on Learning Representations, 2018.
2.Ahmed, O., Träuble, F., Goyal, A., Neitz, A., Bengio, Y., Schölkopf, B., Wüthrich, M., and Bauer, S. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. In International Conference on Learning Representations, 2020.
3.Alemi, A. A., Fischer, I., Dillon, J. V., and Murphy, K. Deep variational information bottleneck. In International Conference on Learning Representations, 2016.
4.Allen, J. F. and Koomen, J. A. Planning using a temporal world model. In Proceedings of the Eighth International Joint Conference on Artificial Intelligence, pp. 741–747, 1983.
5.Arjona-Medina, J. A., Gillhofer, M., Widrich, M., Unterthiner, T., Brandstetter, J., and Hochreiter, S. RUDDER: Return decomposition for delayed rewards. In Advances in Neural Information Processing Systems, volume 32, 2019.
6.Banino, A., Badia, A. P., Walker, J. C., Scholtes, T., Mitrovic, J., and Blundell, C. CoBERL: Contrastive BERT for reinforcement learning. In ICML 2021 Workshop on Unsupervised Reinforcement Learning, 2021.
7.Battaglia, P. W., Pascanu, R., Lai, M., Rezende, D., and Kavukcuoglu, K. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 29, 2016.
8.Battaglia, P. W., Hamrick, J. B., Bapst, V., Sanchez-Gonzalez, A., Zambaldi, V., Malinowski, M., Tacchetti, A., Raposo, D., Santoro, A., Faulkner, R., et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.01261, 2018.
9.Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253–279, 2013.
10.Blundell, C., Uria, B., Pritzel, A., Li, Y., Ruderman, A., Leibo, J. Z., Rae, J., Wierstra, D., and Hassabis, D. Model-free episodic control. arXiv preprint arXiv:1606.04460, 2016.
11.Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Rutherford, E., Millican, K., Driessche, G. v. d., Lespiau, J.-B., Damoc, B., Clark, A., et al. Improving language models by retrieving from trillions of tokens. arXiv preprint arXiv:2112.04426, 2021.
12.Chen, M., Radford, A., Child, R., Wu, J., Jun, H., Luan, D., and Sutskever, I. Generative pretraining from pixels. In International Conference on Machine Learning, pp. 1691–1703. PMLR, 2020a.
13.Chen, T., Kornblith, S., Swersky, K., Norouzi, M., and Hinton, G. Big self-supervised models are strong semi-supervised learners. arXiv preprint arXiv:2006.10029, 2020b.
14.Chevalier-Boisvert, M., Bahdanau, D., Lahlou, S., Willems, L., Saharia, C., Nguyen, T. H., and Bengio, Y. Babyai: A platform to study the sample efficiency of grounded language learning. In International Conference on Learning Representations, 2018.
15.Cover, T. M. Elements of information theory. John Wiley & Sons, 1999.
16.Czarnecki, W. M., Pascanu, R., Osindero, S., Jayakumar, S., Swirszcz, G., and Jaderberg, M. Distilling policy distillation. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1331–1340. PMLR, 2019.
17.Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, Ł. Universal transformers. arXiv preprint arXiv:1807.03819, 2018.
18.Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
19.Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
20.Ernst, D., Geurts, P., and Wehenkel, L. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6:503–556, 2005.
21.Eysenbach, B., Salakhutdinov, R., and Levine, S. Search on the replay buffer: Bridging planning and reinforcement learning. arXiv preprint arXiv:1906.05253, 2019.
22.Florensa, C., Duan, Y., and Abbeel, P. Stochastic neural networks for hierarchical reinforcement learning. arXiv preprint arXiv:1704.03012, 2017.
23.Fortunato, M., Tan, M., Faulkner, R., Hansen, S., Badia, A. P., Buttimore, G., Deck, C., Leibo, J. Z., and Blundell, C. Generalization of reinforcement learners with working and episodic memory. arXiv preprint arXiv:1910.13406, 2019.
24.Frans, K., Ho, J., Chen, X., Abbeel, P., and Schulman, J. Meta learning shared hierarchies. arXiv preprint arXiv:1710.09767, 2017.
25.Galashov, A., Jayakumar, S. M., Hasenclever, L., Tirumala, D., Schwarz, J., Desjardins, G., Czarnecki, W. M., Teh, Y. W., Pascanu, R., and Heess, N. Information asymmetry in KL-regularized RL. arXiv preprint arXiv:1905.01240, 2019.
26.Goyal, A., Brakel, P., Fedus, W., Singhal, S., Lillicrap, T., Levine, S., Larochelle, H., and Bengio, Y. Recall traces: Backtracking models for efficient reinforcement learning. arXiv preprint arXiv:1804.00379, 2018.
27.Goyal, A., Islam, R., Strouse, D., Ahmed, Z., Botvinick, M., Larochelle, H., Bengio, Y., and Levine, S. Infobot: Transfer and exploration via the information bottleneck. arXiv preprint arXiv:1901.10902, 2019a.
28.Goyal, A., Lamb, A., Hoffmann, J., Sodhani, S., Levine, S., Bengio, Y., and Schölkopf, B. Recurrent independent mechanisms. arXiv preprint arXiv:1909.10893, 2019b.
29.Goyal, A., Sodhani, S., Binas, J., Peng, X. B., Levine, S., and Bengio, Y. Reinforcement learning with competitive ensembles of information-constrained primitives. arXiv preprint arXiv:1906.10667, 2019c.
30.Goyal, A., Bengio, Y., Botvinick, M., and Levine, S. The variational bandwidth bottleneck: Stochastic evaluation on an information budget. arXiv preprint arXiv:2004.11935, 2020a.
31.Goyal, A., Lamb, A., Gampa, P., Beaudoin, P., Levine, S., Blundell, C., Bengio, Y., and Mozer, M. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv preprint arXiv:2006.16225, 2020b.
32.Gulcehre, C., Wang, Z., Novikov, A., Paine, T., Gómez, S., Zolna, K., Agarwal, R., Merel, J. S., Mankowitz, D. J., Paduraru, C., et al. Rl unplugged: A collection of benchmarks for offline reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020a.
33.Gulcehre, C., Wang, Z., Novikov, A., Paine, T. L., Colmenarejo, S. G., Zolna, K., Agarwal, R., Merel, J., Mankowitz, D., Paduraru, C., et al. Rl unplugged: A suite of benchmarks for offline reinforcement learning. arXiv preprint arXiv:2006.13888, 2020b.
34.Guu, K., Lee, K., Tung, Z., Pasupat, P., and Chang, M.-W. REALM: Retrieval-augmented language model pre-training. arXiv preprint arXiv:2002.08909, 2020.
35.Hansen, S., Sprechmann, P., Pritzel, A., Barreto, A., and Blundell, C. Fast deep reinforcement learning using online adjustments from the past. arXiv preprint arXiv:1810.08163, 2018.
36.Hausman, K., Springenberg, J. T., Wang, Z., Heess, N., and Riedmiller, M. Learning an embedding space for transferable robot skills. In International Conference on Learning Representations, 2018.
37.Heess, N., Wayne, G., Silver, D., Lillicrap, T., Tassa, Y., and Erez, T. Learning continuous control policies by stochastic value gradients, 2015.
38.Heess, N., Wayne, G., Tassa, Y., Lillicrap, T., Riedmiller, M., and Silver, D. Learning and transfer of modulated locomotor controllers. arXiv preprint arXiv:1610.05182, 2016.
39.Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural computation, 9(8):1735–1780, 1997.
40.Jaderberg, M., Mnih, V., Czarnecki, W. M., Schaul, T., Leibo, J. Z., Silver, D., and Kavukcuoglu, K. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05397, 2016.
41.Kapturowski, S., Ostrovski, G., Quan, J., Munos, R., and Dabney, W. Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations, 2018.
42.Ke, N. R., Goyal, A., Bilaniuk, O., Binas, J., Mozer, M. C., Pal, C., and Bengio, Y. Sparse attentive backtracking: Temporal credit assignment through reminding. In Advances in Neural Information Processing Systems, pp. 7640–7651, 2018.
43.Ke, N. R., Singh, A., Touati, A., Goyal, A., Bengio, Y., Parikh, D., and Batra, D. Learning dynamics model in reinforcement learning by incorporating the long term future. arXiv preprint arXiv:1903.01599, 2019.
44.Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. arXiv preprint arXiv:1312.6114, 2013.
45.Kolodner, J. L. An introduction to case-based reasoning. Artificial intelligence review, 6(1):3–34, 1992.
46.Lake, B. M., Ullman, T. D., Tenenbaum, J. B., and Gershman, S. J. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017.
47.Leake, D. B. Case-based reasoning: experiences, lessons, and future directions. 1996.
48.Lee, K., Chang, M.-W., and Toutanova, K. Latent retrieval for weakly supervised open domain question answering. arXiv preprint arXiv:1906.00300, 2019.
49.Lewis, P., Perez, E., Piktus, A., Petroni, F., Karpukhin, V., Goyal, N., Küttler, H., Lewis, M., Yih, W.-t., Rocktäschel, T., et al. Retrieval-augmented generation for knowledge-intensive nlp tasks. arXiv preprint arXiv:2005.11401, 2020.
50.Lillicrap, T. P., Hunt, J. J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., Silver, D., and Wierstra, D. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971, 2015.
51.Lin, L.-J. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine learning, 8(3-4):293–321, 1992.
52.Liu, S., Lever, G., Wang, Z., Merel, J., Eslami, S., Hennes, D., Czarnecki, W. M., Tassa, Y., Omidshafiei, S., Abdolmaleki, A., et al. From motor control to team play in simulated humanoid football. arXiv preprint arXiv:2105.12196, 2021.
53.Liu, Z., Li, X., Kang, B., and Darrell, T. Regularization matters in policy optimization. 2019.
54.Logan, G. D., Cox, G. E., Annis, J., and Lindsey, D. R. The episodic flanker effect: Memory retrieval as attention turned inward. Psychological Review, 128(3):397, 2021.
55.Mazoure, B., Combes, R. T. d., Doan, T., Bachman, P., and Hjelm, R. D. Deep reinforcement and infomax learning. arXiv preprint arXiv:2006.07217, 2020.
56.Merel, J., Botvinick, M., and Wayne, G. Hierarchical motor control in mammals and machines. Nature communications, 10(1):1–12, 2019.
57.Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., Petersen, S., Beattie, C., Sadik, A., Antonoglou, I., King, H., Kumaran, D., Wierstra, D., Legg, S., and Hassabis, D. Human-level control through deep reinforcement learning. Nature, 518(7540):529–533, February 2015a.
58.Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. nature, 518(7540):529–533, 2015b.
59.Nair, A., Dalal, M., Gupta, A., and Levine, S. AWAC: Accelerating online reinforcement learning with offline datasets. 2020.
60.Pascanu, R., Li, Y., Vinyals, O., Heess, N., Buesing, L., Racanière, S., Reichert, D., Weber, T., Wierstra, D., and Battaglia, P. Learning model-based planning from scratch. arXiv preprint arXiv:1707.06170, 2017.
61.Pertsch, K., Lee, Y., and Lim, J. J. Accelerating reinforcement learning with learned skill priors. arXiv preprint arXiv:2010.11944, 2020.
62.Pritzel, A., Uria, B., Srinivasan, S., Badia, A. P., Vinyals, O., Hassabis, D., Wierstra, D., and Blundell, C. Neural episodic control. In International Conference on Machine Learning, pp. 2827–2836. PMLR, 2017.
63.Racanière, S., Weber, T., Reichert, D. P., Buesing, L., Guez, A., Rezende, D., Badia, A. P., Vinyals, O., Heess, N., Li, Y., et al. Imagination-augmented agents for deep reinforcement learning. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 5694–5705, 2017.
64.Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
65.Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, pp. 1278–1286. PMLR, 2014.
66.Riedmiller, M. Neural fitted Q iteration–first experiences with a data efficient neural reinforcement learning method. In European conference on machine learning, pp. 317–328. Springer, 2005.
67.Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. arXiv preprint arXiv:1911.08265, 2019.
68.Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604–609, 2020.
69.Siegel, N. Y., Springenberg, J. T., Berkenkamp, F., Abdolmaleki, A., Neunert, M., Lampe, T., Hafner, R., Heess, N., and Riedmiller, M. Keep doing what worked: Behavioral modelling priors for offline reinforcement learning. arXiv preprint arXiv:2002.08396, 2020.
70.Silver, D. Reinforcement learning and simulation-based search in computer go. 2009.
71.Silver, D., Sutton, R. S., and Müller, M. Sample-based learning and search with permanent and transient memories. In Proceedings of the 25th international conference on Machine learning, pp. 968–975, 2008.
72.Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484–489, 2016.
73.Silver, D., Hasselt, H., Hessel, M., Schaul, T., Guez, A., Harley, T., Dulac-Arnold, G., Reichert, D., Rabinowitz, N., Barreto, A., and Degris, T. The predictron: End-to-end learning and planning. In International Conference on Machine Learning, pp. 3191–3199. PMLR, 2017.
74.Silver, D., Hubert, T., Schrittwieser, J., Antonoglou, I., Lai, M., Guez, A., Lanctot, M., Sifre, L., Kumaran, D., Graepel, T., et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419):1140–1144, 2018.
75.Springenberg, J. T., Heess, N., Mankowitz, D., Merel, J., Byravan, A., Abdolmaleki, A., Kay, J., Degrave, J., Schrittwieser, J., Tassa, Y., et al. Local search for policy iteration in continuous control. arXiv preprint arXiv:2010.05545, 2020.
76.Stepleton, T., Delalleau, O., Ostrovski, G., and Wenke, S. URL https://github.com/deepmind/pycolab.
77.Strouse, D., Kleiman-Weiner, M., Tenenbaum, J., Botvinick, M., and Schwab, D. Learning to share and hide intentions using information regularization. arXiv preprint arXiv:1808.02093, 2018.
78.Sun, Y., Wang, S., Feng, S., Ding, S., Pang, C., Shang, J., Liu, J., Chen, X., Zhao, Y., Lu, Y., et al. Ernie 3.0: Large-scale knowledge enhanced pre-training for language understanding and generation. arXiv preprint arXiv:2107.02137, 2021.
79.Sutton, R. S. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160–163, 1991.
80.Teh, Y., Bapst, V., Czarnecki, W. M., Quan, J., Kirkpatrick, J., Hadsell, R., Heess, N., and Pascanu, R. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems, pp. 4496–4506, 2017.
81.Tishby, N., Pereira, F. C., and Bialek, W. The information bottleneck method. arXiv preprint physics/0004057, 2000.
82.Van Hasselt, H., Guez, A., and Silver, D. Deep reinforcement learning with double q-learning. In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
83.van Hasselt, H. P., Hessel, M., and Aslanides, J. When to use parametric models in reinforcement learning? Advances in Neural Information Processing Systems, 32:14322–14333, 2019.
84.Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. In Advances in neural information processing systems, pp. 5998–6008, 2017.
85.Vecerik, M., Hester, T., Scholz, J., Wang, F., Pietquin, O., Piot, B., Heess, N., Rothörl, T., Lampe, T., and Riedmiller, M. Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards. arXiv preprint arXiv:1707.08817, 2017.
86.Veerapaneni, R., Co-Reyes, J. D., Chang, M., Janner, M., Finn, C., Wu, J., Tenenbaum, J., and Levine, S. Entity abstraction in visual model-based reinforcement learning. In Conference on Robot Learning, pp. 1439–1456. PMLR, 2020.
87.Vezhnevets, A. S., Osindero, S., Schaul, T., Heess, N., Jaderberg, M., Silver, D., and Kavukcuoglu, K. Feudal networks for hierarchical reinforcement learning. In International Conference on Machine Learning, pp. 3540–3549. PMLR, 2017.
88.Vinyals, O., Babuschkin, I., Czarnecki, W. M., Mathieu, M., Dudzik, A., Chung, J., Choi, D. H., Powell, R., Ewalds, T., Georgiev, P., et al. Grandmaster level in starcraft II using multi-agent reinforcement learning. Nature, 575 (7782):350–354, 2019.
89.Watters, N., Matthey, L., Bosnjak, M., Burgess, C. P., and Lerchner, A. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration. arXiv preprint arXiv:1905.09275, 2019.
90.Wayne, G., Hung, C.-C., Amos, D., Mirza, M., Ahuja, A., Grabska-Barwinska, A., Rae, J., Mirowski, P., Leibo, J. Z., Santoro, A., et al. Unsupervised predictive memory in a goal-directed agent. arXiv preprint arXiv:1803.10760, 2018.
91.Weisz, G., Amortila, P., and Szepesvári, C. Exponential lower bounds for planning in mdps with linearly-realizable optimal action-value functions. In Algorithmic Learning Theory, pp. 1237–1264. PMLR, 2021.
92.Zambaldi, V., Raposo, D., Santoro, A., Bapst, V., Li, Y., Babuschkin, I., Tuyls, K., Reichert, D., Lillicrap, T., Lockhart, E., et al. Deep reinforcement learning with relational inductive biases. In International Conference on Learning Representations, 2018.
93.Zhu, G., Lin, Z., Yang, G., and Zhang, C. Episodic reinforcement learning with associative memory. 2020.

Citation

MLA

Goyal, A., et al. “Retrieval-Augmented Reinforcement Learning”. International Conference on Machine Learning, vol. 162, 2022, pp. 7740–65, https://proceedings.mlr.press/v162/goyal22a.html.

APA

Goyal, A., Friesen, A., Banino, A., Weber, T., Ke, N. R., Badia, A. P., Guez, A., Mirza, M., Humphreys, P. C., Konyushova, K., Valko, M., Osindero, S., Lillicrap, T., Heess, N., & Blundell, C. (2022). Retrieval-Augmented Reinforcement Learning. International Conference on Machine Learning, 162, 7740–7765. https://proceedings.mlr.press/v162/goyal22a.html

Chicago

Goyal, A., A. Friesen, A. Banino, et al. 2022. “Retrieval-Augmented Reinforcement Learning”. International Conference on Machine Learning 162: 7740–65. https://proceedings.mlr.press/v162/goyal22a.html.

Harvard

Goyal, A. et al. (2022) “Retrieval-Augmented Reinforcement Learning”, International Conference on Machine Learning. PMLR, pp. 7740–7765. Available at: https://proceedings.mlr.press/v162/goyal22a.html.

Vancouver

1. Goyal A, Friesen A, Banino A, et al (2022) Retrieval-Augmented Reinforcement Learning. In: International Conference on Machine Learning. PMLR, pp 7740–7765

BibTeX

@InProceedings{pmlr-v162-goyal22a,
  title = 	 {Retrieval-Augmented Reinforcement Learning},
  author =       {Goyal, Anirudh and Friesen, Abram and Banino, Andrea and Weber, Theophane and Ke, Nan Rosemary and Badia, Adri{\`a} Puigdom{\`e}nech and Guez, Arthur and Mirza, Mehdi and Humphreys, Peter C and Konyushova, Ksenia and Valko, Michal and Osindero, Simon and Lillicrap, Timothy and Heess, Nicolas and Blundell, Charles},
  booktitle = 	 {Proceedings of the 39th International Conference on Machine Learning},
  pages = 	 {7740--7765},
  year = 	 {2022},
  editor = 	 {Chaudhuri, Kamalika and Jegelka, Stefanie and Song, Le and Szepesvari, Csaba and Niu, Gang and Sabato, Sivan},
  volume = 	 {162},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {17--23 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v162/goyal22a/goyal22a.pdf},
  url = 	 {https://proceedings.mlr.press/v162/goyal22a.html},
  abstract = 	 {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. The 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.}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Method	Success Rate (50K)
BC (behavior cloning)	$61\%\pm10\%$
RA-BC (single-task retrieval buffer)	$71\%\pm7\%$
RA-BC (multi-task retrieval buffer)	$82\%\pm5\%$

Retrieval-Augmented Reinforcement Learning

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Retrieval-augmented agents keep experience directly available to behavior

Knowl 2 — Bidirectional trajectory summaries make past and future information retrievable

Knowl 3 — Slot-based attention retrieves selected trajectory contexts and updates the agent

Knowl 4 — An information bottleneck limits dependence on retrieved context

Knowl 5 — Retrieval improves Atari R2D2 performance and sample efficiency

Knowl 6 — Atari ablations identify the contribution of retrieval design choices

Knowl 7 — Task-specific retrieval helps offline multi-task Gridroboman learning