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

2024-12-06

arXiv.org

AI Summary

To address the challenge of rapid zero-shot adaptation of generalist agents to novel environments, this paper proposes a retrieval-driven semi-parametric agent architecture. Methodologically, it abandons large-scale pretraining and gradient-based fine-tuning, instead employing 1-nearest-neighbor (1-NN) retrieval as its core inductive bias. The agent integrates a Transformer-based sequential policy model with retrieval-augmented generation (RAG)-style action selection, enabling in-context few-shot adaptation. This work provides the first empirical evidence that lightweight retrieval can substitute for massive parameter counts and extensive training data to deliver strong generalization priors. Experiments demonstrate that the proposed agent surpasses state-of-the-art generalist agents across diverse unseen environments—including robotic control and game-playing—while reducing parameter count by 2/3 and cutting pretraining data requirements by an order of magnitude. Crucially, it achieves significantly improved generalization performance.

Technology Category

• Search and Optimization: Learning to Search
• Multiagent Systems: Adversarial Agents
• Cognitive Modeling & Cognitive Systems: Agent Architectures

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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

Problem

• Adapts agents to new environments efficiently
• Uses retrieval for fast adaptation
• Reduces parameters and data for pre-training

Innovation

• Retrieval-augmented adaptation
• In-context learning
• Transformer-based policy training