Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper mainly experimented on Go. It leverages fast, approximate nearest neighbor techniques in order to retrieve relevant data from a
set 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.
strengths_and_weaknesses: 1. The paper is well organized and easy to follow.

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


questions: 1. In figure 5a, it shows the retrieval networks using varying numbers of retrieved neighbors and for a baseline non-retrieval
network. 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
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: 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. 
To 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. 
strengths_and_weaknesses: Strengths:
- 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. 
- Evaluation against a human competitor and analysis over different retrieval data size demonstrate a strong benefit.

Main concern: Evaluation does not support some of the claims of the paper
- 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. 
- 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. 

Minor concerns:
- 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. 
questions: How would one need to extend the presented approach to incorporate non-trajectory auxiliary data?
limitations: yes
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.

strengths_and_weaknesses: ## Strengths

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

## Weaknesses

While I am generally positive about this work there are some weaknesses, detailed below, that keep me from a more enthusiastic recommendation.

- Why not use AlphaZero?

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

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

- Grand aims but limited scope

This 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:

1. Using on-policy training data.
2. Using multi-model data or, at the very least as noted above, not artificially generated expert trajectories.
3. Showing applicability to other games / environments.

## Line-by-line comments

- Fig. 1
  * 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. ---| |---> ).

- Lines 142-143
  * End-to-end learning of the query vector seems challenging, can you expand on how you intend to accomplish this?

- Lines 198-199
  * 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?

- Lines 224-225
  * 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.
 
- Fig. 4
  * If you evaluate the AlphaZero agent on this dataset, what is its Top-1 accuracy?

- Lines 271-273
  * 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.

- Lines 275-291
  * 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.
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.

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.
review_title: null

final_score: rating: 4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.
review_content: 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).
strengths_and_weaknesses: ## Strengths
The 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.

As 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!

## Weaknesses
My 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.

As 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?

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

Finally, 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.”

## Minor issues / typos
Figure 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.

## I Request Title and Figure Changes
Update 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.**
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?

Why 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.
review_title: null

Meta-review: metareview: This paper uses nearest neighbor methods to retrieve and exploit information from similar games during planning, whilst playing the game of go (although the method is extensible to other environments which support muzero-style agents). The reviewers found this approach interesting and ultimately worth publishing, although there was a range of scores. However, the emerging consensus during the review phase seemed to lean towards acceptance, a recommendation I am happy to support from having read the discussion.
confidence: Certain
Decision: Accept