PaperlayerView exact revision on arXiv

Source-linked AI summary

Human-Timescale Adaptation in an Open-Ended Task Space

Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmieg, Michael Chang, Natalie Clay, Adrian Collister, Vibhavari Dasagi, Lucy Gonzalez, Karol Gregor, Edward Hughes, Sheleem Kashem, Maria Loks-Thompson, Hannah Openshaw, Jack Parker-Holder, Shreya Pathak, Nicolas Perez-Nieves, Nemanja Rakicevic, Tim Rocktäschel, Yannick Schroecker, Jakub Sygnowski, Karl Tuyls, Sarah York, Alexander Zacherl, Lei Zhang

arXiv:2301.07608v1cs.LGcs.AIcs.NE

TL;DR

The paper asks whether RL can produce generally adaptive agents without relying on the large datasets used by supervised and self-supervised foundation models. It trains AdA with large-scale meta-RL, attention-based memory, and an automated curriculum, finding rapid adaptation across open-ended 3D tasks at a timescale similar to humans, with benefits from first-person prompting and scaling.

  • Problem

    The central gap is whether large-scale RL can train generally adaptive models that learn rapidly across open-ended tasks without relying on large supervised or self-supervised datasets.

  • Method

    AdA is trained with meta-RL in XLand 2.0 using an automated curriculum, large-scale attention-based memory, and distillation from a pretrained teacher.

  • Results

    AdA adapts rapidly across a vast open-ended task space at a timescale similar to human players, while first-person prompting and increases in parameters, context length, and training-task richness improve performance.

  • Takeaways & Limitations

    The results support RL as a route to foundation-like agents that perform rapid in-context adaptation across increasingly broad embodied task spaces.

  • Takeaways & Limitations

    Human demonstrations cannot overcome inherent limitations of AdA’s task distribution, and the largest model may not be the best choice when compute cost is considered.

Abstract

from arXiv · show

Foundation models have shown impressive adaptation and scalability in supervised and self-supervised learning problems, but so far these successes have not fully translated to reinforcement learning (RL). In this work, we demonstrate that training an RL agent at scale leads to a general in-context learning algorithm that can adapt to open-ended novel embodied 3D problems as quickly as humans. In a vast space of held-out environment dynamics, our adaptive agent (AdA) displays on-the-fly hypothesis-driven exploration, efficient exploitation of acquired knowledge, and can successfully be prompted with first-person demonstrations. Adaptation emerges from three ingredients: (1) meta-reinforcement learning across a vast, smooth and diverse task distribution, (2) a policy parameterised as a large-scale attention-based memory architecture, and (3) an effective automated curriculum that prioritises tasks at the frontier of an agent's capabilities. We demonstrate characteristic scaling laws with respect to network size, memory length, and richness of the training task distribution. We believe our results lay the foundation for increasingly general and adaptive RL agents that perform well across ever-larger open-ended domains.

1. Introduction

The paper introduces AdA, an RL agent trained across a vast open-ended task space to adapt rapidly to unseen embodied 3D tasks. Its approach combines scalable meta-RL, attention-based memory, and an automated curriculum, enabling human-timescale adaptation and first-person demonstration prompting.

  • Motivation: AdA addresses the challenge of fast adaptation in unseen environments, where trial-and-error learning and continued refinement are required from only a few episodes.The motivation is to move beyond zero-shot policy generalization toward rapid in-context learning from feedback.
  • Contribution: The work targets an RL foundation model that pre-trains on a vast task distribution and adapts few-shot to downstream tasks without prompts, fine-tuning, or offline datasets.AdA is designed for sparse-reward, open-ended embodied tasks.
  • Results: AdA adapts to challenging held-out tasks in minutes, with a human study finding adaptation timescales comparable to trained human players and improved performance from first-person demonstrations.The reported behavior includes hypothesis-driven exploration and refinement toward near-optimal performance.
  • Approach: AdA combines meta-RL at scale, an automated frontier-focused curriculum, model-based RL with large-scale attention-based memory, and distillation.The training approach uses XLand 2.0 and distillation to make large Transformer architectures practical.
  • Task space: XLand 2.0 supplies a diverse 3D task space with production-rule dynamics, partial observability, and trial resets that preserve agent memory across trials.Production rules expand the range of transition functions by conditionally removing and spawning objects.
  • Architecture: Transformer-XL provides longer variable-length context, and it significantly outperforms a vanilla RNN on adaptation tasks.Its output embeddings feed the policy and value components of the agent.

3. Experiments and Results

AdA adapts across diverse held-out tasks, matching human-timescale improvement through hypothesis-driven exploration and later refinement. Its performance also scales with architecture, memory, curriculum, task-pool size and complexity, and distillation.

  • 0.04 to 0.61: AdA’s 20th-percentile normalized score improves from the first trial to the 13th, quantifying its adaptation gap.
  • AdA improves on over 80% of the single-agent task set with additional trials, without degrading zero-shot performance.The strongest gains occur from zero-shot to two trials, but performance continues increasing with more trials.
  • On 30 hand-authored tasks, AdA and humans both improve with more trials, with AdA’s score slopes as steep as or steeper than humans’ for lower trial counts.Some tasks favor humans or AdA, but most show adaptation by both groups.
  • AdA uses hypothesis-driven exploration in early trials, then refines successful strategies by gathering the right objects more efficiently.This behavior explains higher last-trial scores when the agent receives more trials on the Wrong Pair Disappears task.
  • Self-play outperforms random-action co-play on cooperative held-out tasks, consistent with agents dividing labour and coordinating paths after discovering solutions.The multi-agent agent shows adaptation across almost 90% of held-out test tasks.
  • Transformer-XL performs best among matched-size memory architectures, while attention-enhanced RNNs recover most of the Transformer’s performance.

4. Related Work

The paper combines open-ended learning, adaptive curricula, procedural environments, and memory-based meta-RL to train agents that adapt rather than merely generalise zero-shot. Its key distinction is using large Transformer models and on-the-fly hypothesis-driven exploration in vast task spaces.

  • The work combines attention-based meta-learning with automatic task-distribution curation to train an in-context RL algorithm in an open-ended task space.
  • Architectures and adaptation: The paper uses large Transformer models for adaptation in a vast multi-agent task space, whereas closely related work used smaller models and targeted zero-shot performance.
  • Open-ended learning: The approach extends prior open-ended learning and curriculum methods by focusing on adaptation and curating randomly sampled environments with high regret.
  • Adaptation: Unlike related zero-shot generalisation work, AdA discovers hidden environment rules through on-the-fly hypothesis-driven exploration over multiple trials.
  • Meta-RL: Memory-based meta-RL provides a framework in which rewards and terminations can support a learning algorithm when the policy memory persists across trials.

5. Conclusion

The paper asks whether large, generally adaptive models can be trained with RL and answers by presenting AdA, which adapts rapidly across a vast open-ended task space. It attributes this capability to scalable meta-RL, automatic curricula, attention-based memory, and distillation.

  • The paper addresses whether large-scale, generally adaptive models can be trained through reinforcement learning rather than relying mainly on supervised or self-supervised datasets.
  • AdA demonstrates rapid in-context adaptation across a vast open-ended task space at a timescale similar to human players.
  • AdA explores held-out tasks structurally, refines its policy toward optimal behaviour from few interactions, accepts first-person prompting, and scales with parameters, context length, and training-task richness.
  • Training and architecture: Automatic curricula shape the data distribution to provide sufficient learning signal in an open-ended space, while attention-based architectures exploit it more effectively than purely recurrent networks.
  • Scaling: Distillation enables the potential of large-scale Transformer architectures, supporting the paper’s recipe for training a 500M-parameter model.

6. Authors and Contributions

The supplied passages list contributors, their roles, and selected XLand 2.0 implementation materials. They do not provide a complete narrative of author contributions beyond those listed assignments.

  • The contribution list assigns technical leadership, curriculum research, infrastructure engineering, task authoring, paper writing, agent research, scaling, and agent analysis across named authors.
  • Other listed contributors are assigned project management, QA testing, evaluation research, XLand development, human experiments, strategic advice, and paper writing.
  • XLand 2.0 materials: The XLand 2.0 materials introduce configuration changes including the touching(a,b) relation, frozen objects, and single-relation predicates.
  • Reference materials: Reference tables list the shapes, colours, and predicates used for objects, goals, and production rules.

A.2. Pre-sampling tasks for training

The task pool combines sampled goals, production rules, worlds, distractors, dead ends, and cooperative games, with held-out evaluation and normalized last-trial scores. Aggregation emphasizes coverage across tasks rather than outliers.

  • Task-space construction: The task space contains at least 10^40 distinct tasks from goals and production rules, despite restrictive bounds on objects and rules.
  • Task sampling: Single-player tasks sample goals, sequential production rules, initial worlds, distractors, dead ends, and hiding masks to create exploration and rule-avoidance demands.
  • Evaluation tasks: Multi-player evaluation is restricted to fully cooperative two-player games, while test tasks use held-out world topologies and exclude games contained in training.
  • Evaluation scope: Quantitative test evaluation covers 1000 held-out tasks, supplemented by human-level capability probes on 30 single-agent and 28 multi-agent tasks.
  • Scoring: Test scores divide total reward by a task-specific normalizer so scores are roughly on [0, 1], although evaluated agents can exceed 1.
  • Scoring: Performance uses the last trial across k ∈ {1, 2, 3, 5, 8, 13}, testing whether additional experience improves adaptation, including held-out trial counts.
  • Aggregation: Across games, a fixed percentile provides a lower-bound view of performance on most tasks and emphasizes broad coverage over outlier improvements.

B.3. Adaptation metric

The adaptation metric ranks agents by combining zero-shot performance with efficient few-shot improvement across probe tasks. It uses median task scores and a Nash-average scalarization of Pareto comparisons.

  • Metric construction: The metric evaluates last-trial total reward across probe tasks and trial counts k∈{1, 2, 3, 5, 8, 13}.Per-task scores are normalized, aggregated by the median, and then compared across trial counts.
  • Metric construction: Agents are preferred only when their task-aggregated scores Pareto-improve over another agent across trial counts.This favors both high-quality zero-shot generalization at k=1 and efficient improvement when k>1.
  • Metric construction: The Nash average converts Pareto comparisons into a scalar adaptation metric using a maximum-entropy Nash equilibrium.Agents that do not Pareto-dominate one another receive equal Nash payoffs.
  • Use: The metric was used to select hyperparameters for the best-performing agent in Section 3.1.

C.1. Agent Architecture

The agent encodes visual, goal, rule, and scalar observations into a Transformer-XL memory, then uses Muesli heads to produce policy and value outputs. Its persistent memory spans 300 cached activations, or 1800 effective timesteps.

  • Observation encoding: RGB observations are encoded by a ResNet, while goals and production rules use separate embedding networks.The RGB encoder outputs size 256; goal embeddings have size 8 and production-rule embeddings size 256.
  • Observation encoding: Encoded observations are concatenated with reward, action, proprioception, and timing information before projection to the Transformer memory dimension.
  • Transformer memory: Transformer-XL provides causal memory with 300 cached activations, corresponding to 1800 effective timesteps in the default configuration.The learner uses an 80-step rollout context, while the actor uses a context window of 1.
  • Policy and value heads: The Transformer output initializes a Muesli LSTM whose policy and value MLP heads operate over six decomposed action groups.
  • Adaptation mechanism: Meta-RL episodes contain multiple trials, and memory persists across trial boundaries so the agent can experiment and reuse discovered information.

D.2. Single-agent training

Single-agent training uses a task pool generated from an earlier XLand 2.0 version and trains a distillation teacher for the single-agent experiments.

  • Training setup: The single-agent setup uses a task pool generated according to the procedure described in Section 2.1.
  • Environment scope: Single-agent training used an earlier XLand 2.0 version without frozen objects.Frozen objects were consequently excluded from the single-agent test and hand-authored probe task sets.
  • Training setup: The experimental setup for the single-agent distillation teacher is summarized in Table D.1.

D.3. Multi-agent training

Multi-agent training builds a large mixed task pool, introduces potentially frozen objects, and samples increasingly capable co-players through fictitious self-play. PLR then curates tasks using TD-error fitness.

  • Task pool: The multi-agent task pool matches the single-agent pool size and combines single-player and two-player tasks in equal proportions.
  • Task variation: For each task, objects may be frozen with 50% probability, and each spawned object type is frozen with 20% probability while preserving solvability.
  • Co-player training: Co-player policies are sampled from a fictitious-self-play pool initialized with random actions and expanded every 500M training frames.Each added snapshot increases the pool's capability over training.
  • Curriculum: PLR curates worlds, games, and co-players using the agent's TD-error-based fitness.
  • Training setup: The multi-agent distillation teacher setup is summarized in Table D.2.

D.4. Architecture experiments

This section documents the experimental setup used to compare memory architectures, as reported in Table D.3.

  • Table D.3 specifies the experimental setup for comparing different memory architectures.

D.5. Auto-curriculum learning

The auto-curriculum experiments evaluate task filtering, replay fitness metrics, and curriculum efficiency, showing that adaptive curricula improve training efficiency over uniform sampling.

  • No-op filtering: No-op filtering admits tasks only when baseline and agent scores, differences from no-op, and score variability satisfy threshold criteria.Admitted tasks are then run for their full trial count across 30 episodes, with the resulting experience sent to the learner.
  • Prioritised Level Replay: Prioritised Level Replay maintains a bounded archive and preferentially replays tasks using fitness and staleness.Proposal tasks enter the archive when their fitness exceeds the current minimum, and the lowest-fitness task is removed when capacity is exceeded.
  • Fitness metric: Last-trial fitness outperforms first- or average-trial fitness because late error better indicates residual suboptimality after adaptation.Earlier error can reflect appropriate exploratory behaviour rather than a deficient policy.
  • Fitness metric: TD-error and action-model fitness outperform value-model fitness, while TD-error is selected for stronger asymptotic zero-shot and few-shot performance.Action-model fitness matches TD-error in few-shot performance but has weaker zero-shot performance.
  • Curriculum efficiency: No-op filtering and PLR both improve sample efficiency over uniform task sampling, with PLR more efficient early in training.The comparison uses few-shot k=13 and zero-shot evaluation scores across learner steps and FLOPs.
  • Emergent curricula: The experiments analyse the task metrics selected by PLR and No-op filtering to characterise their emergent curricula.Figure D.3 reports average metrics over dynamic training sets with standard-error shading and compares them with test and hand-authored task sets.

D.10. Scaling the complexity of the task pool

This section reports experiments on adaptation across multi-agent and probe tasks, including task-pool scaling, trial conditioning, and repeated distillation.

  • Task-pool scaling: The task-pool scaling experiments use a setup whose distillation teachers differ between the compared agents.Consequently, the effects of distillation and task complexity cannot be disentangled, although the results remain indicative of task-complexity importance.
  • Trial adaptation: AdA adapts across a wide range of percentiles on a held-out multi-agent task set when evaluated with varying numbers of trials.The evaluations use cooperative self-play and compare normalised last-trial rewards against fine-tuned test-task performance.
  • Trial adaptation: AdA’s last-trial performance increases when it receives more trials on the Irreversible Production for Two probe task.The evaluation is conducted in self-play, with representative last-trial trajectories shown for different trial counts.
  • Trial conditioning: Trial-1 scores are almost identical between one-trial and eight-trial episodes, indicating no trial-count-dependent early exploratory adjustment.A more exploratory first trial with eight total trials would instead be expected to score lower.
  • Trial conditioning: The authors note that the absence of early exploration may reflect either the optimal policy for XLand 2.0 or a training-procedure limitation.Learning from all episode rewards with discounting may encourage myopic behaviour and discourage delayed-information strategies.

E.3. Scaling complexity of the task pool

This section studies how task-distribution complexity and compute accounting affect scaling, finding that environmental complexity and evaluation cost constrain the benefits of larger models and memories.

  • Task-distribution complexity: The complexity comparison varies world topology while holding goals, production rules, and objects fixed, adding 5,000 initialisation points to the empty-room condition.The empty-room proxy is deliberately less diverse and complex than the main distribution’s 4,000 worlds.
  • Task-distribution complexity: Low environment complexity bottlenecks model scaling: in the empty-room condition, performance begins declining beyond 42M Transformer parameters.The comparison uses median and 20th-percentile performance, whereas the multi-topology distribution continues to support scaling in the reported comparison.
  • Compute accounting: The scaling figures separately evaluate Transformer model size and memory length under compute-controlled settings.Figures E.5–E.8 distinguish total-FLOP and learner-FLOP accounting for these axes.
  • Compute-matched scaling: Under matched total FLOPs, performance decreases beyond a sweet spot around 57M Transformer parameters and 141M total parameters.The compute-matched analysis also reports continued benefit from increasing context length.
  • Compute accounting: Learner-only FLOP accounting still shows performance increasing with model size and memory length, but less steeply than sample-efficiency comparisons.The authors caution that actor-step FLOPs may be poorly optimised for single-step Transformer queries on TPU.
  • Repeated distillation: The third distillation generation provides additional few-shot gains after 25 billion further frames, but with diminishing improvement relative to the second generation.The experiment counts combined experience beginning with the original teacher’s experience.

F.1. Probe tasks

AdA is evaluated on diverse single-agent probe tasks designed to test hidden-rule discovery, navigation, object handling, and multi-step reasoning. Across these tasks, prompting can improve performance when demonstrations fit the agent’s memory and match capabilities, but benefits are task-dependent.

  • Task design: The probe suite uses complex worlds with hidden production rules, testing adaptation under sparse information and challenging object interactions.Unless otherwise noted, tasks contain many objects and multiple production rules fully hidden from players.
  • Task design: Several probes test whether agents discover non-obvious object-handling strategies, such as pushing rather than lifting or using another object as a tool.Other variants require transforming objects before safely transporting them or combining many objects through repeated rules.
  • Task design: Tasks span hierarchical crafting, navigation, distractor-rich exploration, object duplication or destruction, and time-limited choices between alternatives.Examples include a seven-rule crafting chain, growing populations of objects, noisy distractors, and tasks where only one option can be reached per trial.
  • Task design: Controlled variants manipulate rule visibility, hidden shortcuts, task topology, and the number of production rules while holding other task properties similar.Same-signature tasks can look identical initially but require different hidden rule sequences, and workstation variants expose or hide rules.
  • Prompting results: Prompting produces perfect or near-perfect performance on some tasks, especially short navigation and short production-rule problems, but does not consistently improve longer or more complex tasks.A full demonstration may exceed memory context, while unfamiliar object configurations, subtle positioning, or tool-use requirements can prevent effective reuse.
  • Prompting results: Human demonstrations reproduce the fine-tuned-teacher prompting pattern: they help on some tasks but do not unlock capabilities outside AdA’s training distribution.AdA approached optimal performance on Navigation: hold up high and Object permanence: yellow cube, while failing on Navigation: find the cube with teaser, Spacer tool, and Transform to transport.
Loading 2301.07608v1…