Can Retrieval-Augmented Reinforcement Learning Learn Faster Than Parametric Agents?

Why distilling experience into weights loses detail

Most deep reinforcement learning agents compress what they have seen into the weights of a policy or value network through gradient updates. The paper lists four costs of that approach: it is computationally expensive, it can take many updates to absorb an experience, experiences that are not fully absorbed do not properly influence behavior, and behavior is capped by model capacity. The authors also make a sharper claim about timing: because information from individual trajectories must first be distilled into a general rule, an agent may be unable to exploit the specific guidance a handful of past experiences could offer, and may need many replays before novel experience is incorporated.

The paper's alternative is case-based reasoning in the style of recalling a similar past situation. Rather than only reading its own current state, the agent is given a retrieval process with direct access to a dataset of experiences. That dataset can be the agent's own past, expert demonstrations, or another relevant source. The point is that at test time the agent can condition its behavior on the entire dataset, not just the current state or current trajectory.

A separate network that selects trajectories, then states

The retrieval process is parameterized as a neural network with its own internal state, kept separate from the agent's state. At each step it forms a query and matches it against a retrieval batch, a pre-processed sample drawn from the retrieval dataset. Selection happens in two stages: attention scores first rank whole trajectories and keep the top few, then scores are recomputed within those trajectories to pick the most relevant states. Information from the selected states is then integrated and passed to the agent process to inform its policy or value output.

Two design details matter for reading the results. First, each trajectory in the retrieval batch is summarized in both directions, with a forward summary of the past and a backward summary of the future at each state, so retrieval can draw on what followed a state as well as what preceded it. Second, the retrieval state is split into slots that each retrieve independently and interact through self-attention. The authors state that all retrieval hyperparameters except the number of trajectories in the retrieval batch stayed the same across the Atari, gridroboman, and BabyAI experiments.

Retrieval helps most as tasks accumulate

The first test is a multi-task offline setting on gridroboman, a grid-world manipulation environment with 30 tasks involving red, green, and blue objects. A single agent trains on multiple tasks at once, but at evaluation time the retrieval dataset contains only trajectories from the task being evaluated. The baseline is a vanilla DQN agent, used because the task does not need recurrence.

With fewer tasks, the retrieval-augmented DQN and the baseline perform identically. As the number of training tasks grows, the retrieval-augmented agent learns much more effectively. The authors' explanation is that training on more tasks requires either more model capacity or the ability to extract information from fewer relevant samples per task, and querying task-relevant experiences directly improves sample efficiency. They also address an obvious alternative explanation: an ablation reported in the paper indicates that the retrieved information, not the extra capacity the retrieval process adds, is what is crucial to performance.

Faster learning and higher scores on Atari

The second test integrates retrieval into an online R2D2 agent and compares it with baseline R2D2 on Atari. The paper reports that the retrieval-augmented agent learns significantly faster and achieves higher scores. The reported per-game numbers show the comparison is not uniform across games: on Alien the baseline scores 52,268 ± 4,125 against 81,626 ± 5,123 for the augmented agent, and on Amidar 10,976 ± 2,341 against 24,693 ± 1,982, while on Assault the augmented agent is slightly lower at 43,497 ± 6,653 versus 45,521 ± 4,751, and on Asterix and Asteroids the two are close.

That pattern matters for how the headline claim should be read. The paper's summary is about learning speed and aggregate score, and the per-game table shows gains concentrated in some games rather than a uniform improvement. The authors also report extensive ablations measuring the contribution of the method's components, and the discussion credits two things in particular: retrieving information from past experiences, and parameterizing the agent and retrieval process separately instead of giving the agent direct access to the replay buffer.

Multi-task retrieval data can help rather than mislead

A natural worry is that retrieving from data covering other tasks could misguide the agent. The paper tests this in BabyAI, a partially observable multi-room grid world where harder tasks are composed from simpler ones and expressed in a subset of a synthetic language. Because observations are partial, the baseline is a recurrent DQN, compared with a retrieval-augmented version. Performance is measured as success rate, the ratio of tasks completed within a fixed number of steps per task.

Retrieval augmentation improves on the baseline with a replay specific to the current task, and improves further with a multi-task replay, at both 50K and 200K offline trajectories per task. The authors attribute this to the compositional structure of BabyAI tasks, where information about a subtask can be more informative than information about the overall task. This is the authors' interpretation of the environment's structure, not a separately measured mechanism.

The cost of querying only part of the dataset

Retrieval adds computation at every step: a slot selects relevant trajectories and relevant states, attention scales with the number of trajectories in the retrieval batch, and the selection changes across time steps and repeats for each slot. The authors note this can be performed efficiently on GPUs or TPUs and describe optimizations used in their experiments. For Atari and gridroboman they skip learned ranking of trajectories and simply use high-return trajectories, since the retrieval data is already specific to the task being trained. They also restrict state selection to states semantically similar to the agent's current state.

The paper's limitations section names a related constraint directly: R2A queries only a subset of the retrieval dataset, which limits the method's generality. The authors also list directions they did not pursue, including training the retrieval and agent processes with different objectives rather than end to end, scaling to complex multi-agent problems such as StarCraft where a retrieval process might be shared between agents, learning an abstract model with abstract internal actions and rewards instead of querying a dataset, and evaluating R2A in a few-shot setting where a pre-trained encoder must help an agent adapt quickly to a new task. No few-shot result is reported.

Where the faster-learning claim holds

The paper's evidence supports a specific priority rather than a blanket claim. The clearest case for retrieval is many-task offline learning: with few gridroboman tasks the augmented and baseline agents are indistinguishable, and the separation appears as tasks accumulate, which is exactly the regime where a parametric model must either grow in capacity or extract more from fewer relevant samples. The ablation pointing to retrieved information rather than added capacity is what makes this a statement about retrieval itself.

On Atari the supported claim is faster learning and higher scores for retrieval-augmented R2D2 overall, with the per-game table showing that the size of the gain depends on the game. In BabyAI, the supported claim is that multi-task retrieval data helped rather than hurt, in an environment whose tasks are compositional. The practical reading is to expect the largest benefit when one agent must cover many tasks from a fixed dataset and relevant past trajectories exist to be found, and to check per-task outcomes rather than assuming a uniform improvement.

Sources used in this answer