PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kaustubh Sridhar, Souradeep Dutta, Dinesh Jayaraman, Insup Lee

arXiv:2412.04759v2cs.AI

TL;DR

Generalist agents need to adapt rapidly to unseen environments, but existing systems often depend on large models and datasets and still struggle with transfer. REGENT uses retrieval-augmented transformer policies and in-context learning, without finetuning, to adapt from a few demonstrations. It outperforms strong generalist baselines across robotics and game-playing settings while using fewer parameters and pre-training datapoints.

  • Problem

    Existing generalist agents struggle to adapt rapidly to unseen environments despite large policies and datasets.

    Source: §1 · p. 1

  • Method

    REGENT combines a transformer policy with retrieved state-reward-action tuples and aggregates its predictions with a nearest-neighbor Retrieve and Play policy for in-context adaptation.

    Source: §1 · p. 1, §4 · p. 4, §4.1 · p. 4

  • Results

    REGENT significantly outperforms JAT/Gato and MTT on unseen environments while using 1.4x to 3x fewer parameters and an order-of-magnitude fewer pre-training datapoints.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    Retrieval provides a powerful bias for rapid adaptation, allowing even a simple nearest-neighbor agent to match or exceed state-of-the-art generalist agents.

    Source: §1 · p. 1, §7 · p. 10

    The evaluation focuses on unseen environments within the same suite as training environments, while very long-horizon Atari tasks and new Mujoco embodiments remain challenging.

    Source: §4.2 · p. 7, §6 · p. 9

Abstract

from arXiv · show

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

1 INTRODUCTION

Generalist agents must adapt rapidly to unseen environments, yet existing approaches rely on large policies and datasets and still struggle to transfer. The paper instead uses retrieval and in-context learning, with REGENT outperforming baselines while using fewer parameters and pre-training datapoints.

Motivation

  • Existing generalist agents are expected to perform across many training environments but often require large datasets and still struggle to generalize to unseen environments.Rapid adaptation is especially important for agents operating in changing digital and real-world environments.Source: §1 · p. 1

Retrieval baseline

  • R&P matches or exceeds state-of-the-art generalist agents using only nearest-neighbor retrieval from demonstrations, with no pre-training environments or neural-network policy.These results motivate retrieval as a strong bias for rapid adaptation.Source: §1 · p. 1

REGENT

  • REGENT combines a transformer policy with retrieved state, reward, and action tuples, aggregating transformer and R&P action predictions for in-context adaptation.This semi-parametric design supports deployment in entirely unseen environments and tasks with only a few demonstrations.Source: §1 · pp. 1–2

Evaluation

  • The evaluation covers robotics and game-playing settings based on Gato/JAT and ProcGen environments.These two settings provide the paper’s main testbeds for unseen-environment generalization.Source: §1 · p. 2

Main results

  • REGENT significantly generalizes to unseen environments without finetuning, outperforming JAT/Gato and MTT while using 1.4x to 3x fewer parameters and an order-of-magnitude fewer pre-training datapoints.It also performs better than baselines within pre-training environments.Source: §1 · p. 2

2 RELATED WORK

Related work shows that many generalist agents struggle with unseen-environment transfer or cannot adapt through in-context learning. REGENT differs by combining retrieval augmentation with in-context adaptation and achieving stronger transfer with smaller data and models.

Generalist agents

  • JAT/Gato struggles to transfer to unseen Atari games even after finetuning, while REGENT outperforms it with fewer parameters and an order-of-magnitude fewer pre-training datapoints.REGENT has 138.6M parameters versus JAT’s 192.7M, while JAT uses up to 5-10x more data.Source: §2 · p. 2

In-context learning

  • REGENT and R&P adapt to new Atari games with as little as 10k transitions and no finetuning, unlike approaches requiring additional fine-tuning or repeated data-collection cycles.Several robotics and vision-language-action agents do not evaluate or demonstrate in-context learning capabilities.Source: §2 · pp. 2–3
  • Existing in-context reinforcement-learning methods generally adapt to new goals or tasks within the same environment but do not handle changes in observations, controls, and dynamics.This distinguishes their setting from REGENT’s unseen-environment transfer.Source: §2 · p. 3

ProcGen comparison

  • In ProcGen, REGENT outperforms MTT while using an order-of-magnitude fewer pre-training transitions and a model about one-third its size.MTT adapts by placing demonstrations into context, whereas REGENT adds retrieval augmentation.Source: §2 · p. 3

Retrieval augmentation

  • Retrieval-augmented generation has enabled language models to perform competitively with fewer parameters and adapt quickly to new or up-to-date data.REGENT applies this retrieval-augmentation perspective to decision-making agents.Source: §2 · p. 3

3 PROBLEM FORMULATION

The paper formulates adaptation as episodic decision making across training and unseen MDPs, where abundant expert demonstrations are available for training environments but only a few for unseen ones. The agent uses in-context learning without finetuning, assuming unseen state and action spaces are known.

Problem setting

  • The agent is pre-trained on datasets from multiple environments and then generalizes to unseen environments from a few expert demonstrations without additional finetuning.The formulation assumes that the unseen environments’ state and action spaces are known.Source: §3 · p. 3

MDP formulation

  • Each environment is modeled as a finite-horizon episodic Markov Decision Process with states, actions, transitions, rewards, discounting, and an initial-state distribution.Policies are evaluated by expected cumulative reward over an episode.Source: §3 · p. 3

Demonstrations

  • Expert demonstrations consist of sequential tuples containing the current state, previous reward, and action for each training or unseen environment.Training environments provide substantially larger demonstration datasets than unseen environments.Source: §3 · p. 4

Training and unseen environments

  • The formulation distinguishes K training environments with sizeable expert datasets from unseen environments indexed K+1 through M with only a handful of demonstrations.This imbalance defines the few-shot adaptation setting addressed by the paper.Source: §3 · p. 4

4 REGENT: A RETRIEVAL-AUGMENTED GENERALIST AGENT

This section develops retrieval-based agents that adapt to unseen environments from demonstrations, first using nearest-neighbor Retrieve and Play (R&P), then a transformer-based semi-parametric REGENT policy. REGENT combines retrieved context with learned predictions and is deployed without finetuning across robotics and game-playing settings.

4.1 RETRIEVE AND PLAY (R&P)

  • R&P selects the action associated with the nearest state retrieved from demonstrations in the current environment, providing adaptation without learned transfer from pre-training.Its simplicity makes it a direct retrieval baseline for few-shot adaptation.Source: §4.1 · p. 4

4.2 RETRIEVAL-AUGMENTED GENERALIST AGENT (REGENT)

  • REGENT trains a causal transformer to combine the current state and previous reward with retrieved (state, reward, action) tuples from demonstrations.The retrieved tuples are ordered by closeness to the query state and supplied as context.Source: §4.2 · p. 4, Figure 3 · p. 5
  • REGENT interpolates between the nearest retrieved action and the transformer policy using a distance-dependent exponential weight.Near a retrieved state, the policy favors the retrieved action; farther away, the parametric policy receives more weight.Source: §4.2 · pp. 4–5
  • REGENT is evaluated on JAT/Gato-derived robotics and game environments and ProcGen, with separate encoders, distance metrics, and action heads for different modalities.The JAT/Gato setting uses ResNet18 image embeddings or proprioceptive vectors with ℓ2 distance, while ProcGen uses raw images with SSIM distance.Source: §4.2 · pp. 5–6
  • At deployment, REGENT retrieves n = 19 nearby states and predicts the query action directly from their context without finetuning.The same retrieval-and-prediction procedure is used in both evaluation settings.Source: §4.2 · p. 6
  • The model is restricted to unseen environments within the same training suites, while game evaluations add sticky stochasticity absent from demonstrations.The study leaves broader generalization beyond those suites for future work.Source: §4.2 · pp. 6–7

5 SUB-OPTIMALITY BOUND FOR REGENT POLICIES

This section bounds REGENT’s sub-optimality relative to the expert policy that generated retrieval demonstrations. The analysis uses the most isolated state to connect demonstration coverage with the policy’s sub-optimality gap.

5 SUB-OPTIMALITY BOUND FOR REGENT POLICIES

  • The bound targets the discrete-action case and measures REGENT’s sub-optimality against the expert policy that generated the retrieval demonstrations.The continuous-action case is left for future work.Source: §5 · p. 7
  • The analysis defines the most isolated state as the state whose closest retrieval demonstration is farthest away.This distance upper-bounds every state’s distance to its closest retrieved state and measures demonstration coverage.Source: §5 · p. 7
  • The sub-optimality gap decreases as demonstrations increase because the distance to the most isolated state decreases.The paper identifies this as the theorem’s main consequence and observes it in the results.Source: §5 · p. 7

6 EXPERIMENTAL EVALUATION

Across unseen Metaworld, Atari, and ProcGen environments, REGENT and the retrieval-only R&P baseline generalize from demonstrations, with REGENT consistently improving on R&P and outperforming larger or finetuned baselines. REGENT also preserves strong performance on training environments, while interpolation, context design, and long-horizon or embodiment changes define important boundaries.

Generalization to Unseen Environments

  • REGENT and R&P generalize across unseen Atari, ProcGen, and Metaworld environments, while REGENT consistently improves on R&P and generally benefits from more demonstrations.The settings span image observations with discrete actions and vector observations with continuous actions.Source: §6 · p. 8
  • REGENT outperforms JAT/Gato and its 5–10x larger-data variants in unseen environments, including variants using retrieval at inference time.JAT/Gato fails to generalize to most unseen environments, whereas retrieval-augmented pretraining provides an advantage beyond inference-time RAG alone.Source: §6 · pp. 7–9
  • REGENT and R&P outperform MTT on unseen ProcGen despite MTT using an order of magnitude more training levels and data and approximately 3x more parameters.The comparison uses the best MTT results reported in [37].Source: §6 · pp. 7–8, Figure 5 · p. 9

Ablations and Discussions

  • Removing interpolation between R&P and the transformer output prevents REGENT from generalizing to unseen environments, making interpolation a key design component.Without interpolation, REGENT performs like a random agent on unseen environments.Source: §6 · p. 8

Effect of Finetuning

  • REGENT further improves with finetuning on few demonstrations, while full or parameter-efficient finetuning of JAT/Gato and training from scratch generally fail to achieve comparable unseen-environment performance.In Atari, the comparison varies demonstration states up to 25k, whereas the multi-game decision transformer [28] requires 1M transitions.Source: §6 · pp. 7, 9

Generalization to Variations of the Training Environments

  • In unseen levels of training ProcGen environments, REGENT performs best and can approach its seen-level upper bound with many demonstrations through retrieval augmentation and in-context learning.The plotted subset covers 4 of 12 training environments; the remaining 8 are reported in Figure 19.Source: §6 · p. 9, Figure 6 · p. 9

Aggregate Performance on Training Environments

  • On training environments, REGENT significantly exceeds JAT/Gato on Metaworld, exceeds it on Atari, matches it on Mujoco, and nearly matches it on BabyAI.Figure 7 aggregates normalized returns across 45 Metaworld, 9 Mujoco, 52 Atari, and 39 BabyAI training environments.Source: §6 · p. 9, Figure 7 · p. 10

Ablations and Discussions

  • A larger context improves performance before saturating, context order matters especially for Atari-like distribution shifts, and ℓ2 distance outperforms cosine distance.These findings come from ablations on context length, context order, and retrieval distance metric.Source: §6 · p. 10

7 CONCLUSIONS AND FUTURE WORK

The paper concludes that retrieval enables rapid adaptation without large models or vast datasets. REGENT combines retrieval augmentation with in-context learning to outperform larger generalist agents, while long-horizon environments and new embodiments remain limitations.

Conclusions

  • Retrieval provides a powerful bias for rapid adaptation, and even the 1-nearest-neighbor R&P agent matches or exceeds state-of-the-art generalist agents in unseen environments.This result motivates REGENT’s semi-parametric design.Source: §7 · p. 10
  • REGENT pre-trains a transformer on query states, rewards, and retrieved context, enabling deployment in unseen environments with only a few demonstrations.The approach uses retrieval augmentation and in-context learning without requiring large models or vast datasets.Source: §7 · p. 10
  • Despite fewer parameters and an order of magnitude fewer pre-training datapoints, REGENT outperforms other generalist agents across unseen environments and after their finetuning.REGENT itself also improves after finetuning on a small number of demonstrations.Source: §7 · p. 10

Future Work

  • REGENT’s main limitations are adaptation to new embodiments and long-horizon environments, motivating more diverse embodiment data and improved retrieval from longer demonstrations.These directions are proposed as future work.Source: §7 · p. 10

A ADDITIONAL DETAILS FOR REGENT AND VARIOUS BASELINES

The additional details specify REGENT’s retrieval, architecture, training, baseline, and evaluation procedures, while documenting comparisons with competing methods and key limitations.

Architecture

  • REGENT uses a 12-layer, 12-head transformer with hidden size 768, 19 retrieved states, and 1 query state in a maximum 40-position context.Its decoders support continuous actions up to size 513 and discrete actions up to 18 classes.Source: A ADDITIONAL DETAILS FOR REGENT AND… · p. 16

Training

  • The training procedure uses early stopping after one epoch because overtraining reduces in-context learning performance.This applies in both JAT/Gato and ProcGen settings, with AdamW optimization and decaying learning rates.Source: A ADDITIONAL DETAILS FOR REGENT AND… · p. 16

Scope

  • The study evaluates unseen environments within the same suites as training environments, leaving generalization to entirely new suites for future work.Held-out environments include Metaworld, Atari, Mujoco, and ProcGen tasks selected within the evaluated suites.Source: A ADDITIONAL DETAILS FOR REGENT AND… · p. 17

Baselines

  • REGENT significantly outperforms JAT/Gato plus retrieval augmentation, showing benefits from retrieval-augmented pretraining rather than inference-time retrieval alone.The retrieval-augmented baselines nevertheless outperform vanilla and fully finetuned JAT/Gato in unseen Metaworld environments.Source: A ADDITIONAL DETAILS FOR REGENT AND… · p. 17
  • REGENT outperforms DRIL in all three unseen Atari environments even though DRIL includes an online reinforcement-learning stage and REGENT is fully offline.REGENT and the other comparison methods never interact with environments during training.Source: A ADDITIONAL DETAILS FOR REGENT AND… · p. 18

Comparisons

  • REGENT significantly outperforms JAT/Gato and JAT/Gato Finetuned despite using 1.5x fewer parameters and 5–10x less data, while outperforming MTT with a 3x smaller model and an order-of-magnitude less pretraining data.The MTT comparison reports the best values from its paper because its weights are closed-source.Source: A ADDITIONAL DETAILS FOR REGENT AND… · p. 18

B PROOF OF THEOREM 5.2

The proof bounds REGENT’s policy variation and connects imitation-learning risk to a sub-optimality guarantee. The resulting gap decreases as more demonstrations reduce retrieval distances.

Proof

  • For fixed state, reward, and retrieved context, the R&P component is unchanged across transformer parameters, so parameter-dependent policy differences arise from the learned component.The proof interprets the remaining inequality through Softmax behavior and distance to the most isolated state.Source: B PROOF OF THEOREM 5.2 · pp. 18–19
  • The proof applies an imitation-learning bound that converts population total-variation risk into a performance gap bounded by min{H, H^2ϵ}.The bound is combined with the total-variation lemma for REGENT’s policy class.Source: B PROOF OF THEOREM 5.2 · p. 19

Theorem consequence

  • The theorem’s main consequence is that REGENT’s sub-optimality gap reduces with more demonstrations because they reduce the distance to the most isolated state.This conclusion is also observed empirically in the paper’s results.Source: B PROOF OF THEOREM 5.2 · p. 19

C ADDITIONAL JAT/GA T O RESULTS

Additional JAT/Gato-setting analyses examine context size, retrieval ordering, runtime, distance metrics, unseen Mujoco adaptation, and qualitative behavior across environments.

Context size

  • 20 states in context produces REGENT’s best overall performance, while larger contexts generally help but often saturate quickly.The context includes the query state, so n states correspond to a context size of 2n − 1.Source: C ADDITIONAL JAT/GA T O RESULTS · p. 20, Figure 9 · p. 20

Context ordering

  • The current ordering, placing retrieved tuples from closest to furthest, performs best overall and significantly outperforms random and reverse orderings on unseen Atari.Metaworld results are sometimes similar across orderings because retrieved states can be highly alike.Source: C ADDITIONAL JAT/GA T O RESULTS · pp. 20–21, Figure 10 · p. 20

Inference runtime

Distance metric

  • L2 distance significantly outperforms cosine distance in most unseen Metaworld environments and in the more challenging unseen Atari environments.The authors speculate that L2 preserves observation magnitude information that cosine distance discards.Source: C ADDITIONAL JAT/GA T O RESULTS · p. 21, Figure 12 · p. 21

Unseen Mujoco

  • After finetuning on 25 demonstrations in two unseen Mujoco tasks, only REGENT significantly improves beyond other methods and continues improving with more demonstrations.R&P remains a strong baseline even when generalizing to new embodiments.Source: C ADDITIONAL JAT/GA T O RESULTS · p. 22, Figure 13 · p. 22

Qualitative examples

  • Qualitative metaworld-bin-picking examples show REGENT predicting actions similar to, but not identical with, the first retrieved R&P action.The paper attributes these differences to in-context learning from pretraining data and associates them with better performance.Source: C ADDITIONAL JAT/GA T O RESULTS · p. 22, Figure 14 · p. 22

Additional evaluations

  • Figures 15–18 report normalized returns across 45 Metaworld, 9 Mujoco, 52 Atari, and 45 BabyAI environments, averaged over the stated rollout counts and different seeds.The Metaworld values are provided in Table 3.Source: Figure 15 · p. 23, Figure 16 · p. 23, Figure 17 · p. 23

D ADDITIONAL PROCGEN RESULTS

The ProcGen-related material documents additional evaluation protocols and situates REGENT among retrieval-augmented and in-context learning methods. It emphasizes broad generalization while noting difficulty on new Mujoco embodiments.

D ADDITIONAL PROCGEN RESULTS

  • The ProcGen evaluation reports normalized returns on unseen levels across all 12 training environments, averaging over 10 levels and 5 rollouts per level.Table 8 provides the corresponding pre-normalization values.Source: Figure 19 · p. 28, Table 8 · p. 29

E ADDITIONAL RELATED WORK

  • REGENT differs from RAEA by pretraining a generalist policy that adapts without finetuning to new environments with different observations, dynamics, and rewards.RAEA instead performs behavior cloning in each new environment and shows only initial signs of new-task generalization after training on demonstrations.Source: E ADDITIONAL RELATED WORK · p. 30
  • REGENT is an in-context imitation-learning method that generalizes to unseen Metaworld, Atari, and ProcGen environments through retrieval augmentation and in-context learning.The paper contrasts this with GLAs’ minor improvements over random agents and RADT’s inability to generalize in several listed environments.Source: E ADDITIONAL RELATED WORK · p. 30
  • The paper notes that REGENT still struggles to generalize to new Mujoco embodiments, despite its stronger generalization in the listed Metaworld, Atari, and ProcGen settings.This limitation is shared with GLAs in the comparison discussed.Source: E ADDITIONAL RELATED WORK · p. 30
Loading 2412.04759v2…