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

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. Website: https://kaustubhsridhar.github.io/regent-research

• The paper introduces REGENT, a novel retrieval-augmented agent architecture combining retrieval with transformer policy learning for efficient generalisation in unseen environments.
• REGENT achieves superior zero-shot performance on diverse tasks with significantly fewer parameters and pre-training data compared to existing generalist agents like JAT.
• This retrieval-based approach facilitates in-context learning, allowing REGENT to adapt effectively to new environments using only a few demonstrations.

Retrieval-Augmented Generalist Agents with In-Context Learning

The research paper introduces a novel approach towards developing generalist AI agents capable of rapid adaptation to new environments. Instead of scaling up existing architectures, the authors propose a retrieval-augmented method called REGENT (Retrieval-Augmented Generalist Agent). The paper emphasizes leveraging retrieval as a means to enhance agent adaptability across diverse, previously unseen environments, achieving significant efficiency in terms of parameters and dataset sizes used for pretraining.

Key Contributions and Methodology

1. Retrieve and Play (R) Baseline: The paper starts by exploring a rudimentary retrieval-based agent termed "Retrieve and Play." This agent employs a simple 1-nearest neighbor approach by selecting and mimicking the action from its closest state in a retrieval dataset. Even this simplistic model shows competitive performance against state-of-the-art generalists, underscoring the potential of retrieval-based biases in action selection.
2. REGENT Design: Building on the promising results of the R agent, REGENT combines retrieval with transformer-based policy learning. It operates on sequences comprising queries and retrievals from a previously encountered set of demonstrations. By including both state information and corresponding actions/rewards in the retrieved context, REGENT tunes a transformer model while maintaining generalization capabilities across different robotics and game-playing settings.
3. Data and Parameter Efficiency: REGENT uses significantly fewer parameters (3x less) and pre-training data (an order-of-magnitude less) than typical architectures like Gato and JAT, yet achieves superior performance. This efficiency marks a substantial leap towards deploying capable agents in resource-constrained scenarios.
4. Evaluation on Diverse Environments: The method is validated on two settings: JAT (Metaworld, Atari, Mujoco, BabyAI) and ProcGen environments. It excels by adapting to unseen environments without finetuning, outperforming baselines like JAT, even when JAT has been extensively finetuned.
5. In-Context Learning: REGENT's architecture allows for context-based learning similar to retrieval-augmented LLMs. This facilitates its ability to generalize using only a handful of new environment-specific demonstrations, paralleling advancements in LLMs tailored for in-context learning tasks.

Implications and Discussion

The implications of introducing REGENT are considerable in the context of reinforcement learning and the design of generalist agents. By focusing on retrieval-based augmentation, the approach overcomes barriers related to large model size and extensive dataset requirements. This offers a refreshing perspective on how general capabilities might be harnessed more efficiently across a vast array of environments.

1. Theoretical Insights: The work provides theoretical bounds on REGENT's sub-optimality, illustrating that increasing the number of retrieval contexts progressively ameliorates performance. This aligns practical achievements with theoretical predictions, further substantiating the methodology.
2. Robustness Across Modalities: REGENT manages different observation modalities (image-based, proprioceptive, and text-based inputs) and action spaces (discrete and continuous) exemplified by its successful applications across Metaworld and Atari environments.
3. Opportunities for Future Research: The limitations acknowledged—particularly regarding adaptations to novel embodiments and extremely long-horizon tasks—pave avenues for future research. These might include diversifying the training dataset further or refining retrieval methods to align more closely with the tasks' specifics.

In conclusion, REGENT's retrieval-augmented strategy for training AI generalists represents a forward-thinking step towards efficient and adaptable agent architectures. It challenges the notion that larger, more complex models necessarily equate to better performance in unseen environments. Instead, leveraging retrieval as a bias for in-context learning provides a promising framework for developing robust, versatile agents. These findings are particularly relevant for extending agent functionalities within domains that demand quick adaptation to changes—a haLLMark of real-world applications.