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

TL;DR

This paper introduces AMAGO-2, a Transformer-based meta-reinforcement learning method that addresses multi-task learning challenges by balancing training losses, enabling scalable multi-task adaptation without explicit task labels.

Contribution

It proposes a novel loss balancing technique for multi-task RL using classification objectives, improving scalability and performance in diverse environments.

Findings

Significant improvements in multi-task adaptation across multiple benchmarks.

Effective decoupling of optimization from return scales.

Enhanced scalability without requiring explicit task labels.

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.