Retrieval-Augmented Reinforcement Learning

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 . 1 INTRODUCTION . A host is preparing a holiday meal for friends . They remember that the last time they went to the grocery store during the holiday season , all of the fresh produce was sold out . Thinking back to this past experience , they decide to go early ! The hypothetical host is employing case-based reasoning ( e.g. , Kolodner , 1992 ; Leake , 1996 ) . Here , an agent recalls a situation similar to the current one and uses information from the previous experience to solve the current task . This may involve adapting old solutions to meet new demands , or using previous experiences to make sense of new situations . In contrast , a dominant paradigm in modern reinforcement learning ( RL ) is to learn general purpose behaviour rules from the agent ’ s past experience . These rules are typically represented in the weights of a parametric policy or value function network model . Most deep RL algorithms integrate information across trajectories by iteratively updating network parameters using gradients that are computed along individual trajectories ( collected online or stored in an experience replay dataset , Lin , 1992 ) . For example , many off-policy algorithms reuse past experience by “ replaying ” trajectory snippets in order to compute weight updates for a value function represented by a deep network ( Ernst et al. , 2005 ; Riedmiller , 2005 ; Mnih et al. , 2015b ; Heess et al. , 2015 ; Lillicrap et al. , 2015 ) . This paradigm has clear advantages but at least two interrelated limitations : First , after learning , an agent ’ s past experiences no longer plays a direct role in the agent ’ s behavior , even if it is relevant to the current situation . This occurs because detailed information in the agent ’ s past experience is lost due to practical constraints on network capacity . Second , since the information provided by individual trajectories first needs to be distilled into a general purpose parametric rule , an agent may not be able to exploit the specific guidance that a handful of individual past experiences could provide , nor rapidly incorporate novel experience that becomes available—it may take many replays through related traces in the past experiences for this to occur ( Weisz et al. , 2021 ) . In this work , we develop an algorithm that overcomes these limitations by augmenting a standard reinforcement learning agent with a retrieval process ( parameterized via a neural network ) . The purpose of the retrieval process is to help the agent achieve its objective by providing relevant contextual information . To this end , the retrieval process uses a learned attention mechanism to dynamically access a large pool of past trajectories stored in a dataset ( e.g. , a replay buffer ) , with the aim of integrating information across these . The proposed algorithm ( R2A ) , shown in Figure 1 , enables an agent to retrieve information from a dataset of trajectories . The high-level idea is to have two different processes . First , the retrieval process , makes a “ query ” to search for relevant contextual information in the dataset . Second , the agent process performs inference and learning based on the information provided by the retrieval process . These two processes have different internal states but interact to shape the representations and predictions of each other : the agent process provides the relevant context , and the retrieval process uses the context and its own internal state to generate a query and retrieve relevant information , which is in turn used by the agent process to shape the representation of its policy and value function ( see Fig . 1A ) . Our proposed retrieval-augmented RL paradigm could take several forms . Here , we focus on a particular instantiation to assay and validate our hypothesis that learning a retrieval process can help an RL agent achieve its objectives . Summary of experimental results . We want RL algorithms that are able to adapt to the available data source and usefully ingest any dataset . Hence , we test the performance of the proposed method in three different scenarios . First , we evaluate it on Atari games in a single task setting . We build upon R2D2 ( Kapturowski et al. , 2018 ) , a state-of-the-art off-policy RL algorithm . Second , we evaluate it on a multi-task offline RL environment , using DQN ( Mnih et al. , 2013 ) as the RL algorithm , where the data in the queried dataset belongs to the same task . Third , we evaluate it on a multi-task offline RL environment where the data in the dataset comes from multiple tasks . In all these cases , we show that R2A learns faster and achieves higher reward compared to the baseline . 2 RETRIEVAL-AUGMENTED AGENTS . We now present our method for augmenting an RL agent with a retrieval process , thereby reducing the agent ’ s dependence on its model capacity , and enabling fast and flexible use of past experiences . A retrieval-augmented agent ( R2A ) consists of two main components : ( 1 ) the retrieval process , which takes in the current state of the agent , combines this with its own internal state , and retrieves relevant information from an external dataset of experiences ; and ( 2 ) a standard reward-maximizing RL agent , which uses the retrieved information to improve its value or policy estimates . See Figure 1 for an overview . The retrieval process is trained to retrieve information that the agent can use to improve its performance , without explicit knowledge of the agent ’ s policy . Importantly , the retrieval process has its own internal state , which enables it to integrate and combine information across retrievals . In the following , we focus on value-based methods , such as DQN ( Mnih et al. , 2015a ) and R2D2 ( Kapturowski et al. , 2018 ) , but our approach is equally applicable to policy-based methods . 2.1 RETRIEVAL-AUGMENTED AGENT . Formally , the agent receives an input xt at each timestep t. Each input is processed by a neural encoder ( e.g. , a resnet if the input is an image ) to obtain an abstract internal state for the agent st = f enc θ ( xt ) . For clarity , we focus here on the case of a single vector input , however , each input could also include the history of past observations , actions , and rewards , as is the case when f encθ is a recurrent network . These embeddings are used by the agent and retrieval processes . The retrieval process operates on a dataset B = { ( xt , at , rt ) , . . . , ( xt+T , at+T , rt+T ) } of l-step trajectories , for l ≥ 1 . This dataset could come from other agents or experts , as in offline RL or imitation learning , or consist of the growing set of the agent ’ s own experiences . Then , a retrieval-augmented agent ( R2A ) consists of the retrieval process and the agent process , parameterized by θ = { θenc , θretr , θagent } , Retrieval process f retrθ , B : mt−1 , st 7→mt , ut Agent process f agentθ : st , ut 7→ Qθ ( st , ut , a ) Retrieval Process . The retrieval process is parameterized as a neural network and has an internal statemt . The retrieval process takes in the current abstract state of the agent process st and its own previous internal statemt−1 and uses these to retrieve relevant information from the dataset B , which it then summarizes in a vector ut , and also updates its internal statemt . Agent Process . The current state of the agent st and the information from the retrieval process ut is then passed to the action-value function , itself used to select external actions . The above defines a parameterization for a retrieval-augmented agent . For retrieval to be effective , the retrieval process needs to : ( 1 ) be able to efficiently query a large dataset of trajectories , ( 2 ) learn and employ a similarity function to find relevant trajectories , and ( 3 ) encode and summarize the trajectories in a manner that allows efficient discovery of relevant past and future information . Below , we explain how we achieve these desiderata . At a high-level , to reduce computational complexity given a experience dataset of hundreds of thousands of trajectories , R2A operates on samples from the dataset . R2A then encodes and summarizes the trajectories in these samples using auxiliary losses and bi-directional sequence models to enable efficient retrieval of temporal information . Finally , R2A uses attention to select semantically relevant trajectories . 2.2 RETRIEVAL BATCH SAMPLING AND PRE-PROCESSING .. Sampling a retrieval batch from the retrieval dataset . To reduce the computational complexity , R2A uniformly samples a large batch of past experiences from the retrieval dataset , and only uses the sampled batch for retrieving information . We denote the sampled batch as “ retrieval batch ” and the number of trajectories in the retrieval batch as nretrieval . Encoding and forward-backward summarization of the retrieval dataset and corresponding auxiliary losses . Since the agent ’ s internal state extracts information from observations which relate to the task at hand , we choose to re-encode the raw experiences in the ” retrieval batch ” using the agent encoder module ( i.e. , f encθ ) . However , this representation is a function only of past observations ( i.e. , it ’ s a causal representation ) and may not be fully compatible with the needs of the retrieval operation . For that reason , we propose to further encode the retrieved batch of information , by additionally learning a summarization function , applied on the output of the encoder module , and which captures information about the past and the future within a particular trajectory by using a bi-directional model ( e.g. , parameterized as a bi-directional RNN or a Transformer ) . Forward Summarizer f fwdθ : ( s1 , . . . , st ) 7→ ht Backward Summarizer f bwdθ : ( sT , . . . , st ) 7→ bt For each trajectory in the retrieval batch , we represent each time-step within a trajectory by a set of two vectors hi , t and bi , t ( Figure 6 in the appendix ) where hi , t summarizes the past ( i.e. , from Algorithm 1 One timestep of a retrieval-augmented agent ( R2A ) . Input : Current input xt , previous retrieval process statemt−1 = { mt−1 , k| k ∈ { 1 , . . . , nf } } , dataset of l-step trajectories B = { ( xit , hit , bit , ait , rit ) . . . ( xit+l , hit+l , bit+l , ait+l , rit+l ) } for l ≥ 1 and 1 ≤ i ≤ ntraj , where h and b are the outputs of the forward & backward summarizers . Encode the current input at time-step t. st = f enc θ ( xt ) Step 1 : Compute the query . For all 1 ≤ k ≤ nf , compute m̂kt−1 = GRUθ ( st , m k t−1 ) qkt = fquery ( m̂ k t−1 ) Step 2 : Identify the most relevant trajectories . For all 1 ≤ k ≤ nf , 1 ≤ j ≤ l and 1 ≤ i ≤ ntraj , κi , j = ( h i jW e ret ) T ` ki , j = ( qkt κi , j√ de ) αki , j = softmax ( ` ki , j ) . Given scores α , the top-ktraj trajectories ( resp . top-kstates states ) are selected and denoted by T kt ( resp . Skt ) . Step 3 : Retrieve information from the most relevant trajectories and states . αki , j = softmax ( ` ki , j ) , i ∈ T kt , j ∈ Skt . gkt = ∑ i , j α k i , jvi , j where vi , j = bi , jW v ret Step 4 : Regularize the retrieved information by using information bottleneck . zkt ∼ p ( z|gkt ) Step 5 : Update the states of the slots . Slotwise update using retrieved information : m̃kt ← m̂kt−1 + zkt ∀k ∈ { 1 , . . . , nf } Joint slot update through self-attention : ckt = m̂ k t−1W q SA ∀k { 1 , . . . , nf } βk , k′ = softmaxk′ ( ckt κ k′ t√ de ) where κk ′ t = ( m̃ k′ t W e SA ) T ∀ k , k′ ∈ { 1 , . . . , nf } mkt ← m̃kt + ∑ k′ βk , k′vk′ where vk′ = m̃ k tW v SA ∀k ∈ { 1 , . . . , nf } Step 6 : Update the agent state using the retrieved information . dt = stW q ag κk = ( zktW e ag ) T ∀k ∈ { 1 , . . . , nf } γk = softmaxk ( dtκ k √ de ) ut ← ∑ k γkvk where vk = z k tW v ag ∀k ∈ { 1 , . . . , nf } . s̃t ← st + ut t′ = 0 to t′ = t time-steps of the ith trajectory ) while bi , t summarizes the future ( i.e. , from t′ = t to t′ = ` time-steps ) within the ith trajectory . In addition , taking inspiration from ( Jaderberg et al. , 2016 ; Trinh et al. , 2018 ; Ke et al. , 2019 ; Devlin et al. , 2018 ; Mazoure et al. , 2020 ; Banino et al. , 2021 ) , we use auxiliary losses to improve modeling of long term dependencies when training the parameters of our forward and backward summarizers . The goal of these losses is to force the representation ( hi , t , bi , t ) i , t≥0 to capture meaningful information for the unknown downstream task . For our experiments , we use supervised losses where we have access to actions or rewards in the retrieval batch . For ablations we also experiment with self-supervised losses . For supervised auxiliary losses , we use policy , value and reward prediction ( Silver et al. , 2017 ; Schrittwieser et al. , 2019 ) , and for self-supervised losses , we use a BERT-style masking loss ( Devlin et al. , 2018 ) .

Summary: The paper presents an approach that augments an RL agent with a parameterized retrieval process directly accessing prior experiences. The retrieval process is trained to retrieve from the dataset based on the current context of the agent and thus make a better prediction of the Q-values. The approach is called Retrieval Augmented RL (RARL) and is demonstrated on Atari domains.