Large-Scale Retrieval for Reinforcement Learning
Authors: Peter Conway Humphreys
Arthur Guez
Olivier Tieleman
Laurent Sifre
Theophane Weber
Timothy P Lillicrap
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
retrieval
neural networks
offline RL

Abstract: Effective decision making involves flexibly relating past experiences and relevant contextual information to a novel situation. In deep reinforcement learning (RL), the dominant paradigm is for an agent to amortise information that helps decision-making into its network weights via gradient descent on training losses. Here, we pursue an alternative approach in which agents can utilise large-scale context-sensitive database lookups to support their parametric computations. This allows agents to directly learn in an end-to-end manner to utilise relevant information to inform their outputs. In addition, new information can be attended to by the agent, without retraining, by simply augmenting the retrieval dataset. We study this approach for offline RL in 9x9 Go, a challenging game for which the vast combinatorial state space privileges generalisation over direct matching to past experiences. We leverage fast, approximate nearest neighbor techniques in order to retrieve relevant data from a set of tens of millions of expert demonstration states. Attending to this information provides a significant boost to prediction accuracy and game-play performance over simply using these demonstrations as training trajectories, providing a compelling demonstration of the value of large-scale retrieval in offline RL agents.

Large-Scale Retrieval for Reinforcement Learning
Peter C. Humphreys, Arthur Guez', Olivier Tieleman, Laurent Sifre, Théophane Weber, Timothy Lillicrap Deepmind, London
{peterhumphress, aguez, ....@google.com
Abstract
Effective decision making involves flexibly relating past experiences and relevant contextual information to novel situation. In deep reinforcement learning (RL), the dominant paradigm for an agent to amortise information that helps decision- making into its network weights via gradient descent training losses. Here, we pursue an alternative approach in which agents can utilise large-scale context- sensitive database lookups to support their parametric computations. This allows agents to directly learn in an end-to-end manner to utilise relevant information to inform their outputs. In addition, new information can be attended to by the agent, without retraining, by simply augmenting the retrieval dataset. We study this approach for offline RL in 9x9 Go, challenging game for which the vast combinatorial state space privileges generalisation over direct matching to past experiences. We leverage fast, approximate nearest neighbor techniques in order retrieve relevant data from set of tens of millions of expert demonstration states. Attending to this information provides significant boost to prediction accuracy and game-play performance over simply using these demonstrations as training trajectories, providing compelling demonstration of the value of large-scale retrieval in offline RL agents.
Introduction
How can reinforcement learning (RI.) agents leverage relevant information to inform their decisions? Deep RL agents have typically been represented as a monolithic parametric function, trained to gradually amortise useful information from experience by gradient descent. This has been effec- tive [25] 38), but a slow means integrating experience, with no straightforward way for an agent to incorporate new information without many additional gradient updates. In addition, this requires increasingly massive models (see e.g. 31) as environments become more complex scaling is driven by the dual role of the parametric function, which must support both computation and memorisation. Finally, this approach has a further drawback of particular relevance in RL - the only way in which previously encountered information (that is not contained in working memory) can aid decision making in a novel situation is indirectly through weight changes mediated by network losses. There end-to-end means for agent attend information outside of working memory directly inform its actions. While there has been significant amount of work focused on increasing the infor- mation available from previous experiences within an episode (e.g., recurrent networks, slot-based memory [19 (28), more extensive direct use of more general forms of experience or data has been limited, although some recent works have begun explore utilising inter-episodic information from We seek to drastically expand the scale of information that is accessible to an agent, allowing it attend to tens of millions of pieces of information, while learning in an end-to-end manner how to use this information for decision making. We view this as a first step towards a vision which an agent can flexibly draw on diverse and large-scale information sources,
the same agent 6
29 34
'These authors contributed equally this work.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

including its own (inter-episodic) experiences along with experiences from humans and other agents. In addition, given that retrieved information need not match the agent's observation format, retrieval could enable agents to integrate information sources that have not typically been utilised such as videos 2]. text, and third-person demonstrations.
We investigate a semi-epaaaae cemi--aramettic agent architecture for large-scale retrieval, in which fast and efficient approximate nearest neighbor matching used dynamically retrieve relevant information from dataset of experience, We evaluate this approach in an offline RL setting for an environment with a combinatorial state space - the game of 9x9 Go with = 10t possible games, where generalisation from past data to novel situations is challenging. We equip our agent with a large-scale dataset of ~50 Go board-state observations, finding that retrieval retrieval-based approach utilising this data is able to consistently and significantly outperform a strong non non-retrieval baseline.
Several key advantages of this retrieval approach are worth highlighting: Instead of having amortise all relevant information into its network weights, a retrieval-augee -arrreee-- network can utilise more of its capacity for computation. In addition, this parameeratttric form allows us update the information available to the agent at evaluation time without having to retrain it. Strikingly, we find that this enables improvements to agent performance without further training when games played against the evaluation opponent are added to its knowledge base.
2
Methods
Before introducing the details of our method, let first present its high-level ingredients. We train a model-based agent [36] that pre- diets future policies and values conditioned on future actions in a given state. This semi- parametric model incorporates a retrieval mecha- nism, which allows it to utilise information from large-scale dataset to inform its predictions. We train the agent in supervised offline RL setting, which allows us to directly evaluate the degree to which retrieval improves the quality of the model predictions. We subsequently eval- uate the resulting model, augmented by Monte- Carlo tree search, against reference opponent. This allows to determine how these improve- ments translate to an effective acting policy for out-of-training ciidiinnddistriuttion situations.
ÿ - agent outputs
Forward computation
Gradients
Learn end-to-end how to use retrieved information
x retrieved neighbors and associated metadata
Nearest-neighbor retrieval
O agent observation
D, large-scale dataset
Retrieving information to support deci- sion making. The agent observation o, is used to generate a query to retrieve relevant information from large-scale dataset D,.. This information is used to inform network outputs ÿt- The agent is trained end-to-end to this information to inform its decisions.
Figure I:
To effectively improve model predictions using auxiliary information, we need I) a scalable way to select and retrieve relevant data and 2) ro- bust way of leveraging that data in our model. As with many successful attention mechanisms, we establish relevance for 1) through an inner product an appropriate key-query space and select the top N relevant results. Choosing the relevant key-query space is critical, since perform- ing nearest-neight neirest-ncighbor on the raw observations will only deliver desirable results for the simplest of domains. For example, in the game of Go, single stone difference on the board can dramatically change the outcome and reading the board, while seemingly different board positions might still share high-level characteristics (e.g. influence, partial local sequences, and life-and-death of groups).
At the scale we want to operate, learning the key query embeddings end-to-end in order optimize the final model predictions is challenging. Instead, as in the language modelling work by Borgeaud et al. 77. we learn an embedding function through a surrogate procedure and use the resulting frozen function to describe domain-relevant similarity. Moreover, scale also constrains the nearest-neighhon lookup to be approximate, since getting the true nearest-neighhors too time-consuming at inference time and good approximations can be obtained.

---

To address 2) and make the best use of the relevant data, we provide the data associated with the nearest-neighbor as additional input features to the parametric part of our model, so that the network can learn to interpret it. We provide some light inductive biases in the architecture to ensure permutation invariance in the neighbors and robustness to outliers and distribution shifts (see Secs. 2.1.3 & 2.3) Letting the network interpret the neighor data is essential in large and complex environments, such as Go, because the typically imperfect match between the current situation and retrieved neighbors will require different aspects of this retrieved data to be ignored or emphasized in a highly context specific manner. This is in contrast to episodic RL approaches like 6 [9] which prescribe how to use the retrieved data. Subsequent sections describe the model, retrieval process, and the full algorithm in more detail.
2.1 Retrieval-augmented agents
We consider a setting in which we wish to train a neural network model ma on a set of environment trajectories 7 E 'D (for example, the dataset D> of trajectories produced by a policy TT). For each timestep t of r, the model takes an observation of and must produce predictions ÿ of targets ye-
In our retrieval paradigm, in addition to having access to or, the model also has access an auxiliary dataset of knowledge or experience D,,. The auxiliary set D, could be the same as D. but it can also contain more information, or in the most general case, information from a completely different distribution and in a different format. If there is overlap between DD and D, it is important for robustness ensure that the model cannot retrieve the same trajectory as it is being trained on (otherwise the network will tend to overly trust information retrieved from D,). Common information sources in RL are trajectories from other agents or experts (offline RL) or the agent's own previous experiences (episodic memory, replay buffer). Note that, in contrast to offline RL, we do not assume that we should directly use this auxiliary data as trajectories train on F
We wish to use D, to inform the model predictions ft. such that C, = =mm (ot D,). This is challenging, as D, is typically far too large directly be consumed as model input. One solution, shown in Fig. 1 is to adopt a retrieval-base the model my is provided with informative subset of data (z - a² conditioned on oc:
model, wherein the parametric and differentiable portion of
retrieved from D,,
Ue ma Te Le
(1)
This approach requires a number of design choices relating to the retrieval mechanism, which we explore further below.
2.1.1 Scalable nearest-neighhor retrieval using SCaNN
Inspired by previous work in language modelling. we chose leverage the SCaNN architecture []] for fast approximate nearest-neighhor retrieval. This requires each entry o in the dataset D, to be associated with key vector h; E IRd and a given observation o, be mapped to query vector q: E Rdd During retrieval, the squared Euclidean distances between q and the dataset keys k; are used determine which neighbors to retrieve reminiscent of neural attention mechanisms, This retrieval process is very efficient and can be scaled to datasets with billions of items.
We will typically want to retrieve further associated meta-data, or context, for each neighbor. For example, if a neighbor is part of a trajectory, we would like to retrieve information about action choices and their consequences. The neighbor observation 0z, together with its meta-data, forms the auxiliary input oi to the model.
The nearest-neighhor retrieval process is non-differentiable, which means that the query and key mappings cannot be trained end-to-end directly. Instead, in this first study, we pre-train a non-retrieval prediction network m² on experience dataset D (details in Sec. A.4). We then use this network to generate embedding corresponding to given observation o, by retrieving the network activations from a specified layer of m We use principal component analysis to compress these activations to a d dimensional vector representing this observation state. The embedding and projection step together form our key (& query network k; = (o:).
For simplicity of notation, omit the dependence on past observations for partiall domains. For example, D, needs not contain actions, rewards, or the full context used to select actions
g
In this study, retrieval dataset entries and have the same format, but this need not be the case in general.
3

---

00
Forward computation Gradients
hg recurrent forward model
a
fe observation & neighbor encoding
White
White witn
White wins
retrieved neighbors, includes: board, next 10 actions, game outcome, final board
SCaNN fast approximate
nearest-neighbor retrieval
qt query
g pretrained embedding network & PCA projection
Dp.-5MM game states keys &
O2 Go game state
Figure 2: Details of the architecture used for retrieval-augmentee Go playing agent. A pre-trained network is used to generate a query qe corresponding to the current Go game state oz. This query is used for fast approximate nearest-neighbor retrieval using SCaNN. Retrieved neighbors x processed using an invariant architecture, and used to inform an action-conditional recurrent forward model that outputs game outcome predictions eh and distributions over next actions a
We preprocess all of the observations in D, using Is to produce corresponding keys. The resulting dataset of keys and observations then used by ScaNN to retrieve neighbors given query vector. For offline RL training with a fixed training dataset, we can also preprocess the nearest neighbor lookups for all training dataset observations, i.c., q: go Vo, D. However, acr during evaluation, must do this searest-neighhor lookup online for each new observation. In future experiments, we intend incorporate end-to-end learning of the query vector [12 to improve agent performance. future online RL pipeline will also require us to dynamically retrieve neighbors during training.
2.1.2 Policy optimization with retrieval
The objective in our offline RL setting is to leverage the available data in order to optimize a policy for acting. Our proposed semi-parametric model (Eqn. [ compatible with several methods for policy optimization in this offline setting. Indeed, it can represent the actor æ (o(o;, D,,) and/or the critie Qo a, D,) in an offline actor-critic method [2]] or the Q-value
value iteration
approach [32] Here we focus
model-based approach, inspired by MuZero [36), where the learned model is employed in a search procedure based on Monte-Carlo Tree Search (MCTS)
This is motivated by the efficacy of MuZero in the offline RL setting 37 24].
Model-based search with retrieval Following MuZero, our prediction model my is conditioned on the current observation o, but also on sequence future actions d, = a++1, at+2; at+K The model my is therefore redefined as Ú, my (o,, ae _ .?.... ,N) The architecture for this retrieval prediction model is illustrated in Fig. P The model ms can decomposed into encoding step 8: fo _t, followed by iterative action-conditioned inference s;
..... ,,"), ......,??) which incorporates the observation and neighbor information,
A+1 h a++x) produce embeddings

---

of corresponding to subsequent time steps (where de = st). The model output de is composed of K + value and policy distributions for the current and next K time-steps: {or # ko=o...K- These outputs are obtained from their respective embeddings sf. The target for value predictions is the game outcome and the targets for the policy predictions are the actions taken by the expert az, a1+1........ a++K- The sample loss [(0, 9t. yt) is obtained by summing the K + 1 individual loss terms for value and policy predictions (see detail in App. A.2 The training procedure for this
retrieval prediction model is summarized in Alg
Algorithm 1 Training semi-parametric action-conditional model
Input: Training dataset and retrieval dataset DD,
1: Initialize parameter vectors 0, Φb
2: Train key/query network g using D.
3: Pre-compute k; = (o(o,) for all 0; D,..
4: for each gradient step do
5: Sample mini-batches from D of (e..e).
6: For each element On of the batch, compute = go (ot).
7: Fetch N neighbor keys with (approx.) smallest distance |lue koll2 for all ky E D,..
8: Gather meta-data associated with these keys as s.
Compute model output Út my (or, a. z...... .........] M ....
9:
10: Compute and sum losses C(0, Vt, y1).
11: Update ! based on VgL.
12:
end for
13:
output 9
In order to using the trained model, online search is used generate improved policy Ts Search(ma,, Dr). We use MCTS with a PUCT rule for internal action selection [35] 40).
carrying out ssss simulations per time step. Model-based search is implemented follows: after retrieving the observation's 's neighbors, the encoded state 9; is computed. Search is carried out from 8, by varying the input action sequence ā for each simulation and collecting the model outputs to update search statistics. Details of this process can be found in [36]. At the end of the search, the resulting policy x (alo:) can be sampled to select the next action.
Due to the semi-parametric nature of the model supporting the search, introducing changes to D, will have immediate effect the acting policy T., even if the model parameters & remain fixed. As explore in Sec. 3.4. this provide a mechanism for fast adaption to new information or experiences.
2.1.3 Using neighbor information
Several choices possible for the network fe that processes the observation and neighbors. Since this is not the main focus of this work, we chose a straightforward approach. We first compute an embedding of of the observation o₁. We then compute embedding for each neighbor ef po of)), using the same network po for all neighbors. These streams are combined in a permutation invariant way through sum, and then concatenated with of produce the final embedding st,
N
St = folor, .I...., IN)= [of.
N
2.2 Evaluating retrieval in the domain of Go
In order to test the utility of retrieval in RL, we wish to evaluate whether agents can effectively generalise between related but distinct experiences, opposed simply retrieving the outcome of a previous instance of an identical situation. This motivates our choice of Go as a domain. We focus on 9x9 Go, instead of full-scale 19x19 Go, the less demanding computational costs of 9x9 experiments enable a more thorough analysis. Even for 9x9 Go. the state space of 1038 possible games vastly larger than the number of positions that we could hope to query.
We collected a dataset of 3.55 expert 9x9 Go self-play games from an AlphaZero-style agent [40] We randomly subsampled 15% of the positions from these games, leaving us with ~50M board
In cases where non-terminal rewards exist, also output reward estimates for each model transition.
5

---

state observations. These, along with metadata on the game outcome, final game board, and future actions (all encoded as input planes), form our retrieval dataset D,- We chose subsample as we hypothesised that this would reduce the chance of retrieving multiple neighbors from the same trajectory, and therefore boost the retrieved neighbor diversity. However, we did not perform ablations on this choice. A future solution would be to use filtering mechanism to reject neighbors from the same trajectory. In this initial work, training and retrieval datasets are the same at least during training. During training, we split D, in two halves such that each game's observations one the datasets. We retrieve neighbors for an observation or from the half is not contained in. This is simply to avoid retrieving the same position as the query. Other ways to obtain the same effect could be devised.
only in
2.3 Regularisation
We wish to make our network robust to poor quality neighbors, in order to ensure that the network can perform well in settings for which there lower overlap with the retrieval dataset than encountered in training. We therefore explore several techniques to improve network robustness to irrelevant neighbors. We randomly zero-out a subset retrieved neighbors during training ("neighbor dropout"), and/or more ad adversarially, randomly replace subset of retrieved neighbors with the neighbors of different observation (" ("neighbor randomisation" "). Inspired by [10]. we also explore using loss to regularise the embedding produced by the neighbor retrieval towards the embedding produced with the observation alone ("neighbor regularisation") Further details are given Sec. A.5
We carried out ablations of these techniques (Appendix Fig. 77. which show that neighbor randomi- sation is important in some contexts for maintaining performance, but that the others do not seem to have significant effect in our final configuration. The results reported in this study utilise all of these augmentations, as during the development process they had been found to slightly benefit per- formance. Interestingly, as we explore in later sections, MCTS with enough simulations compensates for the harmful effect of low-quality neighbors, and can perform effectively without these training augmentations.
Results
3.1
Qualitative examination of retrieved neighbors
Neighbors
Query
Example
Example 2
Example 3
Figure 3: Visualisation of N W approximate nearest-neighhors retrieved using learned Go-specific distance function for 3 query positions (one per row). Red stone indicates the next action(s) - light red stones indicate white moves, dark red stones indicate black moves. The color of the board border indicates whether the current player play (for each board) won the game.

---

Changing a single stone in a Go board can dramatically affect the game, and the number of Go positions enormous. We first wanted to assess whether, despite the combinatorial aspect of the domain, meaningful nearest-neighhon could be retrieved from the dataset using our learned keys/queries. Examples of retrieved positions in Go are shown in Fig- 3 where we observed relevant matches in terms of both local and global structure. In some cases, especially in the early game, the retrieved position is an exact match even though the rest of the game (and therefore the associated nearest-neighhor meta-data) differs. Row in Fig. is an example of this in this case, the retrieved data effectively provides sample rollouts from the query position akin to those derived from MCTS.
0.48
retrieval baseline
Top accuracy
0.46
accuracy
0.48
Top-1 accuracy
0.48
retrieval baseline
accuracy
0.48
0.44
0.47
0.47
0.40
0.46
retrieval baseline
0.46
0.46
neighbors
0.38
0.45
0.45
0.300
10
base lnn
0.185
0.275
Value MSE
0.250
0.800
MSE
0.180
MSE
0.180
0.175
0.175
0.225
0,200
0.170
0.175
0.165
0.165
0.165
250
500
10
20
10a
10'
Train steps (/10%)
Num. neighbors
Relative mode
Retrieva dataset size
Figure 4: a) Test-set top-1 policy prior # accuracy (top) and value 6 mean-squared error (MSE) (bottom) over the course of training, for the (non-retrieval) baseline and retrieval models with
N = 2, 10. Results are averaged over seeds (error typically too small see). Note that the sharp transitions are due to the learning rate schedule. b) Final performance after training plotted as function of number of neighbors N retrieved neighbors. c) Final performance as a function of model size relative to the size used elsewhere in this study, for networks with N 10 retrieved neighbors and the non-retrieval baseline. d) Final performance function of the evaluation retrieval dataset size, for networks with N = 10 retrieved neighbors. The equivalent baseline performance is shown as dashed line for reference. Note that the networks are trained with dataset fraction 0.5 and simply evaluated at other dataset fractions.
a different network trained for each number of
3.2 Impact of retrieval on supervised training
We next evaluated the extent to which the model can learn exploit retrieved information for better model predictions and better decision- mmaiing. First, we evaluated the impact of retrieval on supervised learning losses. As baseline comparison, we trained the same network architectures 0 for all
but setting and useful capacity in the network, while removing access to the retrieval data. We evaluated losses for retrieval networks trained with different N (Fig. 4D) and model sizes (Fig. D. see I for details). Across conditions, we consistently observed a significant boost in test-set accuracy for all metrics and over the course of training. This improvement to predictions is further observed across game trajectories, and is not limited, for example, to opening play positions.
this has the effect of maintaining the number of parameters, flops,
One advantage of the camipparaetrrc approach we outlined is that we can modify the retrieval dataset and potentially see immediate effects on the prediction, without changing the parameters We first verified this by evaluating our model when allowed access to varying fractions of the full dataset D, (Fig. ft). We observed clear gains in prediction accuracy from increasing the size of the retrieval set (only half is used at train time). A further important observation is that large-scale datasets are clearly important- the evaluation metrics drop below the baseline level for dataset D, that is 1% the size of the full dataset.
7

---

3.3 Evaluation against a reference opponent
While the results observed on the offiine dataset suggest strong positive effect from learning to leverage the retrieval data, this is in the context of fixed test data distribution that matches the retrieval data distribution. When deployed. games played against different opponents will likely diverge from that distribution. We evaluated the performance of our search policy (Msiss 200) by playing against fixed reference opponent - the Pachi program (5), which can perform beyond strong amateur level of play in 9x9 Go given sufficient simulation budget (we evaluate against 400k simulations). We observed significant boost performance for retrieval-based networks over the equivalent baseline network of the same capacity (Fig. 5 Interestingly, as shown in Appendix Fig. g playing using only the base policy prior 7 shows a much smaller boost over the baseline network. As we explore further below, the quality of retrieved neighbors available during play is much lower than for training we hypothesise that the search policy more robust to this distributional shift than #0
b)
0.80
0.85
0.78
Win rate
Win rate
0.80
0.75
0.74
0.70
retrieval
baseline
0.72
Num, neighbors
Relative mode size
Figure Win rate against fixed reference opponent (Pachi) when playing using the MCTS policy T, for (a) retrieval networks using varying numbers of retrieved neighbors and for a baseline non-retiee non-retrieval network. (b) Win rate as function of model size relative to the size used elsewhere in this study. Retrieval leads to a clear performance boost compared to non-retrieval baselines of the same capacity,
3.4 Augmenting the retrieval dataset
b)
0.00175
Tuot delened, Trai
Q opponent byyy for RachC) R. Traie
0.00150
O opporsstt bh txi Rachal Ri trase
0.80
Density
0.00125
Win rate
0.78
0.00100
0.00075
0.76
0.00050
from Train
rectieval from Train and Pachi baseline
0.00025
0.00000
500
1000
1500
10
15
20
Simi arity (Squared L2)
neighbors
Figure 6: (a) Distribution of similarity distances from an encoded observation to its retrieved nearest neighbors. Retrieving from the original retrieval dataset using positions from the test dataset (orange) gives neighbors with a markedly different distance distribution to positions queried during play against a Pachi opponent (blue). Augmenting the retrieval dataset with ~600k recorded agent-Pachi game states improves the similarity distribution (green). (b) For play with the MCTS policy n. this augmentation also leads to a consistent win-rate boost (green) over using the train retrieval dataset alone (bluc).
We observed a significant distribution shift between game positions observed in play against the Pachi opponent versus positions in the retrieval dataset (Fig. (o),, as measured by the empirical distance distribution of game observations to their approximate nearest neighbors. Changing or augmenting the retrieval dataset D, modifies this distributional shift. A particularly interesting modification in this setting is to augment the dataset with play recorded between our agents and the Pachi opponent. As can be seen in Fig. 6h, augmenting the dataset with set of ~600 agent-Pad Pachi game states increases

---

the similarity between positions observed in play and their retrieved neighbors. Strikingly, we found that integrating these games into our retrieval dataset leads to improvements in performance without further training for subsequent play against the Pachi opponent (play with the MCTS policy n, is shown in Fig- Gb. play with the policy prior # is shown in Appendix Fig. 9 This highlights the potential of a semi eari-paraettiic model to rapidly integrate recent experience without further training.
As a related intervention, we instead inputted randomly retrieved neighbors into the model at evaluation time (see Appendix Fig. 6 This significantly impaired performance for play using the policy prior #", but somewhat unexpectedly, for play with the MCTS policy 7ss performance did not fully regress to the level of the baseline network. This suggests that our retrieval network is able to amortise some information from neighbors during training, leading to better performance even when no relevant neighbors are available.
4 Related work
The idea of supporting decision making by directly attending to a cache of related events has been visited many times in different contexts, under the names of case-based reasoning [30,, non- parametric RL [16 4 27). or episodic memory (20 23 The aim of some methods is to better support the working memory of agent during given episode 26 4
series of successive related episodes 33 Other methods, including ours, aim to leverage a broader class of relevant experience or persistent knowledge (including across episodes, and from other agents) better support reasoning and planning. One differentiiaingg factor in our work to these past approaches is that we do not prescribe how to process the information from the available data (e.g. through specifying the agent's action-value directly in terms of previously generated value estimates 66] 13 15 [2], or a model from observed transitions (39] but rather learn end-to-end how the data can support better predictions within the parametric model. A recent approach by Goyal et al. [10] has considered an attention mechanism to select where and what to use from available trajectories. but over a small retrieval batch of data rather than the full available experience data. Another class of method to leverage a transition dataset to replay the data at training time in order to perform more gradient steps per experience, this is widespread technique in modern RL algorithms [2] 22 25 36 but it does not benefit the agent at test time, requires additional learning steps to adapt to new data, and does not allow end-to-end learning of how to relate past experience to new situations.
5 Discussion
Our approach and empirical results highlight how reinforcement learning agents can benefit from direct access to a large collection of raw interaction data at inference time, through retrieval mechanism, in addition to their already effective parametric representation. We showed this was the case even 1) when the domain is large enough require generalisation in how to interpret past data, 2) when there is significant distribution shift when acting, and 3) at scale where only approximate nearest-neighbors can be retrieved. We believe this already demonstrates the potential of this approach for many possible scenarios and applications.
We show that retrieval can be effectively combined with model-based search. We find that the benefits from retrieval and search are synergistic, with increasing numbers of retrieved neighbors and increasing simulations both leading to performance increases in almost all contexts we investigated. Furthermore, empirical evidence suggests that search significantly improves agent robustness to distributional shift compared playing with the policy prior. In Appendix Fig. TO compare the boost in performance as a function of parametrie-model compute cost for increasing MCTS simu- lations versus increasing the number of retrieved neighbors processed by the model. This provides a tentative indication that retrieval is also compute efficient means of improving performance.
A key future direction is to investigate the online learning scenario in which recent experience of the agent is rapidly made available for retrieval, hence progressively growing the retrieval dataset over the course of training. While there are additional challenges associated with the online paradigm (e.g., it may be desirable to update the queries and/or keys during training [2],, the fast adaptation effect highlighted in this work may have even more impact there.
There are many potentially relevant sources of information beyond an agent's own experience, or that of other humans agents. For example, it has been shown that YouTube videos are useful

---

for learning to play the Atari game Montezuma'ss 's revenge (2] Training an embedding network on sufficiently diverse data may enable retrieval of information from wide range of contexts 4 including third-person demonstrations, videos and perhaps even books.
Disclosure of Funding
The authors received no specific funding for this work.
References
[1] Christopher G. Atkeson, Andrew W. Moore, and Stefan Schaal. Locally weighted learning for control. Artificial Intelligence Review, 11:75-113, 2004.
[2] Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando De Fre- itas. Playing hard exploration games by watching youtube. Advances in neural information processing systems, 31, 2018.
[3] Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Claudio Fantacci, Jonathan Godwin, Chris Jones, Tom Hennigan, Matteo Hessel, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Lena Martens, Vladimir Mikulik, Tamara Norman, John Quan, George Papa- makarios, Roman Ring. Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojeiech Stokowiec, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. /ttp: com/ denedee
[4] André MS Barreto, Doina Precup, and Joelle Pineau. Practical kernel-based reinforcement learning. The Journal of Machine Learning Research, 17(1):2372-2411, 2016.
[5] Petr Baudiš and Jean-loup Gailly. Pachi: State of the art open source Go program. In Advances in computer games, pages 24-38. Springer, 2011.
[6] Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control. arXiv preprint 1601v:1606.04460, 2016.
[7] Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katic Millican, George van den Driessche, Jean-Baptin Sran-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, al. Improving language models by retrieving from trillions of tokens. arXiv preprint 12X1v:2112.044426, 2021.
[8] Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International conference computers and games, pages 72-83. Springer, 2006.
[9] Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Advances in Neurul Information Processing Systems, 32, 2019.
[10] Anirudh Goyal, Abram L Friesen, Andrea Banino, Theophane Weber, Nan Rosemary Ke, Adria Puigdomenech Badia, Arthur Guez, Mehdi Mirza, Ksenia Konyushkova, Michal Valko, al. Regmened reinforcement learning. arXiv preprint arXiv:2202.084/7 2022.
[11] Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simeha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, pages 3887-3896. PMLR, 2020.
[12] Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-We Wei Chang. Realm: Retrieval-augmented language model pre-training. arXiv preprint arXiv:2002.0888 08909, 2020.
[13] Steven Hansen, Alexander Pritzel, Pablo Sprechmann, André Barreto, and Charles Blundell. Fast deep reinforcement learning using online adjustments from the past. Advances in Neural Information Processing Systems, 31, 2018
[14) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference computer vision, pages 630-645. Springer, 2016.
10

---

[15] Mika Sarkin Jain and Jack W Lindsey. Semiparametrie reinforcement learning. In ICLR, 2018.
[16] H. JoséAntonioMartín, Javier de Lope Asiaín, and Darío Maravall Gómez-Allende. The knn-td reinforcement learning algorithm. In /WINAC, 2009.
[17] Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th anmual international symposium on computer architecture, pages 1-12, 2017.
[18] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[19] Andrew Lampinen, Stephanie Chan, Andrea Banino, and Felix Hill. Towards mental time travel: a hierarchical memory for reinforcement learning agents. Advances Neural Information Processing Systems, 34, 2021.
[20] Máté Lengyel and Peter Dayan. Hippocampal contributions to control: the third way. Advances in neural information processing systems, 20, 2007.
[21] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 2020.
[22] Long-Ji Lin. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine learning, 8(3):293-32) 3211, 1992.
[23] Zichuan Lin, Tianqi Zhao. Guangwen Yang, and Lintao Zhang. Episodic memory deep q- networks. arXiv preprint arXiv: 1805.08805.06603 2018.
[24] Michael Mathieu, Sherjil Ozair, Srivatsan Srinivasan, Caglar Gulcehre, Shangtong Zhang, Ray Jiang, Tom Le Paine, Konrad Zolna, Richard Powell, Julian Schrittwieser, al. Starcraft ii unplugged: Large scale offline reinforcement learning. In Deep RI. Workshop NeurIPS 2021, 2021.
[25] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland. Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 318(7540):529-533 2015.
[26] Junhyuk Oh, Valliappa Chockalingam, Honglak Lee, al. Control of memory, active perception, and action in minecraft. In International Conference Machine Learning, pages 2790-2799.
PMLR, 2016.
[27] Dirk Ormoneit and Saunak Sen. Kernel-based reinforcement learning. Machine learning, 49(2): 161-178, 2002.
[28] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International Conference on Machine Learning.
pages 7487-7498. PMLR, 2020.
[29] Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In International Conference Machine Learning, pages 2827-2836. PMLR. 2017.
[30] Ashwin Ram and Juan Carlos Santamaria. Continuous base- reasoning. Artificial Intelli- gence, 90(1-2):25-77 1997.
[31] Scott Reed, Konrad ZoIna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Marth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.061 2022.
[32] Martin Riedmiller. Neural fitted q iteration-first experiences with data efficient neural reinforcement learning method. European conference on machine learning, pages 317-328. Springer, 2005.
11

---

[33] Sam Ritter, Ryan Faulkner, Laurent Sartran, Adam Santoro, Matt Botvinick, and David Raposo. Rapid sosk-no novel environments. arXiv preprint Ntt1:22006003662,, 2020.
[34] Samuel Ritter, Jane Wang, Zeb Kurth-Nelson, Siddhant Jayakumar, Charles Blundell, Razvan Pascanu, and Matthew Botvinick. Been there, done that: Meta-learning with episodic recall. In International conference on machine learning, pages 4354-4363, PMLR, 2018.
[35] Christopher D Rosin. Multi-armed bandits with episode context, Annals of Mathematics and Artificial Intelligence, (3):2033-223 2011.
[36] Julian Schritwwiser,, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon 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-660 -609)6999 2020.
[37] Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with learned model. Advances in Neural Information Processing Systems, 34, 2021.
[38] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learming, pages 1889-1897. PMLR, 2015.
[39] Aayam Shrestha, Stefan Lee, Prasad Tadepalli, and Alan Fern. Deepaveragers: offline rein- forcement learning by solving derived non-parametric mdps. arXiv preprint (rXiv:2010.08891, 2020.
[40] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel. et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self- Science, 10204419)11400-1144 2018
[41] Pablo Sprechmann, Siddhant M Jayakumar, Jack W Rae, Alexander Pritzel, Adria Puig- domenech Badia, Benigno Uria, Oriol Vinyals, Demis Hassabis, Razvan Pascanu, and Charles Blundell. Memory-based parameter adaptation. arXiv preprint arXiv: 10542, 2018.
[42] 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 in goal-directed agent. arXiv preprint a /0760, 2018.
[43] Hu Xu, Gargi Ghosh, Po-Yao Huang, Prahal Arora, Masoumeh Aminzadeh, Christoph Feicht- enhofer, Florian Metze, and Luke Zettlemoyer. Vlm: Task-agnostic video-language model pre-training for video understanding. arXiv preprint arXiv:2105. 100.0099,,
09996, 2021.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A] No theoretical results.
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [No] Instructions and links to relevant libraries were provided.
12

---

(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (c.g., code, data, models) or curating/releusing new assets...
(a) your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A] No people were involved in the data gathering process.
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "oxmQWYza7t", "reviewer": "Reviewer_WzqK", "summary": "This paper mainly experimented on Go. It leverages fast, approximate nearest neighbor techniques in order to retrieve relevant data from a\nset of tens of millions of expert demonstration states. This information provides a significant boost to prediction accuracy and game-play performance over simply using these demonstrations as training trajectories, providing a compelling demonstration of the value of large-scale retrieval in reinforcement learning agents.", "questions": "1. In figure 5a, it shows the retrieval networks using varying numbers of retrieved neighbors and for a baseline non-retrieval\nnetwork. However, I found that a larger number of neighbors (10 vs 20) actually hurts the performance. Do you have a explanation for that?", "limitations": "None", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1. The paper is well organized and easy to follow.\n\n2. The motivation of this paper is really interesting. It show that the retrieval can help RL for training. It also shows that retrieval can be effectively combined with model-based search. The benefits from retrieval and search are synergistic, with increasing numbers of retrieved neighbors and increasing simulations both leading to performance increases.\n\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "iMPS0Xkay-Y", "reviewer": "Reviewer_BXji", "summary": "This work investigates the impact of integrating a large-scale, nearest-neighbors, retrieval mechanism into a MuZero-inspired model-based RL agent with application to playing the game of Go. In particular, a large dataset of 50M Go board states is first embedded by a pretrained network. During training, an input board state (representing the current state of the game) is itself embedded by the same network and used to perform a fast approximate nearest-neighbor (NN) lookup across the retrieval dataset. The top-K NNs are (along with relevant metadata) given as input to the model-based agent who may use the context of these games to both improve its prior policy as well as inform its rollouts during Monte-Carlo tree search. A number of experiments show that this retrieval-based agent (RBA) outperforms related baselines trained without retrieval and that, interestingly, the RBA agent can even adapt, without additional training, to new retrieval datasets.\n", "questions": "My main concerns in this work are outlined in the weaknesses section and largely center around the applicability of this method to other domains and more diverse retrieval datasets. I would be happy to increase my rating if given strong arguments addressing these weaknesses or if new relevant results can be presented in the rebuttal period. I have also listed a number of questions in the \"Line-by-line comments\" section in the above answer.\n", "limitations": "There is some note of limitations throughout the paper but no dedicated section for their discussion. I believe there is room to more emphasize some of the weaknesses I've addressed above (assuming I have not identified them incorrectly). I don't see any direct potential for negative societal impact from this work.\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\nResearch on semi-parametric approaches that enable agents to reuse training data beyond what they can distill into their model weights is an exciting direction that is underdeveloped in the reinforcement learning area. Differing from past work in this direction, this work integrates retrieved experience using a simple, but highly generic, mechanism that can, in principle, be extended to retrieving information beyond training data (e.g. an embodied agent might retrieve YouTube videos related to completing some household task). The empirical results presented are limited in scope but are convincing in that there is clear evidence that the RBA agent makes effective use of the retrieved data and that this leads to measurable performance improvements. Finally, the paper is well-written and all details are well explained.\n\n## Weaknesses\n\nWhile I am generally positive about this work there are some weaknesses, detailed below, that keep me from a more enthusiastic recommendation.\n\n- Why not use AlphaZero?\n\nAs noted on line 189, the training dataset used in this work is generated by an AlphaZero-style agent which raises a somewhat unfair question: if AlphaZero-level performance is what we're attempting to achieve, why not just use the AlphaZero agent? I.e. what is to be gained from training a model to imitate a model that we already have? I assume the rebuttal to this is that this dataset choice was purely practical and that anyone looking to apply these ideas to a new domain where no such expert is available should instead use the likely human-generated, expert demonstrations they have available. My worry is then that the use of this AlphaZero-generated dataset may not be a fair stand-in for multi-modal, error-rich, human-generated expert datasets; because of this, the performance of the retrieval-based agent using this artificially generated dataset may be a \"best-case\" and others may find achieving such gains more difficult.\n\nWhile the results from Sec. 3.4 with the augmented retrieval dataset provides some evidence that the above is not the case, I do believe that this work would be much stronger if it had used a human-generated expert dataset or if there were additional results applying this method to other domains.\n\n- Grand aims but limited scope\n\nThis paper (especially the title and abstract) suggests a grand scope: reinforcement learning augmented with retrieval towards a \"vision in which an agent can flexibly draw on diverse and large-scale information sources.\" As one reads deeper into the paper, however, the scope is iteratively narrowed until we are left with an agent that is trained to play 9x9 Go with expert supervision in an off-policy manner with a uni-modal dataset being used for retrieval (namely the training dataset). This is a failure in expectation management, while I still find the work interesting it is difficult to not feel let down going from the grand vision to the reality. This would work be significantly stronger if any step was made along:\n\n1. Using on-policy training data.\n2. Using multi-model data or, at the very least as noted above, not artificially generated expert trajectories.\n3. Showing applicability to other games / environments.\n\n## Line-by-line comments\n\n- Fig. 1\n  * I find the use of lines to denote gradient flow to be a little confusing. I would suggest instead only showing the forward computation and using a symbol along these lines to denote a stop-gradient (often represented as a line break, e.g. ---| |---> ).\n\n- Lines 142-143\n  * End-to-end learning of the query vector seems challenging, can you expand on how you intend to accomplish this?\n\n- Lines 198-199\n  * To confirm, does this mean that all observations from a particular game will be in the same portion of the dataset split? I.e. if o_t is in split 1 then o_{t+1} will also be?\n\n- Lines 224-225\n  * Is it possible that the gap between the retrieval-based method and the baseline can be closed simply by doing more MCTS rollouts? If so, how many are needed to do so? I believe Fig. 10 begins to get at this question but not completely. It would be interesting to know how many rollouts approximately equal one retrieved neighbor.\n \n- Fig. 4\n  * If you evaluate the AlphaZero agent on this dataset, what is its Top-1 accuracy?\n\n- Lines 271-273\n  * This is quite surprising to me, doesn't this suggest that the majority of your gains over the baseline are actually not related to retrieval at test-time? I would be interested in seeing more analysis of this.\n\n- Lines 275-291\n  * While I am not an expert in this domain, it seems as though this related work is missing of how people use retrieval-based ideas outside of RL. E.g. a quick search on Google Scholar shows a large number of works related to building NLP models with access to large-scale knowledge bases.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "GR6SvKdgZQW", "reviewer": "Reviewer_LtnM", "summary": "The authors propose a method for retrieving relevant training experience and then making decisions using that experience as context. They run experiments in the game of Go, showing that when they equalize model sizes, their model looking at retrieval information performs better than a baseline method that doesn’t see retrieval information. Specifically, the retrieval information helps predict values (Figure 4) and helps win the game (Figure 5).", "questions": "Why do you consider an offline RL version of Go, a game where we can easily simulate environment dynamics and already have superhuman algorithms such as MuZero? Why did you not use an offline RL benchmark such as D4RL or RL Unplugged?\n\nWhy did you not directly use a baseline from prior work? Is MuZero Unplugged not directly applicable to your setting?", "limitations": "While the title of the work makes it seem like a general-purpose RL study, the experiments are only in the game of Go. The authors motivate the work via offline RL, but Go is a setting where it is very easy to simulate environment experience, so we don't need offline RL in Go.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\nThe motivation of this work – conditioning on prior experience – is creative. I am not aware of any state-of-the-art algorithms that condition on prior experience, so it would be neat to see that this helps in offline RL.\n\nAs far as I can tell, the authors propose a method that can easily be generalized to other domains. If the authors’ method helps in diverse settings, that would be significant!\n\n## Weaknesses\nMy major issue with this paper is that I don’t see the significance of the authors’ experimental setup. The authors motivate their work via offline RL, but then they study the game of Go. Go is precisely the type of domain where we _don’t need offline RL_ because we can easily simulate the environment dynamics! If the authors want to demonstrate the relevance of their method for offline RL, they should use an offline RL benchmark such as D4RL or RL Unplugged. This would allow them to compare against other offline RL algorithms in suitable offline RL environments.\n\nAs it stands, the authors only compare against a baseline that they themselves came up with. However, there is little analysis to justify that this is a good baseline. Why did the authors not directly use MuZero Unplugged ([36] in the authors’ list of references) as a baseline? Is this not possible?\n\nI am also unconvinced that the proposed method does that much better than the baseline. Figure 5 would be the most important figure in the paper, showing that the proposed method helps win the game. However, if I understand correctly, the “model size” on the x-axis of Figure 5 is an unfair comparison because it is leaving out the parameter count of $g_\\phi$ which is needed for retrieval. It also seems that the largest model size only shows a 5% increase in winrate over the baseline. This doesn’t seem very impressive to me as the opponent is a fixed program with “strong amateur level of play” (and I’m also not confident in the baseline).\n\nFinally, I think the title of the work is overclaiming because it is too broad. Since the authors only evaluate in Go, I think the title should make this clear, e.g., “Large-Scale Retrieval for Reinforcement Learning in the Game of Go.”\n\n## Minor issues / typos\nFigure 2 shows $D_r$ and describes it as having “keys $k$”. I believe the keys $k$ are generated using $g_\\phi$, but there is no indication of this in the figure. This confuses me.\n\n## I Request Title and Figure Changes\nUpdate after author response: My rating remains \"borderline reject.\" **If this paper gets accepted anyway, then I request that (1) the title be changed to be more modest (eg \"Retrieval for Reinforcement Learning in 9x9 Go\"), and (2) figures 4 and 5 be modified to include the parameter count of the retrieval network.**", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "FNgcnFVhKOT", "reviewer": "Reviewer_C5Cc", "summary": "The goal of the paper is to enable RL agent to retrieve large amount of information from sources using nearest neighbor queries over experience datasets. As testbench the authors rely on the game GO, where the agent can retrieve information from an auxiliary source of 50M board-state observations. The work studies the trade-off between leveraging such external information repositories and storing all necessary information into the network as weight to optimize a policy in an offline RL setting. An extension for online RL is left open for future work. \nTo retrieve information the model first precomputes embeddings for all possible observations based on hidden activations of a surrogate model which are compressed by PCA. The embeddings are used as query and key for NN search over an external dataset. Lastly, the model output is then generated using the NN results and the current observation. ", "questions": "How would one need to extend the presented approach to incorporate non-trajectory auxiliary data?", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- In my opinion, augmentation of RL agents with \"any kind of\" external data is a great idea, since it can deliver benefits for real-world use cases. For instance, by altering the external data an operator can have direct and immediate influence on a deployed RL policy. It would be interesting to further study how the auxiliary data influences the behavior of the RL policy. \n- Evaluation against a human competitor and analysis over different retrieval data size demonstrate a strong benefit.\n\nMain concern: Evaluation does not support some of the claims of the paper\n- The authors make the claims that potential auxiliary sources can be diverse, e.g. text, video and that the approach is versatile regarding the type of auxiliary data for training. However, in the experiments they only use trajectories as auxiliary data. To be fair other kinds of external data should be included in the evaluation to assess if this versatility claim really and how general the approach is. \n- For estimating the impact of the retrieval the paper only compares a retrieval-enabled model against retrieval-disabled model with the same architecture. However, how do we know that the improvements are not only due to an increased parameter space? The evaluation would benefit from a comparison where only the number of parameters in the retrieval-disabled model are increased to study the claim that a retrieval-augmented network can utilise more of its capacity for computation instead of having to amortise all relevant information into its network weights. \n\nMinor concerns:\n- A lot is left for future experiments, such as the online RL setup or learning of query vectors, indicating that the paper is still preliminary. ", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept