Paper: Memory-Consistent Neural Networks for Imitation Learning (ICLR-2024)

Reviewer comment (Novelty):
The proposed method seems equivalent to supervised training of a NN with the codebook entries oversampled. (Edit: this has been clarified by the authors in an additional experiment)

Author response:
We respectfully disagree and note that the MCNN is not equivalent to a neural network with memories oversampled. For all the training data points that are not memories (upto 97.5% of the dataset when memories are 2.5%), the nearest neighbor component and the neural network (both weighted by their corresponding exponential distance factor) contribute to the predicted action. This is particularly effective as seen in Figures 4, 5, and 8 across various tasks, environments, demonstration data, input modalities, and neural network architectures in beating 7 state-of-the-art baselines including various vanilla neural network architectures, all of which have access to the exact same dataset. 

We performed a comparison between MCNN+MLP and MLP-BC with memories oversampled by various amounts and reported these **new results in Appendix F (Table 5)**. Across all Adroit tasks, we see that MCNN+MLP significantly outperforms the oversampled MLP-BC.

We also respectfully note that other reviewers disagree. Reviewer Cf33 describes our method as “an interesting approach for tackling behavior cloning. Even more so since this can be combined with underlying network architectures”, reviewer kyQU describes our method as “simple yet elegant” and reviewer 5EEF describes our method as providing “something that does not exist for vanilla neural networks”.

Reviewer comment (Writing):
It is unclear to me why the clipping function in Fig. 12 should be referred to as "tanh-like".

Author response:
We thank the reviewer for the question. We describe it as tanh-like since it maps inputs to a value in [-1, 1] just like the familiar tanh activation function in the deep learning literature. But, it is actually a combination of ReLUs as given in Algorithm 2, detailed in Appendix E, and depicted in Figure 12.

Reviewer comment (Presentation):
In Fig. 1, what are the absolute returns for ~10, ~100, ~1,000 and ~10,000 trajectories?

Author response:
As requested by the reviewer, in the **newly added Tables 3 and 4 in Appendix F**, we provide the mean and standard deviation of the return values for various number of demos for various tasks (across various runs) that were in the scatter plot in Figure 1.

Reviewer comment (Experiments):
The authors mention that they provide results on 9 tasks across 5 environments. But I only see 5 tasks, 1 per environment. It would be great if the authors could clarify the 9 tasks that they evaluate on and where they have provided the results.

Author response:
> The authors mention that they provide results on 9 tasks across 5 environments. But I only see 5 tasks, 1 per environment. It would be great if the authors could clarify the 9 tasks that they evaluate on and where they have provided the results.

Following notation in prior works [1, 2, 3], we define an imitation learning task as a tuple of (environment, reward function, and demonstration dataset). Since we perform experiments with both "expert" and "human" demonstrations, our experiment list has the following 9 tasks: pen-human-v1, pen-expert-v1, hammer-human-v1, hammer-expert-v1, door-human-v1, door-expert-v1, relocate-human-v1, relocate-expert-v1, and carla-lane-v0. We added this clarification to Appendix E. 

[1] Fu et al, D4RL: Datasets for Deep Data-Driven Reinforcement Learning, 2021.

[2] Kumar et al 2020, Conservative Q-Learning for Offline Reinforcement Learning, 2020.

[3] Florence et al, Implicit Behavioral Cloning, 2021.

Reviewer comment (Experiments):
For CARLA, the images have been embedded using a fixed off-the-shelf ResNet34 encoder. This might not be ideal for more complicated visual scenes such as the Franka Kitchen environment used in BeT and Diffusion Policy. It would be great if the authors could evaluate MCNN in this benchmark and provide some comments on possible pretraining approaches for the encoder when such an off-the-shelf encoder is insufficient.

Author response:
> For CARLA, the images have been embedded using a fixed off-the-shelf ResNet34 encoder. This might not be ideal for more complicated visual scenes such as the Franka Kitchen environment used in BeT and Diffusion Policy. It would be great if the authors could evaluate MCNN in this benchmark and provide some comments on possible pretraining approaches for the encoder when such an off-the-shelf encoder is insufficient.

> “MCNN can even improve the performance of simple MLP architectures to beyond that of more sophisticated recent architectures such as Diffusion models.”: In quite a few cases, MCNN works better with MLP than with BeT or Diffusion Policy. Both BeT and Diffusion Policy were developed to deal with multimodal data distributions. Hence, I wonder if the tasks shown are not multimodal enough. Evaluating MCNN on the Franka Kitchen environment used in BeT and Diffusion Policy would help clarify this.

Thank you very much. As the reviewer requested, we have evaluated MCNN, and in particular MCNN+Diffusion, on the Franka Kitchen environment and have **added these new results to Appendix F (Table 6)**. We would like to clarify that the Franka Kitchen environment used in the BeT and Diffusion Policy papers is state/proprioception based. They do not use image inputs. We obtain performance better than (or equal to) diffusion policy in all metrics from success rates in interacting with 1 object up to 5 objects. For 1 to 4 object interaction, diffusion policy is already a strong baseline achieving close to 100%. Our method (MCNN+Diffusion) achieves 100% in all four metrics. For interacting with 5 objects, we see a 4x improvement with MCNN+diffusion over the performance of the best baseline which is diffusion policy. 

Also, as requested by the reviewer, we recommend end-to-end training both the vision encoder and policy with the behavior cloning loss when an independently trained vision encoder such as [1,2,3,4] is not sufficient.

We would also like to refer the reviewer to videos of our policies in both Adroit and Franka Kitchen environments and videos of demonstrations in Adroit human tasks available at our anonymous website https://sites.google.com/view/mcnn-anon . The multimodality in Adroit, particularly in the human datasets collected by humans wearing VR gloves, is visible in these videos. Some examples include (1) in the hammer task, some demonstrations move the hammer above the bolt before hitting it, some move it below, and some directly hit the bolt and (2) in the pen task, some demonstrations have the little finger above the pen and some have it below the pen throughout the episode.

Finally, through the statement:

*“MCNN can even improve the performance of simple MLP architectures to beyond that of more sophisticated recent architectures such as Diffusion models”*

we only aimed to indicate that MCNN+MLP outperforms plain Diffusion-BC and not that MCNN+Diffusion did not outperform MCNN+MLP.  Indeed, as seen in Figures 3 and 4, in the majority of the tasks (5 of 8), MCNN+Diffusion outperforms (or equals) MCNN+MLP. Among the baselines too, Diffusion-BC outperforms (or equals) MLP-BC in 7 out of 8 tasks. We note that multimodality is not the only challenge in imitation learning. MCNN helps address the significant compounding error challenge which MLP-BC fails to address and hence results in MCNN+MLP having an overall better performance than Diffusion-BC even though Diffusion is better equipped to handle multimodality. Our results in Franka Kitchen highlight this where MCNN+Diffusion improves upon the already strong Diffusion-BC baseline. We added this clarification to Appendix E.

[1] He et al 2019, Momentum Contrast for Unsupervised Visual Representation Learning

[2] Nair et al 2022, R3M: A Universal Visual Representation for Robot Manipulation

[3] Ma et al 2022, VIP: Towards Universal Visual Reward and Representation

[4] Majumdar et al 2023, Where are we in the search for an Artificial Visual Cortex for Embodied Intelligence?

Reviewer comment (Experiments):
*“MCNN can even improve the performance of simple MLP architectures to beyond that of more sophisticated recent architectures such as Diffusion models.”*: In quite a few cases, MCNN works better with MLP than with BeT or Diffusion Policy. Both BeT and Diffusion Policy were developed to deal with multimodal data distributions. Hence, I wonder if the tasks shown are not multimodal enough. Evaluating MCNN on the Franka Kitchen environment used in BeT and Diffusion Policy would help clarify this.

Author response:
> For CARLA, the images have been embedded using a fixed off-the-shelf ResNet34 encoder. This might not be ideal for more complicated visual scenes such as the Franka Kitchen environment used in BeT and Diffusion Policy. It would be great if the authors could evaluate MCNN in this benchmark and provide some comments on possible pretraining approaches for the encoder when such an off-the-shelf encoder is insufficient.

> “MCNN can even improve the performance of simple MLP architectures to beyond that of more sophisticated recent architectures such as Diffusion models.”: In quite a few cases, MCNN works better with MLP than with BeT or Diffusion Policy. Both BeT and Diffusion Policy were developed to deal with multimodal data distributions. Hence, I wonder if the tasks shown are not multimodal enough. Evaluating MCNN on the Franka Kitchen environment used in BeT and Diffusion Policy would help clarify this.

Thank you very much. As the reviewer requested, we have evaluated MCNN, and in particular MCNN+Diffusion, on the Franka Kitchen environment and have **added these new results to Appendix F (Table 6)**. We would like to clarify that the Franka Kitchen environment used in the BeT and Diffusion Policy papers is state/proprioception based. They do not use image inputs. We obtain performance better than (or equal to) diffusion policy in all metrics from success rates in interacting with 1 object up to 5 objects. For 1 to 4 object interaction, diffusion policy is already a strong baseline achieving close to 100%. Our method (MCNN+Diffusion) achieves 100% in all four metrics. For interacting with 5 objects, we see a 4x improvement with MCNN+diffusion over the performance of the best baseline which is diffusion policy. 

Also, as requested by the reviewer, we recommend end-to-end training both the vision encoder and policy with the behavior cloning loss when an independently trained vision encoder such as [1,2,3,4] is not sufficient.

We would also like to refer the reviewer to videos of our policies in both Adroit and Franka Kitchen environments and videos of demonstrations in Adroit human tasks available at our anonymous website https://sites.google.com/view/mcnn-anon . The multimodality in Adroit, particularly in the human datasets collected by humans wearing VR gloves, is visible in these videos. Some examples include (1) in the hammer task, some demonstrations move the hammer above the bolt before hitting it, some move it below, and some directly hit the bolt and (2) in the pen task, some demonstrations have the little finger above the pen and some have it below the pen throughout the episode.

Finally, through the statement:

*“MCNN can even improve the performance of simple MLP architectures to beyond that of more sophisticated recent architectures such as Diffusion models”*

we only aimed to indicate that MCNN+MLP outperforms plain Diffusion-BC and not that MCNN+Diffusion did not outperform MCNN+MLP.  Indeed, as seen in Figures 3 and 4, in the majority of the tasks (5 of 8), MCNN+Diffusion outperforms (or equals) MCNN+MLP. Among the baselines too, Diffusion-BC outperforms (or equals) MLP-BC in 7 out of 8 tasks. We note that multimodality is not the only challenge in imitation learning. MCNN helps address the significant compounding error challenge which MLP-BC fails to address and hence results in MCNN+MLP having an overall better performance than Diffusion-BC even though Diffusion is better equipped to handle multimodality. Our results in Franka Kitchen highlight this where MCNN+Diffusion improves upon the already strong Diffusion-BC baseline. We added this clarification to Appendix E.

[1] He et al 2019, Momentum Contrast for Unsupervised Visual Representation Learning

[2] Nair et al 2022, R3M: A Universal Visual Representation for Robot Manipulation

[3] Ma et al 2022, VIP: Towards Universal Visual Reward and Representation

[4] Majumdar et al 2023, Where are we in the search for an Artificial Visual Cortex for Embodied Intelligence?

Reviewer comment (Experiments):
The main weakness of this work, as is common with BC approaches, is assumption of representative state-action pairs and optimal behavior provided by the demonstrator. There is no insight as to how the method will behave with reasonably suboptimal demonstrations, as it being a BC based approach has no apparent mechanism that enables it to focus on the better demonstrations of a provided set.

Author response:
Thank you, as you have pointed out, our experimental validation has followed standard practice in the literature on an extensive set of tasks, against well-chosen SOTA baselines. While the reviewer is right that the standard evaluation metrics could be better, we make an earnest effort to closely track the community accepted standards. Among our baselines, nearly all have already proven useful for many realistic tasks with suboptimal demonstrations including on real robots with human-teleoperated demonstrations [1, 2, 3]. Furthermore, note that our experiments already include settings with small, noisy demonstration sets. In all four Adroit human tasks, there are a total of only 25 demonstrations collected by humans wearing VR gloves. As shown in Figure 1, even the lowest number of demonstrations (5 demos) leads to large improvements over the baselines.

We would also like to clarify that all policies in adroit environments are **goal-conditioned**.  The environment provides goal position and orientations as part of the observation vector. Every time that the environment is reset (such as for a new evaluation rollout), the goal is randomly chosen. These include random goal orientations of the pen, goal locations for the relocate task, door locations in door tasks, and nut-and-bolt location in the hammer task. This can be seen clearly in the videos of our policy rollouts at this anonymous link https://sites.google.com/view/mcnn-anon .

With only 25 demonstrations or fewer in Adroit human tasks, it is not possible to cover every random goal location sampled during evaluation rollouts for each of the tasks (such as every possible pen orientation) in the 25 or fewer demos. Hence, coverage too is not available! Even under these stressful conditions, with very few demos that do not provide adequate coverage, MCNN is able to generalize well to new randomly sampled evaluation goals (across 20 evaluation trajectories and 3
random seeds) and significantly outperform baselines across tasks and environments. We have also added this information to the experimental evaluation section in the paper and in Appendix E.

All this said, we do acknowledge that there is notable space within the imitation learning literature for methods such as T-REX and Subdomain Minimization that filter demonstration datasets. However, we view such improvements as orthogonal to our method which aims at improving the function class.

[1] Chi et al, Diffusion policy: Visuomotor policy learning via action diffusion, 2023.

[2] Pari et al, The surprising effectiveness of representation learning for visual imitation, 2021.

[3] Florence et al, Implicit Behavioral Cloning, 2021.

Reviewer comment (Novelty):
There is recent interest in work that can discern between useful demonstrations and harmful / irrelevant ones that can inhibit training (i.e TREX, or Subdominance minimization (Ziebart 2022)). While understandably not the direct focus of this work, it can be inhibiting for complex tasks to assume a plethora of very carefully curated demonstrations that are both optimal and extremely descriptive of the entire problem space, hence the aforementioned attempts to augment imitation learning with such capabilities.

Author response:
Thank you, as you have pointed out, our experimental validation has followed standard practice in the literature on an extensive set of tasks, against well-chosen SOTA baselines. While the reviewer is right that the standard evaluation metrics could be better, we make an earnest effort to closely track the community accepted standards. Among our baselines, nearly all have already proven useful for many realistic tasks with suboptimal demonstrations including on real robots with human-teleoperated demonstrations [1, 2, 3]. Furthermore, note that our experiments already include settings with small, noisy demonstration sets. In all four Adroit human tasks, there are a total of only 25 demonstrations collected by humans wearing VR gloves. As shown in Figure 1, even the lowest number of demonstrations (5 demos) leads to large improvements over the baselines.

We would also like to clarify that all policies in adroit environments are **goal-conditioned**.  The environment provides goal position and orientations as part of the observation vector. Every time that the environment is reset (such as for a new evaluation rollout), the goal is randomly chosen. These include random goal orientations of the pen, goal locations for the relocate task, door locations in door tasks, and nut-and-bolt location in the hammer task. This can be seen clearly in the videos of our policy rollouts at this anonymous link https://sites.google.com/view/mcnn-anon .

With only 25 demonstrations or fewer in Adroit human tasks, it is not possible to cover every random goal location sampled during evaluation rollouts for each of the tasks (such as every possible pen orientation) in the 25 or fewer demos. Hence, coverage too is not available! Even under these stressful conditions, with very few demos that do not provide adequate coverage, MCNN is able to generalize well to new randomly sampled evaluation goals (across 20 evaluation trajectories and 3
random seeds) and significantly outperform baselines across tasks and environments. We have also added this information to the experimental evaluation section in the paper and in Appendix E.

All this said, we do acknowledge that there is notable space within the imitation learning literature for methods such as T-REX and Subdomain Minimization that filter demonstration datasets. However, we view such improvements as orthogonal to our method which aims at improving the function class.

[1] Chi et al, Diffusion policy: Visuomotor policy learning via action diffusion, 2023.

[2] Pari et al, The surprising effectiveness of representation learning for visual imitation, 2021.

[3] Florence et al, Implicit Behavioral Cloning, 2021.

Reviewer comment (Experiments):
There is no experimentation for generalization. For example, with demonstration provided for task A, how would the method perform for a slightly altered task A’ where perhaps the goal has changed slightly, the demonstrations are perturbed by noise or simply slightly different. It is not unreasonable to believe that this method would perform poorly in such a setting, and good results would strengthen the method’s appeal considerably.

Author response:
Thank you, as you have pointed out, our experimental validation has followed standard practice in the literature on an extensive set of tasks, against well-chosen SOTA baselines. While the reviewer is right that the standard evaluation metrics could be better, we make an earnest effort to closely track the community accepted standards. Among our baselines, nearly all have already proven useful for many realistic tasks with suboptimal demonstrations including on real robots with human-teleoperated demonstrations [1, 2, 3]. Furthermore, note that our experiments already include settings with small, noisy demonstration sets. In all four Adroit human tasks, there are a total of only 25 demonstrations collected by humans wearing VR gloves. As shown in Figure 1, even the lowest number of demonstrations (5 demos) leads to large improvements over the baselines.

We would also like to clarify that all policies in adroit environments are **goal-conditioned**.  The environment provides goal position and orientations as part of the observation vector. Every time that the environment is reset (such as for a new evaluation rollout), the goal is randomly chosen. These include random goal orientations of the pen, goal locations for the relocate task, door locations in door tasks, and nut-and-bolt location in the hammer task. This can be seen clearly in the videos of our policy rollouts at this anonymous link https://sites.google.com/view/mcnn-anon .

With only 25 demonstrations or fewer in Adroit human tasks, it is not possible to cover every random goal location sampled during evaluation rollouts for each of the tasks (such as every possible pen orientation) in the 25 or fewer demos. Hence, coverage too is not available! Even under these stressful conditions, with very few demos that do not provide adequate coverage, MCNN is able to generalize well to new randomly sampled evaluation goals (across 20 evaluation trajectories and 3
random seeds) and significantly outperform baselines across tasks and environments. We have also added this information to the experimental evaluation section in the paper and in Appendix E.

All this said, we do acknowledge that there is notable space within the imitation learning literature for methods such as T-REX and Subdomain Minimization that filter demonstration datasets. However, we view such improvements as orthogonal to our method which aims at improving the function class.

[1] Chi et al, Diffusion policy: Visuomotor policy learning via action diffusion, 2023.

[2] Pari et al, The surprising effectiveness of representation learning for visual imitation, 2021.

[3] Florence et al, Implicit Behavioral Cloning, 2021.

Reviewer comment (Experiments):
The distance used to select the nearest neighbor is quite critical for this method and problem specific. For example, the distance needed for an image-based problem could be quite different from a control state-based problem. This can potentially deduct from the appeal of a traditional BC approach that doesn't need such an engineered mechanism. (The authors have clarified that the distance used is a generic L2 distance across all problem spaces considered.)

Author response:
Thank you for the questions. We would like to clarify that we did not choose a measure tailored to any environment or task and simply used the general L2 distance everywhere (between two proprioceptive states and between two image embeddings). With this general distance measure, we obtain significant performance improvements across tasks and environments as seen in our results section. We refer the reviewer to Section 5 under “Embedding CARLA images” where we note that the nearest memory neighbor function and neural network operate on the same pre-trained image embeddings or proprioceptive states. 

> What was the motivation behind a static nearest neighbor function? Have you considered training it as well?

Thank you for the question. We would also like to clarify that we learn a neural gas (and memories) following the neural gas clustering algorithm [Fritke 1994] which is an unsupervised learning algorithm. Hence, our nearest neighbor function is not static but learned depending on the demonstration dataset. This nearest neighbor function helps us control the deviation that MCNNs can have from the expert-recommended actions. We refer the reviewer to the detailed explanation provided under the "Neural gas" subheading in Section 4.3.

Reviewer comment (Experiments):
As a (semi-)parametric method, the computational cost of training and inference scales with the number of memories. There is a discussion on computational complexity in the appendix, but some analysis on training time would be appreciated here as it's difficult to tell whether this is a significant factor.

Author response:
Thank you for the question. We refer the reviewer to the "discussion on computation costs" in Appendix E where we note that the wall time required is less than 1 millisecond for finding the memory for a single (or even a batch of up to size 1024) of datapoints by parallel search on a RTX 3090 GPU. This time is not a significant factor during inference or training. Forward propagation through the MLP/transformer/diffusion models also similarly takes less than a millisecond. Further, even on our largest datasets of 1 million transitions for the Adroit expert tasks, with 10% of the dataset or 10k datapoints as memories, nearest memory search uses less than half the GPU VRAM used by the neural network model (diffusion, BeT, or larger MLPs). This results in no additional GPU VRAM usage over the amount used by the model and hence no additional usage when compared to the other baselines.

Reviewer comment (Evaluation):
The performance improvement seems sensitive to the underlying network architecture and the task. E.g. in Fig. 4, different models exhibit different levels of improvement for each task. This could make it difficult to use in cases where we can't sample from the environment.

Author response:
Thank you for the questions.
First, we would like to clarify that all the results in the main paper except Fig 7 are obtained with no online interaction, i.e., no sampling from the environment. The clustering method is also purely offline. We apologize that our description in the paper might have left room for possible misinterpretations on this front. 

We would also like to highlight that all "MCNN + X (Fixed)" methods refer to MCNN with fixed hyperparameters, not tuned with any online interaction, and applied out of the box to new tasks. 

Second, yes, we agree that MCNN yields improvements of varying magnitudes for different tasks and different models, but note that this is not a problem for a practitioner. Indeed, this is generally true for most modeling improvements: for example, in the same Fig 4, IBC outperforms MLP by different margins in different tasks. Rather than consistently sized margins, it is much more important for a candidate modeling improvement to have reliably positive margins. This is very much true in our experiments, as demonstrated tellingly in Fig 1, which aggregates results across many tasks and many base models.

Reviewer comment (Experiments):
It seems that the clustering method requires you to be able to sample from the entire state space, is this the case? If it is not possible to sample from the environment at will, is it possible to do this clustering purely from data? If the method's performance strongly depends on sampling from the environment, then it would only be fair to compare it to interactive imitation learning methods that also sample from the environment like SQIL, GAIL, etc.

Author response:
Thank you for the questions.
First, we would like to clarify that all the results in the main paper except Fig 7 are obtained with no online interaction, i.e., no sampling from the environment. The clustering method is also purely offline. We apologize that our description in the paper might have left room for possible misinterpretations on this front. 

We would also like to highlight that all "MCNN + X (Fixed)" methods refer to MCNN with fixed hyperparameters, not tuned with any online interaction, and applied out of the box to new tasks. 

Second, yes, we agree that MCNN yields improvements of varying magnitudes for different tasks and different models, but note that this is not a problem for a practitioner. Indeed, this is generally true for most modeling improvements: for example, in the same Fig 4, IBC outperforms MLP by different margins in different tasks. Rather than consistently sized margins, it is much more important for a candidate modeling improvement to have reliably positive margins. This is very much true in our experiments, as demonstrated tellingly in Fig 1, which aggregates results across many tasks and many base models.

Reviewer comment (Evaluation):
Given the distance lookups that are necessary for finding neighbors, how does this impact the wallclock time for training and inference?

Author response:
Thank you for the question. We refer the reviewer to the "discussion on computation costs" in Appendix E where we note that the wall time required is less than 1 millisecond for finding the memory for a single (or even a batch of up to size 1024) of datapoints by parallel search on a RTX 3090 GPU. This time is not a significant factor during inference or training. Forward propagation through the MLP/transformer/diffusion models also similarly takes less than a millisecond. Further, even on our largest datasets of 1 million transitions for the Adroit expert tasks, with 10% of the dataset or 10k datapoints as memories, nearest memory search uses less than half the GPU VRAM used by the neural network model (diffusion, BeT, or larger MLPs). This results in no additional GPU VRAM usage over the amount used by the model and hence no additional usage when compared to the other baselines.

Reviewer comment (Evaluation):
It seems like in some tasks the model is very sensitive to the number of memories, e.g. carla-town-v0 in Fig. 15. Do you know why this is the case? It seems unintuitive that 5% would perform worse than both 2.5% and 10%. Do you think that the number of memories needed to attain sufficient performance is related to task/behavior complexity?

Author response:
Thank you for the question. While there is indeed some sensitivity to this parameter, we set it to 10 % in all our experiments in the main paper barring Fig 7 where we explicitly gauge the impact of parameter tuning using online environment interactions. Further, this need not entirely be an empirical choice --- we are guided by the following intuition. The available training dataset must be shared between the nearest neighbors function and the neural network in order for MCNN to work. Since generalization beyond training samples comes from the neural network, we expect it to require more of the training data. This intuition is borne out by the general trend of increasing performance until the number of memories is about 10-20% of the dataset followed by a decrease. It is also difficult to run comprehensive experiments with many seeds, and rare deviations from the general trend on individual tasks may come from the fact that these hyperparameter effects were only tested with 3 seeds per choice. It is statistically likely that there will be minor individual deviations between 3-seed statistics and true trends in a very small fraction of our reported experiments.

Reviewer comment (Experiments):
It seems like there might be a relationship to critical/important states which exhibit a large difference in the expected reward between the best and worst action, i.e. it is costly to recover from a mistake in a critical state. It may be interesting to do an analysis where you identify the critical states in a task and then ablate them from memory. Do you have any thoughts on whether you can do more intelligent memory selection?

Author response:
We thank the reviewer for this insightful comment. Yes, we agree that it would be helpful to explore smarter strategies for memory selection in future work. We speculate that this would need correlating memory location with 'return to go' like metrics. However, in this paper, we used well-known statistical algorithms to place memories which can capture the distribution of points. It turns out that modifying the function class with such data points can have a big impact on the generalization of the learned policy.