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

Large-Scale Retrieval for Reinforcement Learning

“…We randomly zero-out a subset of retrieved neighbors during training ("neighbor dropout"), and/or more adversarially, randomly replace a subset of retrieved neighbors with the neighbors of a different observation ("neighbor randomisation"). Inspired by [10], we also explore using a loss to regularise the embedding produced by the neighbor retrieval towards the embedding produced with the observation alone ("neighbor regularisation"). Further details are given in Sec.…”

“…through specifying the agent's action-value directly in terms of previously generated value estimates [6,13,15,28], or a model from observed transitions [38]) but rather learn end-to-end how the data can support better predictions within the parametric model. A recent approach by Goyal et al [10] has considered an attention mechanism to select where and what to use from available trajectories, but over a small retrieval batch of data rather than the full available experience data. Another class of method to leverage a transition dataset is to replay the data at training time in order to perform more gradient steps per experience, this is a widespread technique in modern RL algorithms [21,22,24,35] but it does not benefit the agent at test time, requires additional learning steps to adapt to new data, and does not allow end-to-end learning of how to relate past experience to new situations.…”

“…Inspired by [10], we also explore using a loss to regularise the embedding produced by the neighbor retrieval towards the embedding produced with the observation alone. In this case, we apply a convolutional layer to the embedding produced from the game state observation, and use this to predict the output of the neighbor processing tower.…”

“…There is no end-to-end means for an agent to attend to information outside of working memory to directly inform its actions. While there has been a significant amount of work focused on increasing the information available from previous experiences within an episode (e.g., recurrent networks, slot-based memory [19,27]), more extensive direct use of more general forms of experience or data has been limited, although some recent works have begun to explore utilising inter-episodic information from the same agent [6,10,28,33,40]. We seek to drastically expand the scale of information that is accessible to an agent, allowing it to attend to tens of millions of pieces of information, while learning in an end-to-end manner how to use this information for decision making.…”

Discrete Key-Value Bottleneck

“…The hallmark of machine learning is to be able to develop models that can quickly adapt to new tasks once trained on sufficiently diverse tasks [5,54]. There are multiple ways to transfer information from one task to another: (1) transfer information via the transfer of the neural network weights (when trained on source tasks); (2) reuse raw data as in retrieval-based methods [6,37,32,33,19,51,17]; or (3), via knowledge distillation [23]. Each approach implies inevitable trade-offs: When directly transferring neural network weights, previous information about the data may be lost in the finetuning process, while transfer via raw data may be prohibitively expensive as there can be hundreds of thousands of past experiences.…”