AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers

TL;DR

This work revisits the idea that multi-task RL is bottlenecked by imbalanced training losses created by uneven return scales across different tasks and evaluates a simple yet scalable solution where both an agent's actor and critic objectives are converted to classification terms that decouple optimization from the current scale of returns.

Abstract

Language models trained on diverse datasets unlock generalization by in-context learning. Reinforcement Learning (RL) policies can achieve a similar effect by meta-learning within the memory of a sequence model. However, meta-RL research primarily focuses on adapting to minor variations of a single task. It is difficult to scale towards more general behavior without confronting challenges in multi-task optimization, and few solutions are compatible with meta-RL's goal of learning from large training sets of unlabeled tasks. To address this challenge, we revisit the idea that multi-task RL is bottlenecked by imbalanced training losses created by uneven return scales across different tasks. We build upon recent advancements in Transformer-based (in-context) meta-RL and evaluate a simple yet scalable solution where both an agent's actor and critic objectives are converted to classification terms that decouple optimization from the current scale of returns. Large-scale comparisons in Meta-World ML45, Multi-Game Procgen, Multi-Task POPGym, Multi-Game Atari, and BabyAI find that this design unlocks significant progress in online multi-task adaptation and memory problems without explicit task labels.

Figures (24)

• Figure 1: Task Spaces in RL Generalization. Meta-RL agentso adapt to dense variations of a core task. Multi-Task RL overcomes optimization challenges of learning from isolated tasks. Scalable ideas from both areas allow us to extend adaptive agents towards increasingly general behavior.
• Figure 2: Transformer-based Actor-Critic Architecture.
• Figure 3: Scale-Resistant Value Regression. We plot the value of the standard critic loss (Eq. \ref{['eq:critic_dep']}) as a function of the relative prediction error of the TD target ($y$) across four orders of magnitude (left). Y-axes are self-normalized according to the largest displayed value. Two-Hot classification (Eq. \ref{['eq:actor_ind']}) maps the same relative error to similar loss values across the different absolute return scales of each task (right).
• Figure 4: Meta-World ML45 Train Task Results.(Left) Coverage of the $45$ manipulation skills measured by an adaptation horizon success rate $\geq 2/3$. (Center) Average success rate over tasks, variants, and $3$-episode rollouts. Reference scores for MuZero and RL$^2$-PPO are gathered from results in yu2020metaanand2022procedural. (Right) Total return over a $3$-episode adaptation horizon, averaged across tasks and variants. All error bars indicate the maximum and minimum metric across three random trials.
• Figure 5: Multi-Task POPGym Results. Returns are normalized by single-task experts trained for $15$M timesteps. Error bars indicate the maximum and minimum returns across three trials.