# AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents

Jake Grigsby1, Linxi “Jim” Fan2, Yuke Zhu1
1The University of Texas at Austin 2NVIDIA Research

###### Abstract

We introduce AMAGO, an in-context Reinforcement Learning (RL) agent that uses
sequence models to tackle the challenges of generalization, long-term memory,
and meta-learning. Recent works have shown that off-policy learning can make
in-context RL with recurrent policies viable. Nonetheless, these approaches
require extensive tuning and limit scalability by creating key bottlenecks in
agents’ memory capacity, planning horizon, and model size. AMAGO revisits and
redesigns the off-policy in-context approach to successfully train long-
sequence Transformers over entire rollouts in parallel with end-to-end RL. Our
agent is uniquely scalable and applicable to a wide range of problems. We
demonstrate its strong performance empirically in meta-RL and long-term memory
domains. AMAGO’s focus on sparse rewards and off-policy data also allows in-
context learning to extend to goal-conditioned problems with challenging
exploration. When combined with a novel hindsight relabeling scheme, AMAGO can
solve a previously difficult category of open-world domains, where agents
complete many possible instructions in procedurally generated environments. We
evaluate our agent on three goal-conditioned domains and study how its
individual improvements connect to create a generalist policy.

##  1 Introduction

Reinforcement Learning (RL) research has created effective methods for
training specialist decision-making agents that maximize one objective in a
single environment through trial-and-error [4, 7, 61, 86, 93]. New efforts are
shifting towards developing generalist agents that can adapt to a variety of
environments that require long-term memory and reasoning under uncertainty
[48, 24, 90, 80]. One of the most promising approaches to generalist RL is a
family of techniques that leverage sequence models to actively adapt to new
situations by learning from experience at test time. These “in-context” agents
use memory to recall information from previous timesteps and update their
understanding of the current environment [95, 20]. In-context RL’s advantage
is its simplicity: it reduces partial observability, generalization, and meta-
learning to a single problem in which agents equipped with memory train in a
collection of related environments [5, 105, 29]. The ability to explore,
identify, and adapt to new environments arises implicitly as agents learn to
maximize returns in deployments that may span multiple trials [105, 69, 8].

Effective in-context agents need to be able to scale across two axes: 1) the
length of their planning horizon, and 2) the length of their effective memory.
The earliest versions of this framework [95, 20, 89] use recurrent networks
and on-policy learning updates that do not scale well across either axis.
Transformers [92] improve memory by removing the need to selectively write
information many timesteps in advance and turning recall into a retrieval
problem [78]. Efforts to replace recurrent policies with Transformers were
successful [59, 60], but long-term adaptation remained challenging, and the
in-context approach was relegated to an unstable baseline for other meta-RL
methods [5]. However, it has been shown that pure in-context RL with recurrent
networks can be a competitive baseline when the original on-policy gradient
updates are replaced by techniques from a long line of work in stable off-
policy RL [66, 23]. A clear next step is to reintroduce the memory benefits of
Transformers. Unfortunately, training Transformers over long sequences with
off-policy RL combines two of the most implementation-dependent areas in the
field, and many of the established best practices do not scale (Section 3).
The difficulty of this challenge is perhaps best highlighted by the fact that
many popular applications of Transformers in off-policy (or offline) RL devise
ways to avoid RL altogether and reformulate the problem as supervised learning
[11, 52, 43, 104, 51].

Our work introduces AMAGO (Adaptive Memory Agent for achieving GOals) — a new
learning algorithm that makes two important contributions to in-context RL.
First, we redesign the off-policy actor-critic update from scratch to support
long-sequence Transformers that learn from entire rollouts in parallel (Figure
1). AMAGO breaks off-policy in-context RL’s biggest bottlenecks by enabling
memory lengths, model sizes, and planning horizons that were previously
impractical or unstable. Our agent is open-source111Code is available here:
https://github.com/UT-Austin-RPL/amago and specifically designed to be
efficient, stable, and applicable to new environments with little tuning. We
empirically demonstrate its power and flexibility in existing meta-RL and
memory benchmarks, including state-of-the-art results in the POPGym suite
[62], where Transformers’ recall dramatically improves performance. Our second
contribution takes in-context RL in a new direction with fresh benchmarks.
AMAGO’s use of off-policy data — as well as its focus on long horizons and
sparse rewards — makes it uniquely capable of extending the in-context RL
framework to goal-conditioned problems [44, 77] with hard exploration. We add
a new hindsight relabeling scheme for multi-step goals that generates
effective exploration plans in multi-task domains. AMAGO can then learn to
complete many possible instructions while using its memory to adapt to
unfamiliar environments. We first study AMAGO in two new benchmarks in this
area before applying it to instruction-following tasks in the procedurally
generated worlds of Crafter [33].

Figure 1: In-context RL techniques solve memory and meta-learning problems by
using sequence models to infer the identity of unknown environments from test-
time experience. AMAGO solves core technical challenges to unify the
performance of end-to-end off-policy RL with long-sequence Transformers in
order to push memory and adaptation to new limits.

##  2 Related Work

##### RL Generalization and Meta-Learning.

General RL agents should be able to make decisions across similar environments
with different layouts, physics, and visuals [15, 31, 87, 70, 103]. Many
techniques are based on learning invariant policies that behave similarly
despite these changes [48]. Meta-RL takes a more active approach and studies
methods for rapid adaptation to new RL problems [83]. Meta-RL is a crowded
field with a complex taxonomy [5]; 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. In practice, these methods are variants of RL2 [20, 95],
where we train a sequence model with standard RL and let memory and meta-
reasoning arise implicitly as a byproduct of maximizing returns [8, 60, 89,
66, 69, 23]. There are many other ways to use historical context to adapt to
an environment. However, in-context RL makes so few assumptions that
alternatives mainly serve as ways to take advantage of problems that do not
need its flexibility. Examples include domains where each environment is fully
observed, can be labeled with ground-truth parameters indicating its unique
characteristics, or does not require meta-adaptation within trials [42, 45,
55, 79, 81, 94].

The core AMAGO agent is conceptually similar to AdA [90] — a recent work that
trains meta-learning Transformers with off-policy RL. AdA performs large-scale
experiments in a closed-source domain that does not share our focus on long-
term recall. AMAGO’s technical details are more closely inspired by open-
source work in off-policy actor-critic algorithms with sequence models
(Section 3) [100, 66]. Algorithm Distillation (AD) [51] is another approach to
training off-policy Transformers that reformulates meta-RL as supervised
sequence modeling. AD aligns the rate of in-context adaptation with the sample
efficiency of the “source” single-task agents that make up its expensive
training datasets. In contrast, AMAGO trains Transformers on longer sequences
with end-to-end RL that can directly optimize adaptation over any horizon we
specify.

##### Goal-Conditioned RL.

Goal-conditioned RL (GCRL) [44, 77] trains multi-task agents that generalize
across tasks by making the current goal an input to the policy. GCRL often
learns from sparse (binary) rewards based solely on task completion, reducing
the need for reward engineering. GCRL methods commonly use hindsight
experience replay (HER) [2, 63, 64, 21] to ease the challenges of sparse-
reward exploration. HER randomly relabels low-quality data with alternative
goals that lead to higher rewards, and this technique is mainly compatible
with off-policy methods where the same trajectory can be recycled many times
with different goals. Goal conditioned supervised learning (GCSL) [28, 58, 57,
49, 82, 88] uses HER to relabel every trajectory and directly imitates
behavior conditioned on future outcomes. GCSL has proven to be an effective
way to train Transformers for decision-making [43, 11, 104, 52, 16, 98] but
has a growing number of theoretical and practical issues [22, 73, 9, 99].
AMAGO avoids these limitations by providing a stable and effective way to
train long-sequence Transformers with pure RL.

##  3 Background and Problem Formulation

##### Solving Goal-Conditioned CMDPs with In-Context RL.

Partially observed Markov decision processes (POMDPs) create long-term memory
problems where the true environment state $s_{t}$ may need to be inferred from
a history of limited observations $(o_{0},\dots,o_{t})$ [10]. Contextual MDPs
(CMDPs) [36, 48, 75] extend POMDPs to create a distribution of related
decision-making problems where everything that makes an environment unique can
be identified by a context parameter. More formally, an environment’s context
parameter, $e$, conditions its reward, transition, observation emission
function, and initial state distribution. CMDPs create zero-shot
generalization problems when $e$ is sampled at the end of every episode [48]
and meta-RL problems when held constant for multiple trials [5]. We slightly
adapt the usual CMDP definition by decoupling a goal parameter $g$ from the
environment parameter $e$ [55]. The goal $g$ only contributes to the reward
function $R_{e,g}(s,a)$, and lets us provide a task description without
assuming we know anything else about the environment. Domains that jointly
sample from a distribution of environments and a distribution of observable
goals in that environment create “multi-task, multi-environment”
generalization problems. However, goals can be ignored to recover classic
generalization and meta-learning benchmarks.

Knowledge of $e$ lets a policy adapt to any given CMDP, but its value is not
observed as part of the state (and may be impossible to represent), creating
an “implicit POMDP” [29, 42]. From this perspective, the only meaningful
difference between POMDPs and CMDPs is the need to reason about transition
dynamics ($o_{i},a_{i},o_{i+1}$), rewards ($r$), and resets ($d$) based on
full trajectories of environment interaction
$\tau_{0:t}=(o_{0},a_{0},r_{1},o_{1},a_{1},r_{2},d_{2},\dots,o_{t})$. In-
context RL agents embrace this similarity by treating CMDPs as POMDPs with
additional inputs and using their memory to update state and environment
estimates at every timestep. Memory-equipped policies $\pi$ take trajectory
sequences as input and output actions that maximize returns over a finite
horizon ($H$) that may span multiple trials:
$\pi^{*}=\mathop{\text{argmax}}_{\pi}\mathop{\mathbb{E}}_{p(e,g),a_{t}\sim\pi(\cdot\mid\tau_{0:t},g)}\left[\sum_{t=0}^{H}R_{e,g}(s_{t},a_{t})\right]$.

##### Challenges in Off-Policy Model-Free In-Context RL.

In-context RL is a simple and powerful idea but is often outperformed by more
complex methods [106, 79, 5]. However, deep RL is a field where implementation
details are critical [40], and recent work has highlighted how the right
combination of off-policy techniques can close this performance gap [66]. Off-
policy in-context agents collect training data in a replay buffer from which
they sample trajectory sequences
$\tau_{t-l:t}=(o_{t-l},a_{t-l},r_{t-l+1},\dots,o_{t})$ up to a fixed maximum
context length of $l$. Training on shorter sequences than used at test-time is
possible but creates out-of-distribution behavior [46, 38, 39]. Therefore, an
agent’s context length is its most important bottleneck because it establishes
an upper bound on memory and in-context adaptation. Model-free RL agents
control their planning horizon with a value learning discount factor
$\gamma\in[0,1)$. Despite context-based learning’s focus on long-term
adaptation, learning instability restricts nearly all methods to
$\gamma\leq.99$ and effective planning horizons of around $100$ timesteps.
Trade-offs between context length, planning horizon, and stability create an
unpredictable hyperparameter landscape [66] that requires extensive tuning.

Agents that are compatible with both discrete and continuous actions use
$\tau_{t-l:t}$ to update actor ($\pi$) and critic ($Q$) networks with two loss
functions optimized in alternating steps [54, 14] that make it difficult to
share sequence model parameters [66, 100, 101]. It is also standard to use
multiple critic networks to reduce overestimation [27, 12] and maintain extra
target networks to stabilize optimization [61]. The end result is several
sequence model forward/backward passes per training step — limiting model size
and context length [76, 66] (Figure 12). This problem is compounded by the use
of recurrent networks that process each timestep sequentially [56].
Transformers parallelize sequence computation and improve recall but can be
unstable in general [97, 65], and the differences between supervised learning
and RL create additional challenges. RL training objectives do not always
follow clear scaling laws, and the correct network size can be unpredictable.
Datasets also grow and shift over time as determined by update-to-data ratios
that impact performance but are difficult to tune [25, 17]. Finally, we
typically do not know when training will end and cannot prevent model collapse
with learning rate cooldowns [92, 91]. Taken together, these challenges make
us likely to optimize a large policy network for far longer than our small and
shifting dataset can support [68, 1]. Addressing these problems requires
careful architectural changes [59, 72] and Transformers are known to be an
inconsistent baseline in our setting [60, 51, 90].

##  4 Method

We aim to extend the limits of off-policy in-context RL’s three main barriers:
1) the memory limit or context length $l$, 2) the value discount factor
$\gamma$, and 3) the size and recall of our sequence model. Transformers are a
strong solution to the last challenge and may be able to address all three
barriers if we were able to learn from the maximum context length $l=H$ and
select actions for $\gamma\geq.999$ at test time. In principle, this would
allow agents to remember and plan for long adaptation windows in order to
scale to more challenging problems while removing the need to tune trade-offs
between stability and context length. AMAGO overcomes several challenges to
make this possible. This section summarizes the most important techniques that
enable AMAGO’s performance and flexibility. More details can be found in
Appendix A and in our open-source code release.

A high-level overview of AMAGO is illustrated in Figure 1. The observations of
every environment are augmented to a unified goal-conditioned CMDP format,
even when some of the extra input information is unused. This unified format
allows AMAGO to be equally applicable to generalization, meta-RL, long-term
memory, and multi-task problems without changes. AMAGO is an off-policy method
that takes advantage of large and diverse datasets; trajectory data is loaded
from disk and can be relabeled with alternative goals and rewards. The shape
and format of trajectories vary across domains, but we use a timestep encoder
network to map each timestep to a fixed-size representation. AMAGO allows the
timestep encoder to be the only architecture change across experiments. A
single Transformer trajectory encoder processes sequences of timestep
representations. AMAGO uses these representations as the inputs to small feed-
forward actor and critic networks, which are optimized simultaneously across
every timestep of the causal sequence. Like all implicit context-based methods
(Sec. 2), this architecture encourages the trajectory encoder outputs to be
latent estimates of the current state and environment $(s_{i},e)$ (Sec. 3). In
short, AMAGO’s learning update looks more like supervised sequence modeling
than an off-policy actor-critic, where each training step involves exactly one
forward pass of one Transformer model with two output heads. This end result
is simple and scalable but is only made possible by important technical
details.

##### Sharing One Sequence Model.

The accepted best practice in off-policy RL is to optimize separate sequence
models for the actor and critic(s) — using additional target models to compute
the critic loss. Sharing actor and critic sequence models is possible but has
repeatedly been shown to be unstable [66, 100, 101]. AMAGO addresses this
instability and restructures the learning update to train its actor and
critics on top of the same sequence representation; there are no target
sequence models, and every parameter is updated simultaneously without
alternating steps or conflicting learning rates. We combine the actor and
critic objectives into a single loss trained by one optimizer. However, we
must put each term on a predictable scale to prevent their weights from
needing to be tuned across every experiment [37, 34]. The key detail enabling
the simultaneous update to work in AMAGO (as opposed to [66]) — without
separating the Transformer from the actor loss [101, 100] — is to carefully
detach the critic from the actor loss (Figure 12). We compute loss terms with
a custom variant of REDQ [12] and CRR [96] for discrete and continuous actions
that removes unintuitive hyperparameters wherever possible.

##### Stable Long-Context Transformers in Off-Policy RL.

While AMAGO lets us optimize shared parameters between the actor and critics,
our experiments show that these changes are not enough to ensure the stability
of Transformers at long sequence lengths. Transformers bring additional
optimization questions that can be more challenging in RL than supervised
learning (Sec. 2). The shared off-policy update may even amplify these issues,
as we were unable to successfully apply existing architectural changes for
Transformers in on-policy RL [72, 59] to our setting. A robust solution would
let us pick model sizes that are safely too large for the problem without
worrying that long training runs or the wrong update-to-data ratio will lead
to collapse. Based on a clear pattern of network activations across hundreds
of trials, and ideas from continual RL [1], network plasticity [68, 53], and
sequence modeling [102, 85], we develop a custom Transformer variant that
stabilizes our shared off-policy update (Figure 15).

##### Long Horizons and Multi-Gamma Learning.

While long-sequence Transformers improve recall, adaptation over extended
rollouts creates a challenging credit assignment and planning problem. An
important barrier in increasing adaption horizons is finding stable ways to
increase the discount factor $\gamma$. Time information in the trajectory
improves stability [71], but we also use a “multi-gamma” update that jointly
optimizes many different values of $\gamma$ in parallel. Each value converges
at different rates and improves representation learning of the shared
Transformer. AMAGO can select the action corresponding to any $\gamma$ at test
time. If we train on $\gamma_{N}$ values of $\gamma$, discrete actor and
critics’ output layers become $\gamma_{N}$ times larger. Similarly, the
effective batch size of continuous-action actor/critic feed-forward networks
is multiplied by $\gamma_{N}$. In either case, the relative cost of this
technique becomes low as the size of the shared Transformer increases, and its
empirical benefit is significant.

##### Hindsight Instruction Relabeling.

AMAGO’s use of off-policy data and high discount factors allows us to relabel
long trajectories in hindsight. This ability provides a unique opportunity to
expand the in-context RL framework to include goal-conditioned problems with
sparse rewards that are far too difficult for on-policy agents with short
planning horizons to explore. AMAGO goes a step further and introduces a
variant of HER (Sec. 3) for multi-step goals where $g=(g^{0},\dots,g^{k})$.
Besides adding the flexibility to create more complex multi-stage tasks, our
experiments show that relabeling multi-step goals effectively creates
automatic exploration plans in open-world domains like Crafter [33].

[capbesideposition=right,top,capbesidewidth=.26]figure[]

Figure 2: The agent (top left) navigates a maze to reach goal locations
($g^{0},g^{1},g^{2}$). Yellow, red, and blue paths with locations (A, …, E)
are examples of trajectories and achieved goals for relabeling. We show how
instruction-relabeling creates a variety of alternative reward sequences.

Our relabeling scheme works in two key steps. First, we do not restrict goals
to represent target states and allow them to be an abstract set of tokens or
strings from a closed vocabulary. During rollouts, we can evaluate the goal
tokens we were instructed to achieve and record alternatives that can be used
for relabeling. The second step extends the HER relabeling scheme to
“instructions”, or sequences of up to $k$ subgoals. If an agent is tasked with
$k$ goals but only completes the first $n\leq k$ steps of the instruction, we
sample $h\in[0,k-n]$ timesteps that provide alternative goals, and then sample
from all the goals achieved at those timesteps. We merge the $h$ alternative
goals into the original instruction in chronological order and replay the
trajectory from the beginning, recomputing the reward for the new instruction,
which leads to a return of $n+h$ when rewards are binary. Figure 2 illustrates
the technique with a maze-navigation example. Our implementation also
considers importance-sampling variants that weight goal selection based on
rarity and improve sample efficiency. Details are discussed in Appendix B.

##  5 Experiments

Our experiments are divided into two main parts. First, we evaluate our agent
in a variety of existing long-term memory, generalization, and meta-learning
environments. We then explore the combination of AMAGO’s adaptive memory and
hindsight instruction relabeling in multi-task domains with procedurally
generated environments. Additional results, details, and discussion for each
of our experiments can be found in Appendix C.

###  5.1 Long-Term Memory, Generalization, and Meta-Learning

Figure 3: Summary of POPGym Suite Results. (Left) Aggregate results based on
raw returns. (Right) Relative performance standardized by scores in each
environment, as reported in [62].

POPGym [62] is a recent benchmark for evaluating long-term memory and
generalization in a range of CMDPs. Performance in POPGym is far from
saturated, and the challenges of sequence-based learning often prevent prior
on-policy results from meaningfully outperforming memory-free policies. We
tune AMAGO on one environment to meet the sample limit and architecture
constraints of the original benchmark before copying these settings across
$38$ additional environments. Figure 3 summarizes results across the suite. We
compare against the best existing baseline (a recurrent GRU-based agent [13])
and the most comparable architecture (an efficient Transformer variant [47]).
We also report the aggregated results of the best baseline in each environment
— equivalent to an exhaustive grid search over $13$ alternative sequence
models trained by on-policy RL. AMAGO is the state-of-the-art when measured by
median, mean, or interquartile range. If we standardize the return in each
environment to $[0,1]$ based on the highest and lowest baseline, as done in
[62], AMAGO achieves an average score of $.95$ across the suite, making it a
remarkably strong default for sequence-based RL. We perform an ablation that
naively applies the shared Transformer learning update without AMAGO’s other
details — such as our Transformer architecture and multi-gamma update. The
naive baseline performs significantly worse, and the metrics in Figure 3 mask
its gradient collapse in $9/39$ environments; AMAGO maintains stability in all
of its $120+$ trials despite using model sizes and update-to-data ratios that
are not tuned. Learning curves and additional experiments are listed in
Appendix C.1.

Figure 4: Memory-Intensive POPGym Environments. AMAGO’s off-policy updates can
improve performance in general, but its Transformer turns hard memory-
intensive environments into a straightforward recall exercise. Appendix C.1
reports results on more than $30$ additional environments.

Much of AMAGO’s performance gain in POPGym appears to be due to its ability to
unlock Transformers’ memory capabilities in off-policy RL: there are $9$
recall-intensive environments where the GRU baseline achieves a normalized
score of $.19$ compared to AMAGO’s $.999$. Figure 4 ablates the impact of
AMAGO’s trajectory encoder from the rest of its off-policy details by
replacing its Transformer with a GRU-based RNN in a sample of these recall-
intensive environments.

Figure 5 measures the scalability of AMAGO relative to off-policy agents with
a separated RNN architecture as well as the more efficient (but previously
unstable) shared sequence model. AMAGO lets us train larger models faster than
RNNs with equivalent memory capacity. More importantly, Transformers make
optimizing these super-long sequences more realistic.

Figure 5: Model Sizes, Sequence Lengths, and Training Throughput. AMAGO
enables larger models and longer context lengths than off-policy RNN variants.

We test the limits of AMAGO’s recall using a “Passive T-Maze” environment from
concurrent work [67], which creates a way to isolate memory capabilities at
any sequence length. Their work finds that recurrent agents cannot learn
effective memory beyond $100$ timesteps while Transformers can extend to
$H=1,500$. AMAGO trains a single RL Transformer to recover the optimal policy
in this environment all the way until we run out of GPU memory at context
length $l=H=10,000$ (Figure 6 top right), turning the algorithmic challenge of
memory into a hardware barrier. One concern is that AMAGO’s maximum context
lengths, large policies, and long horizons may hinder performance on easier
problems where they are unnecessary, leading to the kind of hyperparameter
tuning we would like to avoid. We demonstrate AMAGO on two continuous-action
meta-RL problems from related work with dense (Fig. 6 left) and sparse (Fig. 6
bottom right) rewards where performance is already saturated, with strong and
stable results.

Figure 6: Case Studies in In-Context Adaptation. (Left) Adaptation over three
episodes in Half-Cheetah Velocity [26]. (Center) RNNs (GRU) vs. Transformers
as AMAGO’s Trajectory Encoder in the Dark-Key-To-Door Meta-RL environment
[51]. (Bottom Right) Adaptation to a sparse-reward environment with continuous
actions [66]. (Top Right) AMAGO solves the Passive T-Maze memory experiment
[67] up until the GPU memory limit of $l=H=10,000$ timesteps.

Figure 6 (center) compares a Transformer and RNN version of AMAGO on the Dark-
Key-To-Door meta-RL environment [51] over a meta-test horizon of $500$
timesteps, which corresponds to at least $10$ episodes depending on a policy’s
exploration efficiency. The Transformer variant is far more effective at these
episode lengths. We run a similar experiment against our own replication of
Algorithm Distillation [51] — an approach for meta-RL with Transformers that
reformulates the problem as supervised learning — in Appendix C.2. AMAGO’s
ability to successfully optimize the true RL objective over any horizon $H$
lets us control sample efficiency at test-time, while AD’s adaptation rate is
limited by its dataset and converges several thousand timesteps later in this
case.

Figure 7: Package Delivery Results. Relabeling, long-term memory, and multi-
gamma updates are essential to success in this sparse goal-conditioned
adaptation problem $(H=180)$. Figure 8: MazeRunner Results. Memory improves
navigation in a partially observed maze, but sparse rewards create the most
significant barrier to learning ($15\times 15$ $H=400$, $30\times 30$
$H=1,000$).

###  5.2 Goal-Conditioned Environment Adaptation

We now turn our focus towards generalization over procedurally generated
environments $e\sim p(e)$ and multi-step instructions $(g^{0},\dots,g^{k})\sim
p(g\mid e)$ of up to $k$ goals (Sec. 3). A successful policy is able to adapt
to a new environment in order to achieve sparse rewards of $+1$ for completing
each step of its instruction, creating a general instruction-following agent.
Learning to explore from sparse rewards is difficult, but AMAGO’s long-horizon
off-policy learning update lets us relabel trajectories with alternative
instructions (Sec. 3 Figure 2). Before we scale up to a more complex domain,
we introduce two easily-simulated benchmarks. “Package Delivery” is a toy
problem where agents navigate a sequence of forks in a road to deliver items
to a list of addresses. This task is both sparse and memory-intensive, as
achieving any reward requires recall of previous attempts. Appendix C.3
provides a full environment description, and Figure 7 compares key ablations
of AMAGO. Relabeling, memory, and multi-gamma learning are essential to
performance and highlight the importance of AMAGO’s technical details.

Figure 9: MazeRunner with Random Dynamics. The agent’s action space is
randomized in every environment.

“MazeRunner” is a zero-shot maze navigation problem modeled after the example
in Figure 2 and loosely based on Memory Maze [74]. The agent finds itself in a
familiar spawn location in an otherwise unfamiliar maze and needs to navigate
to a sequence of $k\in[1,3]$ locations. Although the underlying maze map is
generated as an $N\times N$ gridworld, the agent’s observations are continuous
Lidar-like depth sensors. When $N=15$, the problem is sparse but solvable
without relabeling thanks to our implementation’s low-level improvements
(Figure 8 left). However, $N=30$ is impossibly sparse, and relabeling is
crucial to success (Fig. 8 right). We run a similar experiment where action
dynamics are randomly reset every episode. While this variant is more
challenging, our method outperforms other zero-shot meta-RL agents and nearly
recovers the original performance while adapting to the new action space
(Figure 9).

Crafter [33] is a research-friendly simplification of Minecraft designed to
evaluate multi-task capabilities in procedurally generated worlds. Agents
explore their surroundings to gather food and resources while progressing
through a tech tree of advanced tools and avoiding dangerous enemies. We
create an instruction-conditioned version where the agent is only successful
if it completes the specified task. Our instructions are formed from the $22$
original Crafter achievements with added goals for traveling to a grid of
world coordinates and placing blocks in specific locations. Goals are
represented as strings like “make stone pickaxe”, which are tokenized at the
word level with the motivation of enabling knowledge transfer between similar
goals. Any sequence of $k\leq 5$ goal strings forms a valid instruction that
we expect our agent to solve in any randomly generated Crafter environment.

Figure 10: Crafter Instruction Success Rates. Here we highlight specific
instructions that reveal exploration capabilities. “All Procedurally Generated
Instructions” indicates performance across the entire range of multi-step
goals in randomly generated worlds. Appendix C.5 includes more examples of
multi-step instruction success rates.

Figure 10 shows the success of AMAGO on the entire goal distribution along
with several highlighted tasks. Resource-acquisition and tool-making tasks
that require test-time exploration and adaptation to a new world are only
possible when AMAGO uses its Transformer policy for long-term memory. As the
steps in an instruction become less likely to be achieved by random
exploration, relabeling becomes essential to success. However, Crafter reveals
a more interesting capability of our multi-step hindsight relabeling. If we
ablate our relabeling scheme to single-goal HER and evaluate on instructions
with length $k=1$ (Fig. 10 “w/o Multi-Goal…”), we find AMAGO loses the ability
to complete advanced goals. Figure 11 illustrates how relabeling multi-step
instructions creates automatic plans for discovering goals that are locked
behind a progression like Crafter’s achievement tree. However, this effect
only emerges when the entire instruction is observed in advance as opposed to
revealing the task one step at a time (Fig. 10 “w/o Full Instruction…”). When
the agent can plan for the full rollout it can learn to complete the final
achievement without the multi-step instruction. We continue this investigation
with additional Crafter experiments in Appendix C.5. The most encouraging
takeaway is that learning any specific task becomes easier as the goal space
grows larger and the agent becomes more general.

Figure 11: Multi-Goal Learning and Exploration. We compare AMAGO’s performance
on difficult single-goal instructions with another task that forms a useful
exploration plan.

##  6 Conclusion

We have introduced AMAGO, an in-context RL agent for generalization, long-term
memory, and meta-learning. Our work makes important technical contributions by
finding a stable and high-performance way to train off-policy RL agents on top
of one long-context Transformer. With AMAGO, we can train policies to adapt
over more distant planning horizons with longer effective memories. We also
show that the benefits of an off-policy in-context method can go beyond
stability and sample efficiency, as AMAGO lets us relabel multi-step
instructions to discover sparse rewards while adapting to unfamiliar
environments. Our agent’s efficiency on long input sequences creates an
exciting direction for future work, as very few academic-scale meta-RL or RL
generalization benchmarks genuinely require adaptation over thousands of
timesteps. AMAGO is open-source, and the combination of its performance and
flexibility should let it serve as a strong baseline in the development of new
benchmarks in this area. The core AMAGO framework is also compatible with
sequence models besides Transformers and creates an extensible way to research
more experimental long-term memory architectures that can push adaptation
horizons even further.

#### Acknowledgments

This work was partially supported by research funds from Salesforce and JP
Morgan. We would like to thank Yifeng Zhu, Zhenyu Jiang, Soroush Nasiriany,
Huihan Liu, Rutav Shah, Mingyo Seo, and the UT Austin Robot Perception and
Learning Lab for constructive feedback on early drafts of this paper.

## References

  * [1] Zaheer Abbas et al.  “Loss of plasticity in continual deep reinforcement learning”  In _arXiv preprint arXiv:2303.07507_ , 2023
  * [2] Marcin Andrychowicz et al.  “Hindsight experience replay”  In _Advances in neural information processing systems_ 30, 2017
  * [3] Jimmy Lei Ba, Jamie Ryan Kiros and Geoffrey E Hinton  “Layer normalization”  In _arXiv preprint arXiv:1607.06450_ , 2016
  * [4] Adrià Puigdomènech Badia et al.  “Agent57: Outperforming the atari human benchmark”  In _International Conference on Machine Learning_ , 2020, pp. 507–517  PMLR
  * [5] Jacob Beck et al.  “A Survey of Meta-Reinforcement Learning”  In _arXiv preprint arXiv:2301.08028_ , 2023
  * [6] Marc G Bellemare, Yavar Naddaf, Joel Veness and Michael Bowling  “The arcade learning environment: An evaluation platform for general agents”  In _Journal of Artificial Intelligence Research_ 47, 2013, pp. 253–279
  * [7] Christopher Berner et al.  “Dota 2 with large scale deep reinforcement learning”  In _arXiv preprint arXiv:1912.06680_ , 2019
  * [8] Matthew Botvinick et al.  “Reinforcement learning, fast and slow”  In _Trends in cognitive sciences_ 23.5 Elsevier, 2019, pp. 408–422
  * [9] David Brandfonbrener et al.  “When does return-conditioned supervised learning work for offline reinforcement learning?”  In _arXiv preprint arXiv:2206.01079_ , 2022
  * [10] Anthony R Cassandra, Leslie Pack Kaelbling and Michael L Littman  “Acting optimally in partially observable stochastic domains”  In _Aaai_ 94, 1994, pp. 1023–1028
  * [11] Lili Chen et al.  “Decision transformer: Reinforcement learning via sequence modeling”  In _Advances in neural information processing systems_ 34, 2021, pp. 15084–15097
  * [12] Xinyue Chen, Che Wang, Zijian Zhou and Keith Ross  “Randomized ensembled double q-learning: Learning fast without a model”  In _arXiv preprint arXiv:2101.05982_ , 2021
  * [13] Kyunghyun Cho et al.  “Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation”  In _Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , 2014, pp. 1724  Association for Computational Linguistics
  * [14] Petros Christodoulou  “Soft actor-critic for discrete action settings”  In _arXiv preprint arXiv:1910.07207_ , 2019
  * [15] Karl Cobbe, Chris Hesse, Jacob Hilton and John Schulman  “Leveraging procedural generation to benchmark reinforcement learning”  In _International conference on machine learning_ , 2020, pp. 2048–2056  PMLR
  * [16] Zichen Jeff Cui, Yibin Wang, Nur Muhammad and Lerrel Pinto  “From play to policy: Conditional behavior generation from uncurated robot data”  In _arXiv preprint arXiv:2210.10047_ , 2022
  * [17] Pierluca D’Oro et al.  “Sample-efficient reinforcement learning by breaking the replay ratio barrier”  In _Deep Reinforcement Learning Workshop NeurIPS 2022_ , 2022
  * [18] Tri Dao et al.  “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”  In _arXiv preprint arXiv:2205.14135_ , 2022
  * [19] Ron Dorfman, Idan Shenfeld and Aviv Tamar  “Offline meta learning of exploration”  In _arXiv preprint arXiv:2008.02598_ , 2020
  * [20] Yan Duan et al.  “RL 2: Fast reinforcement learning via slow reinforcement learning”  In _arXiv preprint arXiv:1611.02779_ , 2016
  * [21] Ben Eysenbach, Xinyang Geng, Sergey Levine and Russ R Salakhutdinov  “Rewriting history with inverse rl: Hindsight inference for policy improvement”  In _Advances in neural information processing systems_ 33, 2020, pp. 14783–14795
  * [22] Benjamin Eysenbach, Soumith Udatha, Sergey Levine and Ruslan Salakhutdinov  “Imitating Past Successes can be Very Suboptimal”  In _arXiv preprint arXiv:2206.03378_ , 2022
  * [23] Rasool Fakoor, Pratik Chaudhari, Stefano Soatto and Alexander J Smola  “Meta-q-learning”  In _arXiv preprint arXiv:1910.00125_ , 2019
  * [24] Linxi Fan et al.  “Minedojo: Building open-ended embodied agents with internet-scale knowledge”  In _arXiv preprint arXiv:2206.08853_ , 2022
  * [25] William Fedus et al.  “Revisiting fundamentals of experience replay”  In _International Conference on Machine Learning_ , 2020, pp. 3061–3071  PMLR
  * [26] Chelsea Finn, Pieter Abbeel and Sergey Levine  “Model-agnostic meta-learning for fast adaptation of deep networks”  In _International conference on machine learning_ , 2017, pp. 1126–1135  PMLR
  * [27] Scott Fujimoto, Herke Hoof and David Meger  “Addressing function approximation error in actor-critic methods”  In _International conference on machine learning_ , 2018, pp. 1587–1596  PMLR
  * [28] Dibya Ghosh et al.  “Learning to reach goals via iterated supervised learning”  In _arXiv preprint arXiv:1912.06088_ , 2019
  * [29] Dibya Ghosh et al.  “Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability”, 2021  arXiv:2107.06277 [cs.LG]
  * [30] Jake Grigsby and Yanjun Qi  “A Closer Look at Advantage-Filtered Behavioral Cloning in High-Noise Datasets”  In _arXiv preprint arXiv:2110.04698_ , 2021
  * [31] Jake Grigsby and Yanjun Qi  “Measuring visual generalization in continuous control from pixels”  In _arXiv preprint arXiv:2010.06740_ , 2020
  * [32] Tuomas Haarnoja et al.  “Soft actor-critic algorithms and applications”  In _arXiv preprint arXiv:1812.05905_ , 2018
  * [33] Danijar Hafner  “Benchmarking the spectrum of agent capabilities”  In _arXiv preprint arXiv:2109.06780_ , 2021
  * [34] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi and Jimmy Ba  “Mastering atari with discrete world models”  In _arXiv preprint arXiv:2010.02193_ , 2020
  * [35] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba and Timothy Lillicrap  “Mastering Diverse Domains through World Models”  In _arXiv preprint arXiv:2301.04104_ , 2023
  * [36] Assaf Hallak, Dotan Di Castro and Shie Mannor  “Contextual markov decision processes”  In _arXiv preprint arXiv:1502.02259_ , 2015
  * [37] Hado P Hasselt et al.  “Learning values across many orders of magnitude”  In _Advances in neural information processing systems_ 29, 2016
  * [38] Matthew Hausknecht and Peter Stone  “Deep recurrent q-learning for partially observable mdps”  In _2015 aaai fall symposium series_ , 2015
  * [39] Nicolas Heess, Jonathan J Hunt, Timothy P Lillicrap and David Silver  “Memory-based control with recurrent neural networks”  In _arXiv preprint arXiv:1512.04455_ , 2015
  * [40] Peter Henderson et al.  “Deep reinforcement learning that matters”  In _Proceedings of the AAAI conference on artificial intelligence_ 32.1, 2018
  * [41] Matteo Hessel et al.  “Rainbow: Combining improvements in deep reinforcement learning”  In _Proceedings of the AAAI conference on artificial intelligence_ 32.1, 2018
  * [42] Jan Humplik et al.  “Meta reinforcement learning as task inference”  In _arXiv preprint arXiv:1905.06424_ , 2019
  * [43] Michael Janner, Qiyang Li and Sergey Levine  “Reinforcement learning as one big sequence modeling problem”  In _ICML 2021 Workshop on Unsupervised Reinforcement Learning_ , 2021
  * [44] Leslie Pack Kaelbling  “Learning to achieve goals”  In _IJCAI_ 2, 1993, pp. 1094–8  Citeseer
  * [45] Pierre-Alexandre Kamienny et al.  “Learning adaptive exploration strategies in dynamic environments through informed policy regularization”  In _arXiv preprint arXiv:2005.02934_ , 2020
  * [46] Steven Kapturowski et al.  “Recurrent experience replay in distributed reinforcement learning”  In _International conference on learning representations_ , 2018
  * [47] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas and François Fleuret  “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention” ISSN: 2640-3498  In _Proceedings of the 37th International Conference on Machine Learning_ PMLR, 2020, pp. 5156–5165
  * [48] Robert Kirk, Amy Zhang, Edward Grefenstette and Tim Rocktäschel  “A Survey of Generalisation in Deep Reinforcement Learning”, 2022  arXiv:2111.09794 [cs.LG]
  * [49] Aviral Kumar, Xue Bin Peng and Sergey Levine  “Reward-conditioned policies”  In _arXiv preprint arXiv:1912.13465_ , 2019
  * [50] Heinrich Küttler et al.  “The nethack learning environment”  In _Advances in Neural Information Processing Systems_ 33, 2020, pp. 7671–7684
  * [51] Michael Laskin et al.  “In-context Reinforcement Learning with Algorithm Distillation”  In _arXiv preprint arXiv:2210.14215_ , 2022
  * [52] Kuang-Huei Lee et al.  “Multi-game decision transformers”  In _Advances in Neural Information Processing Systems_ 35, 2022, pp. 27921–27936
  * [53] Qiyang Li, Aviral Kumar, Ilya Kostrikov and Sergey Levine  “Efficient Deep Reinforcement Learning Requires Regulating Overfitting”  In _arXiv preprint arXiv:2304.10466_ , 2023
  * [54] Timothy P Lillicrap et al.  “Continuous control with deep reinforcement learning”  In _arXiv preprint arXiv:1509.02971_ , 2015
  * [55] Evan Z Liu, Aditi Raghunathan, Percy Liang and Chelsea Finn  “Decoupling exploration and exploitation for meta-reinforcement learning without sacrifices”  In _International conference on machine learning_ , 2021, pp. 6925–6935  PMLR
  * [56] Chris Lu et al.  “Structured state space models for in-context reinforcement learning”  In _arXiv preprint arXiv:2303.03982_ , 2023
  * [57] Corey Lynch et al.  “Learning latent plans from play”  In _Conference on robot learning_ , 2020, pp. 1113–1132  PMLR
  * [58] Ajay Mandlekar et al.  “Iris: Implicit reinforcement without interaction at scale for learning control from offline robot manipulation data”  In _2020 IEEE International Conference on Robotics and Automation (ICRA)_ , 2020, pp. 4414–4420  IEEE
  * [59] Luckeciano C Melo  “Transformers are meta-reinforcement learners”  In _International Conference on Machine Learning_ , 2022, pp. 15340–15359  PMLR
  * [60] Nikhil Mishra, Mostafa Rohaninejad, Xi Chen and Pieter Abbeel  “A simple neural attentive meta-learner”  In _arXiv preprint arXiv:1707.03141_ , 2017
  * [61] Volodymyr Mnih et al.  “Human-level control through deep reinforcement learning”  In _nature_ 518.7540 Nature Publishing Group, 2015, pp. 529–533
  * [62] Steven Morad et al.  “POPGym: Benchmarking Partially Observable Reinforcement Learning”  In _The Eleventh International Conference on Learning Representations_ , 2023  URL: https://openreview.net/forum?id=chDrutUTs0K
  * [63] Ashvin V Nair et al.  “Visual reinforcement learning with imagined goals”  In _Advances in neural information processing systems_ 31, 2018
  * [64] Soroush Nasiriany et al.  “Disco rl: Distribution-conditioned reinforcement learning for general-purpose policies”  In _2021 IEEE International Conference on Robotics and Automation (ICRA)_ , 2021, pp. 6635–6641  IEEE
  * [65] Toan Q Nguyen and Julian Salazar  “Transformers without tears: Improving the normalization of self-attention”  In _arXiv preprint arXiv:1910.05895_ , 2019
  * [66] Tianwei Ni, Benjamin Eysenbach and Ruslan Salakhutdinov  “Recurrent Model-Free RL can be a Strong Baseline for Many POMDPs”, 2022  arXiv:2110.05038 [cs.LG]
  * [67] Tianwei Ni, Michel Ma, Benjamin Eysenbach and Pierre-Luc Bacon  “When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment”  In _arXiv preprint arXiv:2307.03864_ , 2023
  * [68] Evgenii Nikishin et al.  “The primacy bias in deep reinforcement learning”  In _International conference on machine learning_ , 2022, pp. 16828–16847  PMLR
  * [69] Pedro A Ortega et al.  “Meta-learning of sequential strategies”  In _arXiv preprint arXiv:1905.03030_ , 2019
  * [70] Charles Packer et al.  “Assessing generalization in deep reinforcement learning”  In _arXiv preprint arXiv:1810.12282_ , 2018
  * [71] Fabio Pardo, Arash Tavakoli, Vitaly Levdik and Petar Kormushev  “Time limits in reinforcement learning”  In _International Conference on Machine Learning_ , 2018, pp. 4045–4054  PMLR
  * [72] Emilio Parisotto et al.  “Stabilizing transformers for reinforcement learning”  In _International conference on machine learning_ , 2020, pp. 7487–7498  PMLR
  * [73] Keiran Paster, Sheila McIlraith and Jimmy Ba  “You Can’t Count on Luck: Why Decision Transformers Fail in Stochastic Environments”  In _arXiv preprint arXiv:2205.15967_ , 2022
  * [74] Jurgis Pasukonis, Timothy Lillicrap and Danijar Hafner  “Evaluating Long-Term Memory in 3D Mazes”  In _arXiv preprint arXiv:2210.13383_ , 2022
  * [75] Christian Perez, Felipe Petroski Such and Theofanis Karaletsos  “Generalized hidden parameter mdps: Transferable model-based rl in a handful of trials”  In _Proceedings of the AAAI Conference on Artificial Intelligence_ 34.04, 2020, pp. 5403–5411
  * [76] Marco Pleines, Matthias Pallasch, Frank Zimmer and Mike Preuss  “Generalization, Mayhems and Limits in Recurrent Proximal Policy Optimization”  In _arXiv preprint arXiv:2205.11104_ , 2022
  * [77] Vitchyr Pong  “Goal-Directed Exploration and Skill Reuse”, 2021
  * [78] Jack W Rae and Ali Razavi  “Do transformers need deep long-range memory”  In _arXiv preprint arXiv:2007.03356_ , 2020
  * [79] Kate Rakelly et al.  “Efficient off-policy meta-reinforcement learning via probabilistic context variables”  In _International conference on machine learning_ , 2019, pp. 5331–5340  PMLR
  * [80] Scott Reed et al.  “A generalist agent”  In _arXiv preprint arXiv:2205.06175_ , 2022
  * [81] Hongyu Ren et al.  “OCEAN: Online Task Inference for Compositional Tasks with Context Adaptation”  In _Conference on Uncertainty in Artificial Intelligence_ , 2020, pp. 1378–1387  PMLR
  * [82] Juergen Schmidhuber  “Reinforcement Learning Upside Down: Don’t Predict Rewards–Just Map Them to Actions”  In _arXiv preprint arXiv:1912.02875_ , 2019
  * [83] Jürgen Schmidhuber  “Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-… hook”, 1987
  * [84] Max Schwarzer et al.  “Bigger, Better, Faster: Human-level Atari with human-level efficiency”  In _International Conference on Machine Learning_ , 2023, pp. 30365–30380  PMLR
  * [85] Sam Shleifer, Jason Weston and Myle Ott  “Normformer: Improved transformer pretraining with extra normalization”  In _arXiv preprint arXiv:2110.09456_ , 2021
  * [86] David Silver et al.  “A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play”  In _Science_ 362.6419 American Association for the Advancement of Science, 2018, pp. 1140–1144
  * [87] Xingyou Song et al.  “Observational overfitting in reinforcement learning”  In _arXiv preprint arXiv:1912.02975_ , 2019
  * [88] Rupesh Kumar Srivastava et al.  “Training agents using upside-down reinforcement learning”  In _arXiv preprint arXiv:1912.02877_ , 2019
  * [89] Bradly C Stadie et al.  “Some considerations on learning to explore via meta-reinforcement learning”  In _arXiv preprint arXiv:1803.01118_ , 2018
  * [90] Adaptive Agent Team et al.  “Human-Timescale Adaptation in an Open-Ended Task Space”  In _arXiv preprint arXiv:2301.07608_ , 2023
  * [91] Hugo Touvron et al.  “Llama 2: Open foundation and fine-tuned chat models”  In _arXiv preprint arXiv:2307.09288_ , 2023
  * [92] Ashish Vaswani et al.  “Attention is all you need”  In _Advances in neural information processing systems_ 30, 2017
  * [93] Oriol Vinyals et al.  “Grandmaster level in StarCraft II using multi-agent reinforcement learning”  In _Nature_ 575.7782 Nature Publishing Group UK London, 2019, pp. 350–354
  * [94] Bernie Wang et al.  “Improving context-based meta-reinforcement learning with self-supervised trajectory contrastive learning”  In _arXiv preprint arXiv:2103.06386_ , 2021
  * [95] Jane X Wang et al.  “Learning to reinforcement learn”  In _arXiv preprint arXiv:1611.05763_ , 2016
  * [96] Ziyu Wang et al.  “Critic regularized regression”  In _Advances in Neural Information Processing Systems_ 33, 2020, pp. 7768–7778
  * [97] Ruibin Xiong et al.  “On layer normalization in the transformer architecture”  In _International Conference on Machine Learning_ , 2020, pp. 10524–10533  PMLR
  * [98] Mengdi Xu et al.  “Prompting decision transformer for few-shot policy generalization”  In _International Conference on Machine Learning_ , 2022, pp. 24631–24645  PMLR
  * [99] Rui Yang et al.  “Rethinking goal-conditioned supervised learning and its connection to offline rl”  In _arXiv preprint arXiv:2202.04478_ , 2022
  * [100] Zhihan Yang and Hai Nguyen  “Recurrent off-policy baselines for memory-based continuous control”  In _arXiv preprint arXiv:2110.12628_ , 2021
  * [101] Denis Yarats et al.  “Improving sample efficiency in model-free reinforcement learning from images”  In _Proceedings of the AAAI Conference on Artificial Intelligence_ 35.12, 2021, pp. 10674–10681
  * [102] Shuangfei Zhai et al.  “Stabilizing Transformer Training by Preventing Attention Entropy Collapse”  In _International Conference on Machine Learning_ , 2023, pp. 40770–40803  PMLR
  * [103] Amy Zhang, Nicolas Ballas and Joelle Pineau  “A dissection of overfitting and generalization in continuous reinforcement learning”  In _arXiv preprint arXiv:1806.07937_ , 2018
  * [104] Qinqing Zheng, Amy Zhang and Aditya Grover  “Online decision transformer”  In _arXiv preprint arXiv:2202.05607_ , 2022
  * [105] Luisa Zintgraf  “Fast adaptation via meta reinforcement learning”, 2022
  * [106] Luisa Zintgraf et al.  “VariBAD: Variational Bayes-Adaptive Deep RL via Meta-Learning”  In _Journal of Machine Learning Research_ 22.289, 2021, pp. 1–39
  * [107] Luisa M Zintgraf et al.  “Exploration in approximate hyper-state space for meta reinforcement learning”  In _International Conference on Machine Learning_ , 2021, pp. 12991–13001  PMLR

##  Appendix A AMAGO Details

###  A.1 Sharing a Single Sequence Model in Off-Policy RL

Off-policy actor-critic methods compute loss terms with an ensemble of actor
and critic networks along with a moving-average copy of their parameters that
is used to generate temporal difference targets. Extending the feed-forward
(fully observed) setup to sequence-based learning results in an excessive
number of sequence model forward/backward passes per training step (Figure 12
top left) [54]. This has created several ways to share sequence models and
improve efficiency. Ni et al. [66] share a sequence model between the ensemble
of critic networks, which becomes more important when using REDQ [12] (Fig. 12
top right). Parameters can be shared across the actor and critics, but this
has been shown to be unstable. SAC+AE [101] confronts a similar problem in
pixel-based learning and popularized the solution of detaching the larger base
model (the Transformer in our case) from the actor’s gradients (Fig. 12 bottom
left). This approach has also been demonstrated in sequence-learning [100].
AMAGO removes the target sequence model as well — sharing one Transformer
across every actor, critic, and target network while preserving the actor’s
gradients and training with one optimizer (Fig. 12 bottom right).

The same shared approach was evaluated in [66] but was noted to be unstable
and needed to be tuned across experiments. We find this is caused by the
critic receiving gradients from the actor loss, and correct for this effect.
Concurrent to our work, [67] addressed the same problem with a different
approach, although it is not discussed in the paper.

Figure 12: Evolution of Off-Policy Actor-Critic Agent Architectures. Black
arrows give a high-level overview of how the actor, critic(s), and target
networks combine to compute the training objective(s). Network borders are
color-coded according to the loss function they optimize. Sequence models
(green) are more expensive than feed-forward actors (blue) and critics (gray),
which has motivated several ways to simplify the training process while
maintaining stability. AMAGO removes all excess sequence computation and
achieves the minimum of a single shared sequence model without detaching the
actor’s gradients.

###  A.2 Base Actor-Critic Update

AMAGO’s shared sequence model reduces the RL training process to the standard
feed-forward case where the output of the Transformer becomes the state array
and the batch size is larger by a factor of the context length $l$222In this
paper, the context length $l$ is always equal to the full adaptation horizon
$H$.. As a general rule, AMAGO strives to remove any hyperparameters that do
not follow intuitive trends. For example, our actors and critics are setup as
in SAC [32] for continuous actions and SAC-D [14] for discrete actions, but we
do not use policy entropy constraints. This makes AMAGO closer to a stochastic
version of TD3 [27], but we do not use any target action noise or delayed
updates. We improve stability by training parallel ensembles of $M>2$ critics
and using random subsets of the target critic ensemble to generate temporal
difference targets, as in REDQ [12]. We also add the filtered behavioral
cloning (FBC) term from CRR [96], where actions are imitated only when they
have positive advantage relative to the output of the current policy. The
actor’s standard policy gradient term has strong learning signal when the
value surface is steep, but FBC has strong signal when the value surface is
flat because its scale is only dependent on the sign of the advantage and
behaves more like supervised learning. The large batch sizes generated by
parallel learning across causal sequences also prevents FBC’s tendency to
increase variance by filtering too many actions [30]. These terms are combined
to create a shared loss where $\mathcal{L}_{\text{Critic}}$ and
$\mathcal{L}_{\text{Actor}}$ denote the standard one-step temporal difference
and deterministic policy gradient objectives, respectively:

$\displaystyle\displaystyle\mathcal{L}(\theta)$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}\tau\sim\mathcal{D}\end{subarray}}\bigg{[}\frac{1}{H}\sum_{t=0}^{H}\lambda_{0}\mathcal{L}_{\text{Critic}}(s_{t},a_{t})+\lambda_{1}\mathcal{L}_{\text{Actor}}(s_{t})+\lambda_{2}\mathcal{L}_{\text{FBC}}(s_{t},a_{t})\bigg{]}$
(1)

The weights of each term ($\lambda_{0},\lambda_{1},\lambda_{2}$) are examples
of unavoidable but unintuitive hyperparameters. PopArt [37] is intended to
normalize the scale of value-based loss functions when training one agent
across multiple domains, but we use it to reduce hyperparameter tuning by
putting $\mathcal{L}_{\text{Actor}}$ and $\mathcal{L}_{\text{Critic}}$ on a
predictable scale that only needs to be tuned once. Other ways to achieve this
effect have been shown to reduce tuning in agents like DreamerV3 [35].

AMAGO computes $\mathcal{L}_{\text{Actor}}$, $\mathcal{L}_{\text{Critic}}$,
and $\mathcal{L}_{\text{FBC}}$ in parallel across $\gamma_{N}$ values of the
discount factor $\gamma$. As described in Section 4, this is done by
increasing the output dimension of actors and critics in the discrete update
or by expanding the effective batch dimension of the continuous update while
making $\gamma$ an additional input to the critics. An example of the value
scales learned by different values of $\gamma$ is plotted in Figure 13.

Figure 13: Multi-Gamma Actor-Critic Training. AMAGO optimizes one Transformer
on actor-critic loss terms corresponding to many values of the discount factor
$\gamma$, and uses the actions that maximize the longest horizon at test-time.
We show the average Q-value across trajectory sequences throughout training at
different discounts. Note that values converge quickly in this example because
instruction relabeling is inflating returns, but the distribution of goals
shifts as the agent improves.

AMAGO samples from a stochastic policy during both data collection and
evaluation. Because we do not use entropy constraints, we add exploration
noise during data collection. In discrete domains noise is added as in classic
epsilon-greedy, while continuous domains randomly perturb the action vector as
in TD3 [27]. The level of action randomness is determined by a hyperparameter
$\epsilon$ that is typically annealed over a fixed number of environment
steps. AMAGO adapts this schedule to more closely align with the
exploration/exploitation trade-off that occurs within the adaption horizon of
any given CMDP. $\epsilon$ is annealed over the $H$ timesteps of a rollout,
and the intensity of this schedule decreases over training. This process is
visualized in Figure 14. Because the parameters of this schedule are difficult
to tune, we heavily randomize them across AMAGO’s parallel actors — meaning we
are always collecting data at varying levels of action noise. There are no
experiments where we find tuning this schedule to be essential to achieving
our reported results.

Figure 14: Exploration Noise in Implicit POMDPs. AMAGO adapts the standard
random exploration schedule to more closely align with the
exploration/exploitation trade-off that occurs when acting according to a
Bayes-optimal policy in an implicit POMDP.

###  A.3 AMAGO Architecture

Figure 15: AMAGO Transformer Block. (Left) A standard Pre-LayerNorm (Pre-LN)
Transformer layer [97]. (Right) AMAGO replaces all saturating activations with
Leaky ReLUs and uses additional LayerNorms [3] (as in NormFormer [85]) and a
modified linear layer ($\sigma$Reparam [102]). These strategies limit the
magnitude of activations along the residual block and effectively prevent
attention entropy collapse.

Transformer. We observe performance collapse when using a standard Pre-
LayerNorm [97] Transformer in long training runs. In rare cases, we find that
this is caused by gradient collapse due to saturating ReLU activations. For
this reason we replace every ReLU/GeLU (including those in the actor/critic
MLPs) with a Leaky ReLU that will allow learning to continue. This idea is
also motivated by work in network plasticity and long training runs in
continual RL, where activations other than ReLU can be a simple baseline [1].
We find that this change fixes gradient instability, but does not prevent
performance collapse. Instead, collapse is now caused by saturating
activations in the residual block of AMAGO’s Transformer. We apply two
existing methods that effectively solve this problem. First, we add the extra
LayerNorms and HeadScale attention from Normformer [85]. Second, we replace
the Transformer’s linear layers with $\sigma$Reparam [102]. AMAGO uses Flash
Attention [18] to enable long context lengths on a single GPU (Figure 16).
However, this forces training to be done in half precision. We suspect this is
at least partially to blame for the Transformers’ instability, but is
difficult to ablate and is fixed by our architectural changes.

Instruction-Conditioning. AMAGO uses a small RNN to process the instruction
sequence of goal tokens, and the resulting representation is concatenated to
the CMDP information that forms Transformer input tokens. It would be simpler
to add the instruction to the beginning of the context sequence. The only
reason for the extra complexity of the RNN embedding is to allow for fair
baselines that do not use context sequences (“w/o Memory”).

Figure 16: Training Throughput: AMAGO vs. Off-Policy RNNs. We compare training
speed in terms of gradient updates per second at various context lengths with
a batch size of $24$ sequences on a single NVIDIA A$5000$ GPU. These metrics
do not include time spent in data collection. We use the codebase from [66] to
benchmark the RNN agents, and modify AMAGO’s hyperparameters to create a fair
comparison. One exception is that AMAGO is loading its data from disk while
the RNNs are loading from a replay buffer stored in RAM. AMAGO’s buffer can be
slower at long sequence lengths but allows for larger dataset sizes. “No
Flash” baselines remove Flash Attention [18].

##  Appendix B Relabeling with Goal Importance Sampling

AMAGO generates training data in multi-goal domains by relabeling trajectories
with alternative instructions based on hindsight outcomes. Relabeling improves
reward sparsity for actor-critic training, and greatly amplifies the learning
signal of existing data by recycling the same experience with many different
instructions. This technique works by saving the rewards for the entire goal
space during rollouts, rather than just the rewards for the goals in the
intended instruction (Figure 17 Step 1). While evaluating many different dense
reward terms would be unrealistic, it is more practical in sparse goal-
conditioned domains where success can be evaluated with simple rules.

Figure 17: Relabeling With Prioritized Goal Sampling. Long rollouts in multi-
goal domains lead to an unmanageable number of candidate instructions for
relabeling. AMAGO improves sample efficiency without domain knowledge by
prioritizing rare goals.

Generating a diverse training dataset with relabeled sequences of goals allows
our agents to carry out multi-stage tasks and has important exploration
advantages (Appendix C.5), but creates a practical issue where we have too
many alternative instructions to choose from. Domains like Crafter create
rollouts with dozens of candidate goals over the full length of the
trajectory. There are also goal types that can occur simultaneously and for
many consecutive timesteps (e.g., completing goals near navigational
waypoints). AMAGO relabels by sampling one instruction from the many sub-
sequences of these goals. With so many combinations of goal instructions
available to us, we need a way to focus our learning updates on useful
information.

Our solution is a weighted relabeling scheme that helps sort through the noise
of common outcomes by prioritizing interesting goals. While there could be
opportunities to add domain knowledge in this process, we prefer to avoid this
and sample goals according to their rarity. AMAGO tracks both the frequency
that a particular goal occurs at any given timestep, and the frequency that it
occurs at all in a given episode. The difference between these statistics is
caused by goals that occur rarely, but when they happen once are likely to be
repeated many times in the same trajectory. Long-distance navigation and
repeated mining behavior in Crafter are examples of this effect. We assign a
priority score to goals based on their rarity, which then lets us modify the
relabeling scheme to sample based on these scores (Figure 17). AMAGO’s
technical details are designed to reduce hyperparameter sensitivity and
prevent individual tricks like this from becoming unintuitive points of
failure that require manual tuning. We follow this philosophy here and
automatically randomize over several reasonable approaches. Examples include
sampling from the top-$k$ most rare goals according to either frequency
statistic, or those above either the median or minimum rarity in a trajectory
to filter trivial goals. AMAGO still relabels uniformly with some frequency,
which keeps the full diversity of our dataset available and prevents
information from being lost. Randomization over these implementation details
occurs on a per-trajectory level, meaning every batch has sequences that were
generated with a wide range of strategies. We defer the precise details to our
open-source code release. Appendix C.5 provides a quantitative demonstration
of our method.

##  Appendix C Experimental Details and Further Analysis

This section provides a detailed analysis of AMAGO’s main results. Each
subsection will give a description of our custom learning domains, followed by
a more complete discussion of the results in the main text with additional
experiments. Hyperparameter and compute information is listed in Appendix D.

###  C.1 POPGym

We evaluate on $39$ environments from the POPGym suite [62], leaving out only
the navigation-based environments that require extra dependencies. AMAGO’s
navigation capabilities are well covered by our own MazeRunner domain
(Appendix C.4). We follow the original benchmark in using policy architectures
with a memory hidden state of $256$. Learning curves for each environment are
shown on a full page below. Figure 3 reports results at $15$M timesteps,
though the learning curves extend slightly futher when data is available. We
plot the maximum and minimum value achieved by any seed to highlight the
stability of AMAGO relative to the “naive” off-policy Transformer baseline.
Each environment defaults to $3$ random seeds, as in [62]. The variance of
AMAGO at convergence is extremely low (a rounding error on nearly every
environment). However, there can be significant variance in the timestep where
rapid increases in performance begin. This is caused by many of the POPGym
domains containing strong local optima for nearly random policies. There are
two environments where this variance impacts results because it occurs near
the $15$M sample limit (AutoencodeHard and CountRecallHard), which we address
by running $9$ random seeds.

The “naive” agent maintains AMAGO’s shared actor-critic update but disables
many of our other technical details. These include the modified Transformer
architecture (Appendix A.3) and multi-gamma update. The naive agent also
reduces the REDQ ensemble from AMAGO’s default of $4$ parallel critics to the
standard $2$, and uses a lower discount factor of $\gamma=.99$. This causes
learning to collapse in many seeds. However, collapse generally does not
impact the final scores reported in Figure 3 which indicate the maximum mean
episodic return (MMER) achieved during training. The effects of policy
collapse can be much more damaging in longer goal-conditioned experiments
where it occurs before convergence, but is too expensive to evaluate at this
scale. Therefore the performance gap between AMAGO and the naive agent in
Figure 3 can be attributed to the credit assignment and planning differences
between the two learning updates.

AMAGO was briefly tuned on one environment (ConcentrationEasy) where we
adjusted its learning schedule in order to meet the benchmark’s $15$M timestep
limit. However, our results show that we likely did not push the update-to-
data ratio high enough, as AMAGO is still improving well after $15$M timesteps
in some difficult environments. We explore the upper bound of AMAGO’s
performance in select environments by either leaving the standard agent
running for additional timesteps (“AMAGO-L”), or training a deeper $5$ layer
Transformer for $100$M timesteps (“AMAGO-XL”). We briefly summarize these
results as a reference for future work on this benchmark:

  * •

AMAGO-L consistently converges to the plotted upper bound on CountRecallHard
and AutoencodeHard by $20$M timesteps.

  * •

AMAGO-XL converges to an average return of $.20$ on ConcentrationMedium at
approximately $75$M timesteps.

  * •

AMAGO-XL converges to an average return of $.33$ on MineSweeperMedium at
approximately $80$M timesteps.

  * •

AMAGO-XL does not improve performance on MineSweeperEasy.

  * •

AMAGO-XL converges to an average return of $.49$ on BattleshipEasy at
approximately $35$M timesteps.

It appears that there is significant room for improvement in AMAGO’s POPGym
results, and exploring the limits of Transformers’ sample efficiency in off-
policy RL is an interesting direction for future work.

###  C.2 Additional Memory and Multi-Episodic Meta-RL Results

##### Half-Cheetah Velocity.

In Figure 6 (left) we use a classic meta-RL benchmark from [26] as an example
of a case where AMAGO’s long sequences and high discount factors are clearly
unnecessary but do not need tuning. This task is solvable with short context
lengths of just a few timesteps. We evaluate adaptation over the first $3$
episodes ($H=600$) and report baselines from variBAD [106], BOReL [19], RL2
[20], and the recurrent off-policy implementation in [66].

##### Dark-Key-To-Door.

Dark-Key-To-Door is a sparse-reward multi-episodic meta-RL problem [51].
Figure 6 (center) compares a version of AMAGO with a Transformer and RNN
trajectory encoder on a $9\times 9$ version of the environment with a max
episode length of $50$. An agent that always fails to solve the task will
encounter $10$ episodes of a new environment during a meta-testing horizon of
$H=500$ timesteps. The maximum return per episode is $2$. An adaptive agent
learns to solve the task quickly once it has identified the key and door
location by meta-exploration, and tries to complete as many episodes as
possible in the $H=500$ timestep limit. The Transformer and RNN model
architectures are set to have equal memory hidden state size and layer count,
and every other detail of the training loop is exactly the same.

We replicate Algorithm Distillation [51] (AD) on an $8\times 8$ version of the
task by collecting training histories from $700$ source RL agents. While AD
converges to high performance, it does so on roughly the same timescale as the
single-task agents in its training data (Figure 18 right). AMAGO can directly
optimize for fast adaptation over the given horizon $H$, and learns a much
more efficient strategy over the first few episodes in a new environment (Fig.
18 left).

Figure 18: Fast-Adaptation with Long-Context Transformers. AMAGO uses a long-
context Transformer to directly optimize for fast adaptation to new
environments (Left), while supervised methods like AD converge at
approximately the speed of the single-task agents that generated its training
data (Right).

##### Passive T-Maze.

Ni et al. [67] is a concurrent work that includes a T-maze experiment that
effectively unit-tests the recall ability of in-context agents; any policy
that achieves the maximum return must be able to recall information from the
first timestep at the final timestep $H$. Their results show that recurrent
approaches fail around context lengths of $100$ (matching the results in
[66]), but that Transformers can stretch as high as $H=1,500$. AMAGO is able
to learn the optimal policy in this environment all the way out to the GPU
memory barrier with context lengths $H=10,000$ (Figure 6 top right). This task
is intended to isolate memory from the effects of credit assignment and
exploration. However, at extreme sequence lengths the reward begins to test
the sample efficiency of epsilon-greedy exploration. We make a few changes
that preserve the central memory question but make this environment much
easier to evaluate, which are well documented in our source code.

##### Wind Environment.

Figure 6 (bottom right) reports scores of a sparse-reward meta-RL task from
[66]. This task has continuous actions and is noted to be sensitive to
hyperparameters. We evaluate AMAGO with a Transformer and RNN trajectory
encoder — keeping all other details the same as in the Dark-Key-To-Door
experiment. The RNN encoder is much more effective at the shorter sequence
lengths in this environment ($H=75$) than the Dark-Key-To-Door environment. It
appears that the RNN is slightly more sample efficient than the Transformer
when using the same update-to-data ratio.

###  C.3 Package Delivery

Package Delivery is a toy problem designed to test our central focus of
learning goal-conditioned policies that need to adapt to a distribution of
different environments. The agent travels along a road and is rewarded for
making deliveries at provided locations along the way. However, there are
“forks” in the road where it needs to decide whether to turn left or right. If
the agent picks the wrong direction at any fork, or reaches the end of the
road, it is sent back to the starting position. The correct approach is to
advance to a fork in the road, pick a direction, note whether it was the
correct choice, and then recall that outcome the next time we reach the same
decision. For an added factor of variation, the action that needs to be taken
to be rewarded for dropping off a package is randomized in each environment,
and needs to discovered by trial-and-error and then remembered so we do not
run out of packages.

####  C.3.1 Environment and Task Details

Environment Distribution $p(e)$: we randomly choose $f\sim U(2,6)$ locations
for forks on a road with length $L=30$. Each fork is uniformly assigned a
correct direction of left or right. These hyperparmeters could be tuned to
create more difficult versions of this problem. We also uniformly sample a
correct action for package delivery from $4$ possibilities.

Goal Space $\mathcal{G}$: all of the locations $(0,\dots,L=30)$ where packages
could be delivered.

Observation Space: current position, remaining packages, and a binary
indicator of arriving at a fork in the road.

Action Space: $8$ discrete actions corresponding to moving forwards along the
road, left/right at a fork, standing still (no-op), and the $4$ actions that
could correspond to successfully dropping off a package (depending on the
environment $e$).

Instruction Distribution $p(g\mid e)$: we randomly choose $k\sim U(2,4)$
locations (without replacement) that are not road forks for package delivery.

Task Horizon $H$: This is the shortest problem in our goal-conditioned
experiments with a maximum horizon (and therefore a maximum sequence length)
of $H=180$.

####  C.3.2 Additional Results and Analysis

Figure 7: Delivery Results (reproduced here for convenience).

Figure 7 compares ablations of AMAGO on the Package Delivery domain. The main
result is that AMAGO can successfully learn the memory-based strategy
necessary to adapt to new roads. As expected, removing its context-based meta-
learning capabilities (“w/o Memory”) greatly reduces performance. The memory
challenge of generalizing to new environments makes rewards difficult to find
and relabeling is essential to learning any non-random behavior (“w/o
Relabeling”). This gap between relabeling and not relabeling is a common theme
across all our experiments, and causes most external baselines to fail in an
uninformative way. For this reason our experiments focus heavily on ablations
of our own method that can remove the barrier of reward sparsity and evaluate
other factors. Figure 7 directly compares ablations with and without multi-
gamma learning. Multi-gamma is one of several implementation details in AMAGO
meant to enable stable learning from sparse rewards with Transformers.
Filtered behavior cloning (“w/o Filtered BC”) [96] is another such detail, but
we can see that its impact goes from quite positive (Fig. 7 left) to slightly
negative (Fig. 7 right) when combined with multi-gamma. Filtered BC and multi-
gamma are addressing a similar problem, but multi-gamma does so more
effectively and exposes some of the drawbacks of filtered BC once the sparsity
issue is resolved. For example, filtered BC can lead to a higher-entropy
policy by cloning sub-optimal actions that leak through a noisy filter that
should have been avoided.

AMAGO randomly relabels trajectories to a mixture of returns between the true
outcome and a complete success. The “Relabel All” curves follow the goal-
conditioned supervised learning (GCSL) (Sec. 2) approach of relabeling every
trajectory with a successful instruction. While this may improve data quality,
it turns the value learning problem from predicting if a goal will be achieved
to when it will happen, because from the agent’s perspective every goal always
does [99]. AMAGO’s high discount factors cause this to have low learning
signal, making actor optimization difficult. Multi-gamma improves performance
(Fig. 7 right), because many of the $\gamma$ values used during training have
a short enough horizon to mask the problem. Relabeling all trajectories to be
a success while turning off the BC filter and policy gradient actor loss would
create a GCSL method, which have shown promise with large Transformer
architectures (Sec. 2). However we were unable to achieve competitive
performance with this setup and prefer a more traditional RL learning update.

###  C.4 MazeRunner

MazeRunner is a difficult but easily simulated problem that combines sparse
task-conditioning with long-term exploration. The agent needs to learn to
navigate a randomly generated maze in order to reach a sequence of goal
coordinates. Maze navigation problems highlight the online
exploration/exploitation problem at the center of context-based adaptation
[105]: when the agent has plenty of time left in the task it can make sense to
explore widely, but when time is running out it may need to be greedy and
pursue small variations of the best known route to the current goal.

Figure 19: $\mathbf{30\times 30}$ MazeRunner. Rendered at scale with the agent
starting in the bottom center and goal locations shown as colored trophies.

####  C.4.1 Environment and Task Details

Environment Distribution $p(e)$: we randomly generate an $N\times N$ maze, and
then manually adjust the bottom three rows to have a familiar layout where the
agent starts in a small tunnel with open space on either side. An example
$30\times 30$ environment is rendered in Figure 19. We can optionally generate
a random permutation of the action space, which adjusts the actions that
corresponds to each direction of movement.

Goal Space $\mathcal{G}$: all of the locations $((0,0),\dots,(N,N))$ in the
maze.

Observation Space: $4$ depth sensors from the agent’s location to the nearest
wall in each direction. By default we include the $(x,y)$ coordinates of the
agent’s position, which can be removed for an extra challenge that forces the
agent to self-localize.

Action Space: Either $4$ discrete actions or a $2$-dimensional continuous
space that is mapped to the $4$ cardinal directions. The action directions can
be randomized based on the environment $e$.

Instruction Distribution $p(g\mid e)$: we choose $k\sim U(1,3)$ locations
(without replacement) that are not covered by a wall of the maze.

Task Horizon $H$: The task horizon creates a difficult trade-off between
exploration and exploitation. There can also be worst-case scenarios where the
layout of the random maze and position of the goals make completing the task
impossible in a fixed time limit. We compute difficulty statistics with an
oracle tree-search agent and use the results to pick safe values where the
$15\times 15$ version sets $H=400$ while $30\times 30$ sets $H=1,000$.

####  C.4.2 Additional Results and Analysis

Figure 8: MazeRunner Results (reproduced here for convenience).

#####  $\mathbf{15\times 15}$ MazeRunner.

Figure 8 (left) evaluates the default version of MazeRunner with $N=15$.
Without random actions dynamics and with the current status of the task as an
input to each timestep, a standard Recurrent PPO baseline could learn the same
kind of adaptive generalization as AMAGO. The lowest-feature ablation of
AMAGO’s core agent (“w/o Relabeling, w/o Multi-Gamma”) uses a Transformer to
roughly match the performance of Recurrent PPO on the $15\times 15$ domain.
However, our other technical improvements can almost double that performance,
albeit with high variance (“w/o Relabeling, w/ Multi-Gamma”). Adding
relabeling to create the full AMAGO method leads to nearly perfect performance
in $15\times 15$ with low variance (“AMAGO”). Position information $(x,y)$ was
included to help baseline performance, but we find that AMAGO can achieve a
$91\%$ success rate without it.

#####  $\mathbf{30\times 30}$ MazeRunner.

The larger maze size creates an extremely challenging sparse exploration
problem, and all ablations of the relabeling scheme fail along with the
Recurrent PPO baseline. AMAGO performs well and uses its memory to efficiently
explore its surroundings (“w/o Memory” has about half its final success rate).
We believe there is room to continue to use the MazeRunner environment at
larger maze sizes and task horizons to evaluate long-term-memory architectures
at low simulation cost. We experiment with adding a dynamics modeling
component term (“w/ World Model”) to AMAGO’s training objective. World
modeling is a natural fit for AMAGO’s long-context Transformer backbone but
does not appear to have an impact on this low-dimensional domain. The “w/o
Planning” baseline hides future goal locations from the task-conditioned
input. In theory, this effects AMAGO’s ability to make
exploration/exploitation tradeoffs and take routes that maximize the chance of
finding all $k$ locations, and does make a noticeable but small difference as
the agent’s strategy improves. However, $1,000$ timesteps may be too generous
to highlight this tradeoff, and the MazeRunner domain does not require any
long-term behavior other than solving the task. This gap grows much larger in
more complicated domains like Crafter, where multi-stage tasks require more
planning to gather resources and stay alive long enough to succeed.

##### “Rewired” Action Space Adaptation.

Figure 9 shows the final performance of several methods with the randomly
permuted action space feature enabled. Generalizing to new action dynamics
requires a level of adaptation above standard Recurrent PPO, so our baselines
shift to full meta-RL methods. While some multi-episodic methods can be
modified to work in a zero-shot setting [107], we focus on three techniques
that are more natural fits for this problem: variBAD [106], HyperX [107], and
RL2 [20, 95]. In an effort to use validated external implementations of these
algorithms, we switch to the continuous action space version of MazeRunner. We
choose the hyperparameters from the largest-scale experiment in the original
codebase, with some tuning performed by solving easier versions of the
problem, e.g., $15\times 15$ without permuted actions, and $11\times 11$ with
permuted actions. We would expect HyperX to outperform variBAD here due to its
focus on hard-exploration problems, and we do see this trend on easier
environments. However, only RL2 shows signs of life on $15\times 15$ permuted
actions (Fig. 9), possibly because RL2 is more compatible with explicit
partial observability and the memory that comes from navigating around the
maze333Comparing results across different RL implementations is hard, and it’s
also possible that RL2’s performance is the result of it being the least
sensitive to tuning when transferred to a new domain.. $30\times 30$ is too
sparse for RL2, but AMAGO’s relabeling and long-sequence improvements allow us
to nearly recover the metrics in Figure 8 while adapting to the randomized
action space.

###  C.5 Crafter

Figure 20: Examples of Crafter Terrain Generation (Left). Crafter generates
new landscapes and locations of rare resources every episode (reproduced from
[33]). Crafter Observation (Right). An example of the agent-centric view that
forms observations.

Crafter [33] is a simplification of Minecraft with 2D graphics and a high
framerate that facilitates research on building agents with the variety of
skills necessary to survive and develop more advanced technology in a
procedurally generated environment. The standard evaluation setting provides a
lightly shaped dense reward that encourages survival and avoiding danger,
while giving a sparse reward the first time the agent completes each of the
$22$ “achievements” in a new episode. This represents an undirected kind of
multi-task learning, where the agent receives equal credit for completing any
of its skills whenever the opportunity arises, and cannot be prompted with a
specific objective. We create an instruction-conditioned version that creates
a challenging generalization problem over distributions of procedurally
generated tasks and environments. The advantage of our instruction-conditioned
format is that once an agent is trained we can direct it to complete any task
we would like. However, rewards are much sparser and require more exploration,
because agents are rewarded only when they complete their goal.

####  C.5.1 Environment and Task Details

Environment Distribution $p(e)$: Crafter generates a new world map that
randomizes the locations of caves, lakes, and rare resources upon every reset.
Three examples are shown in Figure 20.

Goal Space $\mathcal{G}$: We create goals from Crafter’s $22$ achievements,
which include those listed in Table 1 as well as “collect sapling” (where all
agents have a near $100\%$ success rate) and more difficult goals like “make
iron sword” and “collect diamond” (where all agents have a near $0\%$ success
rate). We add more complexity by creating goals for traveling to every $(x,y)$
location spaced $5$ units apart, such as “travel to (5, 5)” and “travel to
(60, 15)”. The travel goals are more difficult than they may appear because
mountains, lakes, and hostile enemies can make navigation require long-term
planning and survival skills. An “Expanded” version of Crafter adds goals for
placing blocks at every $(x,y)$ location. The Expanded goal space is primarily
designed to answer research questions about breaking through the sparse
exploration barrier that prevents AMAGO and existing baselines from mastering
all of Crafter’s achievements (Table 1), and is discussed in more detail
below. All goal strings are represented as three tokens – e.g., “$<$make$>$
$<$stone$>$ $<$sword$>$”, “$<$travel$>$ $<50>$ $<5>$”, “$<$collect$>$
$<$coal$>$ $<$PAD$>$” — which is meant to improve generalization across goals
relative to one-hot identifiers.

Observation Space: Crafter observations are egocentric views of the agent and
its surrounding area (Figure 20 right), along with health, thirst, and
inventory data. The graphics have a similar complexity to Atari environments
[6], and this benchmark is more focused on general behavior than complex
representation learning. AMAGO supports pixel observations, but our
experiments default to a version that maps the image (Figure 20 right) to
texture IDs, which can be processed with an embedding layer rather than a CNN.
This simplification shortens training times and lets us run many ablations,
and is inspired by the NetHack Learning Environment [50] — another high-
throughput generalization benchmark. However, we do evaluate the full AMAGO
method on pixel observations with a CNN (Table 2), and the results are very
similar to the simplified version.

Action Space: Crafter has $17$ discrete actions corresponding to player
movement, tool making, and item placement.

Instruction Distribution $p(g\mid e)$: By default, our instructions are
generated by randomly choosing an instruction length $k\sim U(1,5)$ and then
filling that length by weighted sampling (with replacement) from the goal
space. We pick sample weights that ensured non-relabeling baselines could make
some progress by down-weighting the hardest goals (like diamond collection).
We defer the exact weights to our open-source code release. One benefit of
hindsight instruction relabeling is that AMAGO will modify failed trajectories
to include logical progressions of goals that actually occur in the
environment. Therefore it is more important that our instruction distribution
be sufficiently random and lead to a diverse range of experience than for the
instructions to make logical sense from a gameplay perspective. However, we
experiment with adding a “Tech-Tree” option that uses the ground-truth Crafter
achievement progression (see Crafter Figure 4 [33]) to generate instructions.
This experiment is discussed in more detail below.

Task Horizon $H$: We enforce a maximum episode length of $H=2,000$ timesteps.

####  C.5.2 Additional Results and Analysis

Table 1 evaluates agents on single-goal instructions and lets us measure
AMAGO’s knowledge of the $22$ main skills similar to existing Crafter results
(Table 1 top section). We add a baseline version of our core AMAGO agent
trained with Crafter’s default reward function to control for changes in state
representation, sample limit, memory, and maximum episode length. Moving from
the dense-reward undirected reward function to sparse instruction-conditioning
leads to a dramatic decline in performance without relabeling or long-term
memory (Table 1 middle section). The lower section of Table 1 compares
versions of our method with different instruction-generation and relabeling
schemes, which will be explained below. AMAGO has learned to complete
Crafter’s achievement progression up to collecting iron despite the added
difficulty that comes with being task-conditioned.

|  | Collect
---
Wood
| Collect
---
Drink
| Place
---
Plant
| Wake
---
Up
| Place
---
Table
| Wood
---
Sword
| Eat
---
Cow
| Defeat
---
Zombie
| Wood
---
Pickaxe
| Collect
---
Stone
| Place
---
Furnace
| Stone
---
Pickaxe
| Stone
---
Sword
| Collect
---
Coal
| Defeat
---
Skeleton
| Collect
---
Iron
Rainbow (Dense) [41] | 74.9 | 24.0 | 94.2 | 93.3 | 52.3 | 9.8 | 26.1 | 39.6 | 4.8 | 0.2 | 0.0 | 0.0 | 0.0 | 0.0 | 0.7 | 0.0
DreamerV2 (Dense) [34] | 92.7 | 80.0 | 84.4 | 92.8 | 85.7 | 40.2 | 17.1 | 53.1 | 59.6 | 42.7 | 1.8 | 0.2 | 0.3 | 14.7 | 2.6 | 0.0
AMAGO (Dense) | 99.9 | 93.3 | 99.9 | 95.8 | 99.8 | 99.1 | 81.1 | 91.3 | 99.4 | 97.5 | 93.6 | 86.3 | 92.3 | 69.5 | 53.5 | 0.0
w/o Relabeling | 97.3 | 79.2 | 99.9 | 4.5 | 94.4 | 0.0 | 15.2 | 68.9 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0
w/o Memory | 99.8 | 91.3 | 99.9 | 91.5 | 98.9 | 98.4 | 62.2 | 85.7 | 99.2 | 39.3 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0
w/o Multi-Goal | 99.8 | 93.1 | 99.9 | 97.8 | 99.4 | 97.8 | 93.9 | 88.9 | 99.5 | 83.1 | 0.0 | 0.0 | 0.0 | 0.0 | 10.4 | 0.0
w/o Planning | 99.9 | 98.1 | 98.7 | 93.6 | 95.7 | 93.9 | 36.0 | 82.0 | 98.6 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0
w/o Importance Sampling | 99.9 | 98.4 | 99.8 | 93.3 | 99.8 | 99.9 | 98.9 | 95.1 | 99.9 | 96.1 | 92.1 | 98.6 | 97.5 | 81.5 | 57.5 | 0.0
AMAGO | 99.9 | 99.9 | 99.3 | 99.9 | 96.7 | 99.8 | 96.6 | 97.2 | 99.8 | 99.5 | 94.2 | 97.8 | 98.3 | 90.1 | 76.5 | 0.0
AMAGO (w/ Tech Tree) | 99.9 | 98.3 | 99.8 | 96.1 | 99.9 | 99.8 | 98.5 | 97.3 | 99.8 | 97.8 | 97.0 | 99.0 | 99.5 | 84.0 | 58.4 | 0.0
AMAGO (Expanded) | 99.9 | 97.4 | 99.9 | 94.3 | 99.9 | 99.8 | 97.9 | 96.7 | 99.9 | 98.6 | 97.7 | 99.0 | 99.0 | 84.8 | 40.4 | 0.0

Table 1: Crafter Achievement Success Rates (%). We compare undirected shaped
reward agents (top section) with directed (sparse-reward) ablations of our
agent (middle). AMAGO (bottom) recovers all the skills of the undirected
reward function while being steerable/instruction-conditioned.

##### Exploring with Instructions.

Figure 11 highlights AMAGO’s automatic exploration in Crafter, which is an
important reason to use instructions and relabeling even in domains where we
might only need to solve single-goal tasks. AMAGO’s instruction format allows
us to discover new behavior by following sequences of steps from our own goal
space, and turns exploration into a multi-task generalization problem. The
process works in three steps. First, relabeling lets AMAGO master sequences of
goals that are frequently achieved early in training. Second, we follow
randomly generated instructions that happen to be made up of those easier
goals, eventually reaching the frontier of what we have already discovered.
Random exploration then has a realistic chance of finding new outcomes. We can
verify these first two steps by ablating AMAGO’s ability to follow multi-goal
instructions during training (Table 1 “w/o Multi-Goal”), and observing a steep
decline in performance on hard-exploration goals like tool making. In the
third and final step, AMAGO learns to internalize the new discovery so that it
can be achieved without the help of the instruction curriculum. We suspect
this is closely tied with AMAGO’s ability to observe the full instruction that
led to the discovery, and test this by hiding the full task (Table 1 “w/o
Planning”), which fails to discover hard goals despite following multi-goal
instructions. We note that Crafter’s default reward function of giving $+1$
the first time each achievement occurs in an episode leads to this same
exploration behavior. The best way for an agent to maximize the Crafter reward
is to quickly enumerate every skill it has already mastered, and then start
exploring for new rare possibilities. AMAGO’s instruction format and
relabeling scheme create a way to bring this behavior to a goal-conditioned
setting.

Figure 21: Crafter Instruction Success Rates. Bar labels are average returns.
AMAGO generalizes well to user-prompted instructions in new Crafter
environments.

Figure 21 expands on the results in Figure 10 to show more examples of user-
selected multi-goal instructions in addition to “All Random Instructions”,
which corresponds to an expectation over the procedurally generated $p(g\mid
e)$. Qualitatively, the agents show a clear understanding of instructions and
typically fail by rushing to complete tasks during nighttime when they become
surrounded by enemies.

Figure 22: Learning Curves in Crafter. Goal Importance Sampling improves
sample efficiency by shortening the “plateaus” between discoveries. However,
both agents converge to the same exploration barrier (Table 1). “w/o Goal
Importance Sampling” final performance is indicated by a dashed yellow line.
This plot records train-time success rates, which are lower than the
validation scores in Fig. 21 due to random exploration noise. The sample
efficiency gain is consistent but this plot shows one random seed to highlight
the sudden spikes in performance of individual training runs.

##### Impact of Goal Importance Sampling.

Our results show that AMAGO demonstrates similar final performance with and
without our goal importance sampling technique (Appendix B). Unlike many RL
problems where learning has a diminishing rate of return, AMAGO learning
curves in Crafter can have a sigmoid pattern where our agent plateaus for long
periods of time before suddenly discovering a new skill that unlocks success
in a higher percentage of tasks. Goal Importance Sampling shortens these
plateaus by prioritizing new skills during relabeling, and this leads to a
significant increase in sample efficiency (Figure 22). Figure 23 demonstrates
this goal prioritization on the kind of low-performance data AMAGO sees early
in training. Unfortunately, all agents eventually run into the same
exploration barrier of consistently finding the iron needed to keep moving
through Crafter’s tech-tree. Below we discuss several preliminary attempts to
address this problem.

Figure 23: Importance Sampling’s Impact on Alternative Goal Selection. We use
replay buffer data from a low-performance ablation (“w/o Planning”) to
demonstrate how our rarity-based relabeling scheme (Appendix B) improves
sample efficiency by prioritizing challenging goals. This plot is interpreted:
“given that a particular goal is achieved in a trajectory that needs to be
relabled, what is the probability that it will be chosen to create an
alternative instruction?” The rarity-based scheme exaggerates the frequency of
challenging goals like “collect coal” and “make stone sword”, and diminishes
trivial goals like “travel (30m, 30m).” As the agent improves, goals that were
once rare become more common and this effect is less extreme.

##### Following “Tech-Tree” Instructions and “Expanded” Crafter.

While AMAGO learns many of Crafter’s core skills, it still converges without
mastering all of the goal space. Our results lead us to believe there are two
main components driving AMAGO’s exploration, and it is not clear which is the
main bottleneck. First is the instruction distribution $p(g\mid e)$ that forms
exploration plans for new goals (see above). The exploration effect relies on
occasionally generating useful instructions that can lead to new discoveries,
and could be improved if the instructions we provide are more likely to be
helpful. We can use ground-truth knowledge of Crafter’s skill progression to
bias $p(g\mid e)$ towards instructions that form useful exploration plans. In
this “Tech-Tree” version, half of the instructions are generated by sampling
sub-sequences of Crafter’s progression up until the most advanced goal we have
previously achieved. Unfortunately, this does not lead to the discovery of new
skills (Table 1 “w/ Tech-Tree”). The second opportunity is to expand the goal
space that forms exploration curricula; if AMAGO is exploring by generalizing
over many tasks, adding more goals may improve exploration. “Expanded” Crafter
experiments with this idea and creates tech-tree instructions with added goals
for being nearby rare resources like iron and diamond. This does not fully
solve the problem, but shows some signs of improvement, as multi-goal prompts
like “make stone pickaxe, collect iron” now find iron nearly $40\%$ of the
time (Figure 21), while “make stone pickaxe, collect iron, collect iron” finds
enough iron to make tools with $23\%$ success. Finding better ways to improve
instruction generation (preferably without the domain knowledge used here) and
extend the goal space by learning in more complex multi-task domains are
exciting directions for future work.

##### Planning and Long-Horizon Tasks.

AMAGO’s ability to see all $k$ steps of its instruction is meant to allow
value learning to form long-term plans. If we only provide the current step,
our instruction is effectively handling all the agent’s planning. However, if
we provide the full instruction, AMAGO is free to explore its environment with
the future in mind and take actions that are not directly related to
completing the task but will better prepare it for later events. For example,
we might prioritize gathering food and water when a Crafter instruction will
take more than a full day to accomplish. This kind of planning is less
noticeable in simpler domains like MazeRunner, where planning is mostly needed
for dealing with time pressure.

Figure 24: Impact of AMAGO’s (Implicit) Planning on Crafter Instructions.

It is difficult to disentangle long-term planning improvements from the
exploration advantages that come from full-task inputs (see above). Still, we
investigate this in Figure 24 by comparing the success rates of instructions
with increasing length. We can use the fact that our default instruction
distribution samples each step independently to get a rough upper bound on
performance: if AMAGO can solve 85$\%$ of single-goal instructions, the best
success rate we can expect on a two-goal instruction is $.85^{2}$, and so on.
This is not a fair expectation because each additional goal extends the length
of the episode and leaves more time for starvation or dangerous enemies to
cause failure. However, AMAGO holds to this line quite well despite solving
complex goals that make long instructions take hundreds of timesteps.
Qualitatively, AMAGO does demonstrate an ability to go out of its way to
gather extra food, though we are yet to observe the shelter-building noted as
an emergent behavior in the original Crafter results [33]. We would ideally
see the “w/o Planning” ablation fall well below its upper bound because it
would be less careful about resource management. However, it learns to solve
so many fewer complex goals to begin with that looking at long instructions
that succeed reduces us to a sample of mostly trivial tasks that would not
require much planning at all.

##### Instruction-Specific Behavior.

One concern is that because Crafter only has $22$ different core skills, our
agents could learn an uninteresting strategy of ignoring their instruction and
cycling through every behavior they have learned to accomplish. This is one of
the reasons we add travel and block placement goals, which brings the total
goal count high enough where this strategy would be very difficult to learn
and execute. While AMAGO often achieves goals that are not part of its
instruction, this mainly happens as a survival reflex (e.g., defeating an
enemy), or in a way that could be considered long-term planning (e.g., making
more powerful tools when future goals involve mining or digging through
mountains for travel). Some of the manually-generated instructions in Figure
21 were specifically chosen to be unrelated to the natural progression of
Crafter’s achievement system. For example, “travel to (10, 10), place stone,
travel to (50, 50), place stone” requires the agent to traverse nearly the
full length of its world and would not involve advanced resource-gathering.
This task is difficult because long-distance travel attracts dangerous
enemies, but even when AMAGO fails it does so with a clear understanding of
what was being asked and does not waste time on unrelated objectives.

##### Relabeling Corrects for Missing Instructions.

An advantage of hindsight instruction relabeling is that the train-time
$p(g\mid e)$ it generates includes every goal that is accomplished, rather
than just those we assign during rollouts. This lets AMAGO learn about
instructions that could not be generated by the domain’s natural instruction
distribution. We create a simple demonstration of this by removing the “wake
up” goal from our instruction generation. Because this goal will occur
unintentionally, AMAGO learns to complete it when requested $99.9\%$ of the
time (Table 1). However, the “w/o Relabeling” ablation has never seen these
goal tokens and its $4.5\%$ success rate represents the episodes where it is
achieved by chance due to a stochastic policy.

##### Learning from Pixels.

Our main experiments map Crafter image observations to discrete block
textures. This is primarily a compute-saving measure that lets us study more
ablations. However, we do evaluate the full AMAGO method on pixel
observations, with a summary of key results in Table 2. Learning from pixels
performs similarly to texture observations on all single-goal instructions.
However, there is a roughly $9\%$ drop in performance on the full range of
instructions $p(g\mid e)$. Further analysis reveals that this decline is
caused by success rates in $4$ and $5$-goal instructions.

|  | Collect
---
Sapling
| Collect
---
Wood
| Collect
---
Drink
| Place
---
Plant
| Wake
---
Up
| Place
---
Table
| Place
---
Furnace
| Eat
---
Cow
| Defeat
---
Zombie
| Wood
---
Pickaxe
| Wood
---
Sword
| Collect
---
Stone
| Stone
---
Pickaxe
| Stone
---
Sword
| Collect
---
Coal
| Defeat
---
Skeleton
| Collect
---
Iron
| All
---
Random
Instructions
Textures | 99.9 | 99.9 | 99.9 | 99.3 | 99.9 | 99.9 | 94.5 | 96.2 | 96.6 | 99.9 | 99.9 | 99.5 | 97.8 | 98.3 | 90.1 | 76.5 | 0.0 | 62.1
Pixels | 99.9 | 99.9 | 99.4 | 99.9 | 93.7 | 99.5 | 94.0 | 99.3 | 96.7 | 99.9 | 99.9 | 98.3 | 98.8 | 97.7 | 89.8 | 80.0 | 0.0 | 53.9

Table 2: AMAGO Crafter Success Rates (%) From Pixels.

##  Appendix D Policy Architecture and Training Details

##### AMAGO Hyperparameter Information.

Network architecture details for each of our main experimental domains are
provided in Table 3. Table 4 lists the hyperparameters for our RL training
process. Many of AMAGO’s details are designed to reduce hyperparameter
sensitivity, and this allows us to use a consistent configuration across most
experiments.

| Package Delivery | MazeRunner | Crafter | POPGym
---|---|---|---|---
Transformer |  |  |  |
Token Emb. | MLP (128, 128) | MLP (128, 128) |  | Discrete Emb. or Atari-Sized CNN
---
$\rightarrow$ MLP (384, 384)
MLP (512, 512)
Model Dim. | 128 | 128 | 384 | 256
FF Dim. | 512 | 512 | 1536 | 1024
Heads | 8 | 8 | 12 | 8
Layers | 3 | 3 | 6 | 3
Other Modules |  |  |  |
Actor MLP | (256, 256) | (256, 256) | (400, 400) | (256, 256)
Critic MLP | (256, 256) | (256, 256) | (400, 400) | (256, 256)
Instruction RNN | Token Emb. Dim 13 $\rightarrow$ RNN Hid. Dim. 96 $\rightarrow$ MLP 64
Total Parameters | 2M | 2M | 17M | Varies Per Env (4-11M)

Table 3: AMAGO Network Architecture Details.

| Package Delivery | MazeRunner | Crafter | POPGym
---|---|---|---|---
Critics | 4
Critics in Clipped TD Target | 2
Context Length $l$ | $H$
Actor Loss Weight | 1
Filtered BC Loss Weight | .1
Value Loss Weight | 10
Multi-Gamma $\gamma$ Values | $.7,.9,.93,.95,.98,.99,.992,.994,.995,.997,.998,.999,.9991,.9992,\dots,.9995$
Target Update $\tau$ | .003
Learning Rate | 3e-4 | 3e-4 | 1e-4 | 1e-4
L2 Penalty | 1e-4 | 1e-4 | 1e-4 | 1e-3
Batch Size (in Full Trajectories) | 24 | 24 | 18 | 24
Max Buffer Size (in Full Trajectories) | 15,000 | 20,000 | 20,000 | 80,000
Gradient Updates Per Epoch | 1500 | 1000 | 1500 | 1000
Parallel Actors | 24 | 24 | 8 | 24
Timesteps Per Actor Per Epoch | $H$ | $H$ | $H$ | 1000
Epochs | 600 | 600 | 2000 | 625

Table 4: AMAGO Training Hyperparameters.

##### Compute Requirements.

Each AMAGO agent is trained on one A$5000$ GPU. Results are reported as the
average of at least three random seeds unless otherwise noted. Learning curves
default to displaying the mean and standard deviation of trials. A notable
exception to this is the POPGym learning curves (Appendix C.2), which show the
maximum and minimum values in order to study the stability of the shared
Transformer gradient update. Wall-clock training time varies significantly
across experiments. For reference, POPGym training runs take approximately $8$
hours to complete. AMAGO alternates between data collection and learning
updates for consistent logging and comparisons across baselines. However,
these steps could easily be done in parallel or asynchronously if wall-clock
speed was critical.

##### Sample Efficiency.

The AMAGO training loop alternates between collecting trajectory rollouts from
parallel actors and performing learning updates. Sample efficiency in off-
policy RL is primarily determined by the update-to-data ratio between these
two stages [25, 12]. It has been shown that common defaults are often too
conservative and that sample efficiency can be greatly improved by increasing
the update-to-data ratio [84, 17]. AMAGO’s use of Transformers resets this
hyperparameter landscape. It is unclear how many more or less gradient updates
should be taken when training multi-million parameter sequence models,
particularly when sharing parameters between the actor and critic(s). In long-
term memory and meta-RL problems it also becomes unclear whether the update-
to-data ratio should be measured in terms of trajectories or timesteps. Sample
efficiency is an important advantage of off-policy RL, and we think this is an
interesting direction for study in long-context agents. However, due to
compute constraints and the large number of technical details we are already
evaluating, our experiments make little effort to optimize sample efficiency.
Our update-to-data ratios are essentially never tuned, so it would be
surprising if the current results represent the best trade-off between sample
efficiency and performance.