Research map / Multi-agent systems and games / Human-AI and goal-conditioned
REGENT: A Retrieval-Augmented Generalist Agent That Can Act In-Context in New Environments
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes REGENT, a semi-parametric retrieval-augmented transformer policy that generalizes to unseen robotics and game environments via in-context retrieval, with a simple 1-nearest-neighbor Retrieve-and-Play baseline that is surprisingly strong.
Problem
Building generalist RL agents that can rapidly adapt to unseen environments with only a few expert demonstrations, without finetuning, while using fewer parameters and pre-training data than large generalist policies like Gato/JAT and MTT.
Approach
Combine a 1-nearest-neighbor retrieval policy (Retrieve and Play) with a causal transformer that is pre-trained on query states plus retrieved (state, reward, action) contexts, interpolating the transformer output with the retrieval action via a distance-dependent exponential weight so that deployment in unseen environments needs no weight updates.
Key results
- REGENT (138.6M params) generalizes to unseen Metaworld/Atari/Mujoco environments without finetuning and outperforms JAT/Gato variants, including All Data and finetuned versions, using 1.4x to 3x fewer parameters and ~10x less pre-training data.
- On unseen Metaworld bin-picking with 100 demos, REGENT reaches 0.67 normalized return versus 0.0 for JAT/Gato and 0.06 for JAT/Gato + RAG.
- Retrieve-and-Play, a 1-nearest-neighbor method with no learned parameters, matches or exceeds JAT/Gato on many unseen environments, e.g., 0.33 vs 0.0 on metaworld-bin-picking with 25 demos.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add REGENT's transformer policy and retrieval-augmented pretraining to Retrieve and Play | Retrieve and Play · unseen environments in the JAT/Gato and ProcGen settings | normalized return | Helped |
| Add retrieval-augmented pretraining rather than using retrieval only at inference | JAT/Gato with RAG · unseen environments | normalized return | Helped |
| Finetune REGENT on demonstrations from the unseen environment | REGENT · unseen JAT/Gato environments | normalized return | Helped |
| Add interpolation between Retrieve and Play and the transformer output (ablation) | REGENT · unseen environments | normalized return | Helped |
| Increase context length (ablation) | REGENT · unseen environments | normalized return | Helped |
| Order retrieved tuples by closeness to the query instead of random-permuted or reverse order (ablation) | REGENT · unseen environments, especially Atari | normalized return | Helped |
| Use ℓ2 distance instead of cosine distance (ablation) | REGENT · unseen Metaworld and Atari environments | normalized return | Helped |
Surprising findings
- A 1-nearest-neighbor retrieval policy with no neural network and no pre-training environments performs on par with or better than state-of-the-art generalist agents on unseen environments.
- Without the interpolation between the retrieval policy and transformer output, REGENT performs like a random agent on unseen environments.
Limitations
- Struggles to adapt to new embodiments (unseen Mujoco tasks) and long-horizon environments such as atari-spaceinvaders and atari-stargunner.
- Evaluation only covers unseen environments within the same suites as training environments; generalization to new suites is left for future work.
- MTT comparison uses best numbers reported in the original paper since MTT weights are closed-source; MTT rollouts were constrained to 200 steps in the original work.
Experimental setup
Datasets: JAT dataset subset (Metaworld, Mujoco, Atari, BabyAI) ProcGen PPO-generated rollouts (12 training environments)
Benchmarks: Metaworld Atari Mujoco BabyAI ProcGen
Metrics: normalized return IQM mean normalized return
Baselines: JAT/Gato JAT/Gato (All Data) JAT/Gato + RAG JAT/Gato Finetuned JAT/Gato PEFT with IA3 MTT Train From Scratch behavior cloning DRIL
Builds on: Gato JAT Multi-Trajectory Transformer (MTT) Retrieval-Augmented Generation Algorithm Distillation
Closest papers by meaning
- Search More, Think Less: Rethinking Long-Horizon Agentic Search for Efficiency and Generalization (NeurIPS 2026 · similarity 0.87)
- AgentRefine: Enhancing Agent Generalization through Refinement Tuning (ICLR 2025 · similarity 0.87)
- Learn-by-interact: A Data-Centric Framework For Self-Adaptive Agents in Realistic Environments (ICLR 2025 · similarity 0.87)
- Search Self-Play: Pushing the Frontier of Agent Capability without Supervision (ICLR 2026 · similarity 0.87)
- C-3PO: Compact Plug-and-Play Proxy Optimization to Achieve Human-like Retrieval-Augmented Generation (ICML 2025 · similarity 0.87)
- ML-Agent: Reinforcing LLM Agents for Autonomous Machine Learning Engineering (ICML 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.