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

Thank you for the review and for the constructive comments on limitations and writing. We will discuss your suggestions below.

Framing issue: The considered setting is interesting and a fine target for research. However, parametric variation is not the grand challenge in meta-RL and the paper, in my opinion, does not do enough to clearly delineate the subproblem it tackles.

Yes, we fully agree that the goal should be adaptation to new tasks. Our perspective is that non-parametric variation will be possible if we can scale up our training set to keep that variation within distribution. This will likely take far more tasks than are available in benchmarks like Meta-World, and we generally think that attempts to measure or improve non-parametric variation on Meta-World, Atari, Procgen, and similar benchmarks can risk drawing conclusions from too small of a training set. The challenge is that we still struggle to overfit these small task sets and mostly leave this issue to be tackled by MTRL techniques where solutions might not necessarily be transferable to larger benchmarks. A first step might be to address the key MTRL challenges while maintaining the label-free meta-RL perspective. We will add more discussion to the paragraph on lines 42-55 of the introduction.

The test set of Meta-World ML45 benchmark is mentioned in the experiments, and reasonably ignored. However, since ML45 is already mentioned in the abstract, it would be better to mention that the paper only considers the training set.

Yes, we can do that. We are not aware of any results that demonstrate strong transfer to the test set without introducing some restriction on the environment or policy to regularize training or without adding additional knowledge like language, goal states, or demonstrations.

The reviewer is clearly familiar with Meta-World, but for any other readers who may not be, the train set of 45 tasks can still serve as a reasonable test of parametric variation because the environments are carefully designed to make these changes unobservable.

Minor framing issue: Line 27 says meta-RL is about task identification at its core … The paper adopts a broader viewpoint in the scaling beyond multi-task barrier paragraph. It would be better to update the introduction to match this perspective.

We will rephrase to align this with the more detailed explanation given in Section 2. This line was meant to be a fast summary of the meta-RL perspective where adaptation reduces to exploring the environment to infer missing details of the unobserved parameters of the current task, which lets us treat multi-task RL problems as a special case.

This is out of scope for the review, but I'm curious, what do you think is the best bet for a meta-RL benchmark that would enable investigating the non-parametric generalization ml45 fails to provide?

Meta-RL is facing a difficult benchmarking challenge at the moment, and it’s the reason our experiments need to drift into long-term memory and multi-task domains. A more affordable subset of XLand 2 [1] would probably be the best large-scale benchmark if it were accessible. There have been some recent efforts to create an open-source grid world version of this domain. One interesting direction for large-scale experiments (that is open-source) could be Gym Retro, which supports thousands of ALE-like games. Arcade games and platformers are a natural fit because they often give the player multiple lives/attempts before a reset. Meta-learning was the original goal of Retro, but this is obviously challenging and there hasn’t been much progress made.

[1] Human Timescale Adaptation in an Open-Ended Task Space, Adaptive Agent Team, 2023

Thank you for the review and comments on the experimental results. The evaluation domains in our work are computationally expensive, so it was not possible for us to fully address your questions with finalized results during the rebuttal window. We have begun work on these experiments and will update the paper when the results are complete.

For Multi-Game Procgen experimental analysis, especially the result in Figure 6 right, I am not sure that I understood correctly.

The actor loss is performing supervised learning on a filtered subset of (obs, action) pairs from the replay buffer that are estimated to have a positive advantage by the critics. This update would be equivalent to imitation learning if the filter were to give a weight of one to every example, and we are relying on an accurate critic to mask low quality actions from the loss. Figure 6 right measures the percentage of each batch that passes the filter (A(s, a) > 0) and is used to compute the supervised loss. When the critic network is initialized it approves approximately half of the actions. As training continues, the curve changes according to a combination of the rate of experience collection, environment difficulty, gradient steps per environment step, and the total size of the replay buffer.

For Multi-game Atari, can you compare this results with the single task agent results such as DreamerV2 [1] or state-of-the-art results on this environment? I want to compare the relative performance of your model with the single task agent like shown in Figure 5.

Atari is a highly optimized domain and we are not claiming state-of-the-art results here by any means. Our goal with this experiment is to use Atari as an example of a standard multi-task setting without the multi-episode adaptation or partial observability of the preceding results. We are measuring whether the change in training objective is enough to make progress without the established Atari trick of clipping rewards in [-1, 1] and multi-task RL techniques that rely on per-task gradients or network outputs. We want to measure a direct ablation of the same basic method that does not add orthogonal details that have led to much of the sample efficiency gains on the ALE. We will work to aggregate the results of single-task versions of the method using the default scale-dependent update (in green) for additional context.

For results in Figure 7, can you test it with diverse length of context? I am interested what happens if very large number of context is given.

The trade-off between context length and performance depends on the episode length of the environment. The high frame rate of pixel-based domains can create much longer episodes than typical toy meta-RL experiments. Ideally, the context would span the entire length of the 2-episode rollout, but computational constraints (GPU memory) prevent this. It is worth noting that although we are not using as long of a context as we would like, 768 images is still an unusually long input sequence when training online RL from scratch.

their investigation can be seen as being limited in the homogenous environment setting. For example, extending it to train a single agent for multiple environments where action spaces are different as did in [1] … I didn't think it is a critical limitation of their study.

Yes we are padding the observation/action spaces to the same shape across environments when necessary. A variable-length token encoding like Gato would probably be the first step in attempting this kind of heterogeneous setting. We avoided this because it adds significant computational cost in terms of sequence length per timestep of policy memory, and we felt that the main point of return scale variance is best demonstrated in domains where the observation and actions are similar but the reward functions are different.

I didn't understand clearly the Gym Retro results. Why did you do this experiment?

This experiment is meant to be an extension of the Atari setting where rewards are on very different scales. The difference between the games we’ve chosen here and Atari is that the Mario levels are similar enough where some kind of zero-shot transfer might be realistic. The SuperMario names refer to the titles of the Mario video game series that are supported by the Gym Retro environment. We are working to add more context to these results by including the other filtered imitation learning ablation (“Ind. Actor / Dep. Critic” - the red curve in Figure 3) as another reference. The "Dep. Actor / Dep. Critic" comparison is not possible here due to the action space of the environment. These runs began before the reviews were released but are expensive to collect and have not finished. They will be added to Figure 9.

Thank you the authors for their rebuttal. They substantially addressed my concerns, so I maintain my score.

Thank you for your review. We will try to address your questions and would be happy to continue the discussion.

By converting the actor and critic objectives to classification terms, the method effectively decouples optimization from return scales, which is a novel and practical approach.

To clarify, the classification loss terms are an implementation detail that appears in recent work we have cited and discussed in Sections 2 and 3. We do not want to claim the objectives themselves are novel. Instead, we show that while these ideas may have various motivations and minor benefits in other settings, they provide a major improvement to a key problem in the area of meta-learning and generalization without task labels.

The paper focuses heavily on empirical results but lacks a thorough theoretical analysis of why the proposed method outperforms existing approaches.

We would be open to suggestions about how we can expand our analysis. Prior works have studied the challenges of multi-task optimization (Section 2). We do not think our work adds to the understanding of this issue, but it does help us remove an extra challenge where RL methods unintentionally place the objective of each task on a different scale (Section 3). We have added a more thorough explanation of this effect in a new figure attached to the rebuttal material.

While the method is validated on several benchmarks, the paper would benefit from a broader analysis of its generalization capabilities … For example, meta-world to multi-game procgen.

We agree that transfer between domains like meta-world → multi-game procgen would be very interesting. This kind of transfer between different observation/action spaces would require a specialized input/output format for multi-modal trajectories of different shapes. Gato [1] is a good example for how we might go about this, but their method stretches input sequence lengths in a way that makes it expensive to evaluate memory.

Although the paper discusses that return scales across tasks can vary significantly … A more detailed discussion on how the method can be adapted or modified for tasks with highly diverse return scales would be valuable.

In our experience, the extreme limits of value classification can be managed by globally rescaling all the rewards to a more stable range. The more challenging issue is the way variations in task difficulty drive uneven changes in values even when the initial and optimal returns are fairly well-bounded. For example, Multi-task POPGym and Meta-World ML45 both have bounded returns, but the benefits of the classification-style loss are much more clear on ML45 than POPGym. Reducing the number of classification bins may help improve performance in tasks where returns rapidly shift from the lower bound to the upper bound. We have added a demonstration of this on a POPGym task in the extra rebuttal material.

What are the computational requirements and scalability of the proposed method when applied to very large-scale environments and datasets? Are there any practical limitations or considerations?

The learning update itself scales with the forward/backward of a single sequence model policy. It is comparable to imitation learning because the relative cost of the extra critic term decreases as the size of the shared Transformer increases. It does not directly depend on the task count or dataset size (replay buffer). The classification-style critic loss can be slower than standard regression as we have increased the dimension of the critic networks’ outputs to B bins. Relative to typical online RL experiments, we are evaluating large policies (10-20M+ parameters) and large datasets (20M+ timesteps). All the experiments were conducted on NVIDIA A5000 GPUs. The model is trained on one GPU while an extra GPU can be used for data collection to reduce wall-clock time.

Please see our global reply for a discussion of limitations.

a more detailed discussion on potential negative societal impacts and how to mitigate them would be beneficial

Sequence-based RL on off-policy data is a flexible approach to memory, adaptation, and generalization settings. The aim of our work is to make this technique easier to use by making more complex multi-task benchmarks realistic for research. We do not think our method adds negative societal impacts beyond the existing risks of RL systems, but we would be happy to discuss this further if you have concerns.

[1] A Generalist Agent, Reed et al., 2022

Thank you for your review. We will discuss your concerns below.

Multi-task RL using in-context adaptation without task labels is not new. Some recent works [1,2] learn in-context meta-RL policies that adapt to diverse tasks

We are focused on end-to-end online RL while Raparthy et al. belongs to a different category of in-context imitation learning methods that is cited on line 90 of our paper.

To expand on this discussion, there are two main differences between (supervised) “in-context” imitation learning (ICIL) and context-based meta-RL.

1. ICIL typically relies on a dataset of expert demonstrations collected by single-task RL agents. This removes the opportunity for positive transfer between tasks during data collection. Many examples of large-scale generalist policy methods distill single-task experts into a multi-task policy, which highlights the difficulty of online multi-task RL even after many years of study.

2. ICIL approaches in-context learning much like few-shot prompting in language models, where we select a prompt demonstration and ask the model to continue this behavior over the following inputs. Meta-RL agents begin from a blank sequence and collect examples that will maximize the overall return. IL on aggregated single-task datasets can make it challenging to explore a new environment by creating a mismatch between the task prior of the multi-task policy and the single-task policies it is trying to imitate [1]. In other words, the multi-task policy may be imitating behavior that is not adaptive, but it lets us generalize from demonstrations.

AMAGO is not a multi-task method, which we will discuss below.

[1] Offline Meta Reinforcement Learning - Identifiability Challenges and Effective Data Collection Strategies, Dorfman et al., 2021

In the setting of online RL, task labels are actually accessible in the simulator … Therefore, balancing multi-task losses without task labels is not necessary.

This perspective treats current benchmarks as an end goal instead of a step towards a more general method. The ultimate goal of adaptive RL is to automatically generalize to very large numbers of tasks. Methods that rely on task labels scale with the number of tasks (gradient editing, separate task heads, etc.). So while it’s true that we can always find the ground-truth task label in the gym environment of common benchmarks like ML45 and Atari, our work ignores this information in hopes of researching a more general method that can extend past popular benchmarks. We do not want to rely on knowing the identity of a new task at test-time, because this removes chances of generalization to unseen tasks. And we would like to scale to more open-ended domains where the total task set is so diverse it is unclear how to one-hot label / count tasks at train-time. Meta-learning without task labels is capable of these goals in theory while multi-task RL with labels is not.

This argument is covered in the main text, including in lines 52-55, 135-140, and 255-261, but will add additional discussion in Section 2.

AMAGO [2], which proposes the multi-task RL framework using context-based Transformers and already achieves great results on these benchmarks, the contribution of this paper … is incremental.

Following the terminology in this paper, AMAGO is not a multi-task method. It learns from many procedurally generated versions of a single task, and uses its Transformer to identify and adapt to those subtle variations. AMAGO evaluates on single-task POPGym, single-task Meta-World (ML1) and so on. It does not attempt ML45 or the multi-task POPGym setting introduced here. The green curves in our figures (“Dep. / Dep.”) represent an AMAGO-style optimization step. We do not claim our method is fundamentally new, but we make an observation that an increasingly common implementation detail in recent work has a simple justification in an online multi-task setting, which unlocks the (very non-incremental) performance improvement in our results.

the paper lacks comparisons with many baseline methods. The experiments only present ablation studies of the proposed classification loss.

The high-level approach of memory-based meta-RL is simple but does not allow for many changes, so most of the differences between methods in this area are implementation-specific. We are evaluating changes in training objectives and ensure a fair ablation of those changes by holding all other details fixed. Our experiments cover a variety of domains, model sizes, and sequence lengths. We provide external reference scores for context where possible (Figures 3, 4, and 5), but some of these settings are rarely (if ever) attempted. Tuning hyperparameters to measure state-of-the-art performance across different conceptual approaches on these benchmarks is not the goal of our work.