PaperlayerView exact revision on arXiv

Source-linked AI summary

Retrieval-Augmented Reinforcement Learning

Anirudh Goyal, Abram L. Friesen, Andrea Banino, Theophane Weber, Nan Rosemary Ke, Adria Puigdomenech Badia, Arthur Guez, Mehdi Mirza, Peter C. Humphreys, Ksenia Konyushkova, Laurent Sifre, Michal Valko, Simon Osindero, Timothy Lillicrap, Nicolas Heess, Charles Blundell

arXiv:2202.08417v4cs.LGstat.ML

TL;DR

Parametric RL can lose detailed experience and may take many updates to integrate useful trajectories. This paper adds a learned retrieval process that lets an agent query past experience for relevant context. Retrieval augmentation improves results across online Atari and offline multi-task RL, including a higher Atari mean human normalized score than R2D2.

  • Problem

    Parametric RL can lose detailed information from past experience and may need many replays to integrate guidance from individual trajectories or novel experience.

    Source: §1 · p. 1

  • Method

    R2A augments an RL agent with a learned attention-based retrieval process that queries a dataset of past trajectories for context to inform policy or value predictions.

    Source: §1 · p. 2

  • Results

    Across Atari and offline multi-task RL, retrieval augmentation improves learning; on Atari, mean human normalized score rose 11.32 ± 1.2% relative to R2D2 over 2 billion environment steps.

    Source: §1 · p. 2, §3.1 · p. 6

  • Takeaways & Limitations

    In the reported multi-task offline setting, retrieval augmentation compensates for insufficient model capacity as the number of training tasks increases.

    Source: §3.2 · p. 7, §5 · p. 9

    R2A queries only a subset of the retrieval dataset, which limits the method’s generality.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

The introduction contrasts parametric RL, which can lose or slowly integrate useful experience, with retrieval augmentation that lets an agent query past trajectories. Across the reported settings, retrieval-augmented agents learn faster and achieve higher reward.

1. Introduction

  • The introduction uses case-based reasoning as an analogy: recalling a similar past situation can guide a solution to the current task.The example is a host who shops early after remembering that holiday produce was previously sold out.Source: §1 · p. 1
  • Many deep RL algorithms integrate trajectory information through gradient updates to parametric networks, often reusing past experience by replaying trajectory snippets.The introduction frames this as the dominant modern RL paradigm.Source: §1 · p. 1
  • Parametric policies and value functions can lose detailed experience, and may need many replays to integrate guidance from individual trajectories or novel experience.The passage attributes these limitations to practical network-capacity constraints and the need to distill experience into a general-purpose rule.Source: §1 · p. 1
  • R2A uses a learned attention-based retrieval process to query past trajectories for context that informs the agent’s policy or value predictions.The retrieval and agent processes have separate internal states and interact to shape each other’s representations and predictions.Source: §1 · p. 2
  • Retrieval augmentation helps agents learn faster and achieve higher reward across the three offline multi-task environments, while also improving Atari performance and sample efficiency.For online Atari, retrieval is from the agent’s own game experience; offline evaluations also use data from other agents and tasks.Source: §1 · p. 2

2. Retrieval-Augmented Agents

R2A augments a standard RL agent with a retrieval process that accesses past trajectories and returns information for the agent’s value or policy estimates. It samples and summarizes trajectory data, then uses attention-based retrieval to update the agent’s representation.

2. Retrieval-Augmented Agents

  • R2A pairs a reward-maximizing agent with a retrieval process that accesses an experience dataset and supplies information to improve value or policy estimates.The retrieval process has its own internal state and is trained to retrieve useful information without explicit knowledge of the agent’s policy.Source: §2 · p. 2

2.1. Retrieval-augmented agent

  • The retrieval process uses the agent’s current state and its own previous state to query the dataset, summarize retrieved information, and update its internal state.The agent then combines its state with retrieved information for action-value computation and action selection.Source: §2.1 · p. 3

2.2. Retrieval batch sampling and pre-processing.

  • R2A samples a retrieval batch to reduce computation, then uses a learned bidirectional summarizer to encode past and future trajectory information.Auxiliary losses train forward and backward summaries to capture information useful for downstream tasks; experiments use supervised losses when actions or rewards are available and also test self-supervised losses in ablations.Source: §2.2 · p. 3, §2.1 · p. 3

2.3. Retrieving contextual information.

  • Multiple retrieval slots use attention to find relevant trajectories and states, retrieve weighted backward summaries, and combine the retrieved information to update the agent’s representation.R2A regularizes retrieved information with an information bottleneck, while slot updates also allow information sharing between slots.Source: §2.3 · pp. 4–5

3. Experimental Results

Across Atari and offline multi-task settings, retrieval augmentation improved performance over the corresponding baselines. Experiments also examined task scale, retrieval sources, and design ablations.

3.1. Atari: Single-task off-policy RL

  • On ten Atari games, agent-state querying matched baseline; removing retrieval or shortening context hurt, while BERT losses helped only alongside action, reward, and value prediction.Per-game hyperparameter optimization further improved performance in an additional ablation, but was not used in the reported experiments.Source: §3.1.1 · pp. 6–7
  • 11.32 ± 1.2% higher mean human normalized score than R2D2 over 2 billion environment steps was achieved by retrieval-augmented R2D2.Figure 2a shows per-game relative improvement; the average result is across the Atari evaluation.Source: Figure 2 · p. 6, §3.1 · p. 6

3.2. Gridroboman: Multi-task offline RL with a task-specific retrieval dataset

  • With 20 or 30 Gridroboman tasks, retrieval-augmented DQN learned more effectively than DQN; the agents performed identically with 10 tasks.The retrieval dataset contained experiences from the task being evaluated, and the authors attribute improved sample efficiency to querying task-relevant experiences.Source: §3.2 · p. 7, Figure 3 · p. 8

3.3. BabyAI: Multi-task offline RL with a multi-task retrieval dataset

  • In BabyAI, retrieval augmentation improved success rate over the baseline with a task-specific replay, and an all-task replay improved performance further.The comparison covered 40 levels and two offline-data amounts; success rate is the ratio of tasks completed within a fixed number of steps.Source: §3.3 · p. 7, Table 1 · p. 8
  • Other-task retrieval occurred for 54% of compositional tasks versus 21% of non-compositional tasks; removing the information bottleneck reduced performance, though it remained above baseline.The 40-task set included 15 compositional tasks.Source: §3.3 · p. 8

3.4. CausalWorld: Multi-task offline continuous control

  • Retrieval augmentation improved behaviour-cloning performance with both multi-task and single-task retrieval buffers on five CausalWorld manipulation tasks.The evaluation compared retrieval-augmented behaviour cloning with vanilla behaviour cloning.Source: Table 2 · p. 8, §3.4 · p. 8

4. Related Work

The related-work discussion distinguishes R2A from episodic-control methods by its learned, differentiable retrieval pathway. It also connects retrieval to language-model methods while noting their different domains.

4. Related Work

  • Unlike episodic-control methods, R2A uses a parameterized retrieval network and attention to shape value predictions differentiably, with the agent accessing replay indirectly.The cited methods instead build local action values from nearest neighbors and give the agent direct replay access.Source: §4 · p. 8
  • Language-model retrieval methods combine parametric models with non-parametric retrieval, but address different domains from R2A.The passage names question answering, controllable generation, and machine translation as application areas.Source: §4 · pp. 8–9
  • Further related-work discussion is provided in Appendix A.1.Source: §4 · p. 9

5. Conclusion.

R2A augments an RL agent with a separate retrieval process, and the paper reports improved sample efficiency over R2D2 and compensation for limited capacity in multi-task offline RL. The conclusion also identifies limits and directions for future work.

5. Conclusion.

  • R2A improves sample efficiency over R2D2 and compensates for insufficient model capacity in multi-task offline RL environments.Ablations also indicate the importance of retrieving past experience and separately parameterizing the agent and retrieval processes.Source: §5 · p. 9
  • The retrieval process and agent have separate states and shape each other’s representations and predictions through attention.The retrieval process draws useful information from an experience dataset to support the agent’s objective.Source: §5 · p. 9
  • R2A queries only a subset of the retrieval dataset, which the authors say limits the method’s generality.Source: §5 · p. 9
  • Future work proposed includes different training objectives, scaling to complex multi-agent tasks, abstract internal models, and few-shot adaptation, which was not evaluated.Source: §5 · p. 9

A.1. Extended Relevant Work.

The extended related work places R2A alongside approaches that integrate experience through models, offline datasets, hierarchical policies, and temporal credit assignment. It also draws a conceptual connection between retrieval and selective attention.

Model-based RL

  • Unlike model-based RL, which may accumulate errors when unrolled across long intervals, R2A queries information across multiple trajectories directly.The retrieval dataset can contain trajectories separated by hundreds of time-steps.Source: §A.1 · p. 14

Structural Inductive Biases

  • Structural inductive biases such as Transformers and slot-based recurrent architectures are described as improving generalization, model-size scaling, and long-range dependencies.Source: §A.1 · p. 14

Reinforcement Learning with Offline Datasets

  • Task-agnostic offline datasets are commonly used to improve value functions or provide reusable skills, while R2A searches across trajectories for relevant information.The passage notes that this information might otherwise require many replays to encounter coincidentally.Source: §A.1 · p. 14

Separation of concerns

  • R2A resembles hierarchical RL in letting a higher-level retrieval process shape the agent’s representation, but the agent also shapes the retrieval representation.Source: §A.1 · p. 14

Efficient Credit assignment

  • Where many credit-assignment approaches select relevant past states within one trajectory, R2A searches a replay buffer spanning multiple trajectories.Source: §A.1 · p. 14

Memory retrieval as attention turned inward

  • The passage frames R2A as retrieval from memory through selective attention, connecting it to the conjecture that perceiving and remembering share an extraction process.The cited conjecture is attributed to Logan et al. (2021).Source: §A.1 · p. 14

A.2. Information Theoretic Formulation

The formulation treats retrieved information as context for the agent’s policy and penalizes policy dependence on that information. It expresses the policy through an encoder, decoder, and learned prior, and relates the objective to a lower-bound optimization.

A.2. Information Theoretic Formulation

  • The policy conditions on state S and retrieved information G, where G is produced by the retrieval process given the current state.The formulation seeks to maximize reward while minimizing I(A; G | S), the policy’s dependence on retrieval information.Source: §A.2 · p. 15
  • The conditional mutual information penalty can be written as a KL divergence from the retrieval-conditioned policy to a default policy that marginalizes retrieval information.The objective subtracts this divergence, weighted by β, from expected reward.Source: §A.2 · p. 15
  • A lower-bound objective replaces the original objective’s KL term with a divergence between the encoder distribution and the learned prior.Source: §A.2 · p. 15
  • The policy is parameterized with an encoder, decoder, and learned prior, with encoder output Z representing retrieval information considered useful for the current state.Source: §A.2 · p. 15
  • The formulation uses an upper bound on I(Z; G | S) by first bounding it at each state and averaging over states, drawing on the data processing inequality.The passage notes that the information bottleneck has been shown to improve generalization.Source: §A.2 · p. 15

A.3. Atari: Implementation details and raw scores for R2D2 and RA-R2D2.

This section documents the Atari R2D2 and RA-R2D2 setup, including retrieval-batch construction, model components, training losses, and implementation choices. It also points to tables of experiment hyperparameters and raw scores.

A.3. Atari: Implementation details and raw scores for R2D2 and RA-R2D2.

  • Table 3 lists Atari R2D2 hyperparameters, while Table 4 reports raw scores for RA-R2D2 and the R2D2 baseline.The baseline and agent-process setup follow Banino et al. (2021); the remaining R2D2 hyperparameters are taken from Kapturowski et al. (2018).Source: §A.3 · pp. 15–16, Table 3 · p. 16
  • The agent process is trained with the RL loss, while retrieval parameters use the RL loss and an information-bottleneck regularizer fixed at β = 0.1.The state encoder is trained with RL and auxiliary summarizer losses; marginals and conditionals are estimated using a single sample throughout.Source: §A.3 · pp. 15–16
  • Each learner step samples 512 trajectories, uses 64 for Q-learning, and assigns the rest to retrieval; the retrieval dataset is resampled and re-encoded per gradient update.The retrieval process selects 10 trajectories and 10 states, using forward and backward GRUs with 512 hidden units; Atari experiments use 8 slots.Source: §A.3 · p. 16
  • When a network takes two tensors, the implementation cross-attends from the primary input to the privileged input rather than concatenating them.GRU-style gating and layer normalization are used for residual state or slot updates, and gating is also used when adding the current encoded state to the forward summary.Source: §A.3 · p. 17
  • The visual encoder is a ResNet-47 with a stride-2 initial convolution, bottleneck groups of 2, 4, 6, and 2 blocks, and a 512-unit final layer.The last convolution channels are 64, 128, 256, and 512, and group normalization is applied after each group.Source: §A.3 · p. 17
  • Retrieval attention scales with the number of trajectories in the retrieval batch, and high-return trajectories or states semantically similar to the current state are described as ways to reduce computation.Multiple retrieval queries per time-step and retrieval at different time scales are listed as improvements that were not evaluated exhaustively.Source: §A.3 · p. 17

A.3.1. RETRIEVAL PROCESS: RETRIEVING INFORMATION FROM THE RETRIEVAL BUFFER.

The retrieval process builds slot-specific queries, selects relevant trajectories and states, then combines information from the selected states. The section also identifies the accompanying score and relative-performance visuals.

A.3.1. RETRIEVAL PROCESS: RETRIEVING INFORMATION FROM THE RETRIEVAL BUFFER.

  • For each slot, attention scores rank trajectories and then states within the selected trajectories, with Figure 4 depicting both selection stages.The figure's scalars represent attention weights.Source: §A.3.1 · p. 18, Figure 4 · p. 18
  • Each retrieved trajectory state has forward and backward summaries, and attention retrieves future information from selected states using forward summaries as keys and backward summaries as values.Figure 5 describes the paired representations as summaries of the past and future, respectively.Source: §A.3.1 · p. 18, Figure 5 · p. 18
  • The section includes a table of scores across Atari games and a figure of RA-R2D2's relative improvement over R2D2 in human-normalized score.The figure averages three seeds per method and game; the supplied passages do not state specific score values.Source: Table 4 · p. 19, Figure 6 · p. 20

A.3.2. DETAILS ON ABLATIONS

The ablations alter retrieval state, retrieval-batch access, trajectory summarization length, and the auxiliary losses used to train the summarizers.

A.3.2. DETAILS ON ABLATIONS

  • A-1 removes the internal retrieval state and computes the query directly from the agent state; A-2 removes the retrieval batch and uses memory slots instead.A-1 updates the agent state with the retrieval vector, while A-2 omits retrieval steps 2–4 and computes keys and values from memory slots.Source: §A.3.2 · p. 20
  • A-3 shortens trajectories used for summarization from 80 steps to 5 while leaving trajectories used for Q-learning unchanged..Source: §A.3.2 · p. 20
  • A-4 adds a BERT loss to the standard summarizer losses, whereas A-5 uses only the BERT loss.For A-4, tokens are obtained by compressing agent states with VQ-VAE.Source: §A.3.2 · p. 20

A.3.3. ADDITIONAL ATARI ABLATIONS

Varying the numbers of retrieved trajectories and states independently can improve R2A performance over the R2D2 baseline. The additional ablations focus on five games where RA-R2D2 performs worse than baseline R2D2.

A.3.3. ADDITIONAL ATARI ABLATIONS

  • The additional ablations are restricted to the five Atari games on which RA-R2D2 performs worst relative to baseline R2D2.The passage identifies six ablations in total, including the additional ktraj and kstates variation.Source: §A.3.3 · p. 20, Figure 7 · p. 21
  • Varying ktraj and kstates independently over 5, 10, and 20 can improve R2A performance over the R2D2 baseline.Figure 7 presents these results among six ablations on the five games where RA-R2D2 performs worse than baseline R2D2.Source: §A.3.3 · pp. 20–21, Figure 7 · p. 21
  • The reported experiments did not use separate hyperparameter optimization for each game, although the authors say it can greatly improve R2A performance.The additional tuning varies ktraj and kstates independently.Source: §A.3.3 · p. 21

A.3.4. EPISODIC CONTROL BASELINE

These sections report where retrieval augmentation helps or hurts, then describe its auxiliary losses and multi-task evaluation setups.

A.3.4. EPISODIC CONTROL BASELINE

  • Non-parametric access to episodic memory hurts R2D2 performance relative to the ablated retrieval-augmented approach.The authors compare this with a simple episodic-memory ablation.Source: §A.3.4 · p. 21

A.3.5. RETRIEVAL AUGMENTED RL HELPING FROSTBITE

  • On Frostbite, reducing retrieved-trace context from 80 to 5 time-steps significantly decreases performance, consistent with the authors’ temporally extended credit-assignment hypothesis.The authors hypothesize retrieval uses information from states far from the current state.Source: §A.3.5 · p. 21

A.3.6. AUXILIARY LOSSES IMPLEMENTATION.

  • Auxiliary trajectory-summary objectives include action, value, and reward prediction, plus masked visual-token recovery with a Transformer.The prediction setup follows MuZero; the masked-token objective reconstructs randomly masked visual tokens from corrupted tokens.Source: §A.3.6 · pp. 21–22

A.4. Gridroboman environment

  • Gridroboman places a movable robot and three colored objects on a 7 × 7 board, with 30 tasks whose identity cannot be inferred from the initial state.The agent observes an 11-dimensional vector and episodes last 50 time-steps.Source: §A.4 · p. 23

A.4.1. GRIDROBOMAN EXPERIMENT SETUP

  • Gridroboman offline data comes from separate online DQN agents, with 100K episodes recorded per task; experiments use the first 10, 20, or 30 tasks.At each learner step, DQN samples 256 states and retrieval samples 64 trajectories; retrieval selects 10 high-return trajectories and 10 relevant states.Source: §A.4.1 · p. 24, §A.4.2 · p. 24

A.4.3. GRIDROBOMAN EVALUATION CURVES WHEN TRAINING ON ALL 30 TASKS

  • The 30-task gridroboman evaluation plots average episode reward for DQN and RA-DQN on each task, with curves averaged over 3 seeds.The supplied passages describe the plotted comparison but do not state which agent wins.Source: §A.4.3 · p. 25, Figure 9 · p. 26

A.5. BabyAI environment.

  • On BabyAI, retrieval accesses other-task information 54% of the time for compositional tasks versus 21% for single-behavior tasks.The authors report no pattern explaining which tasks retrieve across tasks or which states prompt retrieval.Source: §A.5 · p. 25

A.6. CausalWorld environment.

  • CausalWorld evaluation uses five of its eight tasks, with behavior cloning as the underlying algorithm and retrieval setup matching BabyAI and gridroboman.The selected tasks are Pushing, Picking, Pick and Place, Stacking2, and Stacked Blocks.Source: §A.6 · p. 25
Loading 2202.08417v4…