Paper: AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents (ICLR-2024)

Reviewer comment (Presentation):
This paper is perhaps trying to do too much in the limited space to the detriment of the clarity of its most important contributions. Given that it is in the method name, goal-conditioned RL seems to be an important part of the contribution. However, this paper would be much easier to understand if it focused on the meta-RL part, where it already has strong results, and upon which its success in goal-conditioned RL builds.

Author response:
> "goal-conditioned RL seems to be an important part of the contribution. However, this paper would be much easier to understand if it focused on the meta-RL part"

With this suggestion in mind our revision has prioritized the first experiment section (meta-RL and long-term memory) when dealing with space constraints. We have moved one of the goal-conditioned section’s figures to the Appendix.

Reviewer comment (Reproducibility):
The proposed agent uses a new combination of existing actor critic losses to construct its update function, yet the main paper does not contain almost any detail about how the objective is motivated or what the objective even is. Contrast this, for example, to the DreamerV3 paper (Mastering Diverse Domains through World Models, Hafner et al. 2023), which tackles a similar problem setting, where the majority of the paper is describing the method and motivating the design choices. Even after reading the relevant parts of the appendix, it is not clear to me how exactly is the critic trained.

Author response:
> "the main paper does not contain almost any detail about how the objective is motivated or what the objective even is. Contrast this, for example, to the DreamerV3 paper …, where the majority of the paper is describing the method … "

We have revised the Method section to help clarify the motivation of our learning update. More minor RL engineering details that are not discussed in the main text or Appendix A are much easier explained in code/documentation, and we assure you that our agent is fully open-source. To be fair, DreamerV3 is building upon a series of implementation-focused papers and it is quite the outlier in terms of how much low-level detail is explained in the text.

Reviewer comment (Experiments):
The architecture design is described as being based on "clear pattern of network activations across hundreds of trials" but this pattern is not described anywhere.

Author response:
> "The architecture design is described as being based on "clear pattern of network activations across hundreds of trials" but this pattern is not described anywhere."

Yes, this was too vague of an explanation, and it has now been updated in the main text and expanded in the Appendix alongside a new figure (Fig. 14).

Reviewer comment (Experiments):
Popart is proposed to help keep the different losses on the same scale across different environments, but I didn't find any attempt to validate that it is indeed doing so or helping in any other way.

Author response:
> "Popart is proposed to help keep the different losses on the same scale across different environments, but I didn't find any attempt to validate that it is indeed doing so or helping in any other way."

Appendix A.2 Equation 1 shows the weighted combination of three terms used to create AMAGO's shared loss. The actor and critic terms scale with the learned $Q$ value, but the filtered BC term does not. This means that the relative weight of the three terms is based on the scale of rewards in any given environment and would require constant tuning. Popart creates a simple solution by normalizing $Q$ during loss computation. The loss scales then converge to predictable values and we validate this by using the exact same weights in every experiment. If you think it would be helpful we can demonstrate this by directly scaling rewards in the same environment (this is how the idea was first tested). Popart saves time, compute, and frustration when evaluating many environments but does not meaningfully change the method if we only care about a specific environment. We are *not* using Popart to independently normalize the returns of multiple tasks simultaneously (like how it is used in Atari) because the size of our task distribution is assumed to be unknown and large.

Reviewer comment (Experiments):
A complicated exploration schedule is proposed and used, but at the same time the paper says it is not essential.

Author response:
> "A complicated exploration schedule is proposed and used, but at the same time the paper says it is not essential."

The correct amount of exploration noise does not usually depend on the current timestep. However, in meta-RL, each rollout creates its own exploration-exploitation tradeoff as we try to identify the current environment. We feel the annealed schedule is a stronger default here, especially because it can be configured to be the same as regular $\epsilon$-greedy if necessary. However, tuning this schedule on each environment goes against our goal of creating a baseline that is easy to use. We did not prioritize this and leave it as an Appendix detail. We pick a conservative default schedule randomized over parallel actors so that performance in tasks with easy exploration can be mostly unaffected. This schedule would likely be more important if our experiments were more concerned with sample efficiency than stability and final performance.

Reviewer comment (Evaluation):
The paper explicitly claims state-of-the-art results on PoPGym but refrains from doing so for Crafter, where it also seems to be beating the previous state-of-the-art (Dreamerv3). Why?

Author response:
> "The paper explicitly claims state-of-the-art results on PoPGym but refrains from doing so for Crafter …. Why?"

The difference is that we match the evaluation protocol of existing POPGym baselines as closely as possible but are not doing the same for Crafter. Instead, we are using the Crafter world generation and achievement system to create our own evaluation setting for goal-conditioned agents in an open-world domain.

Reviewer comment (Writing):
Have I understood it correctly that the "dense" Crafter is the vanilla Crafter setting?

Author response:
> "Have I understood it correctly that the "dense" Crafter is the vanilla Crafter setting?"

Yes you are right, and the term “dense” in Appendix Table 1 comes from the Crafter rewards having some shaping aside from the binary achievements. However, this is probably a poor word choice on our part, as these rewards are still sparse by most standards. The text uses the term “undirected,” and we have now updated the table to match. The most important difference is that Crafter usually rewards the agent for completing any achievement, while our version rewards completing a specific task that is provided.

Reviewer comment (Reproducibility):
What is the actor objective used by the method?

Author response:
> "What is the actor objective used by the method?"

It is a mixture of the standard $-Q(s, \pi(s))$ taking the min over 2 critics from an ensemble like TD3/REDQ and filtered behavior cloning. The behavior cloning term is motivated by the same idea as multi-gamma training, where we are trying to prevent learning from becoming unstable or uninformative when the Q-surface is flat. Hopefully the method section writing changes have helped clarify this.

Reviewer comment (Reproducibility):
What is the critic objective used by the method?

Author response:
> "What is the critic objective used by the method?"

It is standard one-step temporal difference error implemented with an ensemble of Q-networks to reduce overestimation like in REDQ. We use REDQ to reduce the risk of picking an update-to-data ratio that is too high, which can happen early in training because we are replaying entire rollouts in every update. It is not uncommon for the agent to perform gradient updates on the entire replay buffer several times between rollouts. There are many open questions here (Appendix D - final paragraph).

Reviewer comment (Writing):
What is "the clear pattern of network activations across hundreds of trials"?

Author response:
> "What is "the clear pattern of network activations across hundreds of trials"?"

Hopefully we have addressed this question in Part 1.

Reviewer comment (Novelty):
This is very subjective so I'm not listing this as a weakness, but I find the "in-context learning" label to make it less clear what the paper is about. On an abstract level, the method is a variant of well-known black-box meta-RL methods like RL2. Black-box meta-RL methods have variously been grouped under other names like context-based meta-RL, memory-based meta-RL, zero-shot meta-RL, etc. Given that this is a long ongoing research direction in the RL community, with an already confusing diversity of labels, it's not clear to me that emphasizing a new label makes this easier to understand. Granted that for non-RL audiences "in-context learning" may be more intuitive than "black-box meta-RL". ADA uses the term similarly, but when they get to the technical details, they switch to talking about black-box meta-RL.

Author response:
> "I find the "in-context learning" label to make it less clear what the paper is about. On an abstract level, the method is a variant of well-known black-box meta-RL methods like RL2. Black-box meta-RL methods have variously been grouped under other names like context-based meta-RL, memory-based meta-RL, zero-shot meta-RL, etc. …, it's not clear to me that emphasizing a new label makes this easier to understand. Granted that for non-RL audiences "in-context learning" may be more intuitive than "black-box meta-RL". ADA uses the term similarly, but when they get to the technical details, they switch to talking about black-box meta-RL."

We completely agree that the terminology is confusing and thought a lot about how we should handle this. There are several names for the method and even several names for the environment type and problem statement. Ultimately, a key advantage of this kind of black-box method is that it works for meta-RL *and* zero-shot generalization *and* long-term memory while other meta-RL methods may not. This term feels too restrictive for our work because our experiments emphasize that flexibility more than most meta-RL papers. We decided to call the environments CMDPs for a similar reason (that term is more common in generalization than meta-RL). So, if we are not going to use "meta-RL",  it seems reasonable to use "in-context RL" because "in-context learning" has become a common term for Transformers outside of RL as you note. And at least we are not adding to the problem by creating yet another name because AdA uses the same term. However, we do make a point of using a more technical term in the related work section for those with a meta-RL background: "Meta-RL is a crowded field with a complex taxonomy; this paper uses the informal term "in-context RL" to refer to a specific subset of implicit context-based methods that treat meta-learning, zero-shot generalization, and partial observability as a single problem" followed by many meta-RL references.

Reviewer comment (Reproducibility):
This method still requires server-grade GPUs to run, limiting this approach to well-off labs or industry.

Author response:
AMAGO is more GPU-dependent than many single-task baselines (for example), but we will note that every experiment is performed on a single GPU. We inherit the same GPU memory bottleneck as most other work on Transformers, but there are many hyperparameter choices that can help address this problem (batch size, context length, model size, …).

Reviewer comment (Experiments):
For the POPGym benchmark, they are comparing their off-policy method to an on-policy method given the same number of sampled timesteps. This is not a fair comparison, but they do label the baselines as on-policy.

Author response:
We agree that off-policy methods have an advantage in any sample-limited benchmark, but this can be considered a strength of the method. The sample limit in POPGym is a bit arbitrary as it was established by on-policy baselines that did not solve many of the tasks. The [POPGym paper](https://openreview.net/pdf?id=chDrutUTs0K) Figures 10-14 show that learning curves on many environments have not improved or converged long before the 15M sample limit. One positive sign of our work is that performance is still increasing on many of the harder environments (Appendix C.1.). This creates a convenient opportunity for research because now that we have a way to reach strong performance at convergence we can start experimenting with how to reach those levels faster.

Reviewer comment (Presentation):
They do not list any shortcomings.

Author response:
One shortcoming of the current method is that the goal relabeling scheme only applies to sparse/binary rewards for goal completion. However, we believe this could be expanded to arbitrary reward functions. Many of the shortcomings of (value-based) RL still apply. For example, long-term credit assignment can still be challenging when the discount factor is not the only underlying problem. We are also beginning to saturate the memory lengths of standard benchmarks, as the sequence lengths used in our experiments are often longer than successful solutions to the same/similar problems in prior work. A key motivation of AMAGO was to create an open-source in-context learner that can help create new benchmarks in this area.

Reviewer comment (Evaluation):
I would like the authors to list any shortcomings of their approach, especially because this paper is mostly empirical.

Author response:
One shortcoming of the current method is that the goal relabeling scheme only applies to sparse/binary rewards for goal completion. However, we believe this could be expanded to arbitrary reward functions. Many of the shortcomings of (value-based) RL still apply. For example, long-term credit assignment can still be challenging when the discount factor is not the only underlying problem. We are also beginning to saturate the memory lengths of standard benchmarks, as the sequence lengths used in our experiments are often longer than successful solutions to the same/similar problems in prior work. A key motivation of AMAGO was to create an open-source in-context learner that can help create new benchmarks in this area.

Reviewer comment (Writing):
It is a bit hard to keep track of what exactly the different components of the proposed method are. I think it would greatly improve the paper's clarity if the authors could include a table outlining exactly what the contribution is (e.g. Transformer architecture change, multi-gamma update, relabeling, etc.). Generally, I think the main improvement the authors can make on this paper is clarity.

Author response:
We have updated the Transformer architecture change and multi-gamma update sections of the Method section, which we hope will improve their clarity.

Reviewer comment (Reproducibility):
Similarly, an algorithm box for the relabeling scheme would be nice for clarity. I know the authors say they defer details of the relabeling to the open source code release, but I think it would be good to have it in the appendix at least.

Author response:
We have taken your suggestion and added an algorithm box in Appendix B, along with an expanded discussion of the method. The relabeling scheme extends Hindsight Experience Replay (HER)  to multiple steps where we select sub-sequences of alternative goals. The main details are 1) how hindsight goals are selected and 2) how they are merged into the original sequence of real goals. The multi-step version gives more room for customization than HER; for example, it creates a spectrum between relabeling every trajectory to be a complete success and not relabeling at all.

Reviewer comment (Reproducibility):
The authors write that "AMAGO can select the action corresponding to any $\gamma$ at test time". How exactly is this done?

Author response:
During training, our actor and critics networks are outputting values for an array of different $\gamma$ values in parallel. At test-time, we can select the index of the actor’s outputs that corresponds to any value of $\gamma$. This may be simpler than it sounds because computing RL losses over long causal sequences involves lots of parallel computation (over multiple timesteps, multiple critics, etc.), and the multi-gamma method adds one more axis to this sequence format.

Reviewer comment (Experiments):
Have the authors tried experimenting with a smaller set of multi-gamma values for the multi-gamma update rather than using all the values in Table 4? I'm curious if/how much performance would suffer.

Author response:
The multi-gamma technique is efficient, so once we identify the optimization of a single $\gamma$ value as an issue, it is convenient to add many more values. The most important requirement for AMAGO’s technical details is that learning remains stable even when we default to somewhat extreme settings for its hyperparameters – this makes hyperparameter selection less risky and is what enables us to evaluate on a wide range of domains. It is quite possible that the Table 4 settings are well into the diminishing rate of returns of this technique for some environments. We tried running this ablation for this response, but in the interest of time, we picked an environment that was too simple to create a clear demonstration. We will continue these experiments and update the paper with more information.

Reviewer comment (Evaluation):
The AMAGO architecture's complexity, particularly due to the use of transformers, can be computationally intense. This may limit its efficiency in resource-limited situations or possibly make it less suitable for complicated applications.

Author response:
Transformers are more expensive than common RL architectures, but this approach to generalization does require some kind of memory, so “expensive” becomes a relative term. We show in Figure 6 (previously Fig. 4)  that Transformers can be an efficient choice vs RNNs in the common RL setting of small model sizes and long sequence lengths.

Reviewer comment (Evaluation):
Although AMAGO has made improvements in performance, it still has a success rate of 0 in collecting some key resources (e.g. iron), so AMAGO still faces limitations in exploration challenges.

Author response:
We do make some progress on iron collection in Appendix C.5 (Figure 23 and the paragraph titled "Following Tech-Tree Instructions and Expanded Crafter" on page 30) although this may be the experiment you are already referring to. Our work tries to focus more on the memory/generalization problem of RL with sequence models than pure exploration. We agree that more improvements in exploration are needed. However, successful adaptive agents create new ways to drive exploration by transferring knowledge between similar tasks.  The Crafter experiments also highlight helpful advantages of hindsight relabeling.

Reviewer comment (Theory):
Considering AMAGO's ability to work with multi-goals, how well does AMAGO prioritize or sequence goals when given multiple objectives? Is there a mechanism to manage potential conflicts between goals?

Author response:
When we use multi-step goals, the goal becomes an instruction or a sequence of steps, and agents receive rewards for completing goals in the assigned order. There can be an element of planning involved here where an agent could take proactive measures to prepare for later goals before it completes earlier ones. We explore this in Appendix C.5 “Planning and Long-Horizon Tasks” and Appendix C.4.2. The main benefits of multi-step goals are that 1) the user can break down a complex task into multiple steps (like how "make stone pickaxe, collect iron" can be successful when "collect iron" alone is not ), and that 2) the experience generated by attempting multi-step goals drives exploration of new goals with intermediate steps.

Reviewer comment (Experiments):
The author mentioned replacing ReLU/GeLU with Leaky ReLU, adding LayerNorms, and modifying linear layers. It may be better to explain the necessity and rationality behind each modification through a more thorough ablation study.

Author response:
Our revision updates the description of these details in the method section (Page 5). These ideas are included in the POPGym ablation, where they prevent sudden collapse in performance. Our architecture decisions are motivated by observing activations in the Transformer layer before and after collapse and justified by noting that many environments require specific memory patterns that encourage instability. Following your suggestion, we have recreated an example of our early experiments that led to these decisions in a new Appendix Figure 15.