Research map / Reinforcement learning / Visual and representation RL
AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers
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.
Converts the actor and critic losses of a Transformer-based off-policy meta-RL agent (AMAGO) into scale-independent classification terms (two-hot critic, filtered IL actor) to fix multi-task optimization imbalance without task labels.
Problem
Multi-task RL with memory-based meta-RL agents is bottlenecked by imbalanced training losses caused by uneven return scales across tasks, and existing fixes require task labels.
Approach
Replace scalar Q regression with two-hot classification over symlog-spaced bins and replace the Q-dependent actor objective with a binary-filtered imitation (CRR-style) loss, so both losses are insensitive to the absolute scale of returns and no task labels are needed.
Key results
- On Meta-World ML45 with a 100M timestep budget, the Ind./Ind. (two-hot critic, filtered IL actor) variant more than doubles the success rate of the original RL2-PPO result at the same sample limit.
- On Multi-Game Atari without reward clipping, the scale-resistant update gives dramatic improvements in 8 of 10 games.
- Scale-resistant updates outperform both PPO-GRU on multi-task POPGym and the Q-dependent update on BabyAI, while the Q-dependent update performs best in only two games with outlier random-initialization returns.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace Q-dependent actor and critic objectives with scale-independent classification updates (ablation) | Transformer actor-critic · Meta-World ML45 | success rate | Helped |
| Replace Q regression with two-hot classification (ablation) | Transformer actor-critic · Meta-World ML45 | success rate | Helped |
| Use filtered imitation learning for the policy update instead of the Q-dependent actor update (ablation) | Transformer actor-critic · Meta-World ML45 | performance and sample efficiency | Helped |
| Use scale-independent actor and critic updates instead of Q-dependent updates (ablation) | Transformer actor-critic · Meta-World ML45, 400M timesteps | success rate relative to RL²: 1.4× improvement | Helped |
| Use scale-independent Transformer updates instead of PPO-GRU | Multi-task POPGym agent · Multi-Task POPGym, one-shot evaluation | normalized return | Helped |
| Use scale-independent rather than Q-dependent Transformer updates (ablation) | Multi-task POPGym agent · Multi-Task POPGym, one-shot evaluation | normalized return | No clear effect |
| Use scale-independent actor and critic updates rather than Q-dependent updates (ablation) | Multi-task Transformer agent · Procgen Easy, five games | performance on unseen test levels | Helped |
| Rescale Coinrun rewards by ×100 and Climber rewards by dividing by 10 | Multi-task Procgen agents · Procgen Easy, five games | performance across games: ×100; by 10 | Mixed |
| Use Ind. / Ind. updates instead of Q-dependent updates under rescaled rewards (ablation) | Multi-task Transformer agent · Procgen Easy, five games with rescaled rewards | performance across games | Helped |
| Use scale-resistant updates instead of Q-dependent updates (ablation) | Multi-game Atari agent · Multi-Game Atari, 10 games without reward clipping | game performance: 8/10 games | Mixed |
| Use scale-independent updates instead of scalar regression (ablation) | Multi-task Transformer agent · BabyAI | multi-task performance | Helped |
| Use Ind. / Ind. updates for online adaptation across attempts | BabyAI Transformer agent · BabyAI, unseen tasks and/or layouts | return | Mixed |
Surprising findings
- Two-hot classification of the critic is the most important change on Meta-World ML45, and the actor change adds further gains.
- Coinrun-rescaled Procgen experiments show the Q-dependent update collapses in the other games when reward scale is changed, while the Ind./Ind. update is much less affected.
Limitations
- Bin count, limits and symlog settings trade-offs are underexplored; the method relies on hyperparameter choices for classification resolution.
- Experiments use a fixed, small number of tasks (e.g., 45, 27, 16, 10, 68) and two-attempt adaptation windows rather than very large task sets.
- Filtered IL actor is best suited to offline-style learning and may lose optimistic exploration relative to the Q-dependent actor.
Experimental setup
Datasets: Meta-World ML45 POPGym Procgen Atari (ALE) BabyAI (Minigrid)
Benchmarks: Meta-World ML45 Multi-Task POPGym Multi-Game Procgen Multi-Game Atari Multi-Task BabyAI
Metrics: success rate skill coverage return normalized to single-task experts human-normalized score episode return
Baselines: RL2-PPO MuZero HTrMRL PPO-GRU Multi-Task PopArt AMAGO
Builds on: AMAGO PopArt CRR C51 DreamerV3
Closest papers by meaning
- MetaCARD: Meta-Reinforcement Learning with Task Uncertainty Feedback via Decoupled Context-Aware Reward and Dynamics Components (AAAI 2024 · similarity 0.91)
- Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement (NeurIPS 2024 · similarity 0.91)
- Meta-Reinforcement Learning with Universal Policy Adaptation: Provable Near-Optimality under All-task Optimum Comparator (NeurIPS 2024 · similarity 0.91)
- Meta-World+: An Improved, Standardized, RL Benchmark (NeurIPS 2025 · similarity 0.91)
- Scrutinize What We Ignore: Reining In Task Representation Shift Of Context-Based Offline Meta Reinforcement Learning (ICLR 2025 · similarity 0.90)
- Meta-Reinforcement Learning via Exploratory Task Clustering (AAAI 2024 · similarity 0.90)
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.