REGENT: A Retrieval-Augmented Generalist Agent That Can Act In-Context in New Environments

Authors: Kaustubh Sridhar, Souradeep Dutta, Dinesh Jayaraman, Insup Lee
Primary area: reinforcement learning
Keywords: generalist agent, retrieval, in-context learning, vla, imitation learning, reinforcement learning

Abstract:
Building generalist agents that can rapidly adapt to new environments is a key challenge for deploying AI in the digital and real worlds. Is scaling current agent architectures the most effective way to build generalist agents? We propose a novel approach to pre-train relatively small policies on relatively small datasets and adapt them to unseen environments via in-context learning, without any finetuning. Our key idea is that retrieval offers a powerful bias for fast adaptation. Indeed, we demonstrate that even a simple retrieval-based 1-nearest neighbor agent offers a surprisingly strong baseline for today's state-of-the-art generalist agents. From this starting point, we construct a semi-parametric agent, REGENT, that trains a transformer-based policy on sequences of queries and retrieved neighbors. REGENT can generalize to unseen robotics and game-playing environments via retrieval augmentation and in-context learning, achieving this with up to 3x fewer parameters and up to an order-of-magnitude fewer pre-training datapoints, significantly outperforming today's state-of-the-art generalist agents.

Decision: accept (oral)
Reviewer scores: 6, 8, 5, 8

Reviews and discussion:
<thread>
creation date:
1730302403923

thread_type:
official_review

<thread_object>
creator: reviewer_8qsk

<content>
summary:
This work proposes to leverage retrieval augmentation to reduce the amount of training data required for training generalist agents.
The authors introduce $\texttt{REGENT}$, a retrieval-augmented architecture that incorporates a non-parametric retrieval component (R&P) with learning of a parametric policy.
Further, they demonstrate theoretical guarantees for performance with respect to the number of expert demonstrations.
They pre-train $\texttt{REGENT}$ on 4 different domains and evaluate on unseen levels and holdout tasks, where it exhibits strong performance compared to baselines.

strengths:
The authors provide theoretical guarantees for performance with respect to the number of demonstrations.

$\texttt{REGENT}$ is a novel architecture that effectively incorporates retrieval augmentation for imitation learning.

Retrieval augmentation enables training on much less data than other generalist agents.

The authors demonstrate the promise of retrieval augmentation for unseen levels and tasks.

weaknesses:
**Significance of results:**

The authors only provide averages in all figures (except for Figure 7), but no variance estimates. How do those look? Are the results on the unseen tasks statistically significant?

The authors claim that $\texttt{REGENT}$ outperforms JAT even after fine-tuning, but there is little information about the fine-tuning protocol, i.e. what fine-tuning technique is used, etc. Therefore it is difficult to determine the significance of the results. Can the authors elaborate a on this?
Also there are fine-tuning techniques specifically designed for few-shot learning, which may be more suitable in this setting [1,2].

On unseen levels of ProcGen the authors only compare $\texttt{REGENT}$ to R&P, but no results for MTT, why is this the case?
BC is a very weak baseline for training from scratch. The authors should compare to other more recent imitation learning techniques [3,4,5] that use the same amount of expert demonstrations as $\texttt{REGENT}$. This would shed light on the significance of the reported results.

Another interesting experiment would be zero-shot evaluation without having expert demonstrations on the unseen tasks. If $\texttt{REGENT}$ performs well in this scenario, this would greatly strengthen the paper.

There is a lack of ablation studies, for example performance on varying the number of examples in the context and ablation on different distance metrics.

Finally, it is not really clear why the authors differentiate between the two different settings. Why not include ProcGen in the first setting?
Why are different baselines used for the different settings? MTT could also be applied to the first setting, right?
At least having a fixed set of all baselines across both settings would strengthen the paper.

**Limitations and claims**

- The authors claim that deployment on unseen tasks is instantaneously, this is not true though, as deployment in an unseen environment requires collection of expert demonstrations which in turn requires training of an expert policy. This should be made more explicit as a limitation of the proposed method.
- Moreover, claims on "handful of examples" should be caveated as thousands of samples are not a handful.
- Authors claim apples-to-apples comparison between JAT and $\texttt{REGENT}$ when being trained on the same amount of data. This is not apples-to-apples though as $\texttt{REGENT}$ uses expert demonstrations of the unseen tasks. The more fair comparison is to fine-tuning JAT.
- Further, the work would benefit from a bit more accurate positioning on the landscape of methods that do or do not rely on demonstrations. This is an important aspect, as there are approaches that aim at leveraging in-context learning in the context of learning-to-learn as in meta-learning [6]. Examples are [7,8], where the latter also relies on retrieval augmentation. 

**Presentation:**

Presentation can be improved, in particular:

* Eq. 1 & 2: $c_t$ is provided as input to the policy, but it is never used, it is only used in Eq. 3, as R&P does not rely on context information
* Figure 3: enumeration in caption, but not in figure, either remove them in caption or add in figure
* Paragraph on $\texttt{REGENT}$ architecture is not really clear as the reader does not have an overview what modalities are represented in the training data, therefore the choice of encodings seem arbitrary.
* Theorem 5.2: $\boldsymbol{H}$ represents the horizon, right? In section 3 it is not boldface, this should be consistent.
* Figure 8: hard to follow what actions are "good" or "bad" as there is no explanation for the action index

**References:**

[1] Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning, Liu et al., NeurIPS 2022

[2] Scaling & Shifting Your Features: A New Baseline for Efficient Model Tuning, Lian et al., NeurIPS 2022

[3] Generative Adversarial Imitation Learning, Ho et al., NeurIPS 2016

[4] Disagreement-Regularized Imitation Learning, Brantley et  al., ICLR 2020

[5] Of Moments and Matching: A Game-Theoretic Framework for Closing the Imitation Gap, Swamy et al., ICML 2021

[6]   Evolutionary principles in self-referential learning, or on learning how to learn, Schmidhuber et al., Diploma Thesis, 1987

[7] Towards General-Purpose In-Context Learning Agents, Kirsch et al., NeurIPS 2023 FMDM Workshop

[8] Retrieval-Augmented Decision Transformer: External Memory for In-context RL, Schmied et al., arXiv 2024

questions:
- Why is the scaling factor for MixedReLU set to 10? This seems like an arbitrary choice, why not remove it entirely?
- Does the number of demonstrations for MTT on ProcGen match the number of demonstrations for $\texttt{REGENT}$?
- Is it always the case that $\texttt{REGENT}$ only receives the current query state in its context? Have you tried providing more sequential information?
- What ProcGen mode is being used? Easy or hard?
- How important is the ordering in the context of $\texttt{REGENT}$? Did you try permuting the order?
- In Figure 15 in Appendix D not all curves show improvement for using more demonstrations, do the authors have an intuition why this is the case?
- From Figures 11, 12, 13, 14 it can be seen that JAT is particularly good on BabyAI, but not good on other domains, while $\texttt{REGENT}$ is very good at Metaworld. Do the authors have an intuition why this is the case?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730650308027

thread_type:
official_review

<thread_object>
creator: reviewer_euma

<content>
summary:
The paper proposes REGENT, a retrieval-augmented generalist agent that can adapt to new environments without fine-tuning. REGENT can outperform other RL  agents with fewer parameters in two settings.

strengths:
1. The methodology is well-explained and easy to understand.

2. REGENT demonstrates better performance than larger models (JAT/Gato) with fewer parameters across multiple environment types.

weaknesses:
1. Limited Comparison Scope:

While the paper draws inspiration from retrieval-augmented generation (RAG) in language models, it lacks comparisons with recent LM embodied agents that use retrieval-augmented methods. Some existing work:
- https://arxiv.org/abs/2308.10144
- https://arxiv.org/abs/2402.03610

2. Insufficient Ablation Studies, some valuable ablations could have included:
- Analyzing performance w/wo image information in the context
- Testing various context sizes
 - Analyzing the REGENT's inference latency per state is crucial for real-world deployment. Given the retrieval mechanism and context processing at each step, runtime performance analysis would be valuable for assessing practical applicability


3. Training Data Requirements:

While REGENT shows improved data efficiency compared to baselines like JAT/Gato, it still requires substantial pretraining data across different embodied settings (14.5M transitions in JAT/Gato setting, 12M in ProcGen). This raises questions about its true few-shot learning capabilities and practical applicability in scenarios where large-scale demonstration data is unavailable across task families.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730708623889

thread_type:
official_review

<thread_object>
creator: reviewer_n77b

<content>
summary:
This paper presents a method that pre-trained on diverse demonstrations across different environments and tasks that can be adapted to unseen environments via in-context learning, with very few expert demonstrations, and without any fine-tuning. The method is a semi-parametric model that leverages a simple retrieval-based 1-nearest neighbor agent and in-context learning based on other retrieved neighbors to generalize to unseen robotics and game-playing environments. The authors show that their approach is more sample-efficient and parameter-efficient than strong baselines such as Gato.

strengths:
+ This paper proposes a novel and interesting idea of combining a 1-nearest neighbor agent with a Transformer pre-trained to perform in-context learning from retrieved N-nearest neighbor demos. This approach seems to work well for both the Gato setup and the ProcGen setup regarding generalization over unseen environments for robotics or game-playing tasks.
+ The results of the simple 1-nearest neighbor agent show that it serves as a surprisingly strong baseline, which provides insights for solving tasks in the unseen environment.

weaknesses:
+ The paper will be stronger if its experiments are also performed on a larger set of manipulation tasks (ManiSkill2, RLBench, etc.)

questions:
+ I am curious what are the results of REGENT without the help of the 1-nearest neighbor agent.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731001142819

thread_type:
official_review

<thread_object>
creator: reviewer_jvnh

<content>
summary:
This work introduces REGENT, a retrieval-augmented generalist agent designed to adapt to unseen environments without finetuning. It leverages in-context learning and achieves significant generalization across different environments, including robotics and game-playing. It proposed a retrieve-and-play method and incorporated a transformer-based policy to train on sequences of queries and retrieved neighbors.

strengths:
1. The work is well-organized and easy to follow.
2. The observation that the simple retrieve-and-play method can match or surpass the performance of state-of-the-art generalist agents in unseen environments is interesting.

weaknesses:
1. This paper uses nearest neighbor retrieval, claiming it is effective for small training datasets. Yet, they use transformer architecture, known as data hungry. These contradict each other. To prove that nearest neighbor retrieval is indeed effective, it should also be compared to other retrieval methods. 
2. There are some places that the author overclaims:

- The author claims multiple times that Gato "struggles to achieve transfer to an unseen Atari game even after finetuning, irrespective of the pretraining data". Note that in Section 5.2 and Figure 9, Gato did a comprehensive study to show that they can transfer to new tasks with few demonstrations. 

- The author claims that their model has the ability to adapt to new tasks via in-context learning. However, the ICL refers to a model's capacity to learn and perform tasks by observing examples "without requiring additional training." The LLM can do ICL **without training via RAG**. This is conceptually different from what the author claimed ICL in their paper. 

- The author claims that methods like decision transformers cannot generalize to new goals caused by changes in visual observation/available control/game dynamics. Nevertheless, the experimental setting in this paper is simple: the simulation benchmark, like Metaworld, does not involve many visual changes between different tasks. To verify that the REGENT can indeed handle these environmental changes, it is recommended that a real robot experiment be conducted, as Gato does in their paper. 

3. Why is JAT/Gato finetuned on new demonstrations, as it is already trained by all data, such as 50 tasks in MetaWorld? This seems to make JAT/Gato overfitting on the target dataset, thus performing extremely badly in experiments in Figure 4. The comparisons are not fair, and the author should provide proper reasons and compare them with the vanilla JAT/Gato. 

4. Missing highly-related reference [1,2].

[1] Retrieval-Augmented Embodied Agents, CVPR24
[2] Re-ViLM: Retrieval-Augmented Visual Language Model for Zero and Few-Shot Image Captioning, EMNLP

questions:
Please see weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732111193797

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Official Global Comment by Authors

comment:
Dear reviewers,

We thank you all for your insightful feedback and helpful suggestions.

We are grateful that reviewers jvNh, N77b, 8qsk rated our soundness and contribution as good (3). We are happy that reviewers N77b and eumA also rated our presentation as good (3). We are happy to see the many positive comments on our novelty, architecture, efficiency, performance, and theoretical guarantees.

We have addressed each weakness and question for all reviewers below their review. We have also made the following major changes to the paper and highlighted all changes (major or minor) in red in the pdf.

* New results in Appendix C on the effect of number of states/rewards/actions in the context on REGENT’s performance [reviewers eumA and 8qsk]
* New results in Appendix C on the effect of the ordering of the context on REGENT’s performance [reviewer 8qsk]
* Inference runtime values in Appendix C [reviewer eumA]
* Standard deviations and means computed over training seeds and evaluation rollout seeds in Figure 4 which required pre-training 3 base models and then training 96 finetuned models [reviewer 8qsk]
* Simpler and clearer exposition in subsections 4.1 and the start of subsection 4.2 [reviewer 8qsk]
* Additional related work in Appendix E [reviewers jvNh, eumA, 8qsk]

Please let us know if you have lingering questions and whether we can provide any additional clarifications during the discussion period to improve your rating of our paper.

Thanks,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper looks at generalist agents that can adapt to new environments via retrieval-augmented policy that retrieves nearby states from demonstrations. This retrieved information is used in-context by the transformer-based policy, and it is shown that retrieval is effective but pre-training with retrieval is even more so. Results are demonstrated across a range of environments including robotics and game-playing. The reviewers appreciated that the paper was well-written and the idea simple and effective. Translation of advancements in large language models such as RAG (retrieval-augmented generation) and in-context learning is indeed an interesting direction. However, a number of concerns were raised about the overall execution, including better comparisons to JAT/Gato, comparison to other retrieval/imitation learning methods, statistical significance, ablations such as context size/information type, and lack of some details (e.g. fine-tuning protocol, etc.). The authors provided an extremely comprehensive rebuttal, providing a large amount of experiments addressing the reviewers, which all mentioned that their concerns have largely been resolved. 

 After considering all of the materials, I recommend acceptance of this paper. The  method is simple but interesting, and provides a nice signal of moving towards more generalist agents that can solve new environments. I highly recommend that the authors include many of the updated results and discussions in the full paper, as they make the claims and results much stronger and more rigorous.