AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers
Authors: Jake Grigsby
Justin Sasek
Samyak Parajuli
Daniel Adebi
Amy Zhang
Yuke Zhu
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformers
Meta-RL
Multi-Task RL
Primary area: reinforcement_learning

Abstract: Language models trained on diverse datasets unlock generalization by in-context learning. Reinforcement Learning (RL) policies can achieve a similar effect by meta-learning within the memory of a sequence model. However, meta-RL research primarily focuses on adapting to minor variations of a single task. It is difficult to scale towards more general behavior without confronting challenges in multi-task optimization, and few solutions are compatible with meta-RL's goal of learning from large training sets of unlabeled tasks. To address this challenge, we revisit the idea that multi-task RL is bottlenecked by imbalanced training losses created by uneven return scales across different tasks. We build upon recent advancements in Transformer-based (in-context) meta-RL and evaluate a simple yet scalable solution where both an agent's actor and critic objectives are converted to classification terms that decouple optimization from the current scale of returns. Large-scale comparisons in Meta-World ML45, Multi-Game Procgen, Multi-Task POPGym, Multi-Game Atari, and BabyAI find that this design unlocks significant progress in online multi-task adaptation and memory problems without explicit task labels.

AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcemeee Learning with Transformers
Jake Grigsby Justin Sasekt Samyak Parajulit Daniel Adebi
Amy Zhang
Yuke Zhu
The University of Texas at Austin
i Equal contribution
{grigsby, yukez) seriiisbyyykkzeeeeee
utexas. edu
Abstract
Language models trained on diverse datasets unlock generalization by in-context learning. Reinforcement Learning (RL) policies can achieve a similar effect by meta-learning within the memory of sequence model. However, meta-RL research primarily focuses on adapting to minor variations of a single task. difficult to scale towards more general behavior without confronting challenges in multi-task optimization, and few solutions are compatible with meta-RL's goal of learning from large training sets of unlabeled tasks. To address this challenge, revisit the idea that multi-task RI is bottlenecked by imbalanced training losses created by uneven return scales across different tasks. We build upon recent advancements Transformer-bast (in-context) meta-RL and evaluate a simple yet scalable solution where both an agent's actor and critic objectives terms that decouple optimization from the current scale of returns. Large-scale comparisons Meta- Woold ML45, Multi-Game Proegen, Multi-Tase POPGym, Almt Atari, and Al find that this design unlocks significant progress in online multi-task adaptation and memory problems without explicit task labels.
converted classification
Introduction
generative models trained to imitate human language and behaviors from web datasets have unlocked unprecedented generality in machine learning [1, 2]. While evaluations of these large language models (LLMss) are organized into discrete benchmarks [3] with specific themes like competitive mathematics [4] or reading comprehension [5], their knowledge transfers to countless problems. LLMs' flexibility is largely due an "in-context learning" effect that emerges when input sequences grow long enough to resemble a dataset of examples for new task [6. 7]. However, predicting the most likely continuation to tasks where high-quality demonstrations are available.
sequence is misaligned with optimal control and limits us
Online Reinforcement Learning (RL) lets models continuously self-improve and directly optimizes performance. The RL equivalent of in-context learning is a subset of meta-RL techniques that use sequence models for memory and adaptation [8, 9, 10]. Meta-RL can be viewed as a task inference problem in which an agent explores its new surroundings to discover information that it can use to improve decision-makiie problems. However, current applications are quite narrow usually focusing minor variations of what could considered a single task. For example, imagine platformer video game with procedurally generated levels. The game has consistent theme and controls, but each level introduces a slightly different reward function, physics, and layout. meta-learning policy trained on this game would learn to identify the relevant changes in new level in order maximize returns. By most formal definitions, any two levels from this game would be considered two different "tasks," but seems unlikely that anyone outside the field would refer to the resulting policy as "multi-task" agent.
[11, 12, 13]. In theory, meta-RL can adapt to wide range of control
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Meta-RL
Multi-Task RL
Generalist Agent
Adaplation unseen variations of a single task
Joint optimization multiple tasks
Adaptation and task poverage
+
Figure I: Task Spaces in RL Generalization. Meta-RL agentso adapt to dense variations of a core task. Multi-Task RL overcomes optimization challenges of learning from isolated tasks. Scalable ideas from both areas allow us to extend adaptive agents towards increasingly general behavior.
A truly general agent would also be able to play entirely different games. We will use the term "task" to refer to this kind of multi-game generality while calling different levels of the same game "task variants" [14, 15, 16]. Each task (game) creates a distribution of variants (levels) we can sample. The distinction between two "tasks" and two "variants" of the same task is arbitrary in theory'. However, it is so relevant in practice that having separate terms will be important for clarity.
Ideally, we would train RL agents many different tasks and many variants of those tasks. Recent works have scaled to heterogeneous mixtures of popular RL benchmarks [19, 20, 21], but rely on offline datasets collected by expert single-task agents. Scaling up online learning is challenging because multi-task distributions can create disconnected datasets; as we transition from a single task to small sets of several tasks, reach a barrier where the gap between tasks becomes so wide that practical optimization challenges hinder learning more than any knowledge transfer between them is helping [15, 22, 23]. This conflict is particularly limiting in meta-RL because multi-task behavior should otherwise be a simple case of task inference. Meta-RL agents learn to infer the identity of task variants based on limited information
which becomes more challenging as those variants become more similar. Adding an entirely new task with clearly distinct visuals or controls does little to the difficulty of this problem, and yet it can dramatically reduce performance because introduces the (mostly unrelated) challenge of optimizing second reward function. Fortunately, this challenge is not unique meta-RL and is a focus of broader research in multi-task RL (MTRL) and supervised learning [24]. MTRL. methods can improve multi-objective training by creating task-specific network parameters or editing the gradients of each task to resolve conflicting updates (Section 2). Unfortunately, these solutions rely on our ability to manually label tasks, scale with the size of our training set, and can abandon meta-RL's goal of adapting to new tasks at test-time. While MTRL might be practical on current benchmarks, it would be more scalable to preserve meta-RL's task inference and improve multi-task optimization without introducing new assumptions (Figure 1).
This paper focuses on training sequence model policies that generalize across hybrid task distributions containing elements of meta-RL, MTRL, and long-term memory. We revisit the idea that MTRL is primarily bottlenecked by imbalanced training losses created by tasks with returns on significantly different scales [25]. However, we approach this problem from a meta-RL. perspective where task inference needs to be learned; we do not allow knowledge of task labels to influence training in any way. We experiment with Transformer [26] policies trained by mixtures of actor and critic objectives that are either directly or indirectly dependent on the scale of returns in any given task. Comparisons on Berad ML45 [17], Multi-Task POPGym [27], Multi-Game Procgen [28], Multi-Game Atari [29]. and Multi-Task Baby AI [30] evaluate the importance of scale-resistant updates. We find that converting value regression to classification [31] and policy improvement to filtered imitation learning [32] leads significant improvements in these hybrid settings. Our results suggest that scale invariance should be a priority when designing multi-task agents with adaptive memory. Our empirical success with classification updates further supports recent observations that sequence-based RL can improve by borrowing technical details from supervised sequence modeling [33, 34, 35] while retaining the online aaffimmprvveeett process that makes RL useful [36, 37].
'The distinction between tasks/variants is also called non- paramettic/paraamttic variation [17, 18). Variants are often generated by randomizing a known set of parameters in a single (simulated) environment.

---

2 Background
Memory-Basee Meta-RL. The objective of a meta-RL agent is to adapt its policy to maximize subset of "black-
returns over horizon of more attempts in new task. We focus
box" approaches that reduce meta-learning to the problem of training sequence model policies with end-to-end RL [8]. We view adaptation as form of partial observability [12], where instead of inferring the current state from a sequence of observations (on...., ot), we try to infer the state and the identity of the task from trajectories (7) of observations, actions (a), and rewards
(r): Top = (oo, ao, To, 01..... Qt-1; 1t-1, 9r) [13]. A task's identity is an unobservable variable describing aspects of its dynamies that vary across the training distribution [38]. Because improved estimates of the current task would increase returns, meta-reasoning can arise implicitly in the latent space of a sequence model that is trained by standard RL updates.
This sequence- -bsed meta- RL framework, which began with RL [9, 10, 39], has two key advantages. First, training policies on trajectory sequences turns standard partial observability and zero-shot generalization [40] into special cases. In practice, this creates flexible agents that blur formal boundaries between memory, generalization, and meta-leaming [41]. Second, end-to-end RL. enables self-improvement without demonstrations. Many meta-RI. algorithms improve performance in applications where task identification could benefit from additional assumptions [11, 42. 43, 44]. Another line of work recovers the flexibility of "in-context" sequence learning when a dataset of demonstrations makes self-improvement unnecessary [45, 46, 47, 48]. Most of the algorithmic variety in meta-RL. RL comes from these two areas. In contrast, the RL2² approach has primarily become an engineering problem. Important developments include increasing the adaptation horizon by enabling long-contex: sequence models in an RL setting [49, 50, 511, and shift away from on-policy updates in favor of off off-policy actor-critics [52, 53, 54J. Combinations of these details allow Transformers with long-term recall to self-improve on large datasets of recycled data [55, 56, 57].
Multi- Optimization. Multi- learning addresses trade-off where conflicting objectives decrease performance relative to single-task learning despite an increase in training data. One approach uses specialized optimizers that edit gradients to balance tasks' local optima [22, 58, 59, 60, 61]. However, it can be as effective (and less expensive) to rescale each task's contribution to the overall loss [23, 25, 62, 63, 64]. These techniques require the ability to identify the current task, which is commonly available in multi- task settings: we can view MTRL as a meta-RL problem where the ground-trute identity of each environment has been one-hot labeled in advance. MTRL often depends on task labels in more subtle ways, including networks with task-specific output heads [65] and datasets that balance sampling [66] or normalize rewards [67] across tasks.
Problem Setting. We focus on hybrid generalization problems that require elements of meta- learning, multi-task learning, and long-term memory. Meta-RL often evaluates performance over k >>> 1 attempts or episodes. A "k-episode" objective maximizes the total return over all k attempts and creates an exploration/exploittiin trade-off at test-time [13]. A "k-shot" objective allows for an exploration window where only the final episode counts towards the agent's score [42, 68]. We measure generality by dividing a set of task variants (V) into train and test sets. Meta-RL. benchmarks (informally) characterized by their relative density meaning |V] is large and composed of subtle changes to a common objective. Multi-task problems are created by training a single policy to maximize N qualitatively distinct objectives, such as learning play N Atari games [29]. More general settings might be described as containing N adaptation problems with their own set of variants Vo...., VN.
Scaling Meta-RL Beyond the Multi-Task Barrier. Early meta-RL experiments focused on adapt- ing to minor variations of multi-armed bandits [9], gridworlds [68], and locomotion benchmarks [43, 69, 70J. Despite significant algorithmic progress, these domains are still representative of recent research [44, 45, 54, 71J. Meta-RL requires generating many variations of a core problem while keeping those changes partially observed; is difficult to design a cohesive benchmark that supports the dense level of variety it takes to induce meta-learning [55, 72]. We can create toy meta-RL problems, but there is a wide gap in complexity before we reach naturally occurring domains where long-term adaptation is necessary. For example, one application of meta-RL. would be to learn to master unfamiliar video games over several minutes or hours. Building such an agent would likely involve scaling to large training set of games. It is difficult to make gradual progress towards
3

---

this goal because it begins as memory problem, becomes a multi-task problem, and then only requires meta-learning at an extreme scale. Individual games often require long-term memory, but as we add games, there will inevitably be isolated groups with unique objectives that introduce challenges from MTRL.. Eventually, we reach a level of training diversity where adaptation becomes critical and meta-learning emerges. The central issue considered by this paper is that multi-task optimization limits train-time performance at task counts far smaller than would be necessary for test-time generalization to new tasks. A promising direction is to take a method that is already capable of both memory and meta-learning and then find a way to break the barrier where progress is limited by multi-task optimization. However, if we want adapt to many tasks, we should avoid solutions that scale with the number of training tasks or depend on our ability to identify them manually.
This "multi-task barrier" impacts both long-term applications and current research benchmarks. For a more practical example, consider Meta- World [17]. World is a suite of 50 robotic manipulation tasks. Each task creates a meta-RI. problem by adjusting an unobservable goal so that agents must interpret incoming reward signals to adapt to the current objective. Meta- Woold MLA5 extends generalization by merging meta-RL tasks into a multi-task training set U, Identifying the current task during training should be a special case of meta-RL because the changes between tasks are less ambiguous than the changes in goal location between variants of the same task. For example, an agent might distinguish a pick-and place task from a door opening task after a single observation of the objects on the table. ML45 reserves five tasks to benchmark meta-RL's long-term goal of adapting to entirely new tasks at test-time. However, learning task adaptation from just N = 45 training tasks is not realistic without additional assumptions [17, 73]. ML45 is clearly bottlenecked by MTRL: a similar setup with the meta-RL. component removed is a challenging MTRL benchmark (MT50) by itself [59, 74]. MTRL methods use task labels balance loss functions, separate network architectures, or edit task gradients. However, these techniques scale with N and require a formal identification of each "task" that meta-RL otherwise should not need. Our goal to overcome the same challenge while preserving the meta-RL perspective that task differences are arbitrary so that we can scale to unstructured domains where N is large [19] or infinite. Therefore, even though our experiments will have a clearly defined N, we will enforce strict constraint where the agent can have zero knowledge of which task is active or how many tasks there are (N, [V,l).
3 Multi-Task Adaptation Without Task Labels
We study multi-task adaptation without task labels by building on the flexible memory- based framework in which sequence models optimize standard RL objectives across trajectory inputs (Section
Astion Datmibtiion
Later
Extrituttonn
Actor
hg
Crific
2). Agents observe trajectory slices from a replay buffer with a context length of up to I timesteps (Tr-I:)- The information re-
Transformer
vealed between timesteps (0,, 01:0;;1: TT-1; d1-1) is merged into a single embedding o The sequence (o .........0.. oft) is then passed through a sereence-t--equenne model to produce (hy-t....,hh hy)
Timestop
Og
Timestop Encoder
which serves as a shared representation for output heads represent- ing the critic(s) Q and stochastic policy æ (Figure 2). We focus on the combination of Transformers and off-policy RL updates, which will let us evaluate large-scale policies trained over millions of timesteps. Let y/ rı + +Q(h++1, a' ~ #(h++1) be the one-step temporal difference target, where Q and मे refer to frozen target networks [75]. We can minimize
Figure 2: Transformer-basedd
Actor-Critic Architecture.
standard off-policy actor-critic loss terms [75, 76] in parallel over the context length:
Lemee (t) = =Q((i,, at) - )²)²
(1)
Lacoor (t) = -hQu (hh)) (2)
Following details in AMAGO [57], we compute Equations (1) and (2) with an ensemble of crities
[77] and over several different values of the discount factor Y in parallel [78]. By preventing the parameters of Q from minimizing LActor
update the full architecture with one step of the
joint loss: E Cacuo (i) + XCCeme where À is hyperparameter that balances the objectives for the shared sequence model [56, 57].
Multi-task optimization challenges are not unique to RL [22, 62], but RL typically adds bottleneck that we can remove: both Equations (1) and (2) directly depend on the scale of returns Q in each of our N tasks. When computed simple expectation over tasks, the gradients of CCriti and LActor

---

can favor small relative improvements in tasks with high absolute returns (Figure left). We can almost never assume tasks have similar return scales. For example, even when we design reward functions to have similar initial and optimal returns, differences in task difficulty lead to uneven learning progress and imbalanced Q-values. This issue was proposed and demonstrated by methods such as Multi- -aas PopArt [25]. PopArt's solution is to compute MTRL loss terms on relative basis by creating a new critic output for each of the N tasks. Each critic adaptively normalizes y for its task and automatically rescales its output to help compensate for shifting targets. Simpler versions of this approach (e.g., without network rescaling) are universal in MTRL. However, we cannot use this technique without ground-truth knowledge of task labels. Although we cannot adaptively balance each task's 's loss, we can reformulate Eqs. (1) and (2) to be less sensitive to the scale of Q-
Critics Distributional RL [79] is motivated by representation learning and risk- aware policies, but methods like C51 [80] have a useful side-effect of turning (1) into a classification problem that scales with the number of output bins instead of Q. For this reason, C51-style critic updates have become a key implementation detail in large-scale (offline) MTRL agents [65, 35]. We can keep the classification side-effect without modeling return distributions by converting a scalar Q to discrete label(s) E {O..... B). Two-hot labels create one-to-one mapping between every scalar in fixed range and a discrete distribution where all non-zero probability is placed in two adjacent bins [81, 82]. If QB denotes a modified critic network that outputs probabilities over B bins, the TD-error can be minimized by multi-label classification:
Liiii Ind (t) ttwohot (y.)² log Q (he, oz)
We can reduce tuning of label spac- ing and limits with an invertible trans- form that maps a wide range of returns to a smaller number la-
Mean Squared Erro
Symlog Two-Hot Clas sificatioo
10-1
Relative critic Loss
10-2
bels [83]: we use symlog demon- strated by DreamerV3 [31] and TD-
10-5
10-1
y-10
Y-10 10
MPC2 [20]. We also experimented with using a global (task-agnostic) Q-normalization for this purpose but found that a fixed mapping is im- portant because stabilizing learning against shifting labels required exten- sive hyperparameter tuning. Figure visualizes why we would expect clas- sification to be an improvement in a multi-task setting: Lemme-Ind maps the same relative error in value pre- diction to a similar loss value over a wide range return scales. In contrast, Leiii can bias optimization towards tasks that happen to have large regression targets. This effect can be dependent on proper tuning of the bin count B, upper/lower return bounds, and use of the symlog transform (Appendix A).
10-4
100
100
1,000
1,000
20:31
10,000
10,000
0.0
40
20 40
Target Prediction Error (%)
Target Prediction Error (5)
Value Regression. We plot the value of the standard critic loss (Eq. 1) as function of the relative prediction error of the TD target (3) across four orders of magnitude (left). Y-axes are self-normalized according to the largest displayed value. Two-Ho Hot classification (Eq. 4) maps the same relative error to similar loss values across the different absolute return scales of each task (right).
Figure 3:
A concurrent work [34] evaluates the two-hot critic loss' ability to unlock scaling laws in RL when used as direct replacement for Leiic (1) (or its algorithm-specific equivalent). They find the classification loss improves performance in a range of domains and attribute this to advantages in representation learning and robustness to noisy value targets (y)). Our experiments focus the PopArt hypothesis that saleeinvariianee improves optimization in an online multi-task setting. We will be exploring the impact of classification losses an actor-critic agent with adaptive memory. where we also need to consider the scale of the policy's 's objective (2).
Scale-Resistann Actors. Many policy updates resemble a weighted regression objective where the importance of each action label is scaled by a function f(r, a) that increases with its value or advantage, such as f(T, a) x (x)(A(r))))² Methods differ in where the candidate actions a are sampled from, how values are estimated, and how f(r, a) clips or rescales the update [82, 84, 85,
WWhere A(7, a) = Q(T, a) Q(T, a' ~ (())) is the improvement of over the value of the current policy estimated by actor and critic networks that take trajectory sequences input.
5

---

86, 87]. In the off- off-py setting where actions are sampled from a dataset of previous experience, slight differences advantage estimates and weight functions f create a dense family of similar setaantage-eiighted regression (AWR) methods [86, 87, 88, 89, 90]. AWR is intuitive and simple to implement but directly depends on the scale of returns unless we intentionally avoid this with a binary filter f(T,, = [{A(T, a) > 0} [32, 91]. We follow details from CRR- MRRRBinnry-MMen [32] for a candidate actor loss that is independent of Q:
Lactor-Ind (t) = [(Q oz) I [Q a')] 0} log a(a, h2)
ah~
Intuitively, CRR performs imitation learning (IL) on action labels that we estimate will improve the current policy. Advantage estimates within IncoorInd rely entirely on one-step dynamic programming (I) so that they never become outdated [32, 92]. Because they do not allow the policy to extrapolate to out-of- de-dstrribution actions, weighted IL updates like CRR are best suited to offline RL [93]. In offline RL, the ability to stabilize optimization on static datasets is well worth losing the optimistic exploration of LAeur (2). Choosing to use LactorInnd online implies that the ability to address issues in multi-task optimization is similarly worthwhile trade-off. There are other RI. engineering factors to consider in the specific case of training long-context Transformers. For example, sampling batches of trajectory sequences can create replay ratios high enough to effectively be offline RL [94]. However, the simplest justification for training Transformers on LActor-Ind may come from its similarities to widely used methods for supervised learning in RL.
Transformers and RL via Supervised Learning, Decision Transformer (DT) is a popular approach to training sequence policies that reformulates RL as IL conditioned on target return [36, 95, 37]. DT's advantage is that it reduces to II. on expert datasets, making it a safe improvement over a technique that is already effective for many problems. In field where engineering is critical. and results can be hard to reproduce, there is valuable simplicity in inheriting technical details from supervised sequence modeling. This simplicity has likely contributed to the method's success despite several disadvantages. For example, DT policies are conditioned on the optimal return in order to imitate the best continuation of the current trajectory, but LActerInd does this without needing to define the optimal return in advance. DT determines the value of an action based on the fixed return of its trajectory, while Anctor continuously updates its estimate with the value of the current network. By extension, DT cannot handle stochastic reward functions [96]. Like DT, L Ind becomes II. on expert datasets, and Eqs. (3) and (4) resemble supervised learning with two classification heads.
4 Experiments
Lxcurr Lemiee LActor-ind and Lemiicied create four interchangeable combinations of learning updates where one or both of the actor and critic loss can be dependent or indirectly dependent the scale of Q. We will compare all four update variants, although we are mainly interested in evaluating the two extremes: "Dep. Actor, Dep. Critic" (Cactae Lcmie) and "Ind. Actor, Ind. Critic" (CActor-Ind- InaiicIma). Our experimental setup allows for a direct ablation where all other details can be held fixed. Appendix A includes additional implementation details. Our experiments focus on three main questions: 1) Do scale-resistant actor-critic objectives offer empirical benefits in challenging generalization problems? 2) If so, can these gains be attributed to multi-task return scaling? and 3) What new applications of memory-based policies might be unlocked by improvements in multi-task optimization without task labels?
Meta-World ML45. ML45 is challenging suite of 45 robotic manipulation meta-RL, tasks. We argue that much of MI 45' difficulty can be attributed to variance in return scales: MI.5's training tasks are designed to have similar optimal returns, but they have different levels of difficulty, which causes the scale of Q-values to diverge early in training. The PopArt method of normalizing Q per task may help balance optimization. Methods built on the original Meta- results normalize the returns of each task separately [67], which is simpler approach to the same idea. We force all task identification to be learned within the context of our memory policy so that one-hot labels and the total task size (45) do not influence training. The training environment samples a new task between meta-rollouts and the task identity is not included in the observation. Every task shares a single critic head, and all experience is added to shared replay buffer sampled uniformly at random.
Figure 4 compares the four combinations of learning updates described in Section 3. Actor and critic loss terms that do not scale directly with the Q-values of each task improve skill coverage.

---

Cumulative Tasks with 22 Successes Per 3 Episode Rollout
Avg. Success Rate (%)
Avg. Total Return (3 Episodes)
40
12000
80
Nuzem
10000
60
25
Mall-World
ALL-PPQ
40
15
4000
20
2000
0
0
0M
20M
BOM
80M
100M
40M
100M
Total Timestaps (Arross All Tasks]
Total Timestops (Across All 45 Tasks)
Total Timestops (Across All Tasks)
Ind. Actor, Dep. Critic
Actor, Dep. Critic
Actor, Ind. Critic
Ind. Actor, Ind. Critic
Figure 4: Meta-Worke MT Train Task Results. (Left) Coverage of the 45 manipulation skills measured by an adaptation horizon success rate 2 2/3. (Center) Average success rate over tasks, variants, and 3-episode rollouts. Reference scores for MuZero and RL RL---PPO are gathered from results in [17, 73]. (Right) Total return over a 3-episode adaptation horizon, averaged across tasks and variants. All error bars indicate the maximum and minimum metric across three random trials.
overall success rate, and cumulative return across 3-episode adaptation horizons. Replacing Q regression with two-hot classification appears be the most important change, although substituting the weighted IL policy update improves performance and sample efficiency. Full learning curves for all 45 tasks plotted in Appendix C and highlight the uneven learning progress (and therefore Q-value scale) of tasks throughout training. With sample budget of 100M timesteps, our off-pooicy Transformers trained by max likelihood losses (Lactoc-lad, CCritic-Ind) more than double the success rate of Meta-World's original RL result over the same sample limit while maintaining a 1.4x improvement on its performance 400M timesteps. Our simple one-step Q-learning matches more complex value updates like MuZero [81] at its reported 50M timestep budget [73]. We use general Transformer architecture that would be equally applicable to any sequence modeling problem, but our results are also comparable recent work such as HTrMRL [97] (success rate 85%) that use architectures specialized for multi-episodic RL.
One- Shot Multi-Task Task POPGyn
Normalized Best Single-Task Baseline (%)
Ind. Actar, Ind. Critic
Multi-Task POPGym. Meta-World is de- signed meta-learning benchmark, but its adaptation does not require long-term memory.
100
Dep. Actor, Dep. Critic
PPO-GRU
80
POPGym [27] is a suite of zero-shot POMDP problems. Many POPGym tasks are explicit tests of long-term recall, but others involve state and variant estimation from previous actions and re- wards. We use simple observation and action space padding approach to create multi-task ver- sion of POPGym: we combine all 27 tasks where both the action space and observation space have dimension 30.
Avg. Return
30-
not clear whether zero- shot performance in this setting can be compared against single-task POPGym results, as a strong agent may need lose valuable time identifying which of the 27 memory games is currently playing. Therefore, we evaluate in one-shot setting where the first trial is a free exploration window used to determine which task is active that the second trial creates fair comparison to single-task results [68]. Importantly, we avoid "spoiling" any memory challenges for the second trial by randomly resetting new variant of the
20M
40M
B0M
100M
Total Timestope ps (Across All Taskes
Figure 5: Multi-Task POPGym Results. Re- turns are normalized by single-task experts trained for 15M timesteps. Error bars indicate the maxi- mum and minimum returns across three trials.
task.
compare Transformers trained by the "indirect" and Q-dependent update against the on-policy PPO agent from the POPGym codebase [27] using the best sequence model from the original
7

---

benchmark (a GRU RNN 198]). Figure 5 aggregates the results across all 27 tasks with returns scaled between random policy and the best single-task result (of any sequence model) that appears in Morad et al. [99]. Single-task reference scores are given budget of 15M timesteps, while our multi-task agents see total of 100M across all 27 tasks and train on reward signals for approximately half of those timesteps. Both Transformer learning updates outperform PPO-GRU but are relatively comparable to each other. This may not be surprising because POPGym tasks have similar learning curves and returns bounded in [-1, 1]. However, many of these tasks have sharp short-term credit assignment where the advantage surface an accurate value function would be steep. We expect the standard CActor update to be significant improvement in online exploration, it is encouraging that the offline-stlle Ind can exceed its performance after 15M total timesteps. Learning curves for each task can be found in Appendix C.
Coinrun
Leaper
Ninja
Jumper
Clinber
Avg. Total Return 12 Epiaad))
19.0
18.0
18.0
23.0
25.2-
14.4
24.4
15.4
11.4
10.8
13.2
2.6-
Ind. / Ind. IRowards Scaled)
Niha
Dop. Dep. (Rowards Scateo)
(nd. / /ind. [1 532)
Singls-TaskPPe Referance
Ind. lnd.
Dep./ Dep.
Figure 6: Return Scales in Multi-Game Procgen. We evaluate over 2 episodes in unseen levels after optimizing default or rescaled reward functions. Scores are averaged over two model checkpoints and converted to the default scale. Error bars indicate the difference between two 275M timestep trials.
Multi-Game Procgen. Procgen [28] creates 16 video games with an infinite variety of procedurally generated levels. If memory-based actor-critic learning bottlenecked by the scale of returns across multiple tasks, we should be able to decrease performance by simply rescaling rewards. We begin by training multi-task agents across 2k levels of five games in easy mode. Solid bars in Figure 6 show the performance on unseen test levels, where the scalc-resistant update offers a slight improvement.
5. 5-Gane Procoon Test Levels
Default Rewards
Avg Norma lize Return ber Epsode
0.8
Most Active Critic Label Indices
110
0.10
Actions Approved iniration (%)
0.7
80
82
0.6
0.05
0.5
70-
60
Procgen
0.4
0.00
Procgen (Rescaled)
0.3
Dop./ Dep.
Ind. / Ind.
Rescaled Rewards
0:10
40
POPGym
Meta-World
0.2
(llescaled)
19c
DVD (Rescaled)
90-
0.2
80
0.0
70
0.00
0.2M
0.0M 0.4M 0.8M
2M
1.6M 2.0M
40
80
100
Gradient Steps
Gradient Steps
Training Progress (%)
Figure 7: Learning from Rescaled Returns. (Left) Procgen test level returns normalized across games according to the standard benchmark scale [28]. (Center) Frequencies of the critic classifica- tion label index with the highest probability throughout training. We trim the y-axis to the portion of symlog og [31] space that is relevant to Procgen. (Right) The value of LetttrrIed binary advantage weights create an estimate of the percentage of the replay buffer currently being imitated.
Our multi-task Procgen environment randomly selects new game between resets, which does not account for how episode lengths depend on the agent's skill (Figure 14). We retrain new policies from scratch but scale the rewards of the easiest task (Coinrun) by x 100 while dividing the rewards of the task the agent spends the most timesteps interacting with (Climber) by 10. We would expect the scale of LActor and Lerue to be dominated by Coinrun so that optimization variance overshadows the objectives of the other environments even after learning in Coinrun converges. Hatched bars in Figure 6 show this is indeed the case, with methods recovering most of their original performance in Coinrun while failing to improve in the other four games. The "Ind. Ind." update is far less impacted by the new return scales even though they decrease efficiency in terms of gradient steps (Figure 7 Left). Fig.
8

---

7 (Center) reveals simple explanation where the broader range of returns approximately doubles the number of frequently used labels our critics need learn classify. Multi-task Transformer policies are commonly trained by imitation learning, and Fig. 7 (Right) highlights why LActor Ind is so effective for training multi-task Transformer policies with online RI:: the policy update reduces to IL on dynamic percentage of the replay buffer, but the ability to ignore a fraction of the dataset automatically allows for self-improvement in a way that standard IL does not.
Normalized Score fby Hard Mode)
1.00
0.75
Memory Mode
1338
395
2-28%
Hard Mode
First
-7%
0.50
43% -18%
-2155
0.25
-156
0.00
dottgbbal Cauetyy helt
jumper
maze
mine frdkboo glunder chaser
tapper
bigftsh
darpil bossightt phi cointun dmmbe
Figure 8: Multi-Task Procgen in Memory-Hard Mode. We measure policy performance across the two episodes of its adaptation window. Results are averaged over 30M frames in unseen test levels.
Inspired by the success of the "Ind. / Ind." learning update in the easy distribution of 5 games, train a larger 24M parameter Transformer with a context length = 768 frames on all 16 Procgen games simultaneously. We default to the standard "hard" level distribution but increase the difficulty to the (rarely attempted) "memory" mode in the games it is available. Results after 2.7B frames are highlighted in Figure 8. The agent does improve its performance the second attempt of unseen levels particularly in memory games where partial observability makes adaptation most useful.
Ball Zoro
Boxing
Breakout
Chopperlomnee
Crazy Lirber Lirber
1044
11005
Human Hormatge
2.4
0.-
Tam
2.0
401
500
15m
Copher
NsPacran
Pong
Obart
Seaquest
3.2
Humton Nay Narmaltoc
0.5
0.30
02
SAM
0.1
0.2
0.05
3.0
0.00
20M
Timostops 20e
5M
J8M 20M
25e
Dep./ Dep. [Clip]
1,3
Ind./Ind.
Dep. Dap.
Figure 9: Multi-Game Atari Without Reward Clipping. We train single policy on 10 games simultaneously. Results are plotted relative to human performance (left axis) and the raw scale of returns (right axis). Error bars denote the maximum and minimum four random trials.
Multi-Game Atari. If RL generalization problems lie axes between meta-RL and MTRL (Section 2), the ALE [29] would represent the multi-task extreme. Super-human performance on Atari does not require adaptation or memory but leads to per-game returns separated by orders of magnitude [100], which motivates the infamous trick of clipping rewards in [-1, 1] [1OÎ]; clipping stabilizes learning but misaligns agents' training objective from the raw game score used for evaluation.
PopArt [25] showed that multi-task training on unclipped rewards is possible when we split critic networks into game-specific output heads with independently normalized targets. Figure revisits this experiment on a smaller scale. However, rather than normalizing a multi-task architecture, we train a single critic on 10 games with unclipped rewards and let our shared Transformer recover state and task labels implicitly from context sequences (of RGB images). In other words, we are treating the widely studied MTRL Atari setting as a special case of meta-learning in order isolate the challenge of optimizing returns on different scales. Results are shown in Figure 9. The cealsaresistant update leads to dramatic improvements in 8/10 games; the default Q-dependent update performs best in the two games where returns at random initialization are significant outliers. The standard approach of clipping rewards [-1, (grey) helps mask the multi-task optimization challenge created by scale-dependent learning updates.

---

BabyAI. BabyAI is suite of partially observed gridworld tasks with simple language instructions [30]. We create hybrid multi-task meta-RL problem with the potential for task generalization by randomly generating a train/test split over 68 of BabyAl's task configurations as provided by the Minigrid simulator [102]. Our agents observe a text description of the their task and have two attempts to adapt to a procedurally generated level layout. Learning curves for all 68 tasks are provided in Appendix C. Figure 10 summarizes our results in the Baby Al domain and reveals another case where scale-independent updates outperform scalar regression. Figure 11 evaluates Baby Al as multi-episodic adaptation problem by comparing the returns of the first and second attempt in a new level layout. Despite BabyAl's 's partial observabilty, only small set of tasks benefit from in-context adaptation over a context length / = 512. Transformers are clearly capable of learning effective recall over sequences of this length via RL [56, 57], and we have demonstrated that ccale-invariant updates can unlock significant improvement in terms of multi-task training. Therefore we argue that the next step is to train this agent in a domain with sufficient task diversity to evaluate meta- veta--lerningg over k > 2 episodes. XLannd Minigrid [103] is a concurrent effort to augment gridworlds with millions of unique tasks and may be promising domain for future study.
Mulki-Task BabyAl: Average Total Return Per 2-Episode Rollout I (0))
50 Train Tasks, Test Varients
18 Test Tasks
1.75
1.75
1:50
1.50
1.25
1.25
1:00
1.00
0.75
0.75
Ind. Taak Mealan
Ind. Taak Mean
Dop. Tisk Mean
Dep. Dask intequrrtte
Medlan
0.25
0.00
0.25
Dap. Dep. Tase intarquarti
20M
40M
50M
R0M
100M
120M
0M
20M
40M
60M
BBM
100M
120M
Total
Training Timesteps
Total
Training Timostops
Figure 10: Multli-Task BabyAL Results are the average over four training seeds and plotted according to the median, mean, and interquartile range over the task set.
Episodic Return (E(o. 11)
2- 2-Epis In-Context Improvement in Multi-Task BabyAl
1.0
0.8
First Episode
Second Episode
0.4
Train Tasks (Test Variants)
Test Tasks
Co tea for de dortn
d0 ohaah gaa to Se wre
get and brag duturna tovo of
ope Sing
of Co
Figure 11: In-Context BI We measure the average return of IInd. / Ind." agents by attempt in unseen tasks and/or layouts (variants). Results are arbitrarily sorted in order of increasing second- episode return and are the average of four training seeds and 10M evaluation timesteps.
5 Conclusion
Memory-based RL methods are already capable of long-term recall and test-time adaptation to variations of familiar tasks. However, optimization challenges from multi-task learning limit our ability to expand to more general applications. We attribute some, but not all, of this difficulty to the unpredictable way our agents balance the optimization of each task according to their current performance. We show that this bottleneck can be removed without explicit task labels additional assumptions by simply replacing both the actor and critic objectives with classification terms that do not depend on the current scale of returns. Experiments across meta-learning, generalization, and long-term memory benchmarks find that this update improves multi-task performance without sacrificing sample efficiency. More broadly, these results justify a uniquely accessible approach to adaptive RI. Training a single Transformer on arbitrary recycled data with two loss functions that resemble supervised learning is simple way to unlock the eeff-mmproeement of online RL. while sharing intuition and technical details with widely familiar topics like sequence modeling.
10

---

Acknowledgments and Disclosure of Funding
This work was partially supported by the National Science Foundation (EFRI-2318065), the Office of Naval Research (N00014-24-1-2550)) the DARPA TIAMAT program (HR0011-24-9-0422) JP Morgan, and Sony.
References
[1] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint vrXXiv2337.099888 2023.
[2] Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean- -an-tBpttstte Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11800 2023.
[3] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song. and Jacob Steinhardt. Mcasuring massive multitask language understanding. In International Conference on Learning Representations, 2020.
[4] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Xiaodong Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. ArXiv, abs/2103.0322 2021. URL htttps://ap semanticscholar. org/ Org/CorpusID. 232134851.
[5] Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs. Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 2368- 2378, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics. doi: 10. 18653/v1/119/ 1246. URL https: aalaanhhllogyyy http://aatttooloooorr//lll/----/
[6] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranay Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. Gerprsl- in-context learning by meta-learning transformers. arXiv preprint arXiv:2212.04458. 2022.
[8] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. survey of meta-reinforcemmet learning. arXiv preprint arXiv:2301.08023 2023.
[9] Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning reinforcement learn. arXiv preprint arXiv:1611.0574 2016.
[10] Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI 2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint 1rX.v:1611.00279, 2016.
[1I] Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint MXXiv:9905006244, 2019.
[12] Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan P Adams, and Sergey Levine. Why generalization in rl is difficult: Epistemic pomdps and implicit partial observability. Advances in neural information processing systems, 34:25502-25515, 2021.
[13] Luisa Zintgraf. Fast adaptation via meta reinforcement learning. PhD thesis, University of Oxford, 2022.
11

---

[14] Mandi Zhao, Pieter Abbeel, and Stephen James. On the effectiveness of versus meta-reinforcement learning, Advances in Neural Information Processing Systems, 35:26519- 26531, 2022.
[15] Adrien Ali Taiga, Rishabh Agarwal, Jesse Farebrother, Aaron Courville, and Mare G Bellemare. Investigating multi-task pretraining and generalization in reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. URL https://
openreview.. wttprviwwwwwwefo.o.. net:foorunid=sStffRRSSS
[16] Andrei Alex Rusu, Sebastian Flennerhag, Dushyant Rao, Razvan Pascanu, and Raia Hadsell. Probing transfer in deep reinforcement learning without task engineering. In Conference on Lifelong Learning Agents, pages 1231-1254. PMLR, 2022.
[17] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1100. PMLR, 2020.
[18) Suyoung Lee, Myungsik Cho, and Youngchul Sung. Parameterizing non-parametric meta- reinforcement learning tasks via subtask decomposition. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL id=Jx6U1 oWrmE.
wttp:/rrwwiieeeeeeoroeerrem net/forum?
[19] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-maron, Mai Giménez, Yury Suisky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals. Mahyar Bordbar, and Nando de Freitas. A generalist agent. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https: //wp///eee//eeeee///e/e/ee / forum?id= 1ikiiikkHiiK Featured Certification, Outstanding Certification.
[20] Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpc2: Scalable, robust world models for continuous control. In The Twelfth International Conference on Learning Representations, 2023.
[21] Quentin Gallouédec, Edward Beeching, Clément Romac, and Emmanuel Dellandréa. Jack of all trades, master of some, a multi-purpose transformer agent. arXiv preprinl arXiv:2402.099844, 2024.
[22] Uri Shaham, Maha Elbayad, Vedanuj Goswami, Omer Levy, and Shruti Bhosale. Causes and cures for interference in multilingual translation. In The 61st Annual Meeting Of The Association For Computational Linguistics, 2023.
[23] Derrick Xin, Behrooz Ghorbani, Justin Gilmer, Ankush Garg, and Orhan Firat. Do current multi-task optimization methods in deep learning even help? Advances in neural information processing systems, 1111133773600 2022.
[24] Yu Zhang and Qiang Yang. A survey multi-task learning. IEEE Transactions on Knowledge and Data Engineering, 34(12):5586-560 2021.
[25] Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojciech Czarnecki, Simon Schmitt, and Hado Van Hasselt. Multi-task deep reinforcement learning with popart. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3796-3803 2019.
[26] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[27] Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POP- Gym: Benchmarking partially observable reinforcement learning. In The Eleventh Interna- tional Conference on Learning Representations, 2023. URL https: //ppnrevieeeeeeee// net/ f forum?id=-cDrrutTTsoK
[28] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, pages 2048-2055 PMLR, 2020.
12

---

[29] Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling, The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
[30] Maxime Chessei Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Sa- haria, Thien Huu Nguyen, and Yoshua Bengio. BabyAI: First steps towards grounded language learning with human in the loop. In International Conference Learning Representations, 2019. URL https:/ /ttp::s///ppeevveee www.eeeeeeeeeeeeeeeeeeeeeeeeeeee www//eetouuune id=rJeXCooccc
[31] Danijar Hafner, Jurgis Pasukonis. Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 04104, 2023.
[32] Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems, 13:7768-7778, 2020.
[33] Yevgen Chebotar, Quan Vuong, Karol Hausman, Fei Xia, Yao Lu, Alex Irpan, Aviral Ku- mar, Tianhe Yu, Alexander Herzog, Karl Pertsch, et al. Q-transformer: Scalable offline reinforcement learning via autoregressive q-functions. In Conference on Robot Learning. pages -9288. PMLR, 2023.
[34] Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taiga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, et al. Stop re- gressing: Training value functions via classification for scalable deep rl. arXiv preprint arXiv:2403.03950 2024.
[35] Jost Tobias Springenberg, Abbas Abdolmaleki. Jingwei Zhang, Oliver Groth, Michael Bloesch, Thomas Lampe, Philemon Brakel, Sarah Bechtle, Steven Kapturowski, Roland Hafner, et al. Offline actor-critic reinforcement learning scales to large models. arXiv preprint arXiv:2402.055A4 2024.
[36] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 15044, 2021.
[37] Kuang-Huei Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. Advances Neural Information Processing Systems, 35:27921-27936. 2022.
[38] Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv: 1502.02259 2015.
[39] Matthew Botvinick, Sam Ritter, Jane X Wang, Zeb Kurth-Nelson, Charles Blundell, and Demis Hassabis. Reinforcement learning, fast and slow. Trends in cognitive sciences, 23(5):408 422, 2019.
[40] Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of generalisa- tion in deep reinforcement learning. 2022.
[41] Carolin Benjamins, Theresa Eimer, Frederik Schubert, Aditya Mohan, Sebastian Döhler, André Biedenkapp, Bodo Rosenhahn, Frank Hutter, and Marius Lindauer. Contextualize me-the case for context in reinforcement learning. arXiv preprint arXiv:2202.04500 2022.
[42] Evan 7. Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and exploitation for meta-reiinorremmen learning without sacrifices. In International conference on machine learning, pages 6925-6935. PMLR, 2021.
[43] Kate Rakelly, Aurick Zhou, Chelsea Fino, Sergey Levine, and Deirdre Quillen. Efficient off-policy reta-reiinoreemment learning via probabilistic context variables. In International conference machine learning, pages 5331-5344 PMLR, 2019.
13

---

[44] Jacob Beck, Matthew Jackson, Risto Vuorio, Zheng Xiong, and Shimon Whiteson. Splagger: Split aggregation for meta-reinforcement learning. arXiv preprint arXiv:2403.03020, 2024.
[45] Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. Xrviv preprint arXiv:2210.144: 142155, 2022.
[46] Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning, Advances in Neural Information Processing Systems, 36, 2024.
[47] Sharath Chandra Raparthy, Eric Hambro, Robert Kirk, Mikael Henaff, and Roberta Raileanu. Generalization to new sequential decision making tasks with in-context learning. arXiv preprint arXiv:2312.0380/ 2023.
[48] Lucy Xiaoyang Shi, Yunfan Jiang, Jake Grigsby, Linxi Fan, and Yuke Zhu. Cross-episodic curriculum for transformer agents. Advances in Neural Information Processing Systems, 36, 2024.
[49] Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta-learnrr arXiv preprint 10X.v:1777.0344, 2017.
[50] Luckeciano C Melo. Transformers are meia-reiiforreement learners, In International Confer- ence Machine Learning, pages 15340-15359, PMLR, 2022.
[51] Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. arXiv preprint arXiv:2303.0398/ 2023.
[52] Rasool Fakoor, Pratik Chaudhari. Stefano Soatto, and Alexander J Smola. Meta-q-leaming. arXiv preprint 1910.00151 2019.
[53] Zhihan Yang and Hai Nguyen. Recurrent off-policy baselines for memory-based continuous control. arXiv preprint arXiv:21 10. /2628, 2021.
[54] Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhuidinov. Recurrent model- -reeel-feeeen rl can be a strong baseline for many pomdps, 2022.
[55] Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmieg, Michael Chang, Natalie Clay, Adrian Collister, et al. Human-timescale adaptation in an open-ended task space. arXiv preprint (3t::23001.7600 2023.
[56] Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in rI? decoupling memory from credit assignment. arXiv preprint arXiv:2307.0386. 2023.
[57] Jake Grigsby, Linxi Fan, and Yuke Zhu. AMAGO: Scalable in-context reinforcement learning for adaptive agents. In The Twelfth International Conference Learning Representations, 2024. URL /tttps:://pppeeeeeeeee
net/ffrumm Ed=Mf
[58] Yunshu Du, Wojeiech M Czarnecki, Siddhant M Jayakumar, Mehrdad Farajtabar, Razvan Pascanu, and Balaji Lakshminarayanan. Adapting auxiliary losses using gradient similarity. arXiv preprint arXiv: /812.022422 2018.
[59] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi- learning. Advances in Neural Information Processing Systems, 33:5824-588 5836, 2020.
[60] Xian Li and Hongyu Gong. Robust optimization for multilingual translation with imbalanced data. Advances in Neural Information Processing Systems, 14-25086-25099) 2021.
[61] Zirui Wang, Yulia Tsvetkov, Orhan Firat, and Yuan Cao. Gradient vaccine: Investigating and improving malti-task optimization in massively multilingual models. In International Conference on Learning Representations, 2020.
14

---

[62] Amelie Royer, Tijmen Blankevoort, and Babak Ehteshami Bejnordi. Scalarization for multi- task and multi-domain learning at scale, Advances Neural Information Processing Systems, 36, 2024.
[63] Vitaly Kurin, Alessandro De Palma, Ilya Kostrikov, Shimon Whiteson, and Pawan K Mudigonda. In defense of the unitary scalarization for deep multi-task learning. Advances in Neural Information Processing Systems, 35:12169-122 -91169283,, 2022.
[64] Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context-based representations. International Conference on Machine Learning, pages 9767-9779. PMLR, 2021.
[65] Aviral Kumar, Rishabh Agarwal, Xinyang Geng, George Tucker, and Sergey Levine. Offline q-learning on diverse multi-task data both scales and generalizes. In The Eleventh International Conference on Learning Representations, 2022.
[66] Carlo D) D'Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters. Sharing knowledge in multi-task deep reinforcement learning. arXiv preprint arXiv:2401.0955 0.9561,,,,6,, 2024.
[67] The garage contributors. Garage: A toolkit for reproducible reinforcement learning research. ht https:///iithhi wow/rlworkgroup/ httproorrrrooogggggge...e garage, 2019.
[68] Bradly C Stadie, Ge Yang. Rein Houthooft, Xi Chen, Yan Duan, Yuhuai Wu, Pieter Abbeel, and Ilya Sutskever. Some considerations on learning to explore via meta-reinforceement learning. arXiv preprint 1803.v8803.000 /8, 2018.
[69] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1 1126- 1135. PMLR, 2017.
[70] Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta-policy search. arXiv preprint 1810:0881000678. 2018.
[71] Zohar Rimon, Tom Jurgenson, Orr Krupnik, Gilad Adler, and Aviv Tamar. MAMBA: an effective world model approach for meta-reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024. URL https: /tttps////peeevviee 1d=1REOHGmU7M.
[72] Jane X Wang, Michael King, Nicolas Pierre Mickael Porcel, Zeb Kurth-Nelson, Tina Zhu, Charlie Deck, Peter Choy, Mary Cassin, Malcolm Reynolds, Francis Song. et al. Alchemy: A benchmark and analysis toolkit for meta-reinforcement learning agents. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Truck (Round 2), 2021.
[73] Ankesh Anand, Jacob Walker, Yazhe Li, Eszter Vértes, Julian Schritwiiserr, Sherjil Ozair, Theophane Weber, and Jessica B Hamrick. Procedural generalization by planning with self- supervised world models. International Conference on Learning Representations, 2022. URL httpps://ppenreeeiie
net/forum? fooormmidd-Fiittt Fd-FnBegXJToY.
[74] Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Conflict-averse gradient descent for multi-task learning. Advances in Neural Information Processing Systems, 34: 18878-18890, 2021.
[75] Timothy P Lillicrap. Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971. 2015.
[76] Petros Christodoulou. Soft actor-critic for discrete action settings. arXiv preprint
arXiv:1910.07207; 2019.
[77] Xinyue Chen, Che Wang, Zijian Zhou, and Keith Ross. Randomized ensembled double q-learning: Learning fast without model. arXiv preprint arXiv:2101.05982, 2021.
15

---

[78] William Fedus, Carles Gelada, Yoshua Bengio, Mare G Bellemare, and Hugo Larochelle. Hyperbolic discounting and learning over multiple horizons. arXiv preprint arXiv:1902.006: 0666666.0666 2019.
[79] Marc G Bellemare, Will Dabney, and Mark Rowland. Distributional reinforcement learning. MIT Press, 2023.
[80] Mare G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective reinforce- ment learning. In International conference on machine learning, pages 449-458. PMLR, 2017.
[81] Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Master- ing atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-609, 2020.
[82] Matteo Hessel, Ivo Danihelka, Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre, Theophane Weber, David Silver, and Hado Van Hasselt. Muesli: Combining improvements in policy optimization. In International conference on machine learning, pages 4214-4226.
PMLR, 2021.
[83] Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations, 2018.
[84] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: arXiv:/707.0000 1ee00070777777 2017.
[85] Abbas Abdolmaleki, Jost Tobias Springenberg. Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. arXiv preprint 2018.
[86] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. ddvantage-weighter regres- sion: Simple and scalable off-policy reinforcement learning. arXiv preprint 19110:191..00177. 2019.
[87] Qing Wang, Jiechao Xiong, Lei Han, Han Liu, Tong Zhang, et al. Exponentially weighted imitation learning for batched historical data, Advances in Neural Information Processing Systems, 31, 2018.
[88] Xinyue Chen, Zijian Zhou, Zheng Wang, Che Wang, Yanqiu Wu, and Keith Ross. Bail: Best-action imitation learning for batch deep reinforcement learning. Advances in Neural Information Processing Systems, 33:1 1333-11338, 2020.
[89] Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavioral modelling priors for offline reinforcement learning. arXiv preprint urXiv:2002.08399 2020.
[90] Jake Grigsby and Yanjun Qi A closer look at advantage-filtes behavioral cloning in high- noise datasets. arXiv preprint arXiv:21 04698, 2021.
[91] Ashvin Nair, Bob McGrew, Marcin Andrychowicz, Wojeiech Zaremba, and Pieter Abbeel. Overcoming exploration in reinforcement learning with demonstrations. In 2018 IEEE inter- national conference on robotics and automation (ICRA), pages 0292-6299, IEEE, 2018
[92] Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.0993 2020.
[93] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.01643, 2020.
[94] William Fedus, Prajit Ramachandran, Rishabh Agarwal, Yoshua Bengio, Hugo Larochelle, Mark Rowland, and Will Dabney. Revisiting fundamentals of experience replay. In Interna- tional Conference Machine Learning, pages 3061-3071. PMLR, 2020.
16

---

[95] Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as one big sequence modeling problem. In ICML 2021 Workshop on Unsupervised Reinforcement Learning, 2021.
[96] Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You can't count on luck: Why decision transformers fail stochastic environments. arXiv preprint 220.::2220.....9.... 2022.
[97] Gresa Shala, André Biedenkapp, and Josif Grabocka. Hierarchical transformers are efficient meta-reinforcemenn learners. arXiv preprini urXiv:2402.06402. 2024.
[98] Kyunghyun Cho, Bart Van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using mn encoder- decoder for statistical machine translation. arXiv preprint arXiv: 1078, 2014.
[99] Steven Morad, Ryan Kortvelesy, Stephan Liwicki, and Amanda Prorok. Reinforcement learning with fast and forgetful memory. Advances Neural Information Processing Systems, 36, 2024.
[100] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[101] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. nature, 518(7540):529-53) 2015.
[102] Maxime Chesseler-Boisvertt, Bolun Dai, Mark Towers, Rodrigo VicentV Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular customizable reinforcement learning environments for goal-oriented tasks. In Advances in Neural Information Processing Systems 36, New Orleans, LA. USA, December 2023.
[103] Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey Kolesnikov. Xland-minigrid: Scalable meta-reinforcement learning environments in jax. arXiv preprint arXi::2312..2044. 2023.
[104] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101. 2017.
[105] Sam Shleifer, Jason Weston, and Myle Ott. Normformer: Improved transformer pretraining with extra normalization. arXiv preprint arXiv:21 2021.
[106] Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M Susskind. Stabilizing transformer training by preventing attention entropy collapse. In International Conference on Machine Learning, pages 40770- 40803. PMLR, 2023.
[107] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International conference on machine learning, pages 1407-1416. PMLR, 2018
[108] Yuxin Wu and Kaiming He, Group normalization. In Proceedings of the European conference computer vision (ECCV), pages 3-19, 2018
[109] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data lata-augmented reinforeement learning. arXiv preprint arXiv:2107.09645, 2021.
[110] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 7rXiv:16070064335 2016.
[111] Mark Towers, Jordan K. Terry, Ariel Kwiatkowski, John Balis, Gianluca de Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Arjun KG, Markus Krimmel, Rodrigo Perez- Vicente, Andrea Pierré, Sander Schulhoff, Jun Jet Tai, Andrew Tan Jin Shen, and Omar G. Younis. Gymnasium, March 2023. URL https: zenooo. org/ arg/record/812700
17

---

A Implementation Details
Code for the agent and multi-task environments used in our experiments is available on GitHub at A-A.uuin-PPLLaaaaaee
Base RL Details. We focus on evaluating changes to the training objective of long-term memory policies trained by off-policy actor-critic updates. Our implementation makes use of many orthogonal technical details from AMAGO [57]. In the context of this work, the most important detail is the use of a global (task-agnostic) PopArt [25] layer to normalize Q-values of the "dependent" update (Lactor: LCiiie). which brings loss functions to predictable absolute scale but does not change the fact that each task has different relative values. Trying to compare learning updates while skipping this technique would make the Q-dependent baselines far too reliant on accurate grid-searches of optimization hyperparameters across domains. Table provides the list of hyperparameters used in our main experiments. All of the results in this paper were completed on NVIDIA A5000 GPUs. We train each agent one GPU whenever possible but add second GPU for Procgen Memory- Merd (Figure 8) where model size and context length use all available memory.
Meta- World
POPGym 1e44
Proegen Easy
Proegen Memory-Hard
Atari
BabyAI
Learning Rate
Se 4
1e4
le.4
le.4
1e
Batch Size (In Sequences)
24
24
24
16
32
24
Replay Buffer Size (in Timesteps)
8M
15,000 Full Trajectories
18M
80M
SM
10M
L2 Penalty
le-4
1e-4
Se-3
5e-3
Se-3
1o:3
Grad Clip (Norm)
Learning Update
1, 96, .7. 99, 905, 900
Rollout ?
Critic Ensemble Size
Target Network Update
Table 1: Learning Hyperparameter Details.
Model Architectures. All of our experiments use causal Transformers trained with the A optimizer [104]. We apply Normformer [105] and oReparam [106] with Leaky ReLU activations to all architectures. These details are intended to stabilize the policy against long training runs of millions of gradient steps but are applicable to any sequence modeling problem: our results did not require any changes that are specific to RI. We use fixed (sinusoidal) position embeddings [26]. Table lists Transformer architectural details for each of our main experiments. The "timestep encoder" maps the potentially multi-modal meta-RI. inputs of observations, actions, rewards, and reset signals to fixed-size vector (Figure 2). State-based domains (Meta-Wane World and POPGym) can concatenate and merge this information with small feed-forward network. In pixel-based environments, we first embed the image observation with a CNN. We evaluated both the small "Nature" architecture [101] and the residual IMPALA CNN [107]. Our results default to IMPALA with additional Group Normalization [108, 65]. We apply the random pad and crop data augmentation from DrQV2 [109] to Proegen and Atari experiments. Image features extracted by the CNN are normalized before being added to action, reward, and terminal data. For consistency with the rest of the architecture we use LayerNorm [1 10]. Preliminary experiments found many of these vision details to be flexible
Mora-Workd World
POPGyn
Proegen Easy
Prorgin Menary-Haard
Atari
RabyA1
Timesty Eacode:
BF (512 256)
EF (512, 200)
[16, 33, 33) Nock depths
IMFALA CNN
[DI,
IMPALA CNN bllck degtts
[16. 32, 32 block depris
IMPALA CNN
CNN (Gind) RNN (Langsage), HP (Cibbe)
Moli-Midaa
Trestomer Din,
Din.
1250
1024
2045
1024
Trndiomer
Trndirerr Heats
K
K
Conteat Letgt (Timestaps)
256
600
128 (256, 255)
765 (25n 256)
512 (256, 218)
Acter and MLL
(256, 256)
(256, 255)
(256 156)
Cittc Oupor Bias (lad)
128
128
128
Table 2: Agent Architecture Details.
Value Classification Details. We convert value regression (Equation 1) to classification (Equation 3) by creating labels for B return bins b = [bu. bi......bpr bgl. Bins are typically spaced at fix intervals between pre-defined upper and lower bounds on the return (by = Raow, critic network (Qn) outputs (softmax) probabilities over these bins, and its value prediction can be recovered by Q(h,, dt) = Qn az) Tb. A scalar temporal difference target y is mapped to a classification label by twohotg E + [0, 1 function that outputs zeroes everywhere aside
Rligh). The
18

---

Mcon Squared Error
Wide Limit Symilg 8 256
Wide Limk Symiag B 128
Wide Linit Symlog B
10
0.8
Renstvo Crte Loos
00.5
1
Relstrr Crice
0.0
0.0
00
0.4
0.
205
10
1.002
2,000
15,000
0.0
-42
30
-23
-22
20
Target
Target
Target
Tured Symlng 128
Tured Limst Syming di
Tuned Limit win Syming 128
Tuned Limit Syming & 32
Reaavee CPRIc Leto
Reaatv CPtC LoOS
0.6
0.6
0.0
0..
Readive
0.4
0.4
0.4
0.
0.2
0.0
Target Prediction
Prediction
Target Prediction
Target Prudiction
Figure 12: Scale-Resistant Regression. We plot the critic loss as function of the relative prediction error of the value target (z) across several orders of magnitude. Standard MSE (top left) assigns significantly different weights to the same relative inaccuracy across tasks according to the absolute value of their current targets. Converting regression classification to similar loss, but this effect can benefit from tuning. We show different strategies. "Wide Limit" refers to the DreamerV3 approach of setting extreme (Rjo Ruihh) defaults that do not need tuning. "Tuned Limit" bins set the upper and lower bound close to the expected range for the domain (in this example: [-100, 15, 000)..
map wide range of values
from two consecutive indices corresponding to the closest bin below y (b,-1) and above (be), which filled based on their distance from y:
bi-1 b;
twohot (v)]i b y
(wohota(y)e 1]
(5)
(6)
b; b;-1
Therefore there are two main settings to configure for each experiment: 1) the bin limits [Row Reigh). and 2) the number of (evenly spaced) bins between the limits, B. We can often set bin limits using prior knowledge of the maximum and minimum return a given domain, but it can be challenging to application reward function. The number of bins B defines
establish correct settings for the "resolution" the critic network. Using fewer bins puts more emphasis on outputting correct label weights while adding more bins shifts focus outputting correct label indices. This trade- off is interesting because it is somewhat unusual for the label space classification problem to be tunable hyperparameter; for example, the dataset determines the number of labels in object recognition or language modeling. Figure 13 provides two single-task examples where lower label counts are more sample efficient, and Figure 7 in the main supports similar conclusion in multi-task Procgen. These settings can be easily adjusted in our open-source code, but their trade-off's underexplored in the results of this work. Instead, details in DreamerV3 [3]]. We set [Roow Rnigh extremely wide such that they capture the full range of every experiment and do need to be tuned. We transform y with "symlog" before label creation: twohot (v) + twohot (symlog(y)). Scalar values are recovered from bin probabilities by inverting this transformation with "symexp": Qh+, oz) = symexp (he, oz) Tb), where:
focus on reducing tuning by following
symlog(v) = sign(y) In (1yl 1)
(7) symexp(b) (b) sign (c)((())))())))
(8)
Bins are spaced evenly between by = symlog( (Row) and by symlog( Reigh). This symlog trick skews labels to give more resolution to smaller return values while retaining our ability represent extreme outliers at a lower resolution. In general, we find that setting wide limit (R Bnigh).
19

---

increasing the label count B, and compressing labels with symlog/symexp is the safest way to achieve scale-invariance across domains (Section 3). An example is illustrated by Figure 12. We always use symlog and wide limits of ( (Rew - 1e5, Rhigh 1e5), and only tune the bin count B (Table 2). However, it is likely that sample efficiency can be improved by tuning bin counts, limits, and the use of the symlog transform for each individual domain.
Autoencode Mediun
Battloship Easy
0.2
0.8
0.4
0.0-
-0.4
-0.2
B = 256
-0.4-
B = 32
-0.8
10
12
12
Tlieesteps (Millions)
Timesteps (Millions)
Figure 13: Bin Counts in POPGym. We compare two label space sizes (B) with the same Riow and Hhigh limits on two single-task POPGym environments [27].
B Additional Environment Details
Multi-Task POPGym. The environment samples uniformly between 27 POPGym [27] tasks with discrete actions where both the action space and observation space have dimension < 30. A full list is provided as part of Figure 15. The cutoff of dimension 30 is arbitrary but some limit is necessary because POPGym has several tasks with unusually large action spaces. We are then able create unified input and output spaces by padding observations and actions to dimensions of 26. If the agent selects an action index that not valid for the current task, the environment samples valid action uniformly at random. An additional observation feature indicates whether the previous action was valid. The one-shot evaluation setting is implemented as in E-RL² [68], where rewards for the exploration episode are inputs to the policy but are not allowed to become true reward outputs that are optimized by RL updates.
Multi-Game Procgen. We merge individual Procgen [28] games into a multi-task environment by randomly sampling new game between meta-rollouts. The same procedurally generated level is then played twice before a new game and level are sampled. In order to allow policies without long-term memory to form accurate value estimates, we draw small black box in the top left corner of the screen throughout the second episode. We enforce a maximum rollout length (across both episodes) of H- 2500 for the 5-game easy mode (Figure 6) and H 5000 for the memory- -hadd mode (Figure 8). In memory-hard mode, we sample levels from the "memory" distribution when available and "hard" otherwise. The memory distribution extends the hard difficulty by introducing partial observability. Memory mode was not evaluated in the original Procgen results and has been rarely used since. Therefore the results in Figure 8 are normalized according to the established scale for hard mode. With the exception of the "Miner" game, these upper bound scores transfer to memory mode as it does not adjust the problem difficulty but does hide information about the layout of the level.
Start of Traindng
End Traintng
End of Training
One detail is that the random game between resets does not account for the length of each rollout in terms of timesteps or the way episode length in- creases or decreases as the agent improves. Our multi- task Procgen experiments create imbalanced datasets where each game makes up an uneven amount of in- coming training data (Figure 14). We consider this to be an interesting additional challenge that off-policy methods with large replay buffers are well suited to address.
chncer
Calnca
22%
LAUGW
LEANT
Figure 14: Imbalanced Procgen Datasets. We measure the inflow of experience to the replay buffer by game. Climber accounts for much of our early training data but the buffer is rebalanced as policies improve.
20

---

Atari. We select 10 games that do not involve challenging exploration but that have human-level returns on several different orders of magnitude. The list games is included in Figure We use the standardized v5 variants of the Atari environments in Gymnasium [111]. In addition to removing reward clipping, we do not use frame stacking or greyscaling. Instead, agents can learn short-term state estimation and task identification from context sequences of RGB frames.
BabyAI. BabyAl provides agent with a x 7 partial view of its surroundings in a procedurally generated gridworld. Goals are communicated with relatively simple language instructions, which we represent as sequences from vocabulary of 29 tokens. Each reset randomly samples new goal type and gridworld generation strategy from 50 tasks in the Minigrid registry during training, and 18 tasks during testing. A full list of task names is included as part of Figure 11.
C
Additional Results
Large figures displaying results from each task in POPGym (Figure 15) and BabyAl (Figure 16) experiments are listed on the following two pages. These figures are followed by learning curves for each Meta-
Meta-World World ML45 task.
21

---

Critic
Dep. Actor, Dep. Critic
AutoencodeEasy
AutoencodeMMdiin
AutoencodeHard
AMAGO
Single-Task
00.55
Timesteps)
Morad et al.
00.49
Best Single-lask
(15M Timestops)
4..54
PPO-GRU
Timesfaps
Timestaps
Repea PPreviousEasy
RepeatPreviou
RepestPrevionne Previousttadd
Repea tFirstEss
1.00
1.00
1.00
1.00
0.48
0.04
40.55
25M
50M
75M
100M
25M
50M
100M
25
100M
25M
100M
Nuttiaakk
Mutiiask
RepestF Lun
Repea tFirsthaad
CountRecallEasy
CountRecal Wediin
1.00
1.00
0.92
0.90
0.46
0.61
00.93
25M
500
75M
100M
25M
50M
75M
100M
23M4
50M
75M
200M
25M
50M
75M
100M
Total Mali-eak Timsndaps
Total Nuli-Task Timeslups
Total Mull-lask Timestaps
Total Multifask Timesiaps
Posi
Position@nlyCarted
EezaitionnnlyCCo MawioonnnnyCarrtlettta.o
0.447
1.00
1.00
0.61
0.70
0.68
0.67
0.39
25M
50M
75M
1.0M
25M
50M
75M
100M
25.
50M
75M4
100M
25M
50M
75M
100M
Total Mal-lask Timettes
Total Multi-lask Timeestps
Total Mall-Tesk Timesleps
Total Multi-lakk Timestegs
Velooityonny@Grrtoeeeeerrr
NoisyPosOniyyarru
0.69
0.68
0.67
0.89
0.39
0.36
0.34
0.37
25M
50M
75M
100M
25M
50M
75M
100M
25M
50M
75M
100M
25M
50M
75M
100M
Total Mul--Tak Timesteps
Total MuttiTosk Timesteps
Total Muh-TTsk Timesteps
Total MuttiTask Timesteps
Mo..sy/PsonntCGrrtPPteeeddd
Noisy/osonnyGrrtPPolllrr
inttiarnedeanditee
WbltiarnedBand UMediin
0.45
0.29
0.42
0.25
0.17
0.22
0.00
0.04
50M
75M
Timentaps
Tineslops
Mali-ask Timestaps
Total Multifaak Timesilega
Hi Easy
HigherLove Mediiu
Highe FLowerlard
NineSweeperEasy
0.54
0.51
0.70
0.34
0.33
0.33
0.16
0.05
4.008
0.02
0.42
5001
25M
25:4
25M
100M
Figure 15: One-Shot Multi-Task POPGym. Plotted on the raw scale of returns in each environment. Error bars denote the maximum and minimum results over three trials. Single- Task reference scores 15M timesteps indicate the best of 14 sequence model backbones trained by PPO in the results of Morad et al. [99]. AMAGO [57] single-task reference scores are taken at 4.5M timesteps, which is a safe upper-bound the timesteps encountered for each task during multi-task training.
22

---

Gold
KayCoriid 5382
KoyCanite 53
3.65
85
Te:a0 traii Tiancies
Tetal Traie Tinestras
T... Train Tikeiies
Tutal Trale Time simettes
T:tal trair Tiancing
Tatal trois Tiansteps
Teial Traiin Tiaediis
Untedev0
KayComcar5413w22
Untoctiocalsd
HolupLoca0
DnliccP--uup-d
KoyCur 55
2.-
15-
2.5
15
2.5
2.0
Openiug
GultsOpen-00
2.5
Cof:Doore0
OperRuddDor-v0
Deb Mru
15-
2.5
05
65
8.--
FidusStie-vd
Unflecd cdap-p
OpenBoor-vd
Cmeteom512-40
2.0
8 E
Picep-et
Ge
OperDoorCatoruu
2.0
2.0
2009
1304
Prtupsaaseew3
Bre
Unteckthurlockea
OnelkonS166v
Go Go SANGu3
OeeRoorr
2.0
26
2.
28-
2.0
200n
200M
1064
6/3 amnu
200
26
vo
2.0
Go 716.
2.
Ge Eoreyeo
do dolloooca Sy
KeyCCo idu 40
22:0
2013
SSM
11:M
Art B0
UniocePriaptievo
OperDortoocal
2.3
944
200m
200M
20e
Lock
Unlodiooca Okstivo
R-
40
vo
2.0
24
2.0
2.0
Dep. Actor/ Dep. Cntlcc Actor Ind. Critic
14-
35
8.-
2230
221M
LOM
S5M
1154
100M
Figure 16: Two-Episode Multi-Task BabyAL. Held-out test tasks are highlighted red. Error bars denote the maximum and minimum value of four random trials.
23

---

Dep. Critic
Dep. Ind. Critic
Ind. Actor, Dep. Critic
Ind. Actor, Ind. Critic
assembly-v
Avg. Total Return
Avg. Total Successes
3.0
2.5
6000
4000
1.5
1.0
2000
0.0
0M
20M
40M
60M
BOM
100M
120M
0M
20M
40M
60M
80M
100M
120M
Total Timestape (Across All Tesks!
Total Timesteps (Across All Tasks)
basketball-v2
Avg. Total Return
Total Successes
12500
10000
2.5
2.0
5800
0.5
20M
40M
60M
80M
100M Tesks!
120M
20M
40M
60M
80M
100M
Timestaps (Acrsss
Total Timesteps (Acrows All Tasks)
but ton- tot--sersss tapdown-v v2
Avg. Total Return
Avg. Total Successes
3.0
10000
7500
2.0
5800
1.5
2500
0.0
0M
20M
40M
60M
8OM
100M
120M
0M
20M
40M
60M
80M
100M
120M
(Across
(Across
Avg. perss-pesssseee towdowo -wall-w
button
2.5
2.0
1.0
2500
0.5
0
100M
120M
60M
Total Timestess (Across All 45 Tasks!
Tetal Timestops (Across All Tasks)
button-prees-vv
Avg. Tola Return
Aug, Total Successes
3,0
8000
2.0
1.5
4800
0
0.0
0M
20M
40M
G0M
80M
100M
120M
0M
20M
40M
60M
80M
100M 120M
(Across All
button- Pirtoo--pee---e wass-wallv
Avg.
10000
2.5
5000
2500
1.0
100M
120M
20M
Total Timesteps (Across All 05 Tasks!
Total Timesteps (Across All Tasks)
coffee-button-v2
Avg. Total Return
Avg. Total Successes
12500
3.0
2.0
7500
1.5
5000
2500
0.5
0.0
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M 120M
messttgs
Timesteps
Figure 17: Three-Episode Meta-World MI.45 Learning Curves (Part 1/7)

---

Dep. Critic
Dep. Ind. Critic
Ind. Actor, Dep. Critic
Ind. Actor, Ind. Critic
foffee- pull-v2
Avg. Total Return
Avg. Total Successes
12500
3.0
10000
2.5
1.5
5000
1.0
2500
0.0
0M
20M
40M
60M
BOM
100M
120M
0M
20M
40M
60M
80M
100M
120M
Total T'mest (Across All Tesks!
Total Timesteps (Across All 45 Tasks)
faffee- laffee-push-ve
Avg. Return
Total Successes
12500
2.5
7500
2.0
5000
0.5
20M
40M
60M
80M
100M
120M
20M
40M
60M
80M
100M
120M
Timestaps (Aemss 45 Tesks!
Total Timesteps (Acrows All Tasks)
dial-turn-v
Avg. Total Return
Avg. Total Successes
3.0
12500
10000
7500
1.5
2500
0.0
0M
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
(Across
(Across
Avg. disassemble
12500
10000
2.5
2.0
5000
1.0
2500
0.5
100M 120M
60M
Total Timestess (Across All 45 Tasks!
Total Timestops (Across All Tasks)
door-close-v2
Avg. Tola Return
Aug. Total Successes
3.0
12500
10000
2.0
1.5
5000
2500
0.0
0M
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
All
Avg. door-open-v2
12500
2.5
10000
7500
1.0
2500
100M 120M
20M
Total Timestepe (Across All 05 Taskst
Total Timesteps (Across All Tasks)
drawer-close-22
Avg. Total Return
Total Successes
15000
3.0
12000
2.0
6000
1.5
3000
0.5
0.0
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
mestege
Timesteps
Figure 18: Three-Episode Morld MI.45 Learning Curves (Part 2/7)
25

---

Dep. Critic
Dep. Ind. Critic Ind. Actor, Ind. Critic
Ind. Actor, Dep. Critic
di prawer-oper Swver-open-e
Avg. Total Return
Aug. Total Successes
3.0
12500
2.5
10000
1.5
1.0
2500
0.0
0M
20M
40M
60M
BOM
100M
120M
0M
20M
40M
60M
80M
100M
120M
Total Timestaps (Across All Tesks!
Total Timesteps (Across All Tasks)
Paucet-open-y
Avg. Return
Total Successes
12500
2.5
10000
2.0
5000
0.5
20M
40M
60M
80M
100M
120M
20M
40M
60M
80M
100M
120M
Timestaps (Across
Tests!
Total Timesteps (Acrows All Tasks)
faucet close-vv
Avg. Total Return
Avg. Total Successes
3.0
12500
2.0
7500
1.5
5000
0.0
0M
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
(Across
(Across
hamner-v2
Avg.
Successes
12500
10000
2.5
2.0
7500
5000
1.0
2500
0.5
100M
120M
60M
Total Timestess (Across All 45 Tasks!
Total Timestops (Across All Tasks)
Avg. Tola Return
Aug. Total Successes
15000
3.0
2.0
6000
1.5
3000
0.0
0M
20M
40M
G0M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
All
Avg. handle-press-u
Successes
12000
2.5
9000
6000
1.0
100M
120M
20M
Total T T'mesteps (Across All 65 Taskst
Total Timesteps (Across All Tasks)
handle- pull-uide-ve
Avg. Total Return
Total Successes
3.0
12000
9000
2.0
6000
1.5
3000
0.5
0.0
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
mesttogs
Timesteps
Figure 19; Three-Episode World ML.45 Learning Curves (Part 3/7)
26

---

Dep. Critic
Dep. Ind. Critic Ind. Actor, Ind. Critic
Ind. Actor, Dep. Critic
handle- puleepppllie
Avg. Total Return
Avg. Total Successes
3.0
12500
2.5
7500
1.5
1.0
2500
0.0
0M
20M
40M
60M
BOM
100M
120M
0M
20M
40M
60M
80M
100M
120M
Total Timest (Across All Tisks!
Total Timesteps (Across All Tasks)
lever-pull-
Avg. Return
Total Successes
2.5
10000
2.0
5000
0.5
20M
40M
60M
80M
100M
120M
20M
40M
60M
80M
100M
120M
Timestaps (Acrnss
Tests!
Total Timesteps (Acrows All Tasks)
peg- iee--nnsertsiee vide-v2
Avg. Total Return
Avg. Total Successes
12500
3.0
10000
2.0
7500
1.5
2500
0.0
0M
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
(Across
(Across
pick- Jick-place-war placc-aal-ww
Avg.
Successes
12500
2.5
7500
2.0
5000
1.0
0.5
0
100M 120M
60M
Total Timestess (Across All 45 Tasks!
Total Timestops (Across All 45 Tasks)
pick-out-of-holeev2
Avg. Tola Return
Aug. Total Successes
12500
3.0
10000
2.0
1.5
5000
2500
0.0
0M
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
Avg.
2.5
10000
5000
1.0
100M
120M
20M
Total T mesttee (Across All 65 Taskst
Total Timesteps (Across All Tasks)
push-back-vy
Avg. Total Return
Avg. Total Successes
12500
3.0
10000
7500
2.0
1.5
2500
0.5
0.0
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M 120M
misettgss
Timesteps
Figure 20: Three-Episode Meta-World MI.45 Learning Curves (Part 4/7)
27

---

Dep. Critic
Dep. Ind. Critic
Ind. Actor, Dep. Critic
Ind. Actor, Ind. Critic
push
Avg. Total Return
Avg. Total Successes
3.0
12000
2.5
9000
1.5
1.0
0.0
0M
20M
40M
60M
BOM
100M
120M
0M
20M
40M
60M
80M
100M
120M
Total Timestage (Across All 45 Tesks!
Total Timesteps (Across All Tasks)
pick place-v2
Avg. Return
Total Successes
10000
2.5
2.0
2500
0.5
20M
40M
60M
80M
100M
120M
20M
40M
60M
80M
100M
120M
Timestaps (Acmsss
Tesks!
Total Timesteps (Acrows All Tasks)
plate-slide-Ve slide-v2
Avg. Total Return
Avg. Total Successes
3.0
12500
10000
2.0
1.5
5000
2500
0.0
0M
20M
40M
60M
8OM
100M
120M
0M
20M
40M
60M
80M
100M
120M
(Across
(Across
Avg. plate-slide-side-Ve
2.5
2,0
1.0
2500
0.5
0
100M
120M
60M
Total Timestess (Across All 45 Tasks!
Total Timestops (Across All Tasks)
plate-slide-back-22
Avg. Tola Return
Aug. Total Successes
15000
3.0
2.0
6000
1.5
3800
0.0
0M
20M
40M
G0M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
Avg. plate-slide-bar backddee-b-k--eeee side-v2
12000
2.5
9000
1.0
60M
100M
120M
20M
Total Timesteps (Across All 65 Taskst
Total Timesteps (Across All Tasks)
ung-unplug-side..
Avg. Total Return
Avg. Total Successes
12500
10000
7500
2.0
1.5
2500
0.5
0.0
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M 120M
mestttgs
Timesteps
Figure 21: Three-Episode Meta-Worlh MI.45 Learning Curves (Part 5/7)
28

---

Dep. Critic
Dep. Ind. Critic
Ind. Actor, Dep. Critic
Ind. Actor, Ind. Critic
soccer-v2
Avg. Total Return
Avg. Total Successes
12500
3.0
10000
2.5
5000
1.5
1.0
2500
0.0
0M
20M
40M
60M
BOM
100M
120M
0M
20M
40M
60M
80M
100M
120M
Total Timestapa (Across All Tesks!
Total Timesteps (Across All Tasks)
stick stick-push-v2
Avg. Return
Total Successes
12500
10000
2.5
7500
2.0
5000
0.5
20M
40M
60M
80M
100M
120
20M
40M
60M
80M
120M
Timestaps (Acrsss
Tesks!
Total Timesteps (Acrows All Tasks)
stick-pull.
Avg. Total Return
Avg. Total Successes
3.0
10000
7500
5000
1.5
0.0
0M
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
(Across
(Across
Avg. push-wall-v2
2.5
2.0
6000
1.0
3000
0.5
0
100M
120M
60M
Total Timestess (Across All 45 Tasks!
Total Timestops (Across All Tasks)
reach-wall-v2
Avg. Tola Return
Aug. Total Successes
3.0
12500
10000
2.0
7500
1.5
2500
0.0
0M
20M
40M
G0M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
(Across
shel 1-pcacev
Avg.
10000
2.5
7500
1.0
2500
60M
100M
120M
20M
Total T'mestepe (Across All 65 Tasks!
Total Timesteps (Across All Tasks)
sweep-inter
-into-22
Avg. Total Return
Total Successes
3.0
12000
9000
2.0
6000
1.5
3800
0.5
0.0
20M
40M
60M
80M
100M
120M
0M
20M
40M
60M
80M
100M
120M
mimestogs
Timesteps
Figure 22: Three-Episode World ML.45 Learning Curves (Part 6/7)
29

---

Dep. Actor, Dep. Critic
Dep. Actor, Ind. Critic
Ind. Actor, Dep. Critic
Ind. Actor, Ind. Critic
sweep-v2
Avg. Total Return
Avg. Total Successes
12500
3.0
2.5
10000
2.0
7500
1.5
5000
2500
0.5
DM
20M
40M
60M
80M
100M
120M
20M
40M
60M
8DM
100M
120M
Total
wi window open
Avg. Total Return
Avg. Total Successes
2.5
10000
7500
1.5
5000
2500
0.5
0.0
0M
20M
40M
60M
80M
100M
120M
QM
20M
40M
60M
80M
100M
120M
(Across
Tasks)
window-close-ve
close-v2
Avg. Total Return
Total Successes
12500
3.0
2.5
2.0
1.5
5000
1.0
2500
0.0
0M
20M
40M
60M
80M
100M
120M
20M
40M
60M
80M
100M 120M
Total Timesteps (Across 45
Tesks)
Total Timesteps (Across All 45 Tesks)
Figure 23: Three-Episode Meta- World ML45 Learning Curves (Part 7/7)
30

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction reflect the paper's contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations our approach and results are discussed in the main text paper.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA].
31

---

Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes].
Justification:
Our code is open source and hyperparameter details are provided in Appendix
A.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
32

---

Answer: [Yes].
Justification:
Our code is open source allowing our results to be reproduced.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips) cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /Codeeuomiisseeeeeeee
/nips. wttpccccpppcccaeecccccccccccccat / assssnoliccyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes].
Justification: Training and testing details can be found in Appendix A or in our open-source code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Errorbars are included and explained.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
33

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The compute resources used in our experiments are discussed in Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http///eerrppp.eeeee http://wrrr...cpuull.aaee htte/public/eeeee ctwuidelllelllleecccccccc de ines?
Answer: [Yes]
Justification:
All research conforms to the Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
This research is foundational and not tied to particular applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
34

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not introduce any new datasests or models with a high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
35

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
36


Reviews:
[{"review_id": "6p0bXooiSW", "reviewer": "Reviewer_Ypn5", "summary": "This paper studies multi-task reinforcement learning using a context-based Transformer policy without task labels. To address optimization difficulties caused by imbalanced losses across different tasks, it proposed replacing actor and critic losses with classification losses. Ablation studies on several benchmarks show that the proposed classification-based actor-critic losses outperform the original algorithm.", "questions": "Please see the Weaknesses.", "limitations": "I cannot find a discussion of limitations.", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "1. The paper is well-focused on the problem of balancing losses in multi-task RL. The proposed classification losses are well-motivated and make sense. \n2. Diverse benchmarks are used to evaluate the method.", "weaknesses": "1. Multi-task RL using in-context adaptation without task labels is not new. Some recent works [1,2] learn in-context meta-RL policies that adapt to diverse tasks, rather than adapting to variations within a single task.\n2. In the setting of online RL, task labels are actually accessible in the simulator. All environments used in the experiments, such as ML45 and multi-game Atari, can provide task labels for each task. Therefore, balancing multi-task losses without task labels is not necessary. Additionally, there are methods that first learn a multi-task label-conditioned policy online and then distill it into a context-based policy [3].\n3. The technical contribution is not strong. Compared to AMAGO [2], which proposes the multi-task RL framework using context-based Transformers and already achieves great results on these benchmarks, the contribution of this paper (classification-based actor-critic loss, which exists in the multi-task RL literature) is incremental.\n4. According to the literature, the paper lacks comparisons with many baseline methods. The experiments only present ablation studies of the proposed classification loss.\n5. Open access to the code is claimed in the checklist. But there is no supplementary or external link provided to access the code. I also cannot find a discussion of limitations, which is claimed in the checklist.\n\n[1] Generalization to New Sequential Decision Making Tasks with In-Context Learning, 2023\n[2] AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents, 2024\n[3] In-Hand Object Rotation via Rapid Motor Adaptation, 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We will discuss your concerns below.\n\n\n> Multi-task RL using in-context adaptation without task labels is not new. Some recent works [1,2] learn in-context meta-RL policies that adapt to diverse tasks\n\nWe are focused on end-to-end online RL while Raparthy et al. belongs to a different category of in-context imitation learning methods that is cited on line 90 of our paper.\n\nTo expand on this discussion, there are two main differences between (supervised) “in-context” imitation learning (ICIL) and context-based meta-RL. \n\n1) ICIL typically relies on a dataset of expert demonstrations collected by single-task RL agents. This removes the opportunity for positive transfer between tasks during data collection. Many examples of large-scale generalist policy methods distill single-task experts into a multi-task policy, which highlights the difficulty of online multi-task RL even after many years of study.\n\n2) ICIL approaches in-context learning much like few-shot prompting in language models, where we select a prompt demonstration and ask the model to continue this behavior over the following inputs. Meta-RL agents begin from a blank sequence and collect examples that will maximize the overall return. IL on aggregated single-task datasets can make it challenging to explore a new environment by creating a mismatch between the task prior of the multi-task policy and the single-task policies it is trying to imitate [1]. In other words, the multi-task policy may be imitating behavior that is not adaptive, but it lets us generalize from demonstrations.\n\nAMAGO is not a multi-task method, which we will discuss below.\n\n[1] *Offline Meta Reinforcement Learning - Identifiability Challenges and Effective Data Collection Strategies*, Dorfman et al., 2021\n\n\n> In the setting of online RL, task labels are actually accessible in the simulator … Therefore, balancing multi-task losses without task labels is not necessary. \n\nThis perspective treats current benchmarks as an end goal instead of a step towards a more general method. The ultimate goal of adaptive RL is to automatically generalize to very large numbers of tasks. Methods that rely on task labels scale with the number of tasks (gradient editing, separate task heads, etc.). So while it’s true that we can always find the ground-truth task label in the gym environment of common benchmarks like ML45 and Atari, our work ignores this information in hopes of researching a more general method that can extend past popular benchmarks. We do not want to rely on knowing the identity of a new task at test-time, because this removes chances of generalization to unseen tasks. And we would like to scale to more open-ended domains where the total task set is so diverse it is unclear how to one-hot label / count tasks at train-time. Meta-learning without task labels is capable of these goals in theory while multi-task RL with labels is not.\n\nThis argument is covered in the main text, including in lines 52-55, 135-140, and 255-261, but will add additional discussion in Section 2.\n \n\n> AMAGO [2], which proposes the multi-task RL framework using context-based Transformers and already achieves great results on these benchmarks, the contribution of this paper … is incremental.\n\n**Following the terminology in this paper, AMAGO is not a multi-task method**. It learns from many procedurally generated versions of a single task, and uses its Transformer to identify and adapt to those subtle variations. AMAGO evaluates on single-task POPGym, single-task Meta-World (ML1) and so on. It does not attempt ML45 or the multi-task POPGym setting introduced here. The green curves in our figures (“Dep. / Dep.”) represent an AMAGO-style optimization step. We do not claim our method is fundamentally new, but we make an observation that an increasingly common implementation detail in recent work has a simple justification in an online multi-task setting, which unlocks the (very non-incremental) performance improvement in our results.\n\n\n> the paper lacks comparisons with many baseline methods. The experiments only present ablation studies of the proposed classification loss.\n\nThe high-level approach of memory-based meta-RL is simple but does not allow for many changes, so most of the differences between methods in this area are implementation-specific. We are evaluating changes in training objectives and ensure a fair ablation of those changes by holding all other details fixed. Our experiments cover a variety of domains, model sizes, and sequence lengths. We provide external reference scores for context where possible (Figures 3, 4, and 5), but some of these settings are rarely (if ever) attempted. Tuning hyperparameters to measure state-of-the-art performance across different conceptual approaches on these benchmarks is not the goal of our work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "e7OlQeAqCA", "reviewer": "Reviewer_1L6X", "summary": "This paper addresses the challenge of scaling meta-reinforcement learning (meta-RL) to handle multiple tasks without explicit task labels. It introduces a method where both the actor and critic objectives are converted to classification terms, decoupling optimization from the scale of returns. This approach builds upon recent Transformer-based meta-RL techniques. The effectiveness of this method is demonstrated through large-scale comparisons on several benchmarks, including Meta-World ML45, Multi-Game Procgen, and Multi-Task POPGym. The results show significant improvements in online multi-task adaptation and memory problems.", "questions": "### Q1. Scalability\n\nWhat are the computational requirements and scalability of the proposed method when applied to very large-scale environments and datasets? Are there any practical limitations or considerations?", "limitations": "The authors have acknowledged the limitations related to the variability of return scales and the focus on specific benchmarks. However, a more detailed discussion on potential negative societal impacts and how to mitigate them would be beneficial.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "### S1. Novel Approach and Problem Identification\n\nThe paper addresses the critical issue of imbalanced training losses in multi-task RL, which often arise due to uneven return scales across tasks. By converting the actor and critic objectives to classification terms, the method effectively decouples optimization from return scales, which is a novel and practical approach.\n\n### S2. Empirical Validation\n\nThe method is validated through comprehensive experiments on multiple benchmarks, demonstrating significant improvements in task adaptation and memory problems. The results are clearly presented, with detailed comparisons to existing methods, highlighting the advantages of the proposed approach.\n\n### S3. Clarity of Presentation\n\nThe paper is well-structured, with clear and concise explanations of the proposed method. Figures and tables are effectively used to illustrate results, enhancing the clarity of the presentation.", "weaknesses": "### W1. Lack of Theoretical Motivation and Analysis\n\nThe paper focuses heavily on empirical results but lacks a thorough theoretical analysis of why the proposed method outperforms existing approaches. A deeper exploration of the theoretical foundations and a comparison with existing theoretical frameworks would strengthen the paper.\n\n### W2. Limited Generalization Analysis\n\nWhile the method is validated on several benchmarks, the paper would benefit from a broader analysis of its generalization capabilities. Additional experiments on more diverse datasets and environments would provide a more comprehensive assessment of the method's applicability. For example, meta-world to multi-game procgen.\n\n### W3. Assumption of Return Scale Variability\n\nAlthough the paper discusses that return scales across tasks can vary significantly, the proposed method does not explicitly address how to handle extreme variations in return scales beyond the classification transformation. A more detailed discussion on how the method can be adapted or modified for tasks with highly diverse return scales would be valuable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We will try to address your questions and would be happy to continue the discussion.\n\n> By converting the actor and critic objectives to classification terms, the method effectively decouples optimization from return scales, which is a novel and practical approach.\n\nTo clarify, the classification loss terms are an implementation detail that appears in recent work we have cited and discussed in Sections 2 and 3. We do not want to claim the objectives themselves are novel. Instead, we show that while these ideas may have various motivations and minor benefits in other settings, they provide a major improvement to a key problem in the area of meta-learning and generalization without task labels.\n\n> The paper focuses heavily on empirical results but lacks a thorough theoretical analysis of why the proposed method outperforms existing approaches. \n\nWe would be open to suggestions about how we can expand our analysis. Prior works have studied the challenges of multi-task optimization (Section 2). We do not think our work adds to the understanding of this issue, but it does help us remove an extra challenge where RL methods unintentionally place the objective of each task on a different scale (Section 3). We have added a more thorough explanation of this effect in a new figure attached to the rebuttal material. \n\n> While the method is validated on several benchmarks, the paper would benefit from a broader analysis of its generalization capabilities … For example, meta-world to multi-game procgen.\n\nWe agree that transfer between domains like meta-world → multi-game procgen would be very interesting. This kind of transfer between different observation/action spaces would require a specialized input/output format for multi-modal trajectories of different shapes. Gato [1] is a good example for how we might go about this, but their method stretches input sequence lengths in a way that makes it expensive to evaluate memory. \n\n> Although the paper discusses that return scales across tasks can vary significantly … A more detailed discussion on how the method can be adapted or modified for tasks with highly diverse return scales would be valuable.\n\nIn our experience, the extreme limits of value classification can be managed by globally rescaling all the rewards to a more stable range. The more challenging issue is the way variations in task difficulty drive uneven changes in values even when the initial and optimal returns are fairly well-bounded. For example, Multi-task POPGym and Meta-World ML45 both have bounded returns, but the benefits of the classification-style loss are much more clear on ML45 than POPGym. Reducing the number of classification bins may help improve performance in tasks where returns rapidly shift from the lower bound to the upper bound. We have added a demonstration of this on a POPGym task in the extra rebuttal material.\n\n> What are the computational requirements and scalability of the proposed method when applied to very large-scale environments and datasets? Are there any practical limitations or considerations?\n\nThe learning update itself scales with the forward/backward of a single sequence model policy. It is comparable to imitation learning because the relative cost of the extra critic term decreases as the size of the shared Transformer increases. It does not directly depend on the task count or dataset size (replay buffer). The classification-style critic loss can be slower than standard regression as we have increased the dimension of the critic networks’ outputs to B bins. Relative to typical online RL experiments, we are evaluating large policies (10-20M+ parameters) and large datasets (20M+ timesteps). All the experiments were conducted on NVIDIA A5000 GPUs. The model is trained on one GPU while an extra GPU can be used for data collection to reduce wall-clock time.\n\nPlease see our global reply for a discussion of limitations.\n\n> a more detailed discussion on potential negative societal impacts and how to mitigate them would be beneficial\n\nSequence-based RL on off-policy data is a flexible approach to memory, adaptation, and generalization settings. The aim of our work is to make this technique easier to use by making more complex multi-task benchmarks realistic for research. We do not think our method adds negative societal impacts beyond the existing risks of RL systems, but we would be happy to discuss this further if you have concerns. \n\n\n[1] *A Generalist Agent*, Reed et al., 2022", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PADTmmR0Em", "reviewer": "Reviewer_NjDj", "summary": "In this paper, the authors investigate to utilize the Transformer architecture for multi-task RL tasks by giving the trajectories from multiple previous episodes as input tokens. More specifically and technically, they addressed the issue of the imbalanced rewards from different task (e.g., For Atari games, some games can give very high valued rewards while it is not for other games). It has been hypothetically resolved by utilizing additional knowledge such as the task label or normalizing rewards. They tried to address this without the additional knowledge for the task different from the previous works. Through this, they showed their modeling is effectively working well for diverse multi-task environments, especially when the reward distribution for each task is very different.", "questions": "- Questions on the experimental results\n\t- For Multi-Game Procgen experimental analysis, especially the result in Figure 6 right, I am not sure that I understood correctly. You analyzed as \"the policy update reduces to IL on a dynamic percentage of the replay buffer, but the ability to ignore a fraction of the dataset automatically allows for self-improvement in a way that standard IL does not.\". What I understood from this analysis is that your agent can reduce to train the policy by imitating the given action, then the y-axis value of the graph should be increased as training progress goes on. But it is different for different environments, which is that what you mentioned as the \"ignore a fraction of the dataset automatically\"?\n\t- I didn't understand clearly the Gym Retro results. Why did you do this experiment? As I understood, it is to evaluate the larger scale of difference between tasks, why Gym Retro gives the setting for this testing? How did you set to test this? What means the SuperMario1 and 3? What means SuperMario2japan and SuperMarioWorld? How could I analyze the experimental results?\n- For Multi-game Atari, can you compare this results with the single task agent results such as DreamerV2 [1] or state-of-the-art results on this environment? I want to compare the relative performance of your model with the single task agent like shown in Figure 5.\n- For results in Figure 7, can you test it with diverse length of context? I am interested what happens if very large number of context is given.\n\n[1] Hafner, Danijar, et al. \"Mastering atari with discrete world models.\" arXiv preprint arXiv:2010.02193 (2020).", "limitations": "Toward generalist agent, their investigation can be seen as being limited in the homogenous environment setting. For example, extending it to train a single agent for multiple environments where action spaces are different as did in [1]. However, it can be seen as over their scope, I didn't think it is a critical limitation of their study.\n\n[1] Reed, Scott, et al. \"A generalist agent.\" arXiv preprint arXiv:2205.06175 (2022).", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Their problem setting is well positioned. As they pointed out, lots of meta-RL environments are mainly designed to evaluate that the agent can work when there are the small different in the tasks such as different positions in a single task, but we should aim to more extensive range of tasks to build the generalist agent. To do that, the imbalanced reward issue is importantly addressed even it has been handled with additional knowledge such as task label.\n- This paper is well written for their motivation and proposal. Their writting for their motivation and proposal is very clear and good enough to understand for me.\n- They reported strong experimental analyses. They did experiments on diverse environments such as Meta-World, POPGym, Procgen, Multi-Game Atari and Gym Retro for their modeling by comparing it to the ablation versions with and without some components of their modeling. Those results show that the strength of their modeling (better handling the imbalanced rewards without task label). Additionally, they also analyzed the in-context learning performance on Procgen environment.", "weaknesses": "- Some experimental analyses are hard to understand. I will ask it in Questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and comments on the experimental results. The evaluation domains in our work are computationally expensive, so it was not possible for us to fully address your questions with finalized results during the rebuttal window. We have begun work on these experiments and will update the paper when the results are complete.\n\n\n> For Multi-Game Procgen experimental analysis, especially the result in Figure 6 right, I am not sure that I understood correctly. \n\n\nThe actor loss is performing supervised learning on a filtered subset of (obs, action) pairs from the replay buffer that are estimated to have a positive advantage by the critics. This update would be equivalent to imitation learning if the filter were to give a weight of one to every example, and we are relying on an accurate critic to mask low quality actions from the loss. Figure 6 right measures the percentage of each batch that passes the filter (A(s, a) > 0) and is used to compute the supervised loss. When the critic network is initialized it approves approximately half of the actions. As training continues, the curve changes according to a combination of the rate of experience collection, environment difficulty, gradient steps per environment step, and the total size of the replay buffer. \n\n> For Multi-game Atari, can you compare this results with the single task agent results such as DreamerV2 [1] or state-of-the-art results on this environment? I want to compare the relative performance of your model with the single task agent like shown in Figure 5.\n\nAtari is a highly optimized domain and we are not claiming state-of-the-art results here by any means. Our goal with this experiment is to use Atari as an example of a standard multi-task setting without the multi-episode adaptation or partial observability of the preceding results. We are measuring whether the change in training objective is enough to make progress without the established Atari trick of clipping rewards in [-1, 1] and multi-task RL techniques that rely on per-task gradients or network outputs. We want to measure a direct ablation of the same basic method that does not add orthogonal details that have led to much of the sample efficiency gains on the ALE. We will work to aggregate the results of single-task versions of the method using the default scale-dependent update (in green) for additional context. \n\n> For results in Figure 7, can you test it with diverse length of context? I am interested what happens if very large number of context is given.\n\nThe trade-off between context length and performance depends on the episode length of the environment. The high frame rate of pixel-based domains can create much longer episodes than typical toy meta-RL experiments. Ideally, the context would span the entire length of the 2-episode rollout, but computational constraints (GPU memory) prevent this. It is worth noting that although we are not using as long of a context as we would like, 768 images is still an unusually long input sequence when training online RL from scratch. \n\n> their investigation can be seen as being limited in the homogenous environment setting. For example, extending it to train a single agent for multiple environments where action spaces are different as did in [1] … I didn't think it is a critical limitation of their study.\n\nYes we are padding the observation/action spaces to the same shape across environments when necessary. A variable-length token encoding like Gato would probably be the first step in attempting this kind of heterogeneous setting. We avoided this because it adds significant computational cost in terms of sequence length per timestep of policy memory, and we felt that the main point of return scale variance is best demonstrated in domains where the observation and actions are similar but the reward functions are different. \n\n> I didn't understand clearly the Gym Retro results. Why did you do this experiment? \n\nThis experiment is meant to be an extension of the Atari setting where rewards are on very different scales. The difference between the games we’ve chosen here and Atari is that the Mario levels are similar enough where some kind of zero-shot transfer might be realistic. The SuperMario names refer to the titles of the Mario video game series that are supported by the Gym Retro environment. We are working to add more context to these results by including the other filtered imitation learning ablation (“Ind. Actor / Dep. Critic” - the red curve in Figure 3) as another reference. The \"Dep. Actor / Dep. Critic\" comparison is not possible here due to the action space of the environment. These runs began before the reviews were released but are expensive to collect and have not finished. They will be added to Figure 9.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ckdLFnnL6D", "reviewer": "Reviewer_TXBg", "summary": "This paper proposes a method for training transformer-based policies for multi-task meta-RL settings, where the task distribution has multiple tasks each of which has parametric variation and the training and test sets both include all tasks. Focus is on developing a method that handles the scale variation in returns between different tasks. Strong empirical results.", "questions": "- This is out of scope for the review, but I'm curious, what do you think is the best bet for a meta-RL benchmark that would enable investigating the non-parametric generalization ml45 fails to provide?", "limitations": "Limitations are discussed briefly in the conclusion. Especially limitations concerning the non-parametric generalization mentioned in weaknesses should be discussed more extensively.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The method description is clear and well motivated.\n- The empirical results are strong and clearly presented.", "weaknesses": "- Framing issue: The considered setting is interesting and a fine target for research. However, parametric variation is not the grand challenge in meta-RL and the paper, in my opinion, does not do enough to clearly delineate the subproblem it tackles.\n    - The ultimate goal of meta-RL is training an adaptive agent that can learn any new task it is presented with, not just parametric variation of existing tasks. This paper does not explore generalization to unseen non-parametric task variants. This should be more clearly stated early in the paper.\n    - The test set of Meta-World ML45 benchmark is mentioned in the experiments, and reasonably ignored. However, since ML45 is already mentioned in the abstract, it would be better to mention that the paper only considers the training set.\n- Minor framing issue: Line 27 says meta-RL is about task identification at its core. While many meta-RL tasks can be solved via methods that reduce to identifying the task and then doing multi-task learning, meta-RL as a whole isn't limited to that. The paper adopts a broader viewpoint in the scaling beyond multi-task barrier paragraph. It would be better to update the introduction to match this perspective.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and for the constructive comments on limitations and writing. We will discuss your suggestions below.\n\n> Framing issue: The considered setting is interesting and a fine target for research. However, parametric variation is not the grand challenge in meta-RL and the paper, in my opinion, does not do enough to clearly delineate the subproblem it tackles.\n\nYes, we fully agree that the goal should be adaptation to new tasks. Our perspective is that non-parametric variation will be possible if we can scale up our training set to keep that variation within distribution. This will likely take far more tasks than are available in benchmarks like Meta-World, and we generally think that attempts to measure or improve non-parametric variation on Meta-World, Atari, Procgen, and similar benchmarks can risk drawing conclusions from too small of a training set. The challenge is that we still struggle to overfit these small task sets and mostly leave this issue to be tackled by MTRL techniques where solutions might not necessarily be transferable to larger benchmarks. A first step might be to address the key MTRL challenges while maintaining the label-free meta-RL perspective. We will add more discussion to the paragraph on lines 42-55 of the introduction. \n\n>The test set of Meta-World ML45 benchmark is mentioned in the experiments, and reasonably ignored. However, since ML45 is already mentioned in the abstract, it would be better to mention that the paper only considers the training set.\n\nYes, we can do that. We are not aware of any results that demonstrate strong transfer to the test set without introducing some restriction on the environment or policy to regularize training or without adding additional knowledge like language, goal states, or demonstrations.\n\nThe reviewer is clearly familiar with Meta-World, but for any other readers who may not be, the train set of 45 tasks can still serve as a reasonable test of parametric variation because the environments are carefully designed to make these changes unobservable.\n\n> Minor framing issue: Line 27 says meta-RL is about task identification at its core … The paper adopts a broader viewpoint in the scaling beyond multi-task barrier paragraph. It would be better to update the introduction to match this perspective.\n\nWe will rephrase to align this with the more detailed explanation given in Section 2. This line was meant to be a fast summary of the meta-RL perspective where adaptation reduces to exploring the environment to infer missing details of the unobserved parameters of the current task, which lets us treat multi-task RL problems as a special case.\n\n> This is out of scope for the review, but I'm curious, what do you think is the best bet for a meta-RL benchmark that would enable investigating the non-parametric generalization ml45 fails to provide?\n\nMeta-RL is facing a difficult benchmarking challenge at the moment, and it’s the reason our experiments need to drift into long-term memory and multi-task domains. A more affordable subset of XLand 2 [1] would probably be the best large-scale benchmark if it were accessible. There have been some recent efforts to create an open-source grid world version of this domain. One interesting direction for large-scale experiments (that is open-source) could be Gym Retro, which supports thousands of ALE-like games. Arcade games and platformers are a natural fit because they often give the player multiple lives/attempts before a reset. Meta-learning was the original goal of Retro, but this is obviously challenging and there hasn’t been much progress made.\n\n[1] *Human Timescale Adaptation in an Open-Ended Task Space*, Adaptive Agent Team, 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all the reviewers for their comments, and we will respond to individual questions and concerns below. 


Several reviewers asked for an expanded discussion of our method’s limitations, which we will add to our conclusion. The main technical limitation of our technique is that it does not address all of the challenges of optimizing distinct task objectives. These challenges are not unique to RL and are the focus of many prior works, as discussed in Section 2. Instead, we are minimizing an additional challenge that value-based RL algorithms can introduce, which is that we are unintentionally rescaling our learning objective(s) throughout training according to a schedule that is difficult to predict or control. We demonstrate a simple approach that lets us manage this issue without introducing explicit task knowledge that would make our agents less applicable to general settings.

Another limitation is that our experiments evaluate the ability to generalize across unseen variants of many different tasks but do not evaluate the ability to generalize to entirely new tasks.  This kind of non-parametric task variation requires large training sets of unique tasks that are beyond the reach of many current benchmarks. We think the path towards this goal requires the ability to learn from the small disjoint task sets currently available. Existing methods struggle to do this without multi-task optimization techniques that rely on our ability to label tasks. We aim to provide a simple way to scale a standard meta-learning approach to increasingly diverse training sets and enable new research on complex domains beyond typical toy problems.

Meta-review: This paper tackles the challenge of scaling meta-reinforcement learning (meta-RL) to manage multiple tasks without explicit task labels. Building upon recent meta-RL techniques, it presents a method that transforms both the actor and critic objectives into classification problems, separating optimization from the scale of returns. The method's effectiveness is validated on various benchmarks, including Meta-World ML45, Multi-Game Procgen, and Multi-Task POPGym, demonstrating notable improvements.

All reviewers indicated high confidence in their reviews, with 3 out of 4 reviewers recommended to accept the paper, and reviewer Ypn5 leaning rejection. However, I find that most of Ypn5's concerns are adequately addressed during the rebuttal. In particular, I agree with the authors that further extending experiments with more baseline comparisons requires a large amount of compute and non-trivial modifications to existing methods that are likely beyond the scope of a conference paper. The method as it stands represents a good contribution that future works may build upon.

To authors: from my reading of the paper, the presentation could be further improved. The paper assumes a high-level of familiarity with previous works and is dense at times. For instance, there are multiple under-defined terms such as "two-hot", and $d_{i-1}$ in $ (o_i, a_{i−1}, r_{i−1}, d_{i−1})$ is only defined in the figure but not the main text. Additionally, please carefully revise the final version to include the discussion with the reviewers.
Decision: Accept (poster)