[Paper Review] Human-Timescale Adaptation in an Open-Ended Task Space

Adaptive Agent Team, Jakob Bauer+26 more|arXiv (Cornell University)|Jan 18, 2023

Reinforcement Learning in Robotics22 citations

TL;DR

The paper presents Adaptive Agent (AdA), a large-scale memory-augmented RL agent trained with meta-RL in an open-ended 3D task space, achieving human-timescale adaptation and zero-/few-shot generalization, including first-person demonstrations.

ABSTRACT

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.

Motivation & Objective

• Investigate whether scale-trained RL agents can exhibit human-timescale, in-context adaptation in a vast open-ended task space.
• Develop AdA, an agent with memory-based meta-RL that can adapt from few test-time trials in unseen tasks.
• Explore the roles of memory architecture, auto-curriculum, and task distribution richness in enabling fast adaptation.
• Demonstrate zero-shot prompting via demonstrations and assess adaptation in single- and multi-agent settings.

Proposed method

• Train a large Transformer-based agent with meta-RL on XLand 2.0 tasks sampled from a frontier-rich pool.
• Use auto-curriculum (no-op filtering and PLR) to select tasks at the frontier of the agent's capabilities.
• Employ memory architectures (Transformer-XL and RNN with attention) to enable rapid adaptation across trials.
• Optimize with Muesli RL algorithm plus a distillation (kickstarting) loss from a teacher model to scale training.
• Evaluate zero-shot and few-shot adaptation on held-out test and hand-authored probe tasks, including multi-agent scenarios.

Experimental results

Research questions

• RQ1 Can a large-scale RL agent trained on a vast open-ended task distribution adapt to unseen tasks within minutes of test-time experience?
• RQ2 What is the impact of memory architecture, curriculum, and task distribution on rapid adaptation performance?
• RQ3 Does the agent exhibit human-timescale adaptation in single- and multi-agent settings?
• RQ4 Can zero-shot prompts via first-person demonstrations further improve adaptation?
• RQ5 How do scaling factors (model size, memory length) affect adaptation performance?

Key findings

• AdA demonstrates human-timescale adaptation, improving performance with more test trials across most held-out tasks.
• Transformer-XL memory yields the best adaptation performance, outperforming RNN-based memories.
• Auto-curricula (no-op filtering and PLR) improve zero-shot generalization and few-shot adaptation over uniform task sampling.
• Zero-shot prompting with first-person demonstrations can further boost performance.
• Scaling model size and memory length yields growth in adaptation capabilities, with higher gains as more trials are provided.
• Emergent cooperative behaviors and division of labor arise in multi-agent tasks during test-time adaptation.