REGENT: A Retrieval-Augmented Generalist Agent That Can Act In-Context in New Environments
Authors: Kaustubh Sridhar
Souradeep Dutta
Dinesh Jayaraman
Insup Lee
Venue: ICLR 2025 Oral (2025)
Keywords: Generalist Agent
Retrieval
In-Context Learning
VLA
Imitation Learning
Reinforcement Learning
Primary area: reinforcement learning
TL;DR: We propose a retrieval-augmented generalist agent that can adapt to new environments via in-context 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.

Published as a conference paper at ICLR 2025
REGENT: A RETRIEVAL-AUGMENTER GENERALIST AGENT THAT CAN ACT IN-CONTEXT IN NEW ENVI- RONMENTS
Kaustubh Sridhar', Souradeep Dutta Dinesh Jayaraman', Insup Lee University of Pennsylvania, UUnivesity of British Columbia ceridhar@seee upenn. edu
ABSTRACT
Building generalist agents that can rapidly adapt to new environments 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 relatively small datasets and adapt them to unseen environments via in-context learning, without any fine- tuning. Our key idea is that retrieval offers powerful bias for fast adaptation. Indeed, we demonstrate that even simple retrieval- bbaeed I-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-parametryy agent, REGENT, that trains a wrasformer-basee policy on sequences of queries and retrieved neigh- bors. 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 onder-of-mmagiuus fewer pre-training datapoints, significantly outperforming today's 's statc-of-the-an generalist agents. Website: wtt./kuuuuuuuuuauuuaa ubshsrchhhhhhe gi github. io/ regent resgercceeeeechh
INTRODUCTION
Al agents, both in the digital [38, 19, 37, 28, 53] and real world [5, 7, 63, 33, 48, 24], constantly face changing environments that require rapid or even instantaneous adaptation. True generalist agents must not only be capable of performing well on large numbers training environments, but arguably more importantly, they must be capable of adapting rapidly to new environments. While this goal has been of considerable interest to the reinforcement learning research community, it has proven elusive. The most promising results so far have all been attributed to large policies [38, 19, 37, 28, 5], pre-trained on large datasets across many environments, and even these models still struggle to generalize to unseen environments without many new environment-sppee sppeciffc demonstrations.
In this work, we take different approach to the problem of constructing such generalist agents. We start by asking: Is scaling current agent architectures the most effective way to build generalist agents? Observing that retrieval offers a powerful bias for fast adaptation, we first evaluate a simple --nearest neighbor method: "Retrieve and Play (R&P)". To determine the action the current state, R&P simply retrieves the closest state from few demonstrations in the target environment and plays its corresponding action. Tested on wide range of environments, both robotics and game-playing, R&P performs on-par or better than the saate-of-the-arl generalist agents. Note that these results involve no pre-training environments, and not even neural network policy: it is clear that larger model and pre-training dataset sizes are not the only roadblock to developing generalist agents.
Having thus established the utility of retrieval for fast adaptation of sequential decision making agents, Agent" (REGENT). REGENT architecture: pre-trains a transformer policy whose inputs are not only the
we proceed to incorporate it into the design of Snetrievll-Augmenn a semi-parametio current state and previous reward, but also retrieved tuples of (state, previous reward, action) from a set of demonstrations for each pre-training task, drawing inspiration from the recent successes of retrieval augmentation in language modeling [20]. At each "query" state, REGENT is trained to prescribe an action through aggregating the action predictions of R&P and the transformer policy. By exploiting retrieval-augmentati as well in-context learning, REGENT permits direct deployment

---

Published as a conference paper at ICLR 2025
45 MetaWorid Eris
62 Alari Enws
Melaworld
bin-picking
ailen
docr-lock
Mujoco Envs
docr-unlock
BabrAN Em
fand hand-insert -ineert
stargunner
Mujcco Ervs
hattcheetah
hopper
Retrieval-augmenner training on datasets from many envs
testing on unseen envs
Figure I: Problem setting in JAT/Scto envirvaments.
entirely unseen environments and tasks with only few demonstrations. REGENT only one of two models developed so far that can adapt to new environments via in-context learning: the other model the multi-trajectory transformer (MTT) [37].
We train and evaluate REGENT two problem settings in this paper, shown in Figures 1 and 2. The first setting is based on the environments used in Ga to [38] (and its open source reproduction JAT [19]) and the second setting is based on the ProcGen environments used in MTT [37].
In both settings, REGENT demonstrates significant generalization to unseen environments without any finetuning. In the JAT/Gato setting, REGENT outperforms JAT/Gato even when the baseline finetuned on demonstrations from the unseen environments. In the ProeGen setting, REGENT significantly surpasses MTT. Moreover, in both settings, REGENT trains a smaller model with 1.4x to 3x fewer parameters and with an order-off mrder-of-magniuude fewer pre-training datapoints. Finally, while REGENT's design choices aimed generalization, its gains not limited unseen environments: even performs better than baselines when deployed within the pre-training environments.
RELATED WORK
Recent work the reinforcement learning community has been aimed at building foundation models and multi-task generalist agents [38, 19, 5, 28, 7, 63, 33, 48, 24, 10, 22, 3, 17, 54, 40, 41, 55, 60).
Many existing generalist agents struggle to adapt to new environments. Gate [38], a popular generalist agent trained on variety of gameplay and robotic environments, struggles to achieve transfer to an unseen Atari game even after fine-tuning, irrespective of the pretraining data. The authors attribute this difficulty the "pronounced differences in the visuals, controls, and strategy" among Atari games. They also attribute Gat.o's o's lack of in-context adaptation to the limited context length of the transformer not allowing for adequate data to be added in the context. Our method sidesteps this issue by retrieving only limited but relevant parts of demonstration trajectories to include in the context. JAT [19], an open-source version of Gat o, faces similar problems. We compare with and significantly outperform JAT/Gato using fewer parameters and an order-of -itoooff--annuudeee fewer pre-training datapoints. While REGENT is 138.61. 28.6M parameter model, JAT uses 192.7M 7M parameters. The closed-soured Ga to, with similar performance the open-source JAT, reports using 1.2B parameters. JAT is also pre-trained on up to 5-10x the amount data used by our method and yet cannot generalize to unseen environments. Even after finetuning a few demonstrations from an unseen enviromment, JAT fails to meaningfully improve.
Many recent generalist agents cannot leverage in-context learning. In-context learning capabil- ities enable easier and faster adaptation compared to finetuning. Robocat [5], which builds on the Gato model, undergoes many cycles of fine-tuning, data collection, and pre- traitrain from scratch to adapt to new manipulation tasks. The multi salli--game& decision transformer [28], an agent trained on over 50 million Atari gameplay transitions, requires another I million transitions for fine-tuning
2

---

Published as a conference paper at ICLR 2025
Environment
cavetlyer bigfish
bosefight
chaieer coinrun dodgeizal fruitbot
helst
Ieaper
maze
miher
starpliot
climber
Jumper
nirja
plunder
LE!
Level
Retrieval- -aggeeetee training on datasets from many envs/levels
Retrieval-augmentee auggentedd testing on unseen envs
Retrieval-augmented testing on unseen levels in seen envs
Figure 2: Problem setting in ProeGen environments adapted from [37).
held-out game which is not practical in real robot settings. We, on the other hand, show that REGENT (and even P) adapt to new Atari games with as little 10k transitions and finetuning on said transitions. Finally, the RT series of robotics models [7, 63, 33]. recent Vision- Litoon-Langag-AAttiii models like Octo, Open VLA, Mobility VLA [48, 24, 10]. and other generalist agents like BAKU, RoboAgent [22, 3] do not evaluate or are demonstrated to not possess in-context learning capabilities. REGENT on the other hand can adapt simply with in-context learning unseen environments.
Agents that can adapt to new tasks via in-context learning, only do so within the same envi- ronment. Algorithm Distillation [26), Decision Pretrained Transformer [27]. and Prompt Decision Transformer [52] are three in-context reinforcement learning methods proposed to generalize to new goals and tasks within the same environment. None of these can handle changes in visual observations, available controls, and game dynamics.
We also compare with and outperform MTT [37], the only other model that can adapt in-context in the ProcGen setting, with improved data efficiency and a smaller model size. MTT trains transformer on sequences of trajectories from a particular level and environment and adapts to unseen environments by throwing the demonstrations into its context. The ProcGen variant of REGENT use an order-of-matne mfggnniuudd fewer transitions in pre- pre-training and is about one-third the size of MTT.
Retrieval-augmenter generation for training and deployment is language models trained with sttrival--augmetaale such as the original RAG model [29], RETRO [4], and REALM [21] have demonstrated performance on par with vanilla language models with significantly fewer parameters. Moreover,
part of policy. Various
generation [20] with large language models has enabled them to quickly adapt new or up-to-date data. We hope that our work can enable similar capabilities for ciccision-making agents.
3 PROBLEM FORMULATION
We aim to pre-train generalist agent on datasets obtained from different environments, with the goal of generalizing to new unseen environments. The agent has access to a few expert demonstrations in these new environments. In this work, the agent achieves this through in-context learning without any additional finetuning. We assume that the state and action spaces of unseen environments are known.
We model each environment i as a Markov Decision Process (MDP). The i-th Markov Decision Process (MDP) M, is a tuple (S,, P., R,, I.). where S; is the set of states, A; is the set of actions, Ps(s'(s, [s, a) the probability of transitioning from state s to s' when taking action a, R, (8, a) is the reward obtained in state upon taking action a, 7 E [0, 1) is the discount factor, and I, is the initial state distribution. We assume that the MDP's operate over trajectories with finite length H4, in an episodic fashion. Given policy acting on M,, the expected cumulative reward accrued over the duration of episode (i.e., expected return) is given by J(Te) E,, [D R2(st,
3

---

Published as a conference paper at ICLR 2025
We denote the expert demonstration dataset corresponding to the i-th (training or seen) environment consisting of tuples of (state, previous-reward. action) as D. = {ss, O, ag), ($1, To, di)...., (s,,, 7N. xx,)). )}. Let assume that have access to K such training environments for i from 1 through K and unseen environments indexed from K + 1 through expert policy oi which
M. We assume that the datasets D; 's generated by acting according maximizes J (T.). The training environments (i S. K) have a sizeable dataset of expert demonstra- tions, while the unseen environments have only handful of them, i.e [D:] >> [Dy], where K.
REGENT: A RETRIEVAL-AUGMENT
GENERALIST AGENT
Simple nearest neighbor retrieval approaches have long history in few-shot learning [50, 2, 49, 43, 9]. These works have found that, at small training dataset sizes, while parametric models might struggle extract any signal without extensive architecture approaches perform about as well as the data can support. Motivated by these prior results in other domains, we first construct such an approach for an agent that can learn directly in an unseen env- ironment with limited expert demonstrations in Section 4.1. Then, we consider how to improve this agent through access to experience in pre-training environments, so that can transfer some knowled- ge to novel environments that allows it to adapt even more effectively with limited data in Section 4.2.
hyperparameter tuning, nearest neighbor
4.1 RETRIEVE AND PLAY (R&P)
This is arguably one of the simplest decision agents that leverages the retrieval toolset for adaptation. Given state St from environment j, let assume that it is possible to retrieve the n-nearest states (and their corresponding previous rewards, actions) from Dj. We refer to this as the context et E Cj- The set Cj is the set of all such contexts in environment j. We also call the state 8, the query state following terminology from retrieval-augmennte generation for language modeling [20].
The R&P agent takes the state 8, and context e, input, picks the nearest retrieved state s' in ct. and plays the corresponding action a'. That is, TR&e (&t, Gi) = a'. We describe the retrieval process in detail later Section 4.2. Clearly, R&P devoid of any learning components which can transfer capabilities from pre-training to unseen environments.
4.2 SIIIIVALL-AGGENTTEE GENERALIST AGENT (REGENT)
To go beyond R&P, we posit that an agent learns to meaningfully combine relevant context to act in a set of training environments, then this skill should be transferable novel environments as well. We propose exactly such an agent in REGENT. We provide an overview of REGENT in Figure 3 where we depict the reinforcement learning loop with the retrieval mechanism and the REGENT transformer. In the figure, we also include the retrieved context and query inputs to the transformer and its output interpolation with the R&P action. We describe REGENT in detail below.
REGENT consists of a deep neural network policy Tg. which takes as input the state s,, previous reward 1t-1, context Ct, and outputs the action directly for continuous environments and the logits over the actions in discrete environments. In the context ct, the retrieved tuples of (state, previous reward, action) are placed in order of their closesness to the query state s, with the closest retrieved state s' placed first. Let d(st, s') be the distance between and s'. We perform a distance weighted interpolation between the outputs of the neural network policy and R&P as follows.
TRRGENT (St, 1+) = Ad(se.s)) TRaE (St, 6x) + (1 e Ad(se.s)) (TT()) 11 ct)
(1)
if action space is discrete (xx) = { MixedReLU(x)), if actiion space is continuous
[Softmax(1),
(2)
where MixedReLU R + [_ is tanh-like activation function from [46] and is detailed in Appendix A. Further, L E IR a hyperparameter for scaling the output of the neural network for continuous action spaces after the MixedReLU. We simply set both L and À to 10 everywhere following a similar choice in [46]. The function Te is a causal transformer, which is adept at modeling relatively long sequences of contextual information and predicting the optimal action [29, 4, 8, 28). All distances are normalized and clipped to [0, 1] as detailed in Appendix A. For discrete action spaces, where the transformer outputs distribution over actions, we softened version of Ma.
4

---

Published as a conference paper at ICLR 2025
a'
dist weighted interpolation
(s
dist (s,.*) weighted interpolation
weighted interpolation
linear
linear
linoar
transformer
r r-11
(inear)ockup
linearlookup
liner/lookup
linee/lookup
linerrlookup
rs)
a
(s rr
a
Spp
1-1
St
Image Encoder
S
raw
Retrieved 1,
Retrieval demos
Figure 3: The RECENT architecture and overview, (1) query state (from the unseon environment during deployment or from training envirooments"
during pre-training) processed retrivva. The nearest states few demonstrations envireument from designated retrieval subset of pre training enviromments' dalusets are retrieved. These states, and their corresponding previous rewards actions, added
context order reward. (3) The predictions from the REGENT tranformer are combined with the first retrieved action. (4) At deployment, only the predicted query action is used. During pro-training. the loss from prodicting actions is used to train the transformer.
closeness
query state, followed by the query state previous
described in Appendix A. As shown in Figure 3, for predicting each action, either for retrieved state in the context or for the query state, we set d(s,, s') in Equations (1) and (3) to the distance between that state s and the first (closest) retrieved state in the context s'.
Further, the exponential weights in Equation (1) allows us to smoothly transition between R&P and the parametric policy Ty. When the state 8; is close enough to first retrieved state in the context ct TREGENT simply plays the retrieved action. However, as it moves further away, policy Taae becomes uniform distribution and the parametric policy T. takes more precedence. We also hypothesize that this interpolation allows the transformer to more reudily generalize to unseen environments, since given the easier task of predicting the residual to the R&P action rather than predicting the complete action.
Pre-training REGENT and Loss Function: We train REGENT by minimizing the total cross-entropy loss discrete actions and total mean-squared error continuous actions for all predictions (n in the context and 1 query). We also follow the JAT recipe in ensuring that each training batch consists only of data from a single environment's dataset. We provide details about the training hyperparameters in Appendix A.
action
REGENT Training Data and Environment Suites: In the JAT/Gato setting, pre-train 100k transitions in each of the 45 Metaworld training environments, 9 Mujoco training environments, 52 Atari training environments, and 39 Baby AI training environments. This adds up total of 14 14.5M transitions used to pre-train REGENT. We obtain these transitions by taking a subset of the open-source JAT dataset [19]. The complete JAT dataset consists of 5-10x the amount of data we use in each environment. The Atari environments have discrete actions and four consecutive grayscale images as states. However, we embed each combined four-stacked image into vector of 512 dimensions and use this image embedding everywhere (for R&P and REGENT). The details of embedding images are discussed later in this Section. The Metaworld and Mujoco environments have proprioceptive vector observations and continuous action spaces. The Baby AI environments have discrete observations, text specifying a goal in natural language, and discrete actions. We emphasize that the same generalist policy handles completely different observation and action modalities and dimensions across environments. The dimensions of observations & actions are given in Appendix A.
In the ProcGen setting that borrows from MTT [37], generate 20 rollouts with PPO policies each the 12 training environments for various levels depending on the environment. We do so to ensure that the total number of transitions in each environment 1M. The total size of the pre-training
5

---

Published as a conference paper at ICLR 2025
dataset is 12M. All environments have an image-based state space and discrete actions. The number of levels varies from 63 at the smallest in bigfish to 1565 at the largest in maze. This variation arises from the difference in rollout horizon in each environment. We note that, even the largest value of 1565 is an order-of-magnitude fewer than the 10k levels in each game used by MTT for pre-training.
REGENT Architecture: REGENT adapts the JAT architecture in the JAT/Gato setting. It consists of a causal transformer trunk. has a shared linear encoder for image embeddings, vector observations, and continuous actions. This encoder takes maximum input vector of length 513, accommodating the largest input vectors-image embeddings of size 512-plus a single reward value. When the length of an imput vector is less than that, it is cyclically repeated and padded to the maximum length. It has a large shared lookup table with the GPT2 [35] vocabulary size for encoding discrete actions, discrete observations, and text observations. Another linear layer is used to combine multiple discrete values and text tokens present in a single observation vector. It has a shared linear head for predicting continuous actions. A shared linear head is used for predicting distributions over discrete actions. All of the above components are shared across environments with but only a subset of input encoders and output heads may be triggered during a forward and backward pass depending on the input modalities in an environment and output action space. We note that states and rewards are concatenated together before encoding following the JAT recipe [19]. REGENT has a total of 138.6M parameters, including the frozen ResNet18 image encoder, compared to JAT'S 192.7M parameters.
We simplify REGENT for the ProcGen setting keeping the transformer trunk with only a convolution encoder for direct image inputs; lookup table for encoding discrete actions; and a linear head for predicting distributions over discrete actions. Here, following the MTT recipe, we do not use any rewards. We only use the states and actions. Here, REGENT has a total of 166M parameters compared to MTT's 310M parameters. We detail all architectural hyperparameters and the differences between JAT, MTT, and REGENT in Appendix A.
Processing Raw Observations, Distance Metrics, Retrieval Mechanism, and Preprocessing Training Data: Recall that we are evaluating in two settings, depicted in Figures and 2. In the JAT/Gato setting, if the raw observations consist of images, we embed the image with an off-the- shelf R encoder [13] trained only on ImageNet data. If instead, the raw observations are proprioceptive vectors, we use them directly without modification. In the ProcGen setting, the raw observations consist of images that we use directly without modification.
In the JAT/Gato setting, we use the la distance metric to compute distances between pairs of observations, regardless of whether they are image embeddings or proprioceptive vectors. To significantly speed up the search for nearest states to a query state, we leverage similarity search indices [15]. In the ProcGen setting, we utilize the SSIM distance metric [51] to obtain distances between two images. We perform naive search for the nearest images a query image, however all SSIM distance calculations are parallelized on GPUs to speed up the retrieval process.
R&P simply performs the retrieval process described above at evaluation time, obtains the closest state to a query state, and plays the corresponding action. REGENT on the other hand has to setup its pre-training dataset as follows. It first designates certain number of randomly chosen demonstrations per environment as the retrieval set in that environment. This is described in Appendix A. Then, for each state in each environment's 's training dataset, we retrieve the n 19 closest states from the designated retrieval subset for that dataset. We ensure that none of the retrieved states are from the same demonstration as the query state. In this process, we convert our dataset of transitions to a dataset of (context, query state, query reward) datapoints where the context consists of 19 retrieved (state, reward, action) tuples. Now, we can begin pre-training REGENT on this dataset.
After pre-training, in each of the two settings, we have one model that can be deployed directly, without finetuning, on unseen environments. During deployment an environment, at each step, we retrieve the 19 closest states to the current query state (and their corresponding actions and rewards). We pass this context with the query state, and previous reward into Equation (1) through the architecture described above to predict the query action to take in the environment.
Evaluating REGENT:
In the JAT/Gato setting, for the unseen environments, we hold-out 5 Metaworld, 5 Atari, and two Mujoco environments (see Figure 1). In the ProcGen setting, following MTT [37], we hold-out 5 environments. Unlike MTT, we also evaluate on unseen levels in training environments (see Figure 2). We explain these choices in Appendix A. In this work, focus unseen environments in the same suite as training cnvironments and leave further generalization to future work. Finally, we also

---

Published as a conference paper at ICLR 2025
note that in all game-playine environments, add small sticky probability [31] of 0.05 in unseen Atari environments and 0.2 in unseen ProcGen environments following [37]. This is not present in any data demonstration, which induces further stochasticity and tests the ability of both R&P and REGENT to truly generalize under novel and stochastic dynamics.
SUB-OPTIMALITY BOUND FOR REGENT POLICIES
In this section, aim bound the sub-optimality of the REGENT policy. This is measured with respect to the expert policy т.. that generated the retrieval demonstrations Dj. We focus on the discrete action case here and leave the continuous action case for future work. The sub-optimality gap in (training or unseen) environment is given by (J(x7) ((TTEEETT) ). Inspired by the theoretical analysis of Sridhar et al. [46], we define the "most isolated state" and use this definition to bound total variation in the REGENT policy class and hence the sub-optimality gap.
That is, first, given Dj, we wish obtain the maximum value of the distance term d(st, in Equations (1) and (3). To do we define the most isolated state as follows.
Definition 5.1 (Most Isolated State). For a given set of retrieval demonstrations Dj in environment j.
we define the most isolated state DD; == arg max min ((s,s))). and consequently the distance to sES, s'ED, the most isolated state as d2. min d(sb, s'ED;
All distances between state in this enviromment and its closest retrieved state are less than the above value, which also measures state space coverage by the demonstrations available for retrieval. Using above definition, we have the following.
Theorem
sub-optimality ((Tz) ( ( < min/H, 11 bdb;)>
5.2. The
gap
in environment
j
is
Proof: We provide the proof in Appendix B.
The main consequence of this theorem, also observed in our results, is that the sub-optimality gap reduces with more demonstrations in D, because of the reduced distance to the most isolated state.
6 EXPERIMENTAL EVALUATION
In experiments, we aim to answer the following key questions in the settings depicted in Figures and 2. (1) How well can R&P and REGENT generalize to unseen environments? (2) How does finetuning in the new environments improve REGENT? (3) How well can REGENT generalize variations of the training environments and perform in aggregate on training environments? (4) How does REGENT qualitatively compare with R&P using the retrieved context?
Metrics:
We plot the normalized return computed using the return of a random and expert agent (total cettum-randdd retum)
in each environment (expert -random retur). The expert and random returns for all JAT/Ga environments are made available in the original work [19]. The expert and random returns for the ProcGen environments can be found in the original ProcGen paper [11].
We consider different set of baselines for each of the two settings. In the JAT/Gato setting, we compare R&P and REGENT with two variants of JAT. The first is a JAT model trained on the same dataset as REGENT. The second is JAT model trained on all available JAT data which consists of an order of magnitude more datapoints (in particular, 5x more data in Atari environments and 10x more data all other environments). We label the former apples-to-apples comparison as JAT/Gato and the latter as JAT/Gato (A11 Data). . We also compare with JAT/Gato with RAG at inference time. We use the same retrieval mechanism as REGENT for this baseline. Additional details and all hyperparameters for the baselines are in Appendix A.
Baselines:
In the ProcGen setting, we compare with MTT [37]. MTT only reports results on unseen environments and not unseen levels in training environments. We simply take the best MTT result each of the four unseen environments [37], obtained when 4 demonstrations are included the MTT context.
Finetuning and Train-from-scratee Baselines: We fully finetune and parameter-efficient finetune (PEFT with IA3 [30]) both JAT/Gato and JAT/Gato (A11 Data) on various number demo- nstrations in each unseen environment in the JAT/Gato setting and compare with REGENT. We also
7

---

Published as a conference paper at ICLR 2025
metaworld- -iniid--iiiiiii picking
metaworld-box box-clls
metaworid-door-lock door-lock
metaworld uno-ulddoodoo--ulookkkk
Nornalized return
Normalizes return
Normallzed return
0.81
0.5
0.6
0.6
0.4
0.6
Normelzzed
0.4
0.4
0.
0.2
0.0
0.0
0.0
0.0
Number demos
Number of demos
Number of demos
Number demos
metaworld-hand-ineee
atari-pong
atari-alien
atari-mspacman
Normalized return
0.8
Normalized return
0.30
0.25
Normalized return
Normelized return
0.5
0.20
0.06
0.5
0.4
0.15
0.00
0.4
0.30
0.-
0.05
0.02
0.2
0.00
0.00
25
50
75
100
10000
25000 Number of states in demes
15000
20000
10000
15000
20000
25000
20000
15000
20000
25000
Number of deross
Number of states in demos
Nurbeer of states demas
JAT/Gato + RAG
JAT/Gato
JAT/Sato Data)
REGENT
JAT/Cato Rinetuned
JAT/Gato (All Data) RAG
Train From Scratch
REGENT finetuned
JAT/Gato PEFT with IA3
JAT/Gato (All Data) Finetuned
DRIL
JAT/Geto
Figure 4: Normalized retums in the unsenn Metaworld and Atari environments against the number of demonstration trajectories the agent can on. agent evaluated across rollouts different seeds in Metaworid rolloats different seeds (with Panka 0.05) in Atari. We compute the overall mean and standard deviation over three training seeds for REGENT, REGENT Finetuned, the PEFT with baselines, and DRIL. See Table for detailed results.
retrieve from
compare with train-frome www---roommttttee behavior cloning baseline and another imitation learning baseline (DRIL [6]). For completeness, we also finetune REGENT on various number of demonstrations in each unseen environment the JAT/Gato setting. We use the same hyperparameters in each finetuning method across JAT/Gato variants and REGENT, Additional details are in Appendix A.
Generalization to Unseen Environments: We plot the normalized return obtained by all methods in unseen Metaworld and Atari environments from the JAT/Gato setting for various number of demonstrations (25, 50, 75, 100) Figure 4. For fair comparison across Atari environments with marked differences in episode horizons, we plot normalized return against various number of states in the demonstrations (atleast 10k, 15k, 20k, 25k) rather than the number of demonstrations itself. These demonstrations can be used by the different methods for retrieval normalized return obtained by all methods in unseen ProcGen environments against various number of demonstrations (2, 4, 8, 12, 16, 20) retrieve from in Figure 5. As mentioned before in Section 4.2, unseen Atari environments have a sticky probability of 0.05 and unseen ProcGen environments have a sticky probability of 0.2.
fine-tuning. We also plot the
In both Figures 4 and 5, we observe that R&P and REGENT can generalize well to unseen Atari and ProcGen environments with image observations and discrete actions well to unseen Metaworld enviromments with vector observations and continuous actions. R&P is surprisingly strong baseline. but REGENT improves on R&P consistently. In general, both methods appear to steadily improve with more demonstrations (with only few exceptions).
In Figure 4, we observe that JAT/Gato cannot generalize to most unseen environments. REGENT (and even P&P) outperform even the All Data variants of JAT/Gato which were pre-trained on 5-10x the size the REGENT dataset. The sticky probability in unseen Atari environments adds further stochasticity and stress tests true generalization against simply replaying demonstrations. The JAT/Gato RAG baselines actually improve over zero-shot JAT/Gato but are still significantly outperformed by REGENT. This demonstrates the advantage of REGENT's architecture and retrieval augmented pretraining over just performing RAG at inference time.
In Figure 5, REGENT (and even R&P) outperform MTT which was trained order of magnitude more training levels and data than RECENT while also having approximately 3x the parameters.
We also note that without the interpolation between R&P and the output of the transformer, RE- GENT does not generalize to unseen environments (i.c., it performs like a random agent on unseen environments). This interpolation is key to REGENT.
8

---

Published as a conference paper at ICLR 2025
ninja
climber
plunder
jumper
1.0
1.0
1.0
1.0
Nornalized return
Normal zod return
Normalizzd return
Normalized return
0.8
0.8
0.8
0.0
0.0
0.0
0.0
0.0
0.4
0.4
0.4
0.4
0.2
0.2
0.2
0.2
0.0
0.0
2
12
16 20
12
16 20
12
20
4
8
12
16
20
Number of demos
Number of demos
Number of dernos
Number of deros
REGENT
R&P
MTT (Number of demos=4)
Figure S: Normalized returns in unseen ProcGen envirouments against the number of demonstration trajectories the agent can retrieve from. RECENT and R&P agents evaluated across levels with rollouts and Paiky 0.2. We compute the overall and standard deviation over three training seeds for REGENT. r. The values for MIT are the best scores reported in [37]. See Table for detailed results.
bigfish
fruitbot
heist
miner
1.0
1.0
Normalized return
Normalized return
0.8
Normalized
0.7
0.2
0.0
Number demos
Number of demos
Number of demos
Number demos
REGENT
R&aP
Upper limit based on REGENT in seen levels with 20 demos
levels in 4 12) ProcGon training enviromments against the number of demonstration trajectories the agent can retrieve from. Each value is an average across 10 levels with rollouts each with Puidyy = 0.1. The other 8 plots are similar to the above and found Figure Appendix D). Table detailed results.
Figure 6: Normalized
Effect of Finetuning: From Figure 4, we see that JAT/Gato, along with its All Data variant, struggles perform in all unseen environments even after full finetuning. JAT/Gat o performs slightly better after PEFT with IA3 in the unseen metaworld environments. Training from scratch on the few demonstrations in each unseen environment also does not obtain any meaningful performance in most environments. RECENT (and even R&P), without any finetuning, outperforms all finetuned variants of JAT/Gato. Moreover, we can see that REGENT further improves after finetuning, even with only a few demonstrations.
To highlight REGENT'S ability to adapt from very little data, we note that in Figure 4, we only vary the number of states in Atari demonstrations until 25k. Whereas, the closest generalist policy that fine- tunes to new Atari environments, the multi-game decision transformer [28], requires 1M transitions.
Finally, we note that all methods face hurdle in generalizing to the very long-horizon atari- spaceinvaders and atari-stargunner environments, which have horizons about 10x that of atari-pong and hence have not been shown in Figure 4. Generalization the embodiments in the two unseen Mujoco environments proves equally challenging, we discuss this in detail in Appendix C.
Generalization to Variations of the Training Environments: In the ProcGen setting in Figure 2, the unseen levels the 12 training environments provide an avenue to test a middle ground of generalization between training and unseen environments. We plot the normalized returns on unseen levels of 4 (of 12) training environments (with a sticky probability of 0.1) in Figure 6. The remaining be found in Figure 19 (Appendix D) and similar the above 4. In Figure that REGENT performs the best in unseen levels while R&P remains strong baseline. We also depict the performance of REGENT on seen levels of these training environments with dotted line this Figure. This represents an upper bound for REGENT's performance in unseen levels. We observe that with large number of demonstrations, REGENT appears to reach close to this upper bound simply via retrieval-augmmetttioo and inddii-cooo in-context learning in environments.
again observe
Aggregate Performance on Training Environments: We plot the aggregate normalized return, both IQM [1] and mean, on training environments for each of the 4 suites (in the JAT/Gato setting) in Figure We notice that REGENT significantly exceeds JAT/Gato o on Metaworld, exceeds it on Atari, matches it on Mujoco, and is close to matching it on BabyAl. This demonstrates the dual advantage of REGENT which gains the capability to generalize to unseen environments while preserving overall multi-task performance in training environments.

---

Published as a conference paper at ICLR 2025
IQM
Mean
IQM
JAT/Galo REGENT
JAT/Galo REGENT
1.00
1.00
1.00
MelaWorkd: Aggregate Normalized Return
MujoCo: Aggregate Normal zed Return
IQM
Mean
IQM
Meen
JAT/Gato REGENT
JAT/Gatt REGENT
0.00 0.50
0.00 0.25 0.50 0.75 1.00 0.00 0.25 0.50 0.75 1.00 BabyAl: Aggregate Normalized Return
Atari (Paay 0.00): Aggregate Normalized Return
Figure 7: Aggregate normalized returns in the 45 training Metaworid enviromments. training Mujoco environments. 52 training Alari envirvaments, and training BabyAl enviromeents See Figures 15, 16, 18 Appendix C) for performance ench traininy enviromment each suite.
6
Tp-1
[0][11]-
[0] [3]
[0]
[0] [1]
Retrieved S, a
Query s. r. Pred å
Retrieved
Query 8, r, Pred à
Figure 8: Qualitative examples of inguts and outputs of RECENT for atari-pong euvironment, REGENT leverages its in-context learning capabilitiss and interpolation with R&2 to either make simple decision and predict same
rendered rollout the discrete action space.
blue on the right) predict better actions key black overall performance as seen in Figure 4,
that leads
Qualitative Examples: We plot examples of the inputs and outputs of REGENT at various states during rendered rollout in the atari-pong environment Figure highlight REGENT's learned in-context learning capabilities and interpolation with R&P. We also provide qualitative example in a continuous metaworld environment Appendix
We performed ablations on REGENT's context length, context order, and distance metric in Figures 9, 10, and 12 in Appendix C. We found that a larger context length improves performance (but saturates quickly), the order of states in the context is key especially when unseen environments significantly differ from pretraining environments (such as Atari), and that lg outperforms cosine distance. We also discuss inference runtimes in Figure 11 in Appendix C.
Ablations and Discussions:
7 CONCLUSIONS AND FUTURE WORK
In this work, demonstrated that retrieval offers generalist agents a powerful bias for rapid adaptation to new environments, even without large models and vast datasets. We showed that simple retrieval-based 1-nearest neighbor agent, Retrieve and Play (R&P), strong baseline that matches or exceeds the performance of today's state-of-the-art generalist agents in unseen environments. Building on this, we proposed a semi-parametric generalist agent, REGENT, that pre-trains transformer-base REGENT exploits retrieval-an auggrnattagmttiiiionnn and in-context learning for direct deployment in unseen environments with only few demonstrations, Even after pre-training datapoints than other generalist agents and with fewer parameters, REGENT outperforms them across unseen environments and surpasses them even after they have been finetuned on demonstrations from those unseen environments. REGENT, itself, further improves with finetuning on even a small number of demonstrations. We note that REGENT faces couple of limitations: adapting and long-horizon environments, In future work, we believe a larger diversity of embodiments in the training dataset and improved retrieval from longer demonstrations can help REGENT overcome these challenges. We conclude with the conviction that retrieval in general and REGENT in particular redefines the possibilities for developing highly adaptive and efficient generalist agents, even with limited resources.
policy on sequences of query state, reward, and retrieved context.
order of magnitude fewer
embodiments
10

---

Published as a conference paper at ICLR 2025
Acknowledgements: This work was supported in part by ARO MURI W911NF-20-1-00% NSF 2143274, NSF CAREER 2239301, NSF 2331783, and DARPA TIAMAT HR001 12490421. Any opinions, findings, conclusions recommendations expressed in this material are those of the authors and do not necessarily reflect the views the Army Research Office (ARO), the Department of Defense, the United States Government.
To ensure reproducibility, we have released all code at https: /// kaust i Lchar. github. i0/ regent research. We have also described all datasets, environment suites, basclines, and REGENT itself in detail in Section 4, Section 6, and Appendix A.
Reproducibility statement:
REFERENCES
[1] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304 29320, 2021. (Cited on
[2] M Saiful Bari, Batool Haider, and Saab Mansour. Nearest neighbour few-shot learning for cross- cross-lingual classification. arXiv preprint arXiv:2109.02221, 2021. (Cited on
[3] Homanga Bharadhwaj. Jay Vakil, Mohit Sharma, Abhinav Gupta, Shubham Tulsiani, and Vikash Kumar. Roboagent: Generalization and efficiency in robot manipulation via semantic augmentations and action chunking. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 4788-4795, IEEE, 2024. (Cited on 3)
[4] Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In International conference on machine learning, pp. 2206-2240, PMLR, 2022. (Cited 3, 4)
[5] Konstantinos Bousmalis, Giulia Vezzani, Dushyant Rao, Coline Devin, Alex X Lee, Maria Bauza, Todor Davchev, Yuxiang Zhou, Agrim Gupta, Akhil Raju, al. Robocat: A self- improving foundation agent for robotic manipulation. arXiv preprint arXiv:2306. 11706, 2023. (Cited on 2)
[6] Kiante Brantley, Wen Sun, and Mikael Henaff. Disagreement-reguulaiied imitation learning. In International Conference Learning Representations, 2019. (Cited 8. 18)
[7] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-l: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.068117 2022. (Cited on 1,
2, 3)
[8] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 15044-
2021. (Cited on 4)
[9] Wei- Yu Chen, Yen-Cheng Liu, Zsolt Kira, Yu- Yu-Chiang Frank Wang, and Jia Ban Huang. A closer look at few-shot classification. arXiv preprint arXiv: 1904.0004000 044424422, 2019. (Cited on 4)
[10] Hao-Tien Lewis Chiang, Zhuo Xu, Zipeng Fu, Mithun George Jacob, Tingnan Zhang, Tsang- Wei Edward Lee, Wenhao Yu, Connor Schenck. David Rendleman, Dhruv Shah, et al. Mobility vla: Multimodal instruction navigation with long-context vlms and topological graphs. arXiv preprint arXiv:2407.07775, 2024. (Cited on 2, 3)
[11] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, pp. PMLR, 2020. (Cited on
[12] Erwin Coumans and Yunfei Bai. Pybullet, a python module for physics simulation for games, robotics and machine learning. http://pybul let. org, 2016-2019. (Cited 17)
11

---

Published as a conference paper at ICLR 2025
[13] Sudeep Dasari, Mohan Kumar Srirama, Unnat Jain, and Abhinav Gupta. An unbiased look at datasets for visuo-motor pre-training. In Conference on Robot Learning, pp. 1183-1198.
PMLR, 2023. (Cited on 6, 16, 17)
[14] Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. Carla: An open urban driving simulator. In Conference robot learning, pp. 1-16. PMLR, 2017. (Cited on 17)
[15] Mathijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre- Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library, arXiv preprint arXiv:2401.08281, 2024. (Cited on 6)
[16] Souradeep Dutta, Kaustubh Sridhar, Osbert Bastani, Edgar Dobriban, James Weimer, Insup Lee, and Julia Parish-Morris. Exploring with sticky mittens: Reinforcement learning with expert interventions via option templates. In Conference on Robot Learning, pp. 1499-1509. PMLR, 2023. (Cited on 17)
[17] Kiana Ehsani, Tanmay Gupta, Rose Hendrix, Jordi Salvador, Luca Weihs, Kuo-Hao Zeng, Kunal Pratap Singh, Yejin Kim, Winson Han, Alvaro Herrasti, et al. Imitating shortest paths in simulation enables effective navigation and manipulation in the real world. arXiv preprint arXiv:2312.029700 2023. (Cited on 2)
[18] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures, In International conference on machine learning, pp. 1407-1416. PMLR, 2018. (Cited on 18)
[19] Quentin Gallouédec, Edward Beeching. Clément Romac, and Emmanuel Dellandréa. Jack of all trades, master of some, a multi-purpose transformer agent. arXiv preprint arXiv:2402.09844, 2024. (Cited on 1, 2, 5, 6, 7, 17)
[20] Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yi Dai, Jiawei Sun, and Haofen Wang. Retrieval -ugmeeeted generation for large language models: A survey. arXiv preprint arXiv:2312. 10997, 2023. (Cited on 1, 3, 4)
[21] Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training, In International conference on machine learning, pp. 3929-3938. PMLR, 2020. (Cited on 3)
[22] Siddhant Haldar, Zhuoran Peng, and Lerrel Pinto. Baku: An efficient transformer for multi-task policy learning. arXiv preprint (xXx:24406007339 2024. (Cited on 3)
[23] Tomoyuki Kagaya, Thong Jing Yuan, Yuxuan Lou, Jayashree Karlekar, Sugiri Pranata, Akira Kinose, Koki Oguri, Felix Wick, and Yang You. Rap: Retrieval-augmeented planning with contextual memory for multimodal Ilm agents. arXiv preprint arXiv:2402.036110 2024. (Cited on 30)
[24] Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, Grace Lam, Pannag Sanketi, al. Openvla: An open-source vision-language-action model. arXiv preprint arXiv:2406..99246 2024. (Cited on 1, 2, 3)
[25] Louis Kirsch, James Harrison, Daniel Freeman, Jascha Sohl-Dickstein, and Jürgen Schmidhuber. Towards general-purpose in-context learning agents. Workshop on Distribution Shifts, 37th Conference Neural Information ... 2023. (Cited on 30)
[26] Michael Laskin, Luyu Wang, Junbyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv: 14215, 2022. (Cited on 3)
[27] Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. Advances in Neural Information P Processing Systems, 36, 2024. (Cited on 3)
12

---

Published as a conference paper at ICLR 2025
[28] Kuang- Hueng Lee, Ofir Nachum, Mengjiao Sherry Yang. Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang. Henryk Michalewski, et al. Multi-game decision transformers. Advances in Neural Information Processing Systems, 35:27921-27936, (Cited on 2,2,4,9 9, 17)
2022.
[29] Patrick Lewis, Ethan Perez, Alcksandra Piktus. Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-angmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33:9459- 9974,, 2020. (Cited on 3)
[30] Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang. Mohit Bansal, and Colin A Raffel. Few-shot parameter-cfficient fine-tuning is better and cheaper than in-context learning. Advances in Neural Information Processing Systems, 35:1950-1968 2022. (Cited on 17)
[31] Marlos C Machado, Marc G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
(Cited on
[32] Chiara Dalla Man, Francesco Micheletto, Dayu Lv, Mare Breton, Boris Kovatchev, and Claudio Cobelli. The uva/padova type diabetes simulator: new features. Journal of diabetes science and technology, 8(1):26 34, 2014. (Cited 17)
[33] Abhishck Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint 2, 3)
2023. (Cited on 1,
[34] Jean Park, Sydney Pugh, Kaustubh Sridhar, Mengyu Liu, Navish Yarna, Ramneet Kaur, Souradeep Dutta, Elena Bernardis, Oleg Sokolsky, and Insup Lee. Automating weak label generation for data programming with climicians in the loop. In Conference on Connected Health: Applications, Systems and Engineering Technologies (CHASE), 2024. (Cited on 17)
[35] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. (Cited on 6)
[36] Nived Rajaraman, Lin F. Yang. Jiantao Jiao, and Kannan Ramchandran. Toward the fundamental limits of imitation learning. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS'20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546. (Cited on 19)
[37] Sharath Chandra Raparthy, Eric Hambro, Robert Kirk, Mikael Henaff, and Roberta Raileanu. Generalization to new sequential decision making tasks with in-context learning. arXiv preprint arXiv:2312.03801, 2023. (Cited on 1, 2, 3, 5, 6, 7, 9, 17, 28)
[38] Scott Reed, Konrad ZoIna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Marth, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent, arXiv preprint arXiv:2005.0611 2022. (Cited on 1, 2)
[39] Thomas Schmied, Fabian Paischer, Vihang Patil, Markus Hofmarcher, Razvan Pascanu, and Sepp Hochreiter. Retrieval-augmenned decision transformer: External memory for in-context arXiv preprint arXiv:2410.07071, 2024. (Cited on 30)
[40] Dhruv Shah, Ajay Sridhar, Arjun Bhorkar, Noriaki Hirose, and Sergey Levine. Gnm: A general navigation model to drive any robot. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 7226-7233. IEEE, 2023. (Cited on 2)
[41] Dhruv Shah, Ajay Sridhar. Nitish Dashora, Kyle Stachowicz, Kevin Black. Noriaki Hi- rose, and Sergey Levine. Vint: A foundation model for visual navigation. arXiv preprint arXiv:2306. /4846, 2023. (Cited on 2)
13

---

Published as a conference paper at ICLR 2025
[42] Auditya Singh, Stephanie Chan, Ted Moskovitz, Erin Grant, Andrew Saxe, and Felix Hill. The transient nature of emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 36, 2024. (Cited on 16)
[43] Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. Advances in neural information processing systems, 30, 2017. (Cited on 4)
[44] Kaustubh Sridhar and Srikant Sukumar. Finite-time, event-triggered tracking control of quadro- tors. In 5th CEAS Specialist Conference on Guidance, Navigation & Control (EurGNC 19)
Milano, Italy, 2019. (Cited on 17)
[45] Kaustubh Sridhar, Oleg Sokolsky, Insup Lee, and James Weimer. Improving neural network robustness via persistency of excitation. In 2022 American Control Conference (ACC), pp- 1521-1526. IEEE, 2022. (Cited on 17)
[46] Kaustubh Sridhar, Souradeep Dutta, Dinesh Jayaraman, James Weimer, and Insup Lee. Memory- consistent neural networks for imitation learning. arXiv preprint arXi:23110006711, 2023. (Cited
on 4, 7, 16, 17)
[47] Kaustubh Sridhar, Souradeep Dutta, James Weimer, and Insup Lee. Guaranteed conformance of neurosymbolic models to natural constraints. In Learning for Dynamics and Control Conference, pp. 76-89. PMLR, 2023. (Cited on 17)
[48] Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al. Octo: An open-source generalist robot policy. arXiv preprint arXiv:2405. 12213, 2024. (Cited on 1, 3)
[49] Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. Advances in neural information processing systems, 29, 2016. (Cited on 4)
[50] Yan Wang, Wei-Lun Chao, Kilian Q Weinberger, and Laurens Van Der Maaten. Simpleshot: Re- visiting nearest-neighbor classification for few-shot learning. arXiv preprint arXiv:1911.008 2019. (Cited on 4)
[51] Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing, 13(::: 600-612, 2004. (Cited on 6)
[52] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In international conference machine learning, pp. 24631-24645 PMLR, 2022. (Cited on 3)
[53] John Yang, Carlos E Jimenez, Alexander Wettig, Kilian Lieret, Shunyu Yao, Karthik Narasimhan, and Ofir Press. Swe-agent: Agent-computer interfaces enable automated software engineering. arXiv preprint arXiv:2405. 15793, 2024. (Cited on 1)
[54] Jonathan Yang, Dorsa Sadigh, and Chelsea Finn. Polybot: Training one policy across robots while embracing variability. arXiv preprint arXiv:2307.03719 2023. (Cited on 2)
[55] Jonathan Yang, Catherine Glossop, Arjun Bhorkar, Dhruv Shah, Quan Vuong, Chelsea Finn, Dorsa Sadigh, and Sergey Levine. Pushing the limits of cross-embodiment learning for manipu- lation and navigation. arXiv preprint arXiv: 19432, 2024. (Cited on 2)
[56) Yahan Yang, Ramneet Kaur, Souradeep Dutta, and Insup Lee. Interpretable detection of distri- bution shifts in learning enabled cyber-physich systems. In 2022 ACM/IEEE 13th International Conference
Syberr-Phssccal Systems (ICCPS), pp. 225-235. IEEE, 2022. (Cited 17)
[57] Zhuolin Yang, Wei Ping. Zihan Liu, Vijay Korthikanti, Weili Nie, De-An Huang, Linxi Fan, Zhiding Yu, Shiyi Lan, Bo Li, et al. Re-vilm: Retrieval-augmmeted visual language model for zero and few-shot image captioning. arXiv preprint arXiv:2302.044 04858, 2023. (Cited on 30)
[58] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Merd-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100 PMLR, 2020. (Cited on 17)
14

---

Published as a conference paper at ICLR 2025
[59] Andrew Zhao, Daniel Huang, Quentin Xu, Matthieu Lin, Yong-Jin Liu, and Gao Huang. Expel: Llm agents are experiential learners. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 19632-19642, 2024. (Cited 30)
[60] Haoyu Zhen, Xiaowen Qiu, Peihao Chen, Jincheng Yang, Xin Yan, Yilun Du, Yining Hong, and Chuang Gan. 3d-vla: A 3d assionllaggaage-action generative world model. arXiv preprini arXiv:2403.09631 2024. (Cited on 2)
[61] Yichen Zhu, Zhicai Ou, Xiaofeng Mou, and Jian Tang. Retrieval- Retrieval-augmenned embodied agents. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 17985-17995, 2024. (Cited on 30)
[62] Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martín, Abhishek Joshi, Soroush Nasiriany, and Yifeng Zhu. robosuite: A modular simulation framework and benchmark for robot learning. arXiv preprint arXiv:2009. 2020. (Cited on 17)
[63] Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, et al. Rt-2: lisoonlanguaggeacttonn models transfer web knowledge to robotic control. In Conference on Robot Learning, pp. 2165-2183 PMLR, 2023. (Cited on 1, 2, 3)
15

---

Published as a conference paper at ICLR 2025
APPENDIX
A ADDITIONAL DETAILS FOR REGENT AND VARIOUS BASELINES
Additional Details for REGENT
MixedReLU activation: MixedReLU is tanh-like activation function from [46] given by MiXedReLU(z)) (2(ReLU( ( ReLU(==1)) - 1) which simplifies to -1 for -1,
for 11 S x S 1, and 1 for x >
Normalizing distances to [0, 1]: We compute the 95th percentile of all distances d(s, s') between any (retrieved or query) state s and the first (closest) retrieved state s'. This value is computed from the demonstrations D and used to normalize all distances in that environment. This value is calculated for all (training or unseen) environments during the preprocessing stage. If after normalization, distance value is greater than 1, we simply clip to 1.
Softening Ta&P logits for discrete action environments: Let us assume that the discrete action space consists of Nas different actions. We have the following.
L+ D1-ddee
if - d' if + a'
TaEe (alss,)),, (4)
Nat
d(se,s") Na.
The distribution induced by the modified TasP function assigns all probability mass to the action a' when d(s,, s') = 0 and acts like uniform distribution when d(st, 5') 1, thereby preventing bias toward any single logit.
Architectural hyperparameters: In both settings, the transformer trunk consists of 12 layers and 12 heads with hidden size of 768. We set the maximum position encodings to 40 for 20 (state, previous reward)'s and 20 actions. Of these 20, 19 belong to the retrieved context and belongs to the query. In the JA JAT/Gato Go setting, the maximum multi-discrete observation size is 212 (for BabyAl). The maximum continuous observation size as discussed before in Section 4 is set to 513, consequence of the ResNet18 image embedding models's 's [13] embedding size of 512 with an additional I dimension for the reward. All linear encoding layers map from their corresponding input size to the hidden size. Following the JAT model, the linear decoder head for predicting continuous actions maps from the hidden size to the maximum continuous size discussed above (513). On the other hand, the linear decoder head for predicting distributions over discrete actions maps from the hidden size only the maximum number of discrete actions across all discrete environments (18), whereas in the JAT model, they map to the full GPT2 vocab size of 50257.
In the JAT/Gato setting, the original JAT architecture has a much larger context size (that is not required for REGENT), larger image encoder (than the resnet18 used in REGENT), much larger discrete decoder predicting distributions over the entire GPT2 vocab size (instead of just the maximum number of discrete actions in REGENT), and has (optional) decoders for predicting various observations (that is also not required for REGENT).
In the ProcGen setting. MTT has 18 layers, 16 heads, and a hidden size of 1024 all three of which are much larger than REGENT's 's architectural hyperparameters mentioned above.
Training hyperparameters: In the JA JAT/Gato o setting, we use a batch size of 512 and the AdamW optimizer with parameters 31 = 0.9 and 33 = 0.999. The learning rate starts at 5e-5 and decays to zero through epochs over all pre-training data. We early stop after a single epoch because we find that overtraining reduces in-context learning performance. This observation is consistent with similar observations about the transience of in-context learning in literature [42].
In the ProcGen setting, epochs over all the pre-training data with an early stop after the first epoch. We again use the Adam W optimizer with parameters
batch size of 1024, starting learning of 1c-4 also decaying over
0.9 and Bg 0.95.
Designating retrieval demonstrations in each training environment's dataset: In the JAT/Gato setting, designate 100 random demonstrations from the each of training vector observation environments as the retrieval subset for that environment. We also designate 5 random demonstrations from the training image environments as the retrieval subsets for said environments. In the ProcGen
16

---

Published as a conference paper at ICLR 2025
setting, with only 20 demonstrations per level per environment, designate all 20 as the retrieval set for that level in that environment.
We use all states in all demonstrations in each training environment's dataset query states. So, when a query state is from a demonstration in the designated retrieval subset, we simply retrieve from all other demonstrations in the designated retrieval subset except the current one. When query state is from a demonstration not in the designated retrieval subset, we retrieve simply from the designated retrieval subset. This way, we ensure that none of the retrieved states are from the same demonstration as the query state. A possible direction for future work includes intelligently designating and collecting demonstrations for retrieval (perhaps using ideas like persistency of excitation [45], memory-basee methods [56, 46, 47, 34], expert intervention methods [16]).
Choice of held-out environments: In the JAT/Gato setting, for the unseen environments, we hold-out the 5 Metaworld environments recommended in [58], the 5 Atari environments held-out in [28], and finally, we choose two Mujoco environments of mixed difficulty (see Figure 1). In the ProcGen setting, following MTT [37], we hold-out 5 environments.
In this work, we generalize to unseen environments from the same suites as the training environments. In future work, it would be interesting to examine what is needed for generalization to new suites (such as more manipulation suites [62], quadrotor, driving simulators [12, 14, 44), biological simulation [32], etc.).
Additional details on all environments: In the JAT/Gato setting, the Atari environments have 18 discrete actions and four consecutive grayscale images as states. However, embed each combined four-stacked image, with an off-the-shelf ResNet18 encoder [13], into a vector of 512 dimensions and use this image embedding everywhere (for R&P and REGENT). All Metaworld environments have observations with 39 dimensions and actions with 4 dimensions. Mujoco environments have observations with 4 to 376 dimensions and actions with 1 17 dimensions. The Baby environments have up to 7 discrete actions. They have discrete observations and text observations. After tokenizing the BabyAI AI text observations, which specify goal in natural language, into 64 tokens, BabyAl observations have 212 total dimensions and consist only of discrete values.
Additional Details on Various Baselines
Additional information on re-training JAT/Gato: When retraining JAT/Gato, we follow all hypepermaters chosen in [19]. We note that we skip the VQA datasets in our retraining of JAT/Gato, and use a smaller batch size to fit available GPUs. For JAT/Gato (A11 Data) we train over the full JAT dataset for the specified 250k training steps. For JAT/Gato train on the REGENT subset (which is 5- 10x smaller) for 25k training steps. In both variants, the training loss converges to the same low value at about 5000 steps into the training run.
Additional information on JAT/Gato + RAG (at inference): We note that JAT/Gato is like typical LLM- it has a causal transformer architecture and undergoes vanilla pretraining on In both the JAT/Gato + RAG baselines, we use the same
consecutive states-rewards-actionn retrieval mechanism as REGENT. We provide the retrieved states-rewards-actionn context these RAG baselines replacing the context of the past few states-rewards-actionn in vanilla JAT/Gato baselines.
REGENT significantly outperforms both JAT/Gato + RAG baselines in Figure 4. This demonstrates the advantage REGENT's architecture and retrieval augmented pretraining over just performing RAG at inference time. But, the JAT/Cato + RAG baselines outperform vanilla JAT/Gato as well as the JAT/Gat o fully Finetuned baselines the unseen metaworld environments. From this we can conclude that RAG at inference time is also important for performance.
Additional information on JAT/Gato Finetuned and REGENT Finetuned: Starting from a pre- trained checkpoint, finetune the full REGENT and JAT/Gato checkpoints using the same opti- mizer as pre-training time but starting with 1/10th the learning rate (i.e., 5c-6) and for epochs over the finetuning demonstrations. We this same recipe across all environments and for both REGENT and JAT/Cato. We only use 3 epochs to prevent the overwriting of any capabilities learned during pre-training.
Additional information on Parameter Efficient Finetuing (PEFT) of JAT/Gato with IA3 [30]] We finetuned JAT/Gato and JAT/Gato (All Data) with the TA3 PEFT method [30] and plotted the
17

---

Published as a conference paper at ICLR 2025
results in Figure 4. We used the huggingface implementation and the hyperparameters given in the IA3 paper (i.e. 1000 training steps. batch size of 8, and the adafactor optimizer with 3e-3 learning rate and 60 warmup steps in linear decay schedule). We found that finetuning JAI/Gato with IA3 improved performance over our simple full finetuning recipe, especially in metaworld environments. REGENT still outperforms both the IA3 finetuned JAT/Ga to and the fully finetuned JAT/Gato in all unseen environments in Figure 4.
Aditiinaa information on the fein-from--eeatch behavior cloning baseline: For the policy in this baseline, we use the Impala CNN [18] in Atari environments and a MLP with two hidden layers with 256 neurons each in Metaworld and Mujoco environments. We use the same learning rate (5e-5) for the same 3 epochs with the same optimizer (AdamW W with 31 = 0.9 and B2 0.999) as used for REGENT except we use constant learning rate schedule.
We added an imitation learning baseline, DRIL [6], to the atari environments Figure 4. Please note that DRIL consists of two stages. The first stage is pre training with BC loss. In the second stage, the DRIL ensemble policy is allowed to interact in the online environment where it reduces the variance between the members of its ensemble with an online RL algorithm. We run the second online stage for 1M steps around 1000 episodes. REGENT and all the other methods in Figure 4, on the other hand, are completely offline and never interact with any environment for training. While DRIL improves over the BC baseline, REGENT still outperforms DRIL in all three unseen atari environments.
Additional information on DRIL:
Additional information on MTT: The MTT model weights closed-source. In their ICML publication, the authors do not provide results on new levels. MTT is also difficult to train from scratch in any setting since they put multiple full trajectories in the context of the transformer, leading to exploding GPU VRAM requirements beyond academic compute resources. However, we simply report the best values from their paper their exact setting demonstrate REGENT's significantly improved performance even with a 3x smaller model and an order-of -f-mannitud less pretraining data, Similarly. we compare with JAI/Gato in their exact setting and significantly outperform JAT/Gato and JA T/Gato Finetuned, even though JA has 1.5x our parameters and 5- 10x our data.
We note that likely that MTT performance close random on a couple of unseen ProcGen environments only because the authors are forced to constrain MTT rollouts to 200 steps, stopping short of the true horizon of these environments. This simply because the quadratically increasing complexity of attention with increase in context size. MTT also likely does not report results beyond 4 demonstrations in the context for this reason. Yet, this is not an issue for R&P or REGENT which can scale any number of demonstrations to retrieve from since it only uses up to the 19 closest states in the context.
B PROOF OF THEOREM 5.2
To prove theorem 5.2, first state and prove the following lemma.
Let us refer to the family of policies represented by Th TR (als, r, c) in Equation (1) for various 0 as IT a, r, c). The parameters 0 of the transformer are drawn from the space O.
Lemma B.1. (Total Variation in the Policy Class) The total variation of the policy class (als, c) in environment j. V 01,02 E θ fnr some s E S5, r R2, e E Cj, defined as sup sup I 1 matcewr(alss (,)s,
92 "RRcewt(alss, e) is upper bounded by (1 -bbb,).
Proof: We have,
sup sup 41 (als, (als, 8,
sup sup (1 ->d) 01.02
Softmax(Te, (altma
(4)
= (1 e Xd sup sup Softmax(To,(a))) Softmax(To, (als, 8z
(5)
< (1 -xebb
(6)
18

---

Published as a conference paper at ICLR 2025
The first equality holds because for the same tuple of (state, previous reward, context), the R&P component does not change since it not affected by the choice of parameters. The last inequality can be interpreted property of Softmax and the distance to the most isolated state.
Now, we rewrite theorem 5.2 below and prove it.
Theorem
B.2. The
sub-optimality
gap
environment
(aj)) J J (RRCENT) I H bd
Proof: Recall that in imitation learning, the population total variation (TV) risk T( T²) S 6 then, J(x²) J++ I min H. H (See [36] Lemma 4.3). Using this with Lemma B.1 proves the theorem.
The main consequence of this theorem, also observed in our results, is that the sub-optimality gap reduces with more demonstrations Dj because of the reduced distance to the most isolated state.
19

---

Published as a conference paper at ICLR 2025
C ADDITIONAL JAT/GATO RESULTS
Effect of number of states/rewardsaactions the context: We plot the normalized return obtained by REGENT for different context sizes evaluation time in unseen environments in Figure 9. We find that, in general, larger context size leads to better performance, We also notice that the improvement in performance saturates quickly in many environments. Our choice of 20 states in the context (including the query state) leads to the best performance overall.
mettworrd--iin-o www.bbiii--iiggggg
setawwordd-box-cln
metaworid-doorf doo.d--odookkkkkkk
metaworid metaworidddoor-unlook door-unlock
(100 demos)
(1000 demos)
(1000 demoss
(100 demos)
Normaiizd return
Normalized return
Normalizee lizd return
Nomalized return
with REGENT
0.6
REGENT
0.8
REGENT
with REGENT
0.4
0.6
0.6
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0.0
14
17 (including query state)
states
(including query state)
(including query state)
(including query state)
metaworrd-hand-inee
atari-pang
stari-alien
alari-mspacman
(100 demos)
(16 demos)
[6 demos)
(10 demos)
Normalized return
with REGENT
0.8
Normalized return
with REGENT
Normalizeen saliidddd with REGENT
0.075
Normalized retur
with REGENT
0.8
0.6
0.15
0.6
0.4
0.10
0.4
0.025
0.05
0.00
0.000
2 5
14 17 20
11 14 17
8 11
17 20
2
5
11 14 17 20
Num states in context
Num states in context
Num states in context
Num states in context
(including query state)
(including query state)
(including query state)
(including query state)
Figure 9: Normalized returns obtained by RGGGENT in the unseen Metaworld and Alari enviromments with lixed retrieval buffer against the number states given in-context (including the query state). Please that context Each value from single model training seed but evaluated over 100 rollouts of different seeds Metawerld and 15 rollouts of different seeds
context.
(with Panda 0.05) in Atari.
metaworld-bal wywoorddbii-ppiiking
metaworid-bbo
metaworrd-doof
metaworid www.orrdddooouulllcckk
(100 demos)
[100 demos)
[100
demos)
(100 demos)
Normaizzed return with REGENT
Normalized return with REGENT
Normalized retur with REGENT
0.8
Nomalizzd return
with REGENT
0.8
0.6
0.6
0.4
0.4
0.4
0.1
0.2
0.2
0.2
0.0
0.0
Bandom-
Sendars
Dandeen
Bendom.
metawwrrd d-ann-d--aand (100 demos)
atari-pong (16 demos)
atari-alien [6 demos)
atari-mspacman
(10 demos)
Norma Mirmalized return with REGENT
0.8
Normalized return
with REGENT
0.20
REGENT
0.08
Norallzzd return
with REGENT
0.8
0.15
0.6
0.4
0.10
0.04
0.4
0.05
0.02
0.2
0.0
0.00
0.00
0.0
Servirm- Penmatted
Dumeett
Servieer
Durret
Bendimm Reverse Permintdd
Rendom- Recernee Premuited
Dermined
Figure Normalized returns obtained by REGENT the Metaworld and Atari enviromments with fixed retrieval buller against the choice of ordering states the context. Each value is from a single model training seed hut evaluated over 100 rollouts different seeds Metaworld rollouts of diflerent
Parsky 0.05) in Atari.
Effect of the ordering of the context: We plot the normalized return obtained by REGENT for differ- ent choices of context ordering at evaluation time in unseen environments Figure 10. We evaluated three options: "current", mandom-permuter tuples (state, previous reward. action) are placed in order of their closeness to the query state s,. In "ranntom-permuted"" the retrieved tuples placed in randomly chosen order in the context.
and "reverse". In the "current" order, the retrieved
20

---

Published as a conference paper at ICLR 2025
In "reverse", the "current" order of the retrieved tuples is reversed with the furthest state in the 19 retrieved states placed first. Our current choice performs the best overall. It significantly outperforms the other two orderings in the more challenging task of generalizing unseen Atari games which have different image observations, dynamics, and rewards. In some metaworld environments on the other hand, all three orderings have similar performance denoting that all the 19 retrieved states for each query state are very similar to each other.
Inference runtime values: In Figure 11, plot the inference runtime values of REGENT's two main components in a metaworld environment: retrieval from demonstrations and forward propagation through the REGENT architecture. We find that retrieval takes between 2 to 8 milliseconds time (for increasing number of demonstrations to retrieve from) while forward propagation takes around milliseconds uniformly. Both these runtimes are small enough allow real-time deployment of REGENT on real robots.
Runtime (milliseconds)
metaworld-bin-pickin
10
(25 demos)
forward prop (25 demos)
(SO demos)
forward prop (50 demos)
retrieval (75 demos)
forward prop (75 demes)
retrieval 1000 demes)
forward prop 1100 demos)
Figure Average runtime values retrioval and forward propagation through REGENT in metaworld for various number demonstrations the retrieval buffer.
Effect of the distance metric: We compared REGENT with cosine distance everywhere (from pretraining to evaluation) against our current version of REGENT with la distance everywhere in the JAT/Gat o setting in Figure 12. REGENT with cosine distance performs similarly to REGENT with & distance in 3 unseen metaworld environments. In the other unseen metaworld environments as well as the more challenging unseen atari environments, REGENT with &g distance significantly outperforms REGENT with cosine distance. We speculate that this is because the magnitude of observations is important with both image embeddings and proprioceptive states and this is only taken into account by the & distance (not the cosine distance).
metaworld sintrrld--ii-p-ee bin-d-iii---iiiiggg
metaworld-box-clos.
metaworld weotaworld-doorllock
metawortd-ddor-umlldl
Normalized return
6.6
Norallzedd return
0.8
Nomalized return
0.0
0.4
0.4
0.
0.0
100
75
100
50
75
100
25
50
/3
100
Number demos
Number demos
Number of demos
Number of demos
metaawrid-hand-inssrr
atari-pong
atari-alien
atari-mss Seari-mssacman
0.3
Normalized return
Normalized return
Normalized retum
Normalized returm
0.8
0.075
0.050
0.4
0.025
0.2
0.0
0.000
25
50
75
100
10000 15000 20000 25000
10000
15000 20000 25000 Number of states in demos
10000
15000
20000 25000
Number of demos
Number of states in demos
Number of states demos
R.P.P
REGENT
REGENT Finetuned
.... REGENT [zosine distance)
Figure Normalized returns in the unseen Metaworld Atari environments against number of demonstration trajectories agent can retrieve from. The REGGEN (cosine distance)" values are average 100 rollouts of different seeds Metaworld and 15 rollouts dillerent seeds (with Paiky 0.05) in Aturi.
21

---

Published as a conference paper at ICLR 2025
Generalization to unseen Mujoco envirooment/emboodmene In the two unseen Mujoco tasks shown in Figure 13, R&P appears to be a strong baseline even at generalizing to new embodiments! Moreover, after finetuning just 25 demonstrations, only REGENT, not JAI/Gato, significantly improves to outperform other methods and only continues to get better with more demonstrations.
mujoco-halfchestah
mujoco-hopper
0.3
Noraa lizdd return
0.20
0.2
0.15
Normalizes
0.10
0.05
0.00
25
50
75
100
25
30
75
100
Number of demas
Number of demos
R&
JAT/Gato Finetuned
REGENT
JAT/Gato (All Data)
REGENT Finetuned
JAT/Geto (All Data) Finetuned
JAT/Gato
Figure 13: Normalized returs in the unseen Mujoco environments against the number of demonstration trajectories the agent retrieve from or linetume on. Each value an average across 100 rollouts different seeds.
Additional qualitative examples: We plot examples of the inputs and outputs of REGENT at various states during rendered rollout in the metaworld-bin-picke continuous Metaworld example, REGENT can be seen predicting actions that are somewhat similar but the same as the first retrieved (R&P) action. These differences lead better overall performance seen in Figure 4. These differences are also the direct result of RECENT's 's in-context learning capabilities, learned from the preprocessed datasets from the pre- pre-training environments.
environment in Figure 14. In this
(2000]
634
(2584)
0.83
6287
0284
2007)
280
121
066
641
6.11
Retrieved r.
Query s, Pred
Retrieved s,
Query Pred å
Figure 14: Qualitative examples few inputs and outputs REGENT for various states in rendered rollout the continuous action space, melaworld bin picking enviromment.
22

---

Published as a conference paper at ICLR 2025
metaworld
REGENT
JATAGato
2.-
6.0
essemblle
basketball
battonn cois roodomn
buttondii nol
button pioeppoce
corfee buton
coffee pul (o0
da offee-pot
draner souecco
muusetooo tammo
hendee-pesss cresse
handls prezz
handleppol 2id0
na put
peg inpiiteeeee aie
rack oor--ot ootooootoee
picko peceppioc ond
pick-ptoc
pinte si1e DACC
pint site
pu So---ooc-
pushro
button button peoppp
Figure 15: Normalized returns in each of the 45 seen Metaworld envs. Each value is an average across 100 rollouts of different seeds. See Table 3 values.
mujoco
RECENT JAI/taato
Normeiter
0.0
humanalc
Figure Nommalized returns each 9 seen Mujoca enws. all values.
is average across rollouts seeds, Table
atari (Pactyy 0.00)
1.0
REGEN Maaoo
Wormaitoed
0.0
Figure 17: Normalized returns in each the 52 seen Atari cavs. Each value is an average across 15 rollouts of different soods. Sec Table for all values.
babyai
Normalized return
REGENT JAT/Gato
0.8
0.4
0.2
sclion-ti-doon
boocke unteck pickup boss-leveiine pees-leeve-noo-uee unlock
find ob) ds boxs level
go to door go tastoo---mmne -n---noock go-to-foccl obj
go-to (0-
greyl-grey
go to rod ball grey
qe to red ball
gp-io- ree--red-aa Sel--buu---- sall
go -e-00--0
key comdoo go
trint-boss--ed levee
move two stroso sbr9
one roomss door
open doors order 14
opers-ted--dmm
open
pickup-asi above
pirkup-niis diet
pickup-lor loc
pickup
pac-noot
synthsloe hoo
synth oeq
synth
unblock pttkcc pickup
unlock I
unlock
putmext-looc
unlack-to-unie
Figure 18: Normalized returns in each of the 45 seen AabyAl envs. Each value an average across 100 rollouts of different seeds. See Table 6
for all values.
23

---

Published as a conference paper at ICLR 2025
Table I: Values in Figure For particular training seed, each agent is evaluned across 100 rollout of different seeds in Metaworl and 15 rollouns of different seeds (with Paidky 0.05) in Atd. We compute overall mean sanddrd deviation over thiee fraining seeds for RGGONT. ENT Finetuned,, the PEFT with IA3 baselines, and DRIL.. For RCGGENT this means that we trainod and thoroughly ovalunted three REGENT policics Then, we finsttnned each of these dhree RECENT policies 32 times, with enviromeents and 4 points per environment. for total of 96 finetuned policies. We compure the final average over both training seeds and evaluation seeds and the final standard day ievitiion standard devintion training sonds does
DRL
00-000
0.0 1:00
0.04 0.0-2 100000000 100006
00:000
00
0.01000 0:0000
1.0.00
0.1 = 0.00
0.06 1: 0.000
0.17 1: 0.0.5
0.06
0.18
From
Scratch
0.0
0.0
0.0
0.1
0.0
0.0
0.0
0.0
0.03
001
0.03
0.0
0.01
0.0
0.09
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.
0.0
0.0
0.0
0.0
000
PEFT
with
0.1 = 20
0.01 #:00 00
0.01 00
004000
004000
00
00:100
000
039 0:000
04 04 0.06
0.4 042
015 015 #000
0.5 035 0000
05 03 1000
034 0:3420000
039 039 2004
0.8 100000
0163 0163 5:003
061 2-000
00:00
0000000
00000
001000
00:00
00:00
00400
00100
00 0:0
0.01 0.0
0.0 00 10000
0.01
01145
0.01
(AI Daa) Finetuned
0.0
0.01
0.0
0.0
0.25
0.27
000
0.0
0.5
0.0
0.06
0.08
00
0.0
0.
00
0.0
0.0
0.0
0.02
0.01
0.052
(ALL Daaa) RAGG
0.01
0.01
0.02
0.09
0.2
0.23
0.31
0.8
0.24
0:31
0.38
0-11
0-43
0:44
0.5
0.0
0.1
0.1
0.44
0.0
0.0
0,0
0.0
0.0
0.0
0.0
0.01
0.00
0.00
0.00
122
JAT/Ca0 (AUI Data)
0.0
0.0
0.0
0.0
0.0
0.24
0.24
0.06
0.06
0.06
0.06
0.01
0.01
0.01
00
0.0
0.0
0.0
0.0
0.0
0.01
0.01
0.01
0.01
0..4
PIET
with
1x2
0.01 = 0.005
0.01 = 0,0055
000
0.0 - 0.0
0.0 - 0.
012.
0.09 0.00000000 0.021
0.14 811 0.029
0.15
0.22 0.55
0.23 0.23220.00 0:08
0.3 0.023
0,16 0,003
0009
0.60
0.16 0.16::000222
0.15200054
0.0 0/0
0.0 0.0
0.0 0.0
000
0.
0.0.00 0.0
0.0 0.005
0.0 0.000
0.02 = 0.00
0.02 0.021:00000 0.003
0.03 1 0.005
00033
any training porameters.
000
Finclued
0.0
0.01
0.0
0.0
0.0
0.0
0.35
0.17
0.7
0.00
0.08
0.02
0.08
0.07
0.04
0.04
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.01
0.00
0002
+RAG
0.04
0.08
0.06
080
0.2
014
0.56
0053
0.34
0.30
0.36
0-47
0:11
0.07
0.13
0.22
0/0
0.0
0.0
0.0
0.0
0.0
0.01
0.02
0.02
0.165
0.0
0.0
0.0
0.0
0113
013
001
001
000
001
001
0.0
00
00
00
0.0
0.0
0.019
RECENT
Finstuned
=
06 + 01040
0.69
0.021
0.73
0181 4
0.86 + 102
0.99 0.07
0.99
090 1800
08
0.3 0.04
088 1: 0.022
0.52
0. 0.72 0007
084 0.01
0.92 0.000
093
0.54
0.69 0.0
0.79 0.001
0.3
017
0.23 0.0
0.23
0,3
0.09
0.09 0.09
0.09 0.000
0.82 12220000 0.00
0.83 1: 0.016
0.89 19 0.009
0.91:00
0.633
RBOENT
055 + 0.51
0.50-0000
0.097
073 1113
082 0 142
084 0026
0931
073 073
ON 004T
035 0028
063 0054
08 0.09
038 0008
0001
068 001
0.88 0028
0.8 000
0.0 00
0.2 0.00
023 023=0000 0.09
098 0005
0.88 0000
0.0 00
0.05
082 0017
081 0.001
0,8 - 0.89
0,0022
0.59
034
035
062
0,00
04
0.34
0.33
0.58
0.69
0.68
031
0.58
0.55
0.3
0.4
055
0.55
0.78
0.86
0.77
0.54
0.36
0.69
0.08
0.18
0.88
0.88
0.04
0.05
0.05
0.68
0.68
0.82
0.82
0445
demos
50
75
25
50
75
In0
25
50
75
100
25
50
75
100
25
50
75
(1159)
10 (16533)
(68)
((2755)
(L77355)
5 (22269))
(77722)
((11992)
(L4522)
9 (22400)
training We that
www.mull---iieeee.. biteeicking
mac anstoworrl--dee doo-kee
und indeemot www.eeooollllooeeeeeeo deloeooooooaooo
mns wo doll--rlnnneee ookk
meteovell wew.eeddduddddd dinser invert
de fonn insett
what antt-mon mog aeggeennn
wiari mig mopaemmn
Aggregaic Moan
Miart Mtart-gogg gong
sariipoo pong
attr olien
aitec aatsi--llic ollen
sllinn
24

---

Published as a conference paper at ICLR 2025
Num demoss
R&F
REGENT
REGENT Finetumed
LATiGato
JAT/Gato Finelumed
GAto (AII Data)
JATaGao (AII Data)
Train
From Scratch
Finetuned
mujocc halfcheelah
25
0,08
0.01
0.19
0.0
0.0
0.0
0.0
0.04
mujoco- halfcheetah
50
hippco-halfcheetal
75
0.09
0.0
0.31
0.0
0.03
0.0
0.0
0.04
mujoco-halffehoetl
100
0.1
0.0
0.34
0.03
0.0
0.01
0.04
mujoco-hopper
25
0.21
0.01
0.01
0.0
0.01
0.03
0.04
mujocc-hopper
50
0,26
0.01
0.04
0.0
0.0
0.01
0.03
0.02
mujoco-bege hopper
75
mujoce hopper
100
0.25
0.01
0.12
0.0
0.0
0.01
0.05
0.01
Aggrogate Mean
0.164
0.008
0.171
0.009
0.005
0.021
0.031
Table 2: Values in Figure Each value an average across rollouts different seeds.
Eny
Num demos
REGENT
JAT7Gato
metaworld-asseebbe
100
0.83
0.89
metaworld basketbballl
100
metaworld button press topdown wall
100
0.62
0.43
metaworld press topdownn
metaworld button press wall
100
0.94
0.61
metaworld button press
100
0.62
0.43
metaworld -coffee- button
100
0.84
0.33
nntawwell-ccoffss-eee
100
0.62
0.01
metaworld-cae www.worlddddeeeeeeooooo push
metawoeld dial turn
100
0.83
0.33
metaworld www.worrrddddssembbbl
setaworld-door-clen
100
1.0
0.46
metaworld door-open
100
0.98
0.46
metaworld-drawee drawede----eeeeeee
100
1.0
0.89
metaworld-de wwweeeedddddww--ooo oreeer-open
100
0.96
0.91
metaworkd- laucel-close
metaworld Faucet open
100
0.99
0.68
metaworldd hammer
100
0.95
metaworld srrlwweld--aee dindll-press-side
100
0.99
0.83
metaworld handle- press
100
0.99
0.83
metaworld-hamn hid-hhndl---pl--iiiiiii
100
0.48
metaworld- handle-pull
100
0.48
0.0
metaworld pull
metaworld, insert side
100
0.7
0.15
metaworld- pepoorddd-ppppeeoo pog-unplug unplugg-iidd side
100
0.31
metaworid-ple pick-out-oof holee
100
0.01
0.0
metaworld- diid-pick- pick-place-wall
100
0.99
0.01
metawerld wowdwwerl--iicee pickiic--lacc
100
metawurkd plate slide
slide-back back
100
1.0
0.02
metaworld, plate slide back
metaworrdd Id- Mate-slide-side
100
0.99
0.19
metaworld-g wwwpworrdd--pa-siddd-. plate-slide
100
1.0
metaworid push-bbcck back
100
0.84
metaworid- sesttwordd-uss--eaee pall-
100
0.81
0.01
peldd-usst
me melaworld reach wall
100
0.99
1.03
metawarka
0.99
metaworldded siddd---h-ffa -se--s----l-ooe
100
0.96
0.0
metawoeld-ae -edcoer
100
0.61
0.21
metaworld stick pull
100
0.88
0.04
hetaworld-stick-pae
100
0.75
0.02
metaworld sweep
metaworld sweep
100
0.91
0.16
metaworld wew.dd-widowwwweer
0.9
metaworld-wisdowe window-ppon
100
0.97
0.14
Aggregate Mean
0.82
0.281
Table Values in Figure 15, Euch value is un average across 100 rollouts of different seeds.
25

---

Published as a conference paper at ICLR 2025
Fay
Num demos
REGENT
JAT/Gaao
mujoco-ant
100
0.17
0.12
mujoca- doublependullmm
100
mujocio bumanosd
100
0.02
0.02
mujoco- pendulum
mujoce-pushee
100
0.9
0.95
mejoco-reacher
100
0.9
0.91
mujoco-standup
100
0.26
majoco-swimmee
100
0.82
0.97
mujoco-walker
Aggregate Mean
0.356
0.374
Table 4: Values in Figure 16. Fach value is an average across 100 rollouts of different seeds.
Env
Num demos
REGENT
JAT/Gato
atari-amidar
0.96
0.02
atari-assault
0.02
0.02
atari-asteris
0.13
0.03
stari-atlantis
5
0.18
0.03
atari bunkheist
5
0.46
0.18
hattlezone
5
0.47
0.01
stari-beamrider
5
0.04
0.0
starrbbrzzerk
utari bowling
5
1.0
1.0
stari boxing
atari- breakout
5
0.15
0.0
atari centipede
5
0.35
0.07
atari choppercommand
5
0.02
0.0
atari-crazy cani-crazyelimber
5
0.19
0.14
mari-delender
atari demonattack
5
0.0
0.0
doubledunk
atari-caduro
5
0.0
0.0
atari fishingderhy
5
0.62
0.1
stari-freeway
5
0.82
0.6
atari-trestbite
5
0.44
0.01
stari gopher
atari gravitar
5
0.03
0.06
5
0.17
stari-icebockey
5
0.21
0.1
atari-jameshond
5
0.01
0.0
atari-Kangaros
0.0
atari-Krull
5
0.5
0.41
atari kungfumaster
atari montezzuarrevenge
5
0.0
0.0
stari atari-namen namethisgams
5
0.03
atari- phooenxx
5
0.07
0.0
atari-pitfall
5
0.97
1.01
atari- Mari-privatooye
0.79
alari-qbert
5
0.57
0.01
riverriid
atari-madrunnee
5
0.03
0.02
stari-rebotank
5
0.04
atari-seaquest
5
0.48
0.11
atari-skiing
5
0.0
0.0
atari-solaris
aturi surround
5
0.36
0.03
atari-timepilot
5
0.04
0.0
stari-tutankham
5
0.16
0.02
atari upodown
5
0.13
0.01
atari-venture
5
1.0
1.0
atari videopinball
atari wizardofwar
5
0.76
0.01
stari-yarsrevenge
atarizixxon
0.14
0.01
Aggregate Mean
0.293
0.131
Table S: Values in Figure 17, Each value an average acrows 15 rollouts of different seals.
26

---

Published as a conference paper at ICLR 2025
Fny
Num demos
REGENT
JATiGat
tabyal-actiooo-aa-ddo
20
0.39
0.57
babyal- blw.ko--oocookkeeeeeeooeooooo-ee pickup
babyaii boss level no unlock
20
0.49
0.1
hahyai
Iebyai-find-obbe
20
0.61
0.92
babyal-go-to-door
0.49
0.89
babyai-go-Bee paal-gg-tt--mm-nnlee
0.21
0.03
babyal-go-to-locl
20
0.66
0.46
babyai-go
bahyai to ohj
20
0.14
0,76
habyai-ge cool----tt--edd- ball-grey
bolbyai-go-to seaa---g-tt-re--baeee to---o---baall-----e--
20
0.95
0,98
babyai-go soooyaai-------eeeee co-to-red-ball
0.68
0.76
babyai-go-ton go----oooo---bue-eeee
0.71
0.71
babyal-ge -ooya--------e
20
0.46
0.46
babyat- go-to
habyai key corridor
20
0.07
0.47
babyai-mini cill-mnni-boss--eee boss-level
0,28
20
0.0
0.0
babyai com--room--
0.96
0.97
baby ai-open-dor
0,92
babyai-open open-doors dooor-ordder--eeeee order-n4
20
0.54
U.33
baltyai open
habyai-open--t cetrii--pe------ddeee two-ddors
20
0.37
0.37
babyai-open
0,92
babyai waa---iccup-aabee abooe
0.44
0,86
buby un pickup dist
0.4
babyai-picke sabbya-pickkup-loo
bubyai pickup
20
0.44
0.86
habyai put next
bahyai-pan potta---tt--
20
0.12
0.0
syuthl-sytth-loe
0.3
0.59
0,07
bahyai synth
0.5
0.3
habyai-unblock pickoc
0.25
0.23
babyai unlock local
20
0.93
0.99
bahyai wwwiii-noccc--iiiieee
conyyi--nnook-to-ulleo.
20
0.35
0,79
babyal-unlock
0.99
Angregate Menn
0.508
0.577
Table 6: Values Figure Each value average 100 rollouts of different seeds.
27

---

Published as a conference paper at ICLR 2025
D ADDITIONAL PROCGEN RESULTS
bigfish
bossfight
caveflyer
chaser
Normelized return
0.8
Normalized return
0.8
Normalized retumm
0.8
0.0
0.6
0.6
0.6
0.6
0.4
0.4
0.4
0.4
0.2
0.2
0.2
0.2
0.0
12
16
20
12
10
Number of demos
Number of demos
Nurbeer of demos fruitbot
Number of demos
coinrun
dodgeball
heist
Normalized return
0.8
Normalized return
0.8
Normalized retur
0.8
0.0
0.6
0.6
0.6
Normalized
0.6
0.4
0.4
0.4
0.4-
0.2
0.2
2
12
16
20
2
16
20
16 Number of demos
8
12
4
12
16
20
Number of demos
Number of demos
Number of demos
leaper
maze
miner
starpilot
1.0
1.0
Normalized return
0.6
Normalized return
0.8
Normalized return
0.8
0.8
0.6
0.6
0.6
0.6
0.4
0.4
0.4
0.4
0.2
2
12
16
20
20
16
4
8
12
16
20
Number of demos
Number of demos
Number of demos
Number of demos
REGENT
R&P
---. Upper limit based on REGENT in seen levels with 20 demos
Figure 19: Normalized returns in unseen levels all 12 ProeGen training environments against the number of demonstration trajectories the agent retrieve from. Fach value average across 10 levels with 5 rollouts each with Paike 0.1. See Table for all values.
Eny
sticky prob
num demos
(epoch, end batch)
REGENT
R&2
MIT [371
ninja
p=0.2
(0, 11612)
7.134 0.09
5.67
(,, 11612)
7.311 0.34
5.5
4.1% 0.33
(0, IT612)
6.8 0.85
6.4
1=12
(0, 11612)
7.6% 0.28
6.33
(0, I1612)
7.8.
n=:2
(0, (1612)
8.07 0.19
7.6
climber
p=0.2
Dm2
(0, TT612)
4.65
(0, 11612)
6.011 0.52
3.95
1.851 0.41
Dm8
(0, (T612)
5.8911 0.44
3.12
12
(0, I1612)
7.02: 0.13
3.64
n=16
(0, 11612)
6.2: 0.43
3.77
(0, (T612)
6.131
plunder
p=0.2
(0, 11612)
10.48: 0.43
8.40
Dmm
(0, TT612)
12.33 0.48
8.3
2.5
(0, (1612)
12.33: 0.66
8.2
Dm
(0, 11612)
13.51 0.22
9.36
(0, 11612)
13.75:00.2
n=
(0, 11612)
14.13 0.26
1038
Jumper
p=0.2
(0, (T612)
Dm4
(0, 11612)
5.4 0.57
4.0
4.654 0.28
D=N
(0, IT612)
5.62 0.71
(0, (1612)
6.471 1.06
3.25
D= 16
(0, 11612)
6.474 0.47
6.0
(0. 11612)
Table 7: Values in Figure 5 hefore nomalization. We trained three REGENT generalist agents for three training seeds. We evalunted ench of the three agents' performance, as well as the R&2 agent, across levels with rollouts and Paicy 0.2 in each unseen ProeGen environment. We computed the final average over both training seeds evaluation levels/tellouts and the final standard deviation training seeds. The values for MTT are the best seores reported in [37].
28

---

Published as a conference paper at ICLR 2025
Eny
sticky prob
num demos
(epoch, end baich)
REGENT
RGP
REGEN: (ID seed)
biglish
p=0.1
n=2
10. 11612)
11.32
6.75
m=4
10, (1612)
14.32
11.85
n=8
[0. 11612)
18.82
14.7
R=12
(0. 11612)
20.66
14.02
16
[0, 11612)
21.32
14.43
nm
[0. TT612)
22.74
13228
25.7
bosstight
p=0.1
n=2
(0. 1I612)
3.76
1.45
n=4
(0. 11612)
5.52
3.38
n=8
10, (T612)
5:24
3.67
n=12
(0. 11612)
5.12
3.17
16
(0. 11612)
4.32
3.51
m=22
[0, 1T612)
4,44
4.1
5.6
caverlyer
p=0.1
n=2
(0, 1I612)
7.44
2.74
n=4
(0. 11612)
7.34
2.98
8=8
(0. 11612)
3.29
n=12
10, (T612)
6,96
3:27
m=16
(0. 11612)
6.72
3.31
n=200
[0, T1612)
7,08
4.04
7.68
chaser
p=0.1
n=2
(0, 11612)
1.15
1:01
n=4
10, 11612)
1.31
1.31
n=8
(0, 11612)
1.46
1.47
10, 11612)
1.61
n=16 16
10, (T612)
1.47
n=20
(0, 11612)
1.79
1.73
cninrun
p-0.1
n=2
10, TT612)
7.8
6.2
Dm4
(0, 11612)
7.2
5.9
n=8
10. 11612)
m=12
10, 11612)
6.8
6.2
n=16
10, 11612)
nm20
10, (T612)
8.0
5.8
dodgeball
p=0.1
D=2
T0. 11612)
3.36
1.93
10. TT612)
3.44
1
n=8
(0, 11612)
4.04
2.17
n=12
10. 11612)
3.72
2.07
m=16
10, 11612)
3.64
2.35
n=
10, 11612)
3,56
fruitbot
p=0.1
n=2
10, (1612)
3,18
0.98
n=4
T0. 11612)
7.88
1.6
n=8
10, TT612)
11.02
3.67
n=12
(0. 11612)
13.02
4.38
m=16
10. 11612)
15.04
7.45
n==20
10, 11612)
15.62
4.67
16.36
heist
p=0.1
nm2
10, 11612)
4.6
0=4
10, 11612)
5.0
3.7
n=8
TO. 11612)
3.2
4.9
12
(0, 1I612)
3.4
3.4
n=16
(0. 11612)
6.0
5.1
2=
10. (1612)
leaper
p=0.1
m=2
10, 11612)
6.4
3.35
nm4
10, 11612)
7.4
5.85
8=8
10, 11612)
7.6
7.17
1=12
TO, 11612)
8.2
n=16
(0. 11612)
7.83
m=20
10, 11612)
9.2
6.17
9.2
maze
P=0.1
m=2
10, 11612)
nm4
10, 11612)
4.8
3.1
n=8
(0. 11612)
R=12
(0. 11612)
6.0
4.6
16
[0, 11612)
6.0
4.5
n=20
10. 11612)
miner
p=01
m=2
10, 11612)
2.88
2.14
m=4
10, 11612)
4.72
3.12
n=8
10, 11612)
3.68
2.96
n=12
[0. 11612)
4.54
3.09
n=16
(0, 11612)
4.38
3.45
n=20
(0, 1T612)
3.76
312
6.55
starpilot
P=0.1
n=2
(0. T1612)
12.52
n=4
10, 11612)
11.94
n=8
(0. T1612)
12.82
10.93
1=12
(0. 11612)
14.9
10.8
n=16
(0, 11612)
15.34
10.49
n=
(0, (1612)
11.8
10.08
16.1
Table 8: Values in Figures 6 and 19 before normalization. Each value is an average across 10 levels with 5 rollouts cach.
29

---

Published as a conference paper at ICLR 2025
E ADDITIONAL RELATED WORK
In addition to our related work section, we discuss other relevant papers where a retrieval bias has been recognized useful component here [61, 57, 59, 23, 25, 39].
RAEA [61] performs behavior cloning in each new environment with access to "policy memory bank" and focuses on improving performance in training task. In RAEA's appendix, the authors demonstrate very initial signs of generalization to new tasks only after training on few demonstra- tions from the new tasks. REGENT on the other hand pretrains generalist policy that can generalize without any finetuning to completely new environments with different observations, dynamics, and rewards. Re-V [57] trains retrieval-augne -uugente- image captioning model, which while demon- strating the usefulness of retrieval bias beyond robotics and game- game-playing tasks, is not applicable to our settings.
Expel [59] and RAP [23] build LLM agents for high-level planning top of frozen LLMs. Although off-the-shelf LLMs are not yet helpful with low-level fine-grained continuous control, these methods could be key in high-level semantic text-based action spaces.
GLAs [25] and RADT [39] are recent in-context reinforcement learning (RL) methods. We briefly discussed earlier in-context RL methods in related work section. GLAs attempt to generalize to new mujoco embodiments via in-context RL but only show minor improvements over random agents. RADT emphasizes that they unable to generalize via in-context RL unseen procgen, metaworld, or DMControl environments. REGENT, on the other hand, an in-context imitation learning method that generalizes to unseen metaworld, atari, and procgen environments via retrieval augmentation (from a retrieval buffer a few expert demonstrations) and in-context learning. REGENT also struggles with generalization new mujoco embodiments like GLAs.
30


Reviews:
[{"review_id": "N4TqxrXo5T", "reviewer": "Reviewer_jvNh", "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.", "questions": "Please see weaknesses part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The work is well-organized and easy to follow.\n2. 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. \n2. There are some places that the author overclaims:\n\n- 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. \n\n- 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. \n\n- 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. \n\n3. 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. \n\n4. Missing highly-related reference [1,2].\n\n[1] Retrieval-Augmented Embodied Agents, CVPR24\n[2] Re-ViLM: Retrieval-Augmented Visual Language Model for Zero and Few-Shot Image Captioning, EMNLP", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cgqqDsMn9a", "reviewer": "Reviewer_N77b", "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.", "questions": "+ I am curious what are the results of REGENT without the help of the 1-nearest neighbor agent.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n+ 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.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "elzJQvtVCS", "reviewer": "Reviewer_eumA", "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.", "questions": "See weaknesses.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The methodology is well-explained and easy to understand.\n\n2. REGENT demonstrates better performance than larger models (JAT/Gato) with fewer parameters across multiple environment types.", "weaknesses": "1. Limited Comparison Scope:\n\nWhile 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:\n- https://arxiv.org/abs/2308.10144\n- https://arxiv.org/abs/2402.03610\n\n2. Insufficient Ablation Studies, some valuable ablations could have included:\n- Analyzing performance w/wo image information in the context\n- Testing various context sizes\n - 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\n\n\n3. Training Data Requirements:\n\nWhile 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ltytxSDyni", "reviewer": "Reviewer_8qsk", "summary": "This work proposes to leverage retrieval augmentation to reduce the amount of training data required for training generalist agents.\nThe authors introduce $\\texttt{REGENT}$, a retrieval-augmented architecture that incorporates a non-parametric retrieval component (R&P) with learning of a parametric policy.\nFurther, they demonstrate theoretical guarantees for performance with respect to the number of expert demonstrations.\nThey pre-train $\\texttt{REGENT}$ on 4 different domains and evaluate on unseen levels and holdout tasks, where it exhibits strong performance compared to baselines.", "questions": "- Why is the scaling factor for MixedReLU set to 10? This seems like an arbitrary choice, why not remove it entirely?\n- Does the number of demonstrations for MTT on ProcGen match the number of demonstrations for $\\texttt{REGENT}$?\n- Is it always the case that $\\texttt{REGENT}$ only receives the current query state in its context? Have you tried providing more sequential information?\n- What ProcGen mode is being used? Easy or hard?\n- How important is the ordering in the context of $\\texttt{REGENT}$? Did you try permuting the order?\n- 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?\n- 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?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The authors provide theoretical guarantees for performance with respect to the number of demonstrations.\n\n$\\texttt{REGENT}$ is a novel architecture that effectively incorporates retrieval augmentation for imitation learning.\n\nRetrieval augmentation enables training on much less data than other generalist agents.\n\nThe authors demonstrate the promise of retrieval augmentation for unseen levels and tasks.", "weaknesses": "**Significance of results:**\n\nThe 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?\n\nThe 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?\nAlso there are fine-tuning techniques specifically designed for few-shot learning, which may be more suitable in this setting [1,2].\n\nOn unseen levels of ProcGen the authors only compare $\\texttt{REGENT}$ to R&P, but no results for MTT, why is this the case?\nBC 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.\n\nAnother 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.\n\nThere is a lack of ablation studies, for example performance on varying the number of examples in the context and ablation on different distance metrics.\n\nFinally, it is not really clear why the authors differentiate between the two different settings. Why not include ProcGen in the first setting?\nWhy are different baselines used for the different settings? MTT could also be applied to the first setting, right?\nAt least having a fixed set of all baselines across both settings would strengthen the paper.\n\n**Limitations and claims**\n\n- 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.\n- Moreover, claims on \"handful of examples\" should be caveated as thousands of samples are not a handful.\n- 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.\n- 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. \n\n**Presentation:**\n\nPresentation can be improved, in particular:\n\n* 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\n* Figure 3: enumeration in caption, but not in figure, either remove them in caption or add in figure\n* 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.\n* Theorem 5.2: $\\boldsymbol{H}$ represents the horizon, right? In section 3 it is not boldface, this should be consistent.\n* Figure 8: hard to follow what actions are \"good\" or \"bad\" as there is no explanation for the action index\n\n**References:**\n\n[1] Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning, Liu et al., NeurIPS 2022\n\n[2] Scaling & Shifting Your Features: A New Baseline for Efficient Model Tuning, Lian et al., NeurIPS 2022\n\n[3] Generative Adversarial Imitation Learning, Ho et al., NeurIPS 2016\n\n[4] Disagreement-Regularized Imitation Learning, Brantley et  al., ICLR 2020\n\n[5] Of Moments and Matching: A Game-Theoretic Framework for Closing the Imitation Gap, Swamy et al., ICML 2021\n\n[6]   Evolutionary principles in self-referential learning, or on learning how to learn, Schmidhuber et al., Diploma Thesis, 1987\n\n[7] Towards General-Purpose In-Context Learning Agents, Kirsch et al., NeurIPS 2023 FMDM Workshop\n\n[8] Retrieval-Augmented Decision Transformer: External Memory for In-context RL, Schmied et al., arXiv 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)