Retrieval-Augmented Reinforcement Learning
Authors: Anirudh Goyal
Abram L. Friesen
Theophane Weber
Andrea Banino
Nan Rosemary Ke
Adria Puigdomenech Badia
Ksenia Konyushkova
Michal Valko
Simon Osindero
Timothy P Lillicrap
Nicolas Heess
Charles Blundell
Venue: ICLR 2022 Submitted (2022)
Keywords: replay buffer
reinforcement learning
offline RL
attention

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.


Under review as a conference paper at ICLR 2022
RETRIEVAL-AUGMENTE REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Most deep reinforcement learning (RL) algorithms distill experience into para- metric behavior policies 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 dataset of past experiences to optimal behavior. Specifically, we augment an RI. 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 method into two different RL agents: an offline DQN agent and an online R2D2 agent. In offline multi-task problems, 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.
show that the sarrival--augeenee DQN agent
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, dominant paradigm in modern reinforcement learning (RL) is to learn general purpose behaviour rules from the agent's 's past experience. These rules are typically represented in the weights of parametrie 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 stored in an experience replay dataset,
Lin 1992). For example, many off- tff-policy algorithms reuse past experience by "replaying" trajectory snippets 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 al. 2015).
This paradigm has clear advantages but at least two interrelated limitations: First, after learning, an agent's 's past experiences no longer plays a direct role in the agent's '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 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 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 occur (Weisz et al. 2021).
network capacity. Second, since the information provided by

---

Under review as a conference paper at ICLR 2022
B
C
AGENT OUTPUT
RETRIEVAL BATCH
RETRIEVAL BATCH
AGENT PROCESS
FORWARD BACKWARD SUMMARIZER
RETRIEVAL PROCESS
ENCODER
1l--1
102
BATCH
ENCODER
RETRIEVAL PRE-PROCESS
AGENT
SAMPLED BATCH
DATASET B
DATASET
ZI
RETRIEVAL PROCESS
RETRIEVAL PRE-PROCESS
Figure Retrieval- anrreetll-augmeeted agent (R2A) architecture: (A) R2A augments agent retrieval process. retrioval process and the ageni maintain sepurate internal states, ma, and respectively. The retrieval pracess retrieves information relevant to the agent's current internal state s, from the retrieval batch, which a pee pre-processed sample by the agent process to inform its output (o.g., a policy or value function). (B) A batch of raw trajectorics sampled from the retrieval datasct and encoded (using the same encoder as the agent). Each encoded trajectory is then summarized via forward and backward summarization The retrieval process parameterized
retrieval dataset
retrieved information u, is
Junctions (section 2.2 and sent the retrieval process. state m. is partitioned into slots. Each slot independently retrieves information from the retrioval batch, which is used to update the slot's representation
model and the internal
agent process in uy. Slots also interact with each other via self-attention. See section 2.3 for more details.
In this work, we develop an algorithm that overcomes these limitations by augmenting a standard reinforcement learning agent with retrieval process (parameterized via 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 large pool of past trajectories stored in dataset (e.g., replay buffer), with the aim of integrating information across these. The proposed algorithm (R2A), shown in Figure] [] enables agent to retrieve information from a dataset of trajectories. The high-level idea is to have two different processes. First, the retrieval process, makes "query" to search for relevant contextual information in the dataset. Second, the agen! 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 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. [A]. 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 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 on Atari games in single task setting. We build upon R2D2 (Kapturowski et al. 2018), a state-of state-of-the-aa art off- pooicy RL algorithm. Second, we evaluate it on multi-task offflie 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 multi-task offline RL environment where the data the dataset comes from multiple tasks. all these cases, we show that R2A learns faster and achieves higher reward compared to the baseline.
BETRIEVAL-AUGMENTT
AGENTS
We now present our method for augmenting an RL agent with a retrieval process, thereby reducing the agent's dependence on model capacity, and enabling fast and flexible use of past experiences. A retrieval-augmenner 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 | for overview. The retrieval process is trained to retrieve information that the agent improve its performance, without explicit knowledge of the agent's s policy. Importantly, the retrieval process has its own internal state, which enables to integrate and combine information across
to
2

---

Under review as a conference paper at ICLR 2022
retrievals. In the following. focus value- based methods, such as DQN (Mnih al. 2015a and R2D2 (Kapturowski et al. 2018), but our approach is equally applicable to policy-based methods.
2.1 RETRIEVAL-AUGMENTTI AGENT
Formally, the agent receives an input a; at each timestep Each input processed by neural encoder (e.g., resnet if the input is an image) to obtain an abstract internal state for the agent 8, = -x;,),)),,, For clarity, we focus here on the case of single vector input, however, each input could also include the history of past observations, actions, and rewards, is the case when fene is recurrent network. These embeddings are used by the agent and retrieval processes. The retrieval process operates
dataset B {(Ee at; re), (Ii+T; 0t+T, (T+T)) of 1-step trajectories, for 1 2 1. This dataset could come from other agents or experts, as in offline RL or imitation learning, or consist of the growing of the agent's own experiences. Then, retrieval-augmenned agent (R2A) consists of the retrieval process and the agent process, parameterized by = { (reur ( ),
Retrieval process J0,8 m 8t ++ M4e U4
Agent process from 81, + Qo Ue: a)
Retrieval Process. The retrieval process is parameterized as a neural network and has an internal state m4- The retrieval process takes in the current abstract state of the agent process 8, and its previous internal state m and uses these to retrieve relevant information from the dataset B, which then summarizes in vector u,, and also updates its internal state my
Agent Process.
The current state of the agent and the information from the retrieval process u, is then passed the action-value function, itself used select external actions.
The above defines parameterization for a retrieval-augmented agent. For retrieval to be effective, the retrieval process needs to: (1) be able to efficiently query large dataset of trajectories, (2) learn and employ a similarity function to find relevant trajectories, and (3) encode and summarize the trajectories in manner that allows efficient discovery of relevant past and future information.
Below, we explain how we achieve these desiderata. At 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 rretrieval-
summarization of the retrieval dataset and corresponding aux- iliary 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.c., ) However, this representation is a function only of past observations (i.c., it's 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 summarization function, applied 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 Transformer).
Encoding and forward-backe
Forward Summarizer ffnd (#1,.... 81) + hi
Backward Summarizer /bwd (ST,.... 81) + b
For each trajectory in the retrieval batch, we represent each time-step within a trajectory by a set of two vectors hit and b (Figure 6 in the appendix) where h summarizes the past (i.e., from
3

---

Under review as a conference paper at ICLR 2022
Algorithm I One timestep of a retrieval-auggmented agent (R2A).
Input:
Current input æt, previous retrieval process state m4--1 {m- 1 E [1. nl... datasct i-step trajectories {(a), hi, bi, di, ri) ... (#i+1, hitss bes aeti. fi-1)} for I 2 and S S Venj- where h and are the outputs of the forward & backward summarizers.
Encode the current input time-step I.
Sc pone (+t)
Step 1: Compute the query. For all 1 & l S ny, compute mi = GRUg (s., m
q. fquery (m.
Step 2: Identify the most relevant trajectories. For all S k < ny, 1 j S S i S Anjj
K.j (h:)
W
et Ri.i
vd.
h softmax (e%)).
di
Given scores the top-kraj trajectories (resp. top-ksmee states) are selected and denoted by T (resp. S!):
Step 3: Retrieve information from the most relevant trajectories and states. o softmax (ei.), i T,", . E S g. Σi, o., where v., bu, W
Step 4: Regularize the retrieved information by using information bottleneck.
2, p(>(g?)
Step 5: Update the states the slots. Slotwise update using retrieved information:
mk a 1 + zt Vk E (1.....,ny)
Joint slot update through self-attention:
cf m WSA
yk(L....,,.
BNN = softmaxe
where KS." = (mm) WTAT
Vk, N E (1,...,,y)) (u.
Vd.
m +- mi
Ev UK where uy m, W
Step 6: Update the agent state using the retrieved information.
d, W
= WT)"T
wk E fl,,..,,,mp
(d,xt) var
7 softmaxi
Ex nov where Us 2 W Ak (l...,,
S, St ue
= to t' time-steps of the 4th trajectory) while b,, summarizes the future (i.e., from
to time-steps) within the 7h trajectory. In addition, taking inspiration from Jaderberg et al. 2016: Trinh et al.) 2018; Ke et al. 2019; Devlin et al.] 2018 Mazoure al. 2020 Banino 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 b,,) to capture meaningful information for the unknown downstream task. For our experiments, we use supervised losses where have access actions rewards in the retrieval batch. For ablations we also experiment with self-supervise losses. For supervised auxiliary losses, we use policy, value and reward prediction (Silver et al. [011]] Scchittwieser al. 2019). and for self-supervised losses, use a BERT-style masking loss (Devlin et al. 2018).
2.3 RETRIEVING CONTEXTUAL INFORMATION FROM PAST EXPERIENCES.
In this section, explain how the retrieval process, when provided with relevant contextual infor- mation represented by the agent's current state s,, interacts with the summarized information in the retrieval batch to select information u, provide the agent in return.
4

---

Under review as a conference paper at ICLR 2022
Retrieval process state parameterization. We parameterize the process that retrieves information from past experience as a structured parametric model with multiple separate memory slots (or sub-units). The state of the retrieval process is set of nj memory slots denoted by m, = {mm k e ".....,,yy] (indexed by the agent time-step (). Slots are initialized randomly at the beginning of the episode. Each slot independently queries and retrieves relevant information from the pool of data. Subsequently the slots update their values independently based on the retrieved information. This is followed by an integration step during which information is shared between slots. Algorithm ! provides a specification of the six steps of R2A which we now explain in detail below.
Step 1: Query computation. Each slot independently computes its prestate using an GRU on the contextual information from the agent: m = GRU, (s,, m Ak E J....,ny! Then, each slot independently computes a retrieval query which will be matched against information in the retrieval batch: q.¹ queryy (m l E [T, [ where of is the query generated by the Ath slot at timestep t.
(s,,
Identification of most relevant trajectories and states for each slot (Figure 5A) The retrieval mechanism process uses an attention mechanism to match a query produced by the retrieval state associated with each slot m to keys computed on each time step of each trajectory of the retrieval batch. Formally, for each time step and each trajectory in the buffer, we compute a key s.y by using a linear projection with matrix W on the forward summaries h: Ki. = (hi Wee)", Each
Step 2:
query qf is then matched with the set of all keys Ki.j, forming attention logits 2 ck (6 Ki and
vT
corresponding attention weights of. softmax (4t,).. i S Puraj 0 JJ.T.
Intuitively, of,, captures the extent to which the j"h timestep of the 7th trajectory in the buffer will be relevant to memory m through the query of- It follows that E, o,,j is a measure of how relevant the 7th trajectory is as a whole for of. Following previous work (Ke et al. 2018 Goyal et al. 2019b), matching only on the most relevant trajectories will increase the robustness of the retrieval mechanism. We therefore select, for each query, the set T,² of Kujj trajectories with highest aggregated score Note that typically the queries corresponding different slots will select different top-kinj
Σ, att trajectories from the retrieval batch. Following the selection of relevant trajectories, we renormalize the weights a, and use a top-k mechanism ance-again, this time to choose the set of most relevant states S. (i.e. which maximizes Eiers a..)).
Step 3: Information retrieval from the most relevant trajectories and states (Figure 5). The next step of the retrieval mechanism consists in computing the renormalized weights a on the subsets T and St (af,) = softmax (Lt)). i E T j St) and use those weights to compute the final retrieved information. The value retrieved from the buffer for query of is computed as the a-weighted average of linear function of the backward state summaries: g. Ei.s of Vig where Vi, b.. W ret
Step 4: Regularization of the retrieved information via an information bottleneck. We regular- ize the retrieved information of via the of information bottleneck (Tishby et al. |000 Alemi eet al. 2016): intuitively, each query pays a price of to exploit information from the retrieval batch. Formally, we parametrize two Gaussian distributions pl [Z) gf (which has access the retrieved information) and r r(Z\m,-) (which only has access to the memory units). We define z as a single sample from pf 1 2 via the reparameterization trick to ensure differentiability (Kingma & Welling 2013 Rezende et al. 2014), and ensure that 2t does not contain too much information by adding an additional loss (plr) to the overall agent loss. We provide more details in the appendix.
Step 5: Slot update. The representation of each slot is first additively updated as function of the retrieved information m +- m + 2 The final update m consists in an update in which all slots interact through self-attention (as normally done in Transformers) (see Algorithm [ for details).
Step 6: Updating the agent state using retrieved information. The primary goal of the retrieval process is to extract information which may be useful for the agent process. Here, we use the retrieved information to change the state of the agent 8, In the previous step, the retrieved information is used to change the state of the slots. For this step, we similar attention mechanism. Here, the
guuery is parameterized as a neural network. We drop time indexing from attention-related quantities to simplify notation.
5

---

Under review as a conference paper at ICLR 2022
Different games
150
Patselaga Inprovmmett
39
Mering Tea
garres
Figure Relative improvement retrieval augmented R2D2 vanilla R2D2 different Atari games, measured by normalized score. We report the average score from seeds method and per game.
query is function of the state of the agent process d, = sg Wag. which are matched with the keys RK =z, ag
Akk E (1...., ny) }, as a result of retrieving information, forming attention weights
7h. = softmax
ddxc") vd= ) The values generated as a result of retrieved information by different slots and
the attention weights are then used to update the state of the learning agent: ut < Σx Uh where 22 >:W" vk E (1.....ny). U, is the result of the attention over the retrieved information which is then used to change the representation of the agent process 8t + 8; + U2 We also shape the representation of the action-value function Q(s,, 2 az) by conditioning the value function on the retrieved information z (again via a similar attention mechanism).
EXPERIMENTAL RESULTS
To evaluate R2A, we analyze its performance in three different settings. First, we evaluate on the Atari arcade learning environment (ALE) (Bellemare et al. 2013). single-task off- aff-policy setting where the retrieval process extracts relevant information from the agent's 's current replay buffer. We then a series of ablations of R2A better understand the roles and effects of its components. Second, we evaluate on a multi-task, offline environment that we created, called gridroboman. In this environment, single network is trained on data from all tasks and then, at evaluation time, the retrieval process queries a retrieval dataset containing only data from the task being evaluated. Third, we evaluate R2A a multi-task offline version of the Baby Al environment (Chevalier- Boiiierrrro et 2018). Again, a single network is trained on all tasks but now the retrieval process queries retrieval dataset containing data from all tasks.
In our experiments, the retrieval process selects the top kinj 10 most relevant trajectories (step 2. section 2.3), and then retrieves relevant information from the selected trajectories (step 3, section 2.3 using the top kstates 10 most relevant states. To summarize the experiences in the retrieval batch we use a forward and backward GRU with 512 hidden units. To train the representation these, we use auxiliary losses in the form of action, reward, and value prediction (section 2.2). See the section A.3 in appendix for more details about the experimental setup and training losses.
3.1
ATARI: SINGLE-TASE OFFIPYLICY RL
In this first experiment, our goal is to evaluate whether retrieval augmentation improves the perfor- mance and sample efficiency of a strong, recurrent bascline agent on a challenging, visually-comple environment- the Atari 2600 videogame suite (Bellemare et al. 2013). We use recurrent replay distributed DQN (R2D2, Kapturowski et al. (008)) as the baseline agent and compare retrieval- augmented R2D2 (RA-R2D2) to vanilla R2D2. The retrieval dataset is the agent's current replay buffer. The agent process is parameterized as an GRU (Hochreiter Schmidhuber 1997). and the retrieval process is parameterized 8 slots. We randomly sample a retrieval batch consisting of 256 trajectories from the retrieval dataset. In [ we report the relative improvement of RA-R2D2 versus the R2D2 baseline. We observe an increase of 11.32 + 1.2% in the mean human normalized score relative to the R2D2 baseline over 2 billion environment steps, demonstrating that retrieval augmentation is quite beneficial in Atari and
the slot-based recurrent architecture described in Section 2, using

---

Under review as a conference paper at ICLR 2022
that the agent's own replay buffer is useful source for retrieval. Raw scores and training curves presented in appendix A
3.1.1 ABLATIONS AND ANALYSIS
To understand the benefit of different components of retrieval aug- mentation, we ablate RA-R2D2 on the 10 Atari games it performs best relative to R2D2. The ablations are as follows.
80
Relative Improvement
60
Importance of a separate retrieval process (A-1). In R2A, the re- trieval process and the agent process are parameterized separately, i.c. they have their own internal states. Here we examine what happens when the agent's is state is used to query the retrieval batch instead of using the retrieval state my. To implement this we modify
40
Step of Algorithm [ to make the query a direct function state of the agent, Q: [query (.r) The resulting query is used in the
the
same way as above. This is similar to the episodic control baseline
of (Pritzel et al. (017).
RA-R222
A-1
A-2
A-3
A-4
A-5
Importance of retrieving information (A-2). We examine what hap- pens when the retrieval process does not have access the retrieval dataset and hence no information is retrieved, keeping all else the same. In this ablation, the retrieval process updates the state of the slots using a transformer (i.e., in Step we replace GRU with a transformer), and the updated state of the transformer is used by the agent process to shape the representation of its value function.
(Different Ablations)
Figure Relative performance ablated RA-R2D2 versus baseline R2D2 for ab- lations en 10 Atari games. black lines rep- Standard deviation.
Shorter retrieved trajectories (A-3). We decrease the length of the trajectories that are retrieved and summarized during retrieval pre-processing, thus reducing the amount of past and future information the retrieval process can retrieve. By default, the trajectories in the retrieval dataset are of length 80. To perform this ablation, we decrease the length of the effective context to only include information from timesteps.
Importance of auxiliary losses to summarize retrieval batch (A-4, A-5). Here we study the use of self-supervised BERT style masking losses in addition to using action, reward and value prediction. We use these auxiliary losses top of the representation learned by the forward and the backward dynamics model. To implement these losses, we randomly mask 15% of the hidden states a trajectory, and then using the representation of hidden states at other time-steps, we predict the representation of masked hidden states. In A-5, we study using only self-supervised BERT style masking losses for summarizing the trajectories.
Figure 3 shows the performance of RA-R2D2 and each ablation relative to the R2D2 baseline. Ablation A-1 shows that it is crucial parameterize the agent process and retrieval process separately, as using the agent state does no better than the baseline. Ablation A-2 shows that the retrieval process finds relevant information that is useful to the agent. It also shows that the performance gains of R2A are not simply due to an increase in model capacity but depend on retrieving information from the retrieval dataset. Ablation A-3 shows that decreasing the length of the context the retrieval dataset results in worse performance, thus showing the importance of our forward and backward trajectory summarization. Ablation A-4 demonstrates that the performance of R2A can further be improved by incorporating BERT style auxiliary losses but that only using BERT style auxiliary losses results in worse performance (but still better than the baseline R2D2). For completeness, Tigure[TO in the appendix repeats the ablations on the games that RA-R2D2 performs worst (relative to R2D2). The main takeaway that the performance of the R2A can be greatly improved by optimizing hyperparameters for each game separately (A-6), which we did not do for our experiments.
3.2 GRIDROBOMAN: MULTI-TASK OFFLINE RL WITH A TASK-SPECIFIC RETRIEVAL DATASET
Beyond querying the agent's 's own experiences, retrieval can provide helpful information from other sources of experiences, including experts other agents the in offline RL where the agent must learn from a fixed dataset of experiences generated by other agents without interacting with the environment during training. A major challenge in offline RL is distributional shift-the the mismatch
7

---

Under review as a conference paper at ICLR 2022
40
Egisode retur
40
Episode retur
40
-00
30
Episode
20
20
20
10
DQN
10
RA-DON
RA-DON
RA-DON
%
250k
%
250k
500k 750k
8
500K
750k
Learner stops
Leamer stops
Learner stops
(9) Training and testing on 10
(b) Tmining testing on makss
(e) Training and testing on 30 tools.
Figure 4: Gridroboman: Multi-lask oflline RL with tpsk-specific retrieval datasel. Average episode return learner sleps for the malti- task gridrohoman environment when training and evaluating on 10, 20, and 30 tasks. With fewer tasks (n), the huseline DQN agent (blue) and the retricval-augmeett augeeetedd DQN agent (orange) perforn identically; however, when number of tasks increnses c), retrievalanngmented agent learns much more effectively than the baseline DQN agent. Results are the average of 3 seeds for each method. Curves for individual lasks are shown in the appendix (Figure 12
between the distribution of states in the training data and those visited by the agent when acting- which makes it difficult to learn an accurate value function for states and actions rarely seen during training. We hypothesize that the retrieval process by retrieving trajectories (including states, actions, and rewards) relevant to the agent's s current state, particularly for states and actions that rare in the offline dataset. We test this hypothesis multi-task offline RL setup where a single agent is trained on multiple tasks simultancously but at training and evaluation time the retrieval dataset contains only trajectories from the evaluated task.
improve performance in the offline setting
For this experiment, we created a minimalistic grid-world-based robotic manipulation environment (gridroboman) with 30 tasks related the three objects (red, green, and blue) on the board. Gridrobo- man is built on the pycolab game engine (Stepleton et al.). The environment is inspired by the challenges of robotic manipulation, and includes tasks such as "go to object X" and "put object X on object Y". Details of the environment, its tasks, and an example figure are presented in the appendix (section A.4) Here, we incorporate retrieval augmentation into vanilla DQN agent as seent-state-recurrence is not needed for this task. Figure 4 shows the results of training retrieval- augmented DQN (RA-DQN, orange) and DQN (blue) on increasing numbers tasks. With fewer tasks, RA-DQN and DQN perform identically; however, when the number of training tasks increases the retrieval -uggeeetd agent is able to learn much more effectively than the baseline agent. Training on more tasks requires either additional model capacity or the ability extract information from fewer relevant samples for each task. By directly querying task-relevant experiences in the offline dataset, retrieval augmentation improves sample efficiency. Note that while the retrieval process does afford extra model capacity to the agent directly, ablation A-2 in section shows that the retrieved information what crucial to performance, not the increased capacity.
3.3 BABYAI: MULTI-TASK OFFLINE RL WITH MULTI-TASK RETRIEVAL DATASET
In this third experiment, our goal is to evaluate the benefit of retrieval augmentation when data from other tasks is present in the retrieval dataset. Multi-task retrieval data be either harmful if the retrieved information misguides the agent or beneficial if information from the other tasks is relevant the current task. Due the use of attention in the retrieval process, we hypothesize that R2A will be able to ignore distracting information and retrieve relevant information from other tasks.
To test our hypothesis, we use the Alby environment Chheaalier-Boiseer et al. 2018). a partially observable multi-room grid world in which harder tasks are composed of simpler tasks and are formulated using subsets of a synthetic language. At the start of each episode, the agent is placed in random room and must navigate to a randomly located goal. Due to the partial observability, we use a recurrent DQN (RDQN) agent as the baseline and compare its performance to retrievallaugmmented RDQN (RA (RA-RDQN) agent. As is common in this environment, we measure the success rate of each agent, defined as the ratio of tasks the agent was able to accomplish given a fixed number of steps for each task. Table [ shows the performance of RA-RDQN with a multi-task replay, RA-RDQN with a replay specific to the current task, and the baseline for varying amounts offline training data (50K trajectories per task versus 200K trajectories per task). As expected from the previous experiment, retrieval augmentation improves performance over the baseline when using a single-task replay. Performance further improves when using multi-task replay. We believe that this is due to
8

---

Under review as a conference paper at ICLR 2022
Table 1: BabyAl: Multi-task offline RL with multi-task retrieval datasel. Mean success rate of retrieval- recurrntt DQN (RA- RDQN) as compared to a recurrent DQN (RDQN) baseline on the 10 BabyAl lovels as function of the amount of training data. RA-RDQN run twice, once with only the current lask being evaluated in the retrieval dataset and once with all tasks in the retrieval dataset. Results the average of 3 random seeds with standard erross
Method
Success Rate (50K)
Success Rate (200K)
RDQN
32% 1: 4%
45% + 6%
R RA-RDON (single-task retrieval huffer)
48% + 4%
64% + 5%
RA-RDQN (multi-task retricval buffer, without IB)
47% + 3%
59% + 6%
RA-RDQN (multi-task retrieval buffer)
55% 1. 5%
74% 1 3%
the compositional nature of tasks in BabyAl, where retrieving information about the current subtask is more informative than retrieving information about the overall task.
RELATED WORK
Episodic control. The idea of allowing deep RL agents adapt based on past experiences using a non -on-parametric memory is not new (Blundell et al. 2016; Pritzel et al., 2017 Hansen et al.) [2018] Eysenbach et al. 2019 van Hasselt al., 2019 The basic idea is that the agent equipped with an episodic memory system, which is used to remember and recall past experiences to inform decisions. There are two important differences between R2A and these methods. (1) In these methods, a local action-value function is constructed by using information about the nearest neighbors in the replay buffer, and then the agent makes a decision about which action to execute based on both the local value function well the global value function. However, in the proposed work, employ parameterized network (the retrieval process), which has access to the information in the replay buffer, and the agent process uses the retrieved information to shape the predictions of its value function in fully differentiable way (using attention). (2) In these episodic control methods, there is only one process (the agent), which has direct access to the replay buffer. However, in R2A, the agent process only has indirect access to the replay buffer via the retrieval process.
Separation of concerns. In Hierarchical RL (HRL) (Heess et al., 2016 Frans et al. 2017 Vezhnevets et al. 2017 Florensa et al., 2017 Hausman et al. 2018 Goyal et al.
2019c). there's 's separation of concerns among different policies, each policy focuses on different aspect of the task, e.g., giving task relevant information to the high level policy only such that low level policy learns behaviours that are task agnostic. In these methods, the high level policy shapes the behaviour of low level policy by either influencing representations by influencing rewards. It is possible to view our work through an analogous lens: wherein the "retrieval process" is the higher level policy (and has access to the all the information in the replay buffer) and is influencing the representation of the agent process that is interacting with the environment. However, there are also notable differences in our work- for instance, the agent process also directly shapes the representation of the retrieval process, which is generally not the case in HRL (e.g.. in Vezhnevets et al. (017) the manager directly influences the worker, but the worker does not directly influence the manager).
Retrieval in language models. Several retrieval-based methods have recently been developed for question answering. controllable generation, and machine translation (Guu al.) 2020 Lee al.) 2019: Lewis et al. 2020 Sun et al. 2021). The general scheme in such methods is to combine a parametric model (like a BERT-style masked language model or pre-trained seq2seq model) with con-paraettric retrieval system. They show that such models generate diverse and factually correct language than the state of the art parametric only baselines. These methods share some similarities with our proposed model since they all involve retrieval component but focus on different domains.
Additional related work is discussed in the appendix (section A.1
5 CONCLUSION.
In this work, we developed R2A, an algorithm that augments an RL agent with retrieval process. The retrieval process and the agent have separate states and shape the representation and predictions of each other via attention. The goal of the retrieval process is to retrieve useful information from a dataset of experiences to help the agent achieve its objective more efficiently and effectively. We show that R2A improves sample efficiency over R2D2, strong off- off-policy agent, and compensates for insufficient capacity when training in a multi-task offline RL environments. Multiple ablations show the importance of the different components of R2A, including retrieving information from past

---

Under review as a conference paper at ICLR 2022
experiences and parameterizing the agent and retrieval process separately instead of giving the agent process direct access to the replay buffer.
REFERENCES
Abbas Abdolmalcki, Jost Tobias Springenberg. Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. arXiv preprint arXiv: arXiv:1806. 1806.0006222 2018.
Alexander A Alemi, lan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410, 2016.
James F Allen and Johannes Koomen. Planning using a temporal world model. In Proceedings of the Eighth international joint conference Artificial intelligence- Volume 2, pp. 741-747, 1983.
Jose A Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brand- stetter, and Sepp Hochreiter. RUDDER: Return decomposition for delayed rewards. arXiv preprint arXivv1800.008877, 2018.
Andrea Banino, Adria Puigdomenech Badia, Jacob C Walker, Tim Scholtes, Jovana Mitrovic, and Charles Blundell. Coberl: Contrastive bert for reinforcement learning. ICML 202 Workshop Unsupervised Reinforcement Learning. 2021.
Peter W Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, and Koray Kavukcuoglu. Interac- tion networks for learning about objects, relations and physics. arXiv preprint arXiv: 1612:002,, 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: arXiv:/806.00 1806,01261, 2018.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel 7. Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control. arXiv preprint 1XXiv16066004660 2016.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David I Luan, and Ilya Sutskever. Generative pretraining from pixels. In International Conference Machine Learning, pp. 1691- 1703. PMLR, 2020a.
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton. Big self-supervised models are strong semi-supervised learners. arXiv preprint arXiv:2006. 10029, 2020b.
Maxime Chevalier-Boiss Boisvert, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning. In International Conference on Learning Representations, 2018.
Thomas M Cover. Flements of information theory. John Wiley & Sons, 1999.
Wojeiech M Czarnecki, Razvan Pascanu, Simon Osindero, Siddhant Jayakumar, Grzegorz Swirszez, and Max Jaderberg. Distilling policy distillation. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1331-1340. PMLR, 2019.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Fukasz, Kaiser. Universal transformers. arXiv preprint a78.v.1807.03879, 2018.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre- Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: /8/0v 10.0040 2018.
10

---

Under review as a conference paper at ICLR 2022
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6:503-556, 2005.
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. Search on the replay buffer: Bridging planning and reinforcement learning. arXiv preprint arXiv:1906.05255. 2019.
Carlos Florensa, Yan Duan, and Pieter Abbeel. Stochastic neural networks for hierarchical reinforce- ment learning. arXiv preprint arXiv: 1704. 03012, 2017.
Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charlie Deck, Joel Z Leibo, and Charles Blundell. Generalization of reinforcement learners with working and episodic memory. arXiv preprint 1rXiv:1910.13400 2019.
Kevin Frans, Jonathan Ho, Xi Chen, Pieter Abbeel, and John Schulman. Meta learning shared hierarchies. arXiv preprint arXiv:1710.09767, 2017.
Alexandre Galashov, Siddhant M Jayakumar, Leonard Hasenclever, Dhruva Tirumala, Jonathan Schwarz, Guillaume Desjardins, Wojeiech M Czarnecki, Yee Whye Teh, Razvan Pascanu, and Nicolas Heess. Information asymmetry in Kl-rregllrriedd RI. arXiv preprint arXiv: 1905.01240, 2019.
Anirudh Goyal, Philemon Brakel, William Fedus, Soumye Singhal, Timothy Lillicrap, Sergey Levine, Hugo Larochelle, and Yoshua Bengio. Recall traces: Backtracking models for efficient reinforcement learning. arXiv preprint 10Xx::880400037 2018.
Anirudh Goyal, Riashat Islam, Daniel Strouse, Zafarali Ahmed, Matthew Botvinick, Hugo Larochelle, Yoshua Bengio, and Sergey Levine. Infobot: Transfer and exploration via the information bottle- neck. arXiv preprint arXiv: 10902, 2019a.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv:1909.108: 10893, 2019b.
Anirudh Goyal, Shagun Sodhani, Jonathan Binas, Xue Bin Peng, Sergey Levine, and Yoshua Bengio. Reinforcement learning with competitive ensembles of information-constranned primitives. arXiv preprint arXiv:1906. 10667, 2019c.
Anirudh Goyal, Yoshua Bengio, Matthew Botvinick, and Sergey Levine. The variational bandwidth bottleneck: Stochastic evaluation on an information budget. arXiv preprint arXiv:2004. / 1935, 2020a.
Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blundell, Yoshua Bengio, and Michael Mozer. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv preprint arXiv:2006. 16225, 2020b.
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nico- las Heess, Michael Mozer, and Yoshua Bengio. Neural production systems. arXiv preprint arXiv:2103.01937, 2021.
Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel J Mankowitz, Cosmin Paduraru, et al. RI unplugged: collection of benchmarks for offline reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming Wei Chang. REALM: Retrieval- augmented language model pre-training. arXiv preprint arXiv:2002.08800 08909, 2020.
11

---

Under review as a conference paper at ICLR 2022
Steven Hansen, Pablo Sprechmann, Alexander Pritzel, André Barreto, and Charles Blundell. Fast deep reinforcement learning using online adjustments from the past. arXiv preprint arXiv:1810.08163, 2018.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. In International Conference on Learning Representations, 2018.
Nicolas Heess, Greg Wayne, David Silver, Timothy Lillicrap, Yuval Tassa, and Tom Erez. Learning continuous control policies by stochastic value gradients, 2015.
Nicolas Heess, Greg Wayne, Yuval Tassa, Timothy Lillicrap, Martin Riedmiller, and David Silver. Learning and transfer of modulated locomotor controllers. arXiv preprint arXiv:661005182,, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1735-1780, 1997.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv: 1.05397, 2016.
Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations, 2018.
Nan Rosemary Ke, Anirudh Goyal, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding. In Advances in Neural Information Processing Systems. pp. 7640-7651, 2018.
Nan Rosemary Ke, Amanpreet Singh, Ahmed Touati, Anirudh Goyal, Yoshua Bengio, Devi Parikh, and Dhruv Batra. Learning dynamics model in reinforcement learning by incorporating the long term future. arXiv preprint arXiv:1903.0155 5999, 2019.
Diederik Kingma and Max Welling.
Auto-encoding variational Bayes.
arXiv preprint
1rXiv:1312.6114222 2013.
Janet L Kolodner. An introduction case-based reasoning. Artificial intelligence review, 6(1):3-34, -34, 1992.
David B Leake. Case-based reasoning: experiences, lessons, and future directions. 1996.
Kenton Lee, Ming-Wei. Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. arXiv preprint arXiv: /t0::/990.00000 00300, 2019.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, al. Retrieval-augmmente genera- tion for Knowledge-intensive nip tasks. arXiv preprint arXiv:2005. 11401, 2020.
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.0297/, 2015.
Long-Ji Lin. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine learning, 8(3-4):293-33) 2211 1992.
Siqi Liu, Guy Lever, Zhe Wang, Josh Merel, SM Eslami, Daniel Hennes, Wojciech M Czarnecki, Yuval Tassa, Shayegan Omidshafiei. Abbas Abdolmaleki, et al. From motor control to team play in simulated humanoid football. arXiv preprint arXiv:2105.1211 2021.
Zhuang Liu, Xuanlin Li, Bingyi Kang, and Trevor Darrell. Regularization matters in policy optimiza- tion. 2019.
Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and Devon Hjelm. Deep reinforcement and infomax learning. arXiv preprint arXiv:2000607211, 2020.
12

---

Under review as a conference paper at ICLR 2022
Josh Merel, Matthew Botvinick, and Greg Wayne, Hierarchical motor control in mammals and machines. Nature communications, 10(1):1-12, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:13/2.5602, 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529-555 533, February 2015a.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, nature, 518(7540):529-533 2015b.
Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. AWAC: Accelerating online reinforcement learning with offline datasets. 2020.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. arXiv preprint arXiv:2010. 19944, 2020.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In International Conference on Machine Learning, pp. 2827-2836. PMLR, 2017.
Sébastien Racanière, Théophane Weber, David Reichert, Lars Buesing, Arthur Guez, Danilo Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, 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.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference machine learning, pp. 1278-1286. PMLR, 2014.
Martin Riedmiller. Neural fitted Q ittration-fin experiences with a data efficient neural reinforcement learning method. In European conference machine learning, pp. 317-328. Springer, 2005,
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. arXiv preprint arXiv: 191 1.08265, 2019.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature. 588(7839):604 609, 2020.
Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavioral modelling priors for offline reinforcement learning. arXiv preprint arXiv:2002.08.0 083966 2020.
David Silver. Reinforcement learning and simulation -baedd search in computer go. 2009.
David Silver, Richard $ Sutton, and Martin Müller. Sample-based learning and search with permanent and transient memories. In Proceedings of the 25th international conference on Machine learning. pp. 968-975, 2008.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
13

---

Under review as a conference paper at ICLR 2022
David Silver, Hado Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac- Amold, David Reichert, Neil Rabinowitz, Andre Barreto, et al. The predictron: End-to-a End-to-end learning and planning. In International Conference Machine Learning, pp. 3191-3199, PMLR, 2017.
David Silver, Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140-1144, 2018.
Jost Tobias Springenherg. Nicolas Heess, Daniel Mankowitz, Josh Merel, Arunkumar Byravan, Abbas Abdolmaleki, Jackie Kay, Jonas Degrave, Julian Schrittwieser, Yuval Tassa, et al. Local search for policy iteration in continuous control. arXiv preprint 20Xi::0010005555,
2020. 2020
Tom Stepleton, Olivier Delalleau, Georg Ostrovski, and Sam Wenke. URL httpp://githht
com/ http//eemmmm//yyoo....a.
DJ Strouse, Max Kleiman- Weiner, Josh Tenenbaum, Matt Botvinick, and David Schwab. Learning to share and hide intentions using information regularization. arXiv preprint arXiv:/80SS 02093, 2018.
Yu Sun, Shuohuan Wang, Shikun Feng, Siyu Ding, Chao Pang, Junyuan Shang, Jiaxiang Liu, Xuyi Chen, Yanbin Zhao, Yuxiang Lu, et al. Ernie 3.0: Large-scale knowledge enhanced pre-training for language understanding and generation. arXiv preprint Ttl::207.002131 2021.
Richard $ Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 164):160 1991.
Yee Tch, Victor Bapst, Wojeiech M Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems, pp.
2017
Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/0004057, 2000.
Trieu Trinh, Andrew Dai, Thang Luong, and Quoe Le. Learning longer-term dependencies in mns with auxiliary losses. In International Conference on Machine Learning, pp. 49744 PMLR, 2018.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement lcarning with double q- learning. In Proceedings of the AAAI conference artificial intelligence, volume 30, 2016.
Hado P van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in reinforcement learning? Advances in Neural Information Processing Systems, 111432--14333 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Mel Vecerik, Todd Hester, Jonathan Scholz, Fumin Wang, Olivier Pietquin, Bilal Piot, Nicolas Heess, Thomas Rothörl, Thomas Lampe, and Martin Riedmiller. Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards. arXiv preprint arXiv:1707.08877, 2017.
Rishi Veerapaneni, John D Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model- based reinforcement learning. Conference on Robot Learning, pp. 1439-1456. PMLR, 2020.
Alexander Sasha Vezhnevets, Simon Osindero, Tom Schaul, Nicolas Heess, Max Jaderberg, David Silver, and Koray Kavukcuoglu. Feudal networks for hierarchical reinforcement learning. In International Conference Machine Learning, pp. 3540-3549. PMLR, 2017.
14

---

Under review as a conference paper at ICLR 2022
Oriol Vinyals, Igor Babuschkin, Wojciech M Czamecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung. David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft II using multi- agent reinforcement learning. Nature, 575(7782)350-3544 2019.
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration. arXiv preprint arXiv:1905.0920 2019.
Greg Wayne, Chia-Chun Hung. David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack Rae, Piotr Mirowski, Joel Z Leibo, Adam Santoro, et al. Unsupervised predictive memory a goal-directed agent. arXiv preprint arXiv: 1803. 10760, 2018
Gellért Weisz, Philip Amortila, and Csaba Szepesvári. Exponential lower bounds for planning in mdps with linearly-rellizable optimal action-value functions. In Algorithmic Learning Theory, pp- 1237-1264. PMLR, 2021.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, al. Deep reinforcement learning with relational inductive biases. In International Conference on Learning Representations, 2018.
15

---

Under review as a conference paper at ICLR 2022
A
APPENDIX
CONTENTS
A.
Extended Relevant Work.
16
Information Theoretic Formulation
17
A.3
Atari: Implementation details and raw scores for R2D2 and RA-R2D2
17
A.3.I
Retrieval Process: Retrieving information from the retrieval buffer.
20
A.3.2
Details ablations
22
A.3.3
Additional Atari ablations
22
A.4
4
Gridroboman environment
26
A.4.I
Gridroboman experiment setup
27
A.4.2
Gridroboman hyperparameters)
27
A.4.3
Gridroboman evaluation curves when training on 30 tasks
28
A.5
BabyAl envirommett,
28
A.6
Limitations and Future Work.
29
A.7
Social Impact.
29
8
Rebuttal Phase
29
A.1
EXTENDED RELEVANT WORK.
Model-based RL. Their are different ways integrate knowledge across past experiences. One of the most common way is by learning model of the world, and using the predictions from the model improve the policy and the value function (Sutton 1991 Silver et al.) 2008 Silver 2009
Allen & Koomen 1983 Silver et al. 2016 Racanière et al. 2017 Silver et al. 2018 Springenberg
et al. ala.]] 2020 Schrittwieser et al. 2020). To integrate information across different episodes (potentially separated by many time-steps), model may needed be unrolled for many time-steps leading to compounding errors. In R2A, the agent has direct access to the information in the retrieval dataset, and querying across multiple trajectories (in parallel) in the retrieval dataset potentially separated by hundreds time-steps.
Structural Inductive Biases. Deep learning have proposed structural inductive biases such as
Transformers (Vaswani et al. 2017 Dehghani et al. 2018 Radford et al. 2019; Chen et al. [2020a b.
Dosovitskiy al.] (2200) slot based recurrent architectures (Battaglia et al. 2016 2016 Zambald et al.
2018 Battaglia et al. 2018 Goyal et al. [019b] Watters et al. 2019 Goyal et al. 2020bb] Veerapaneni al. 2020 Goyal et 2021) where the induced structure has improved generalization, model-size scaling, and longrange dependencies.
Reinforcement Learning with Offfine Datasets. Recent work in RL has tried exploiting large datasets collected across many tasks to improve the sample efficiency of RL algorithms /Vecerik et al., 2017 Pertsch et al. 2020 Nair et al.) 2020) Siegel et al., 2020). An advantage such large datasets that they can be collected cheaply, and can then be reused for learning many downstream tasks. A general scheme for exploiting information about such task-agnostic datasets either using them to tasks
directly improve the value function, by extracting set of skills or options and learning by recombining them. our work, we try to use information the replay buffer by querying and searching for the relevant information across multiple trajectories which otherwise would take many replays through coincidentally relevant information for this to occur.
Efficient Credit assignment. Learning long term dependencies requires assigning credit to time- steps far back in the past. Common methods for assigning credit in dynamics model like backpropa- gation through time requires information to be propagated backwards through every single step in the
16

---

Under review as a conference paper at ICLR 2022
past. This could become computationally expensive when used with very long sequences. Methods which try to get around this problem only back back-propagate information through a selected time-steps in the past, realized by a learned mechanism that associates current state with relevant past states (Ke et al., 2018; Wayne et al. 2018) Arjona- Medina et al. 2018 Goyal et al. 2018 Footunato et al. 2019). Most of these works consider assigning credit to states within the same trajectory, whereas the proposed model R2A, searches for the relevant information in the replay buffer which includes information from other trajectories also.
A.2 INFORMATION THEORETIC FORMULATION
Let's denote the agent's policy (A S,, G), where S is the agent's state, A the agent's action, G the information from the retrieval process conditioned on the current state of the agent i.e., G = (S)) and 0 the parameters of the neural network representing the policy. We want to train agents that in addition to maximizing reward, minimize the policy dependence on the information from the retrieval process, quantified by the conditional mutual information I(A: G (S).
This approach of minimizing the dependence of the policy on the information from the retrieval process be interpreted as encouraging agents to learn useful behaviours and to follow those behaviours closely, except where diverting from doing so (as result of using information from the retrieval process) leads higher reward (Strouse al. 2018 Teh et al.) 2017 Goyal et al.] 2019a] Galashov et al. 2019; Czarnecki et al. 2019). To see this, note that the conditional mutual information can also be written as G 8) Exg [ (A S, G) (o S)I] where E.,, denotes an expectation over trajectories, [A S) = E, p(g) Ta(A S,g)- is a "default" policy with the information from the retrieval process marginalized out. We maximize the following objective:
((0) = E [r] (T(A; SS
E [r 3Ki. ( | S, mo(A) [S]]],
where 3 > is tradeoff parameter, Ex,, denotes an expectation over trajectories (for ex. generated by the agent's policy) and DxL. refers the Kuhlback-Leinn divergence.
(Goyal et al. 2019a Galashov et al. 2019) proposes to optimize the Eq. by maximizing the lower bound J(0)
J(0) 2 J(0) = Er,(r 3DKL [Pene [Z] S, G) 7( |S]]]-
(2)
We parameterize the policy (A |A|S, G) using an encoder Penc (Z S, S,G.), decoder pdte [A] S, Z)
and the (2 | () is the learned prior such that no (A] S, G) = Σ. Penc(= S, Pdec( (A S, =) The encoder output believes is important to access in the present state S in order to perform well.
meant to represent the information from the retrieval process that the agent
Due to the data processing inequality (DPI) (Cover 1999) (2; G 8) 2 [(A; G S), and hence to obtain an upper bound on I (Z; GJS), we must first obtain an upper bound on [(2; G = s), and then average over p(s). We get the following result:
I(Z:)/S) & Ep(s) Enggssss
(3)
Such an information bottleneck has shown improve generalization (Thh et al.
2017 Goyal et al.
2019a c Galashov et al.) 2019: Merel et al. 2019 Liu et al. 2019: Goyal et al. 2020a Liu et al.
2021
A.3 ATART: IMPLEMENTATION DETAILS AND RAW SCORES FOR R2D2 AND RA-R2D2.
Table 3 shows the raw scores achieved by RA- RA-R2D2 and the R2D2 baseline. For R2D2 baseline and for the parameterization of the agent process, we follow the exact same training setup as in (Banino shows the learning curve for both the proposed model and the R2D2 baseline for different games.
al. 202T). Figure g
WWe ask the reader to refer to Information Theoretic Formulation section in the appendix. *For experiments, we estimate the the marginals and conditionals using single sample throughout.
17

---

Under review as a conference paper at ICLR 2022
Training Losses:
The parameters of the state encoder are trained by the RL loss and the auxiliary losses used to train the forward/backward summarizer. The state encoder encodes information about the past actions, past rewards and the current observation into an abstract state.
The parameters of the retrieval process are trained by the RL loss and the information
bottleneck regularizer.
The parameters of the agent process are trained only by the RL loss.
Hyperparameters:
At each learner step, the agent samples large batch of 512 trajectories from the replay buffer. fixed fraction the batch (64 trajectories) is used for learning the Q-function and the remaining trajectories forms the retrieval batch.
We sample a different retrieval dataset for each gradient update. The re-encoding of the trajectories is performed for each gradient update during training.
The retrieval process selects the top-hinaj 10 most relevant trajectories (in step 2, section
2.3 and then retrieves relevant information from the selected trajectories (step 3, section
2.3 using the top-kistates 10 most relevant states (see 5
To summarize the experiences in the retrieval batch with 512 hidden units.
forward and backward GRU
We use 8 slots for Atari-R2D2 experiments.
We use auxiliary losses in the form of action or policy prediction (section 2.2 We follow the same setup as proposed in (Schrittwieser et al. 2020).
The value of the 3 coefficient for the information bottleneck regularizer is fixed to 0.1.
The rest the hyperparameters for R2D2 are taken from Kapturowski et al. (2018) and are detailed in 2
Hyperparameter
Value
Optimizer
Adam
Learning rate
0.0001
Q's
0.8
Adam epsilon
10-7
Adam betal
0.9
Adam beta2
0.999
Adam clip norm
Q-value transform
h(x)) sign [2]
+
Discount factor
0.997
Trace length (Atari)
Replay period (Atari)
40
Replay capacity
100000 sequences 0.9
Replay priority exponent
Importance sampling exponent
Minimum sequences to start replay
5000
Target Q-network update period
400
Evaluation e
0.01
Target
0.01
Table Hyperparameters used the Atari R2D2 experiments.
18

---

Under review as a conference paper at ICLR 2022
Implementation details:
We use GRU style gating and layernrm whenever changing the state of the agent process or the state of the slots a residual way (Step 5, Algorithm
If the imput neural network fun consists of two tensors I, y, never concatenate the inputs to the neural network. We make use of attention. We assume one of the inputs is the primary input (lets say æ), and other input is privileged input (as result of doing some expensive computation, let's say y). First, we compute (z)(2), and then using the output of this computation, we cross-attend over y using multi-head attention.
When computing the forward summary of the trajectories in the retrieval batch h = h4 + 8t. we also add the information about the current encoded state i.c., St.
Visual Encoder Visual observations are encoded using ResNet-47 encoder. The 47 layers are divided in 4 groups which have the following characteristies:
An initial stride-2 convolution with filter size 3x3 (1 layers).
Number of residual bottleneck blocks (in order): (2, 4, 6, 2). Each block has 3 convolutional layers with ReLU activations, with filter sizes 1x1, 3x3, and 1x1 respectively ((2+4+6+2)
layers).
Number of channels for the last convolution in each block: (64, 128, 256, 512).
Number of channels for the non-last convolutions in each block: (16, 32, 64, 128).
Group norm is applied after each group, with group size S.
After this observation encoding step, a final layer with ReLU activations of sizes 512 is applied.
Computational Complexity. Here discuss the computational complexity of the R2A.
The following computations are happening at each step within a trajectory:
A particular slot selects the relevant trajectories and relevant states each time-step.
We use an attention mechanism to select the set of relevant trajectories and relevant states. This process scales with the number of trajectories in the retrieval batch.
The set relevant trajectories and relevant states changes across different time-steps within trajectories. This process repeated for all the slots independently, but be efficiently performed on GPUs/TPUs.
We can make following optimizations to reduce the computational complexity of the R2A.
For Atari and Gridroboman experiments instead of using learned attention mechanism to rank the trajectories (in step section 2.3 we can also just use those trajectories which are of high return (Abdolmaleki et al. 2018) since we already know that the data in the retrieval dataset is specífic the task on which we are training the agent.
Once we have selected relevant trajectories, we can further reduce the computational cost of selecting most relevant states (in the selected trajectories) by only considering those states which are semantically similar to the agent's current state.
Improvements tried that seem to improve the performance but were not evaluated exhaustively:
Instead of performing only single step of retrieval at each step, we can also query the replay buffer multiple times within time-step at the cost of increased computational complexity.
We can allow different slots to retrieve information at different time scales.
We note that of the hyperparameters for the retrieval process, except for the number of trajectories in the retrieval batch, remain the same across all experiments (Atari, gridroboman, and Al)y
19

---

Under review as a conference paper at ICLR 2022
A.3.1 RETRIEVAL PROCESS: RETRIEVING INFORMATION FROM THE RETRIEVAL BUFFER.
Here discuss the different steps involved in retrieving information from the retrieval buffer:
Query computation. Each slot independently computes its prestate using the contextual information from the agent using a GRU. Then, each slot independently computes a retrieval query which will be matched against information in the retrieval buffer.
Figure Retrieval in step 2, section 2 In the figure scalars represent attention weights. B) Then the retrieval process retrieves relevant information from the selected trajectories by selecting most relevant states from the top-l trajectories as detailed step 3, section
trajectories states. retrieval process top- tpp-keraay most trajectories as specified
Identification of most relevant trajectories and states for each slot. The retrieval mechanism process uses an attention mechanism to match query produced by the the retrieval state associated with each slot to keys computed on each time step of each trajectory of the retrieval batch. This process assigns single scalar each state within trajectory. These scores used to assign scalar to each trajectory and then normalized across trajectories. These normalized scores are then used by the retrieval process to select the top-kıraj most relevant trajectories (see Figure. 5). The process is then again repeated to select most relevant states with the selected trajectories. Here, a single scalar is assigned to each state with in selected trajectories and then normalized across all the states. These scores are then used to select most relevant states within the most relevant trajectories. Backward Summarizer 6
Forward Summarizer he
Figure Each imjectory in the retrieval batch summarized via forward and backward running parameteric structured model (Transformers trajectory represented by (he, by) uplle where represents summary past
based recurrent network). Each and by represents the summary of the future with that trajectory.
Information retrieval from the most relevant states. In the previous step, we selected most relevant states within the most relevant trajectories. Now we integrate information across the selected states. We again use the attention mechanism. Here the query by the slots is matched against the keys (which are the function of the forward running dynamics model), to retrieve information about the future (and values are the function of the backward running dynamics model). See figure. 6 for more details.
20

---

Under review as a conference paper at ICLR 2022
Environment
R2D2
RA-R2D2
Alien
52268 4125
81626 5123
Amidar
10976 2341
24693 1982
Assault
45521 4751
43497 6653
Asterix
997910 51241
997585 89251
Asteroids
262560 94168
254194 72141
Atlantis
1424723 67589
1516463 98140
BankHeist
3923
19397 5931
BattlcZone
590913 + 128715
466783 + 89724
BeamRider
89038
79241 =
Bowling
250.37 +
253.5 1
Boxing
99
99
Breakout
848 1 121
869 51
Centipede
315329 = 151612
497746 = 12515
ChopperComnand
998842 11451
999936 + 3422
Crazy
184132 12912
226469 9251
Defender
538752 18241
554158 21415
DemonAttack
143552 +
143519 +
DoubleDunk
24 = 0
24 =
Enduro
2332 = 32
2359 + 19
FishingDerby
71.357 2.51
73.13 3.12
Freeway
34 =
33,19 +
Frosthite
168225 51891
511159 = 124512
Gopher
123138 + 4121
123841 + 8914
Gravitar
13142 1214
13198 1102
Hero
43257 4121
35431 2412
IceHHockyy
72.74 1.
71.59 T 12
Jamesbond
24873 - 2141
25873 + 2817
Kangaroo
14614
15232
Krull
158509 21415
183921 31415
KungFuMaster
208583 + 2102
224385 41512
MontezumaRevenge
1966.7 + 987
1750 = 1800
MsPacman
33530 1214
44396 4521
NameThisGame
30232 3151
45140 2412
Phoenix
695251 32515
847914 12415
Pitfall
0.0 0
0.0 + O
Pong
21 = 0
21 + 0
PrivateBye
21602 4124
18923 2415
Qbert
242169 51512
241525 32513
Riverraid
32441 5314
44554 7325
RoadRunner
490872 51512
513521 65132
Robotank
128.44 1: 18
102 1. 29
Seaquest
998664 + 1212
999899 1 51
Skiing
-29973
1212 2412
Solaris
3924 + 3415
3853 = 1241
SpaceInvaders
57198 + 14125
62624 4215
StarGunner
256129 =
478231
Surround
9.52 0
10 0
Tennis
7.16 +
0=
TimePilot
168592 21241
260609 41212
Tutankham
390 41
235 + 89
UpNDown
544439 1 5851
561389 21516
Venture
2000 +
2000 +
VidcoPinball
673679
234151 962119 + 24151
WizardOrWor
78431 5159
73500 + 6231
Yars Revenge
674692 = 125116
513549 89151
Zaxxon
53312 +: 15152
62631 5161
Table Scores obtained different atari games. (average over 3 different seeds).
21

---

Under review as a conference paper at ICLR 2022
A.3.2 DETAILS ON ABLATIONS
Ablation A-I: In this ablation, the internal retrieval state, m, is removed. The query qi is computed from 8; directly as qe fquery(8t)) The retrieval vector Jr is computed as per steps 2-4 of the original algorithm. Step 5 does not occur, and step the only update 8, f 8t + gt-
Ablation A-2: In this ablation, there is no retrieval batch. The retrieval process is parametrized as before with ng-slot memory. Step is identical, but steps 2-4 do oceur. Memory slots are updated as per the joint update of step 5. Step 6 is otherwise identical, except that keys and values are computed from the memory slots m instead of the retrieved values 2!"
Ablation A-3: In this ablation, the length of the trajectories used for summarization is reduced from 80 to 5. The trajectories used for Q learning are unchanged.
Ablations A-4 and A-5: In our experiments, we use standard policy, reward and value prediction cross-entropy losses (Schrittwieser al. (0200) to train the forward and backward summarizers. In A-4, we add BERT loss; tokens are obtained by compressing the agent states with VQ-VAE. In A-5, we only use the BERT loss.
A.3.3 ADDITIONAL ATARI ABLATIONS
Here we perform ablations on RA-R2D2 using the Atari games on which RA-R2D2 performs worst relative to baseline R2D2. Figure. I0 shows the relative performance of different ablations compared to the R2D2 baseline. Ablations 1-5 are as described the main text (section 3.1.1 For these results, we also ran sixth ablation.
20
Relative Improvement
10
0
-10
-20
RA-R2D2
A-1
A-2
A-3
A-4
A-5
A-6
(Different Ablations)
Figure Relative performance RA-R2D2 versus worse than baseline R2D2.
ablations on Atari games which RA-R2D2 performs
Role of kaa and kistaes (A-6). In our experiments, the retrieval process selects the top knaj 10 most relevant trajectories (step 2, section 2.3p and then selects the top kstate = 10 most relevant states of these trajectories (step 3, section 2.3) from which to retrieve relevant information. To better understand the role of these hyperparameters, we independently vary these two hyper-parameters (top kuaa and top kstaess over the values {5, 10, 20). Figure 10 shows that after independently varying these two syper-parameten the performance the R2A can be improved as compared to the R2D2 baseline.
22

---

Under review as a conference paper at ICLR 2022
Cal Alien
(b) Amidar
(c) Assaalt
(6) Asterisk
(e) Asteroids
1 Allamis
(g) Bank Heist
(hi Barle7onc
( Bernerk
① Boxiin
(k) Brenkout
@ Centipede
(m) CoeypeCoommadd
(n) CrazyClimber
[o) Defender
(p) DemonAmack
(g) DoubleDunk
Enduro
(6) Fishingerbby
(1 Freeway
Ool Frostbise
Gopber
(w) Graviar
HERO
Figure 7: Comparing R2Awith R2D2: Learning curves for retrieval -agmmeteee R2D2 (RA. R2D2) and the R2D2 buseline acroes different Atari Green RA-R2D2 and orange baseline R2D2.
23

---

Under review as a conference paper at ICLR 2022
(a)
(b)
(2)
Pung
RiverRaid
(it) RoboTank
Seseuest
Skiing
Solaris
Sumoondd
Tennis
TimePilot
(n) Tumnkham
UpNDown
[gJ Video
VideoPinBall
WizardO/War
YarsRevenge
Zaxxon
Figure 8: Comparing R2Awith R2D2: Learning curves for retricval-angmenntd R2D2 (RA-R2D2) and the R2D2 bascline across different
Alari games. Green RA-R2D2 orange baseline R2D2.
24

---

Under review as a conference paper at ICLR 2022
(o) TocHockey
(b) JamesBoud
(e) Knil
KungFuMaster
Municemseeeeege
Figure 9: Comparing R2Awith R2D2: Learning curves for retrieval-nugmented R2D2 (RA-R2D2) and the R2D2 baseline across different
Alari games. Green RA-R2D2 and orange baseline R2D2.
25

---

Under review as a conference paper at ICLR 2022
A.4 GRIDROBOMAN ENVIRONMENT
Figure 1I: Gridroboman enviromment illustration. On hoard there three colored objects and rhot represented black block. The robot can move itself and move the objects. The tasks are motivated by the robotic manipulation.
In order to test our method in multi-task setting we designed minimalistic grid world robotic manipulation (gridroboman) environment with a single embodiment and multiple tasks. It is implemented based on the pycolab (Stepleton et al.) game engine that provides tools for designing customizable grid world games. The environment and its tasks are inspired by the challenges in robotic manipulation. An illustration of state in the environment is shown in Figure 4
Environment semanties, observations and actions There are three colored objects on the 7 x board: red, green and blue. The black block represents the robot. It can move in 4 directions and the action is skipped if it attempts move into wall (gray). Additionally, if the robot is located the same position as any colored object, it can apply lift action that would enforee the object to move together with the robot. Then, put action allows position the object either on the board or on the top of another single object. Additionally, there is option of skipping an action. The initial state of the environment that includes positions of the objects and the robot is sampled randomly at the beginning of each episode. The agent observation is 11 11-dimensional vector: it includes and coordinates three objects and the robot as well as status each three objects. The status of an object is numerical value: 0 if the object is positioned the board, -1 if it is under another object and +1 if is either held by the robot or located on the top of another object. Each episode lasts for 50 timesteps.
Tasks There 30 possible tasks: lift red, green, or blue object (3), touch red, green, blue object (3), move red, green, or blue to center (3) or to corner (3), touch one object with another (6), move two objects close (3) or far from (3) each other, stack one object on the top of another (6). The nature of these tasks is such that it is impossible to identify the task by the initial state.
The reward is binary in each time step: it is zero when the task is not solved and one otherwise. In order to receive positive reward the following condition should met:
Rewards
lift X: the robot should lift and not put back an object of color X:
touch X: the robot should be located at the immediately adjacent cell to the object of X color and not hold any objects;
move X center: the object or color X should be located in the x 3 block at the center of the board;
move X to corner: the object of X color should be located in any of the 2 corners of the board;
touch X with Y: the robot should be located at the immediately adjacent position to the object of color and hold the object color
move X close to Y: the distance between the objects of colors X and should be no more than 1 in both and y direction;
26

---

Under review as a conference paper at ICLR 2022
move X far from Y: the sum of distances between objects of colors X and Y in and y coordinated should be greater than 9;
stack X on Y: the objects of colors X and Y should be located at the same position one on the top of another.
A.4.1 GRIDROBOMAN EXPERIMENT SETUP
The gridroboman offline RL dataset was generated by training single DQN agent online on each task separately and recording the 100K generated episodes. The Q(s, a) network was 3-layer MLP with 256 units in each hidden layer for Q-function. For the experiments with 10, 20, and tasks, we use the first 10, 20, and 30 tasks, as listed here.
10 tasks: touch red, touch green, touch blue, lift red, lift green, lift blue, red touch green, green touch red, and green touch blue.
20 tasks: The above 10 tasks, blue touch red, blue touch green, red to corner, green to corner, blue to corner, red to center, green to center, blue to center, red close to blue, and red close to green.
30 tasks:
The above 20 tasks, blue close to green, red far from blue, red far from green, blue far from green, red blue, red on green, green red, green blue, blue red, and blue green.
A.4.2 GRIDROBOMAN HYPERPARAMETERS
The DQN agent used the same network as used to create the data: a 3-layer MLP with 256 hidden units per layer. The first layers define fod and the final layer predicts Q. The RA-DQN agent used the same base network plus separate retrieval network.
Below we detail hyperparameters specific to gridroboman. Hyperparameters not specified in this section are the same used in Atari.
At each learner step, the agent samples a batch of 256 states from the replay buffer to train the DQN agent and further samples a batch of 64 trajectories from the retrieval dataset to form the retrieval batch.
The retrieval process selects the top-kunj 10 trajectories with the highest return, and then retrieves relevant information from the selected trajectories (step 3, section 2.3) using the top-kssates = 10 most relevant states.
summarize the experiences the retrieval batch we use a forward and backward GRU with 256 hidden units.
We use 4 slots parameterize the retrieval process.
We use auxiliary losses in the form of value, reward, and policy prediction (section 2.2 We follow the same setup as proposed in (Schrittwieser et al. 2020) and weight these auxiliary losses with a coefficient of 0.1.
The value of the 3 coefficient for the information bottleneck regularizer is fixed 0.3.
The hyperparameters and setup for offline DQN are taken from Gulcehre et al. (2020). They are detailed in 5 As in the above, we used double DQN Van Hasselt et 2016
Hyperparameter
Value
Optimizer
Adam
Learning rate
3e-4
Discount factor
0.99
Importance sampling exponent
0.2
Minimum sequences start replay
5000
TD loss function
Huber(1.0) (1.)
Target Q-network update period
2500
Evaluation
6.5e-4
Table 4: Hyperparameters used the gridrohoman DQN experiments.
27

---

Under review as a conference paper at ICLR 2022
A.4.3 GRIDROBOMAN EVALUATION CURVES WHEN TRAINING ON ALL 30 TASKS
Figure 12 we show the average episode reward obtained by the evaluation agent on each individual task when training on all 30 tasks. Curves are shown for the DQN and RA-DQN agents.
Eplecde
DCM
DOS
DDN
D0s DON
RA-00s
2506 5000 13B
UM
250c 3300 750k
2506
2008
75c4
3M
2008
5005 7004
2506
2008 7508
3M
(2) touch red
(b) touah blue
(c) touch green
(d) lift red
(c) lift blue
DON
DCN
DON
D0N
Learmer steps
Leareer
(D) lift green
EgT red touch blue
(b3 red touch groen
O green touch rod
( greess mouch blue
NA-con
RA-OON
FA.OON
NA DON
1A.00 DON
stapa
staps
blue touch
touch green
(an) red conner
[a] green to couner
(0) blue to comer
DON FALON
DCN RADON
DON CCN
D00 CAN
TAL-Om
2028
grezs
rod close blue
(1) cod close grezs
19:2 KA-DOM
022 BA.com
Cprodd
DCN
2008
Un
50%
2506
(u) blue close to green
(-) red far from blue
(w) red far from green
(3) blue far from green
(y) stack red on blue
bos
DCN
Loumerr stepo
Learner steps
Learrer steps
Learner Stage
Learner
(3) stick rod on green
(2a) stack grees on red
(ab) stack green on blue
(oc) stack blue on red
(ad) stack blue on green
Figure 12: Multi-task offline RL with tusk-specific retrieval datasel. Evallation performance for RA-DQN (orange) and haseline DQN (blue) when truining on all 30 gridroboman usks with a single agent. Curves show the performance of each agent (averaged over seeds) when running that agent online in the enviroment specified task.
A.5 BABYAI ENVIRONMENT.
Observations, actions. We use the same setup for the observations and actions as in (Chevalier Boisvert et al. 2018).
Tasks. We use all the 40 available tasks the Baby Al environment.
md
28

---

Under review as a conference paper at ICLR 2022
Hyper-parameters and RL algorithm. We follow the same setup in the Gridroboman experi- ments. Since BabyAI is a partially observable environment, we summarize the history the agent using a recurrent encoder.
Hyperparameter
Value
Optimizer
Adam
Learning rate
3e-4
Discount factor
0.99
Importance sampling exponent
0.2
TD loss function
Huber(1.0)
Target Q-network update period
1000
GRU hidden state
512
Tahle parametrs
BabyAl Recurrent DQN experiments.
Retrieval Process. We use the same areer-parametes for the retrieval process as used in the grid- roboman enviorment. For the case of multi-task retrieval buffer, we use 32 trajectories corresponding to each task to form the retrieval batch.
A.6 LIMITATIONS AND FUTURE WORK.
It would be useful to investigate and extend the proposed idea in these different ways: (a) First, investigate training of the retrieval process and the agent process using different objectives as compared to training them in an end-to-end fashion, (b) Second, scaling R2A to more complex problems multi-agent problems like in Starcraft (Vinyals al., [019], where one may have access to millions of past experiences. In R2A, only query a subset of the retrieval dataset, which limits the generality of the method. (c) Third, even more intriguing would the possibility of learning an abstract model with abstract internal actions, and rewards, rather than learning model which queries for information from the retrieval dataset, and hence avoiding the need for Monte Carlo tree search common in the state of the art planning methods (Schrittwieser et al. 2020 these problems and seek answers in the future work.
We aim formulate
A.7 SOCIAL IMPACT.
The authors do not foresee negative social impact of this work beyond that which could arise from general improvements in ML.
A.8 REBUTTAL PHASE.
Comparison to Episodic Control baselines. We have performed ablations where compared to the scenario where we just keep the simple episodic memory (refer to section, 3.I.I| ablation A-1). The result of ablations shows that for R2D2 accessing the dataset in non-parametric way actually hurts the performance of the RL algorithm.
BERT style unsupervised losses then not used generally in the R2A method BERT style unsu- pervised losses actually improve the learned representations. We did not choose to use BERT style representations as they increase the number of design choices, increase the number of ablations, increase the number of parameters, whereas the relative gains of using BERT style unsupervised losses is not much. Future work should investigate the use of pre-trained encoders as common in NLP literature.
For Frostbite, episodic control is beneficial. As mentioned in "Building Machines That Learn and Think Like People", in Frostbite, players control an agent (Frostbite Bailey) tasked with constructing an igloo within time limit. The igloo is built piece-by-piece as the agent jumps on ice floes in water The challenge is that the ice floes are constant motion (moving either left or right), and ice floes only contribute to the construction of the igloo if they visited in active state (white rather than blue). The agent may also earn extra points by gathering fish while avoiding a number of fatal hazards (falling in the water, snow geese, polar bears, etc.). Success in this game requires
29

---

Under review as a conference paper at ICLR 2022
temporally extended plan to ensure the agent can accomplish sub-goal (such reaching an icefloe) and then safely proceed to the next sub-goal. Ultimately, once all of the pieces of the igloo are in place, the agent must proceed to the igloo and thus complete the level before time expires. One hypothesis for the better performance could be that the retrieval process is able to efficient utilize the information from the states which are very far apart from the current state (i.e., temporally extended credit assignment). This hypothesis is further validated when we decrease the length of the traces in the dataset. In ablation section, if decrease the length of the trajectories that are retrieved and summarized during retrieval preprocessing, thus reducing the amount of past and future information the retrieval process can retrieve. By default, the trajectories in the retrieval dataset are of length 80. To perform this ablation, we decrease the length of the effective context to only include information from time-steps.
Similar Computation baseline. We also compared baseline which does the same amount of "computations" (Importance of retrieving information (A-2)). We examine what happens when the retrieval process does not have access the retrieval dataset and hence no information is retrieved, keeping all else the same. In this ablation, the retrieval process updates the state of the slots using transformer (i.e., in Step I we replace GRU with a transformer), and the updated state of the transformer is used by the agent process to shape the representation of its value function. Since the use of the transformer number of parameters in this baseline more compared to R2D2. Our results show that the proposed method R2A archives better performance as compared to vanilla R2D2 as well as the baseline which does the same amount of computation (but does not retrieve information).
How does runtime compare to vanilla R2D2? How many more parameters are introduced? As compared to vanilla R2D2, a naive implementation of the proposed method is around 3x slow (and somewhat optimized implementation around 2x slower), with the increase of 10% more parameters. In the ablations, we also examine what happens when the retrieval process does not have access to the retrieval dataset and hence no information is retrieved but everything else is kept the same (same computation and same parameter baseline).
Use of hierarchical attention (top-k trajectories and then top-k states). A hierarchical attention where we first filter trajectories, and then filter states reduces the search space for the attention. For example. consider 512 trajectories, and 80 states in each trajectory, then the attention operator has to select 10 most relevant states out of 40960 states. Whereas if we do initial filtering of top-k trajectories, then the attention operator needs to select 10 most relevant trajectories out of 512 trajectories, and then 10 most relevant states out of 800 states. Such an hierarchical attention has also shown to be successful in various NLP tasks also.
Analysis about retrieved information in the multi-task setting? We did some analysis to study the properties about the retrieved information in the multi-task setting in BabyAI. BabyAl setup contains about 40 tasks. Out of these 40 tasks, around 15 tasks are compositional i.e., requires to compose information from or more sub-tasks, and rest of the tasks requires the agent to execute a particular behaviour (ex. going to the door, fetching a key etc).
We study as to how often the agent retrieves information from other tasks. Ideally, for the compo- sitional tasks, the hope would be the agent access information about the other tasks more often as compared to the tasks which require only a particular behaviour. So, during test time, we study the percentage of times agent access information from the same task compared to accessing information from the different tasks. For compositional tasks, the agent access information from other tasks about 54% of times, while the percentage for single tasks is about 21 %. That said, were not able to find any pattern as to find why the agent is retrieving more information from other tasks for some particular tasks as well any structure about the states from which information is being retrieved.
30


Reviews:
[{"review_id": "s8EAGyBdnbR", "reviewer": "Reviewer_yAsc", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\nThe high-level idea of retrieving relevant information from past/expert experiences to help make decisions in the current states is very natural and well-motivated, though not very novel considering the existing works in experience replay and episodic memory.\nThe paper is well-organized and relatively easy to follow (there are several clarification issues, see the weakness part). \nThis work is a good trial to adapt the success of retrieval-based methods from language models to RL problems.\n\nWeakness:\n\nThis work is closely related to episodic control, which also retrieves information from past experiences for value estimation. Could you add the prior work in episode control as a baseline? This will be helpful to confirm that introducing the complicated retrieval process is necessary rather than just keeping the simple episodic memory.\n\nCould we use R2A in continuous control tasks? Based on my understanding the proposed method can be adapted to RL problems with continuous action. It will be great if there are additional experiments to verify this. Because most works in episode control require discrete action space, this could be one advantage of R2A if R2A can be used in continuous control tasks.\n\nThe retrieval process looks complicated (see the dense notations in section 2.2, 2.3), I'm generally concerned about why should we use such a complicated design? what's the importance of each component in the retrieval process? is it possible to simplify the design? The design choice might need additional explanations about the intuition or ablative studies about the effect of each component.\n1. what's the intuition of using independent slots as the internal state of the retrieval process? why not just one slot? how is the information retrieved by each slot different from each other? Is the number of slots a sensitive hyper-parameter? How does not the number of slots influence the final performance?\n2. What does the internal state m_t represent? It is clear that s_t represents the hidden state information in the MDP with partial observation. Then what's the intuition behind m_t? How is it different from s_t?\n3. Information bottleneck is used to regularize the information retrieval. Is the final performance sensitive to the weight of this regularization loss?\n4. Is the retrieval network optimized with temporal error in optimizing Q function? It is not crystal clear in the paper what is the overall loss term for the retrieval process.\n\nThe application of R2A on offline RL settings is interesting. Why not consider the standard offline RL benchmark (e.g. D4RL) and compare R2A with recent offline RL approaches (e.g. CQL, REM)? The experiment on commonly used benchmarks will be more convincing. \n\nAs pointed out in the first paragraph in section 3.2, one big challenge in offline RL is the distributional shift. It is unclear to me why the retrieval process can be helpful \"particularly for states and actions that are rare in the offline dataset\"? For rarely visited state-action pairs, the Q value estimation is hard, does the retrieval process make the Q value estimate more accurate on these state-action pairs?  It will be great to show the error in Q value estimation (i.e. difference between the ground truth Q value and estimated Q value) on simple discrete MDPs (where we can access the ground truth Q value). If the error becomes smaller after using R2A, such a claim will be further verified.\n\nLooking forward to seeing stronger evidence to better support the advantages of the proposed method. \n\n\n", "summary_of_the_paper": "This work proposes a retrieval neural network to help value-based RL algorithms (i.e. DQN and R2D2) retrieve information relevant to the current state from a dataset. The dataset may consist of the agent's own past experiences, expert demonstrations, or experiences from separate behavior policies (i.e. consider the offline RL setting). The retrieval network work is a recurrent model. The internal state in the retrieval network is partitioned into slots. Each slot independently retrieves information from the trajectory dataset. The retrieved information is used to update the internal state of the retrieval network and used as an additional input of the value function.\n\nExperiment on Atari games verifies that the retrieval network improves the performance of R2D2. Also, this paper presents results on grid environments to verify that the retrieval network can be beneficial in offline RL settings.", "summary_of_the_review": "The proposed method involves complicated design choices which are not well justified. Also, some important baselines are missing in the evaluation in the experiment part. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "eX15-G9hAul", "reviewer": "Reviewer_LSKp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* The paper is well written, and despite the relative complexity of the R2A process I was able to understand most of the details with only a few additional passes.\n* The use of an attention mechanism over the sampled replay buffer I believe is novel, and is a worthy addition to the episodic control literature\n* There are a few promising results (e.g., BabyAI), and the ablations appear to show the importance of the introduced components.\n* The creation of the `gridroboman` offline benchmark is a really nice addition, and should be commended.\n\nWeaknesses:\n* The experimental methodology is lacking in several aspects:\n   * There should be significance testing, and there are quite a low number of seeds (3) across all experiments; ideally there should be a minimum of 5. This is further justified by the high variance of the training curves in the appendix.\n   * The paper should compare to other episodic control approaches, such as [1,2]. This is because the idea of augmenting a controller with instantaneous buffer information is not new (i.e., episodic control), so I would like to see the merit to adopting a hierarchical + modern attention mechanism over say a neural dictionary approach.\n  * There should be information about run time and computational resources compared to the vanilla R2D2 agent due to the additional machinery introduced by R2A.\n  * More details about the BabyAI experiment, such as training curves and number of epochs; were the success rates reported when all agents had reached convergence? This is important as viewing the `gridroboman` results shows that although R2D2 learns slower offline, it does also eventually seem to reach the level of performance of R2A.\n  * I'm somewhat concerned by the inclusion of the information bottleneck regularizer; is this necessary for the method to work generally? There should be an ablation with and without this.\n* BERT style unsupervised loss improves the learned representations in Figure 3, but why is this then not used generally in the R2A method? Does it actually harm performance beyond the top 10 Atari games?\n* In some cases, I'm not entirely convinced by the improvements offered by R2A over R2D2. Particular areas of concern are:\n   * In more than 50% of all Atari games, there seems to be no improvement adding the retrieval mechanism. Furthermore the learning curves in the back don't convince me that overall R2A learns quicker either. This would be ok if it could be justified (e.g., is there something intrinsic to `Frostbite` that means episodic control is beneificial v.s. `Tutankham`?).\n   * While R2A learns `gridroboman` quicker, in the offline setting this is less of a concern, since we are not exploring and so the only incurred cost is computational (not data-efficiency). Therefore the results here aren't that convincing to me, as it seems R2D2 can retrieve R2A performance in the limit; is this always the case? Furthermore, if R2A is more expensive to train per step, perhaps if time was placed on the x-axis the improvements of R2A would be not as clear?\n* More details about the process of double filtering of top-k trajectories and states would be welcome. Is it not the case that we could be prefilterting out useful state information in the first round of top-10 trajectory information? (e.g., what if there was actually the most useful state information was in the *11th* most useful trajectory, but we dismissed it because on average it also contains some irrelevant information). Perhaps some experiments to show why we don't just perform attention over the entire set of retrieved states, (e.g., top-10 states over all 512 trajectories, not just top-10 trajectories).\n\nNits:\n* LSTM and GRU appear to be used interchangeably, when they are different models, for instance in **Step 1: Query Computation** section of page 5.\n\nRefs:\n\n[1] Neural Episodic Control, Pritzel et al. arxiv:1703.01988\n\n[2] Fast deep reinforcement learning using online adjustments from the past, Hansen et al., arxiv:1810.08163", "summary_of_the_paper": "In this paper, the authors introduce R2A, an agent that is augmented with an attention-based retrieval mechanism that augments its state with information directly from an experience replay. They show that this mechanism can improve on R2D2 online performance, and can also improve learning in the offline domain.", "summary_of_the_review": "Overall the paper presents a nice idea and is well written, and introduces a novel offline multi-task benchmark. However I believe there are shortcomings in the experimental methodology, and furthermore despite the relatively large increase in complexity introduced by R2A, I don't believe the authors have enough convincing the reader that this increase is worth it.\n\nI would therefore recommend at this point a weak reject, but believe if the authors could address the points made concerning experimental methodology and design choice explanations, I'd be willing to increase this score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "c0X4SEq5-zY", "reviewer": "Reviewer_YWJ5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall the approach is well-motivated and is clear in its presentation. The results are strong to demonstrate the utility of the proposed approach. \n\nThe following are some of my questions and would be great if the authors could clarify some of them: \n\n1. The retrieval process relies on auxiliary losses based on reward, value predictions and on BERT-style losses. Could the authors present the exact loss functions for training these modules in the main text? I was not able to find any details related to how these modules were trained and it is an important detail relevant to the approach. \n\n2. Because the main contribution of the paper is an empirical improvement across domains, would the authors release the code or a very-detailed pseudocode (similar to the one presented in Schrittwieser et al. 2020), which seems necessary for a reimplementation? \n\nSchrittwieser J, Antonoglou I, Hubert T, Simonyan K, Sifre L, Schmitt S, Guez A, Lockhart E, Hassabis D, Graepel T, Lillicrap T. Mastering atari, go, chess and shogi by planning with a learned model. Nature. 2020 Dec;588(7839):604-9.\n\n3. An ablation studying how the retrieval process is working is necessary. Intuitively it makes sense that the retrieval process is considering prior experiences to make an accurate estimate of the Q-values. For example, in the umbrella domain (Harutyunyan et al. 2019), an agent picking up an umbrella much earlier in an episode is useful information to have to make an accurate estimate of Q-values for later parts of the episode. Thus, I would find it interesting to see if the retrieval process makes similar computations.\n\nHarutyunyan, A., Dabney, W., Mesnard, T., Gheshlaghi Azar, M., Piot, B., Heess, N., ... & Munos, R. (2019). Hindsight credit assignment. Advances in neural information processing systems, 32, 12488-12497.\n\n4. Could the authors, in addition Fig 2, present the Atari performance as a learning curve where the median-normalized performance is reported as a function of training steps? This would tell if the agent overall produces a faster learning performance and also tells the final median performance across all atari games.\n", "summary_of_the_paper": "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. \n", "summary_of_the_review": "Overall I think the paper makes an interesting contribution to RL where the agent is able to access prior data in making accurate estimations of Q-values and through that produce a better policy for maximizing rewards from a given task. \n\nSome concerns related to the paper are as follows: Providing detailed pseudocode or code for replicating the results, ablation to understand the retrieval process, and detailed presentation of the proposed approach.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "5oUv19O6xew", "reviewer": "Reviewer_cZ2E", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper proposed a simple approach of end-2-end2 differentiable state-augmentation through retrieval.\n\nStrenght:\n* The approach is well written and the paper is pleasant to read.\n* The modular choices composing the proposed method are reasonably clear and motivated.\n* The experimental setting is clearly explained and reproducible.\n* The experiment's motivations are clearly justified.\n\nWeakness:\n* The contribution is marginal, state augmentation is a reasonably classic topic and the end-2-end retrieval process could have been replaced by an attention model over the dataset of trajectories. We would have hoped to see this baseline studied deeply to justify the approach. In the introduction, the authors claim that a retrieval process allows simplifying the memorization duty of the model. We would have hoped to see experiments defending this claim.\n* We would have hoped for a more exhaustive justification of why certain Atari games seem to improve significantly compared to others. it doesn't seem to have a general pattern of the game occurring where the approach is particularly beneficial.\n* For offline RL, it would have been interesting to compare to SoA approaches, using for example the D4RL dataset which is dedicated to this task.", "summary_of_the_paper": "The paper proposed a state-augmentation strategy based on a fully differentiable model for online and offline reinforcement learning.\nThe authors evaluate the proposal on two algorithms, DQN and R2D2.\nThe augmentation through retrieval approach is decomposed in several steps, namely trajectory, and state retrieval based on the currently perceived state and an internal state, the retriever has been defined with a recurrent neural network type of model.\nThe retrieved information is then aggregated using a so-called summarizing model which is also end-2-end differentiable.\nThe authors claim and show detailed improvement with respect to the state of the art on Atari and multi-task offline RL.\nIn the case of multi-task, the authors show that incorporating multiple tasks into the retrievable dataset is beneficial.", "summary_of_the_review": "The paper proposed a retrieval-based state augmentation approach for online and offline reinforcement learning.\nThe retriever is stateful and end-2-end differentiable.\nThe approach is evaluated in online RL using Atari benchmark and offline RL.\nThe experiments in offline RL missed the main baselines and could have been evaluated on a state-of-the-art benchmark like the D4RL benchmark.\nThe retrieval approach also could have been evaluated against a fully attention-based model over the dataset, like Performer or Reformer Transformer architecture for example.\nIn conclusion, while the paper is reasonably well written, the contribution seems somehow limited and more experiments are needed to assess the full benefit of the proposed approach.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: One of the four reviewers failed to engage in discussion, two acknowledged the author's response and paper revision without changing their scores, and one reviewer engaged in considerable discussion resulting in a score increase to a weak accept.  No reviewer gave the paper a strong endorsement.  I do appreciate the large effort that the authors put into revising their paper and addressing reviewers concerns.  However, major post-submission revision puts an inappropriate burden on reviewers.  In any case, there is not strong support for this paper even from the one heavily engaged reviewer.
Decision: Reject