Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 

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. 

The following are some of my questions and would be great if the authors could clarify some of them: 

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

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

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

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

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

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

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. 

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

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper proposed a state-augmentation strategy based on a fully differentiable model for online and offline reinforcement learning.
The authors evaluate the proposal on two algorithms, DQN and R2D2.
The 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.
The retrieved information is then aggregated using a so-called summarizing model which is also end-2-end differentiable.
The authors claim and show detailed improvement with respect to the state of the art on Atari and multi-task offline RL.
In the case of multi-task, the authors show that incorporating multiple tasks into the retrievable dataset is beneficial.
main_review: The paper proposed a simple approach of end-2-end2 differentiable state-augmentation through retrieval.

Strenght:
* The approach is well written and the paper is pleasant to read.
* The modular choices composing the proposed method are reasonably clear and motivated.
* The experimental setting is clearly explained and reproducible.
* The experiment's motivations are clearly justified.

Weakness:
* 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.
* 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.
* 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_review: The paper proposed a retrieval-based state augmentation approach for online and offline reinforcement learning.
The retriever is stateful and end-2-end differentiable.
The approach is evaluated in online RL using Atari benchmark and offline RL.
The experiments in offline RL missed the main baselines and could have been evaluated on a state-of-the-art benchmark like the D4RL benchmark.
The retrieval approach also could have been evaluated against a fully attention-based model over the dataset, like Performer or Reformer Transformer architecture for example.
In 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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.

Experiment 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.
main_review: Strengths:
The 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.
The paper is well-organized and relatively easy to follow (there are several clarification issues, see the weakness part). 
This work is a good trial to adapt the success of retrieval-based methods from language models to RL problems.

Weakness:

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

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

The 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.
1. 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?
2. 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?
3. Information bottleneck is used to regularize the information retrieval. Is the final performance sensitive to the weight of this regularization loss?
4. 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.

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

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

Looking forward to seeing stronger evidence to better support the advantages of the proposed method. 



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

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
main_review: Strengths:
* 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.
* 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
* There are a few promising results (e.g., BabyAI), and the ablations appear to show the importance of the introduced components.
* The creation of the `gridroboman` offline benchmark is a really nice addition, and should be commended.

Weaknesses:
* The experimental methodology is lacking in several aspects:
   * 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.
   * 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.
  * There should be information about run time and computational resources compared to the vanilla R2D2 agent due to the additional machinery introduced by R2A.
  * 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.
  * 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.
* 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?
* In some cases, I'm not entirely convinced by the improvements offered by R2A over R2D2. Particular areas of concern are:
   * 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`?).
   * 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?
* 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).

Nits:
* LSTM and GRU appear to be used interchangeably, when they are different models, for instance in **Step 1: Query Computation** section of page 5.

Refs:

[1] Neural Episodic Control, Pritzel et al. arxiv:1703.01988

[2] Fast deep reinforcement learning using online adjustments from the past, Hansen et al., arxiv:1810.08163
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.

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

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