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

TL;DR

REGENT is a semi-parametric, retrieval-augmented agent that can adapt to new environments through in-context learning, requiring less data and parameters than existing methods, and outperforming current state-of-the-art generalist agents.

Contribution

The paper introduces REGENT, a novel retrieval-augmented transformer agent that adapts to unseen environments without fine-tuning, using less data and parameters than prior approaches.

Findings

REGENT outperforms state-of-the-art generalist agents in unseen environments.

It requires up to 3x fewer parameters and an order of magnitude less pre-training data.

A simple 1-nearest neighbor retrieval baseline is surprisingly effective.

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