On the Feasibility of Cross-Task Transfer with Model-Based Reinforcement Learning
Authors: Yifan Xu
Nicklas Hansen
Zirui Wang
Yung-Chieh Chan
Hao Su
Zhuowen Tu
Venue: ICLR 2023 poster (2023)
Keywords: model-based reinforcement learning
visual reinforcement learning

Abstract: Reinforcement Learning (RL) algorithms can solve challenging control problems directly from image observations, but they often require millions of environment interactions to do so. Recently, model-based RL algorithms have greatly improved sample-efficiency by concurrently learning an internal model of the world, and supplementing real environment interactions with imagined rollouts for policy improvement. However, learning an effective model of the world from scratch is challenging, and in stark contrast to humans that rely heavily on world understanding and visual cues for learning new skills. In this work, we investigate whether internal models learned by modern model-based RL algorithms can be leveraged to solve new, distinctly different tasks faster. We propose Model-Based Cross-Task Transfer (XTRA), a framework for sample-efficient online RL with scalable pretraining and finetuning of learned world models. By offline multi-task pretraining and online cross-task finetuning, we achieve substantial improvements over a baseline trained from scratch; we improve mean performance of model-based algorithm EfficientZero by 23%, and by as much as 71% in some instances. Project page: https://nicklashansen.github.io/xtra

Published as a conference paper at ICLR 2023
ON THE FEASIBILITY OF CROSS-TASK TRANSFER WITH MODEL-BASSH REINFORCEMENT LEARNING
Yifan Xu, Nicklas Hansen; Zirui Wang, Yung-Chieh Chan, Hao Su, Zhuowen Tu
University of California, San Diego (yix081, ni nansnn,,, ziw029, ychan, has 68, ztu) (ccs)uucdd.ee edu
ABSTRACT
Reinforcement Learning (RL) algorithms can solve challenging control problems directly from image observations, but they often require millions of environment interactions to do so. Recently, model-based RL algorithms have greatly improved sumple-effiiiency by concurrently learning an internal model of the world, and supplementing real environment interactions with imagined rollouts for policy improvement. However, learning an effective model of the world from scratch is challenging, and in stark contrast to humans that rely heavily on world understand- ing and visual cues for learning new skills. In this work, we investigate whether internal models learned by modern model-based RL algorithms can be leveraged solve new, distinctly different tasks faster. We propose Model-Based Cross-Task Transfer (XTRA), a framework for sample-efficient online RI. with scalable pre- training and finetuning of learned world models. By offline multi-task pretraining and online cross-task finetuning, we achieve substantial improvements over a base- line trained from scratch; we improve mean performance of model-based algorithm EfficientZero by 23%, and by as much as 71% in some instances.
INTRODUCTION
Reinforcement Learning (RI.) has achieved great feats across wide range of areas, most notably game-playing (Mnih et al., 2013; Silver et al., 2016; Berner et al., 2019; Cobbe et al., 2020). However, traditional RL algorithms often suffer from poor sample-efficiene
Atari 100k
1.4
1.2
Normalizedd Score
1.0
and require millions (or even billions) of environment interactions solve tasks espe- cially when learning from high-dimensional observations such as images. This is in stark contrast to humans that have a remarkable ability to quickly learn new skills despite very limited exposure (Dubey al., 2018). In an effort reliably benchmark and improve the samplc-efficieny; of image-based RL across a variety of problems, the Arcade Learning Environment (ALE; (Bellemare et al., 2013)) has become long-standing challenge for RL. This task suite has given rise to numerous successful and increasingly sample- efficient algorithms (Mnih et al., 2013; Badia et al., 2020; Kaiser al., 2020; Schrittwieser et al., 2020; Kostrikov et al., 2021; Hafner et al., 2021; Ye et al., 2021), notably most of which are model-based, i.e., they learn model of the environment (Ha & Schmidhuber, 2018).
0.8
0.4
0.0
20
60
100
Environment steps (x10*)
EfficientZero
BC
XTRA (ours)
XTRA (zero-shot)
Figure 1. Atari100k score, normalized by mean EfficientZero performance at 100k environment steps across 10 games. Mean of seeds; shaded area indicates 95% Cls.
Most recently, EfficientZero Ye et al. (2021), a model-based RL algorithm, has demonstrated impres- sive cample-effiieene surpassing numan-level performance with little as 2 hours of real-time game play in select Atari 2600 games from the ALE. This achievement attributed, in part, to the algorithm concurrently learning an internal model of the environment from interaction, and using the learned model to imagine (simulate) further interactions for planning and policy improvement, thus reducing reliance on real environment interactions for skill acquisition. However, current RL algorithms, including EfficientZero, are still predominantly assumed to learn both perception, model, and skills
"Equal contribution. Project page: https: /ni ck1 wwwaaasaaannnnnnneeeeeeee io/xtra

---

Published as a conference paper at ICLR 2023
Action
Z
Pretrained Model
Online Finetuning
Target Tesk
Reward
MuZero Offiie Pretraining
Offine Tasks
Concurrent Cross-Task Learning
Observation
Stage 1: Offline Multi-Task Pretraining
Stage 2: Online Finetuning On Target Task
Figure 2. Model-Basee Cross-Tase Transfer (XTRA): sample-efficient online RL framework with scalable pretraining and finetuning of learned world models using auxiliary data from offline tasks.
tabula rasa (from scratch) for each new task. Conversely, humans rely heavily prior knowledge and visual cues when learning new skills a study found that human players easily identify visual cues about game mechanics when exposed to new game, and that human performance is severely degraded if such cues are removed or conflict with prior experiences (Dubey et al., 2018).
In related areas such computer vision and natural language processing. large-scale pretraining on large-scale datasets (Devlin et al., 2019; Brown al., 2020; Li et al., 2022; Radford et al., 2021; Chowdhery et al., 2022) has emerged as powerful framework for solving numerous downstream tasks with few samples (Alayrac et al., 2022). This pretraining paradigm has recently been extended visuo-motor control various forms, e.g., by leveraging frozen (no finetuning) pretrained representations (Xiao et al., 2022; Parisi et al., 2022) or by finetuning a supervised setting (Reed et al., 2022; Lee et al., 2022). However, the success of finetuning for online RL. has mostly been limited to same-task initialization of model-free policies from offline datasets (Wang et al., 2022; Zheng et al., 2022), or adapting policies to novel instances of given task (Mishra et al., 2017; Julian et al., 2020; Hansen et al., 2021a), with prior work citing high-variance objectives and catastrophical forgetting as the main obstacles to finetuning representations with RL (Bodnar et al., 2020; Xiao et al., 2022).
In this work, we explore whether such positive transfer can be induced with current model-based RL algorithms in online RL setting, and across markedly distinct tasks. Specifically, seek answer the following questions: when and how can a model-based RL algorithm such as EfficientZero benefit from pretraining on a diverse set of tasks? We base our experiments on the ALE due cues that are easily identifiable to humans despite great diversity in tasks, and identify two key ingredients - cross-task finetuning and task alignment- for model-based adaptation that improve sample-efficiency substantially compared to models learned tabula rasa. In comparison, we find that a naïve treatment of the finetuning procedure as commonly used in supervised learning (Pan & Yang, 2010; Doersch et al., 2015; 5; He et al., 2020; Reed et al., 2022; Lee et al., 2022) is found to be unsuccessful or outright harmful in an RL context.
Based on our findings, we propose Model-Based Cross-Task Transfer (XTRA), a framework for sample-efficient online RL with scalable pretraining and finetuning of learned world models using extra, auxiliary data from other tasks (see Figure 2). Concretely, our framework consists of two stages: (i) offline multi-task pretraining of a world model on an offline dataset from m diverse tasks, (ii) finetaning stage where the world model is jointly finetuned target task in addition to m offline tasks. By leveraging offline data both in pretraining and finetuning, XTRA overcomes the challenges of catastrophical forgetting. To prevent harmful interference from certain offline tasks, adaptively re-weight gradient contributions in unsupervised manner based on similarity to target task.
We evaluate method and of strong baselines extensively across 14 Atari 2600 games from the Atari I 100k benchmark (Kaiser et al., 2020) that require algorithms to be extremely sample-efficient. From Figure and Table observe that XTRA improves sample-efficiency substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25%, respectively.
2 BACKGROUND
Problem setting. We model image-based agent-environment interaction as an episodic Partially Observable Markov Decision Process (POMDP; Kaelbling al. (1998)) defined by the tuple M
2

---

Published as a conference paper at ICLR 2023
P. where O is the observation space (pixels), A is the action space, P: Ox O a transition function, p is the initial state distribution, r: x A H IR is a scalar reward function, and Y [0, 1) is discount factor. As is standard practice in ALE (Bellemare al., 2013), we convert M to a fully observable Markov Decision Process (MDP; Bellman (1957)) by approximating state 8, E S time as stack of frames S, = (ot; Ot-1; 0t-2,-.- where o E O (Mnih et al., 2013), and redefine P. P, r to be functions of s. Our goal is then to find a (neural) policy Ne ma(s)) parameterized by 0 that maximizes discounted return E.a [Σi- (s at)] where a ~ (a(as)) and 7 is the episode horizon. For clarity, we denote all parameterization by throughout this work. To obtain a good policy from minimal environment interaction, we learn "world model" from interaction data and use the learned model for action search. Define M as the target task that we aim to solve. Then, seek to first obtain good parameter initialization & that allows us to solve task M using fewer interactions (samples) than training from scratch, i.e., we wish to improve the sample-efficien
~ P(s at), So ~
of online RL. We do by first pretraining the model on offline (fixed) dataset that consists of transitions (s, a, 7, s') collected by unknown behavior policies in m environments {M' [M # M, 1M,1sis i S m}. and then finetune the model by online interaction on the target task.
EfficientZero (Ye et al., 2021) is model-based RL algo- rithm based on MuZero (Schrittwieser et al., 2020) that learns a discrete-action latent dynamics model from en- viromment interactions, and selects actions by lookahead via Monte Carlo Tree Search (MCTS; (Abramson, 1987; Coulom, 2006; Silver et al., 2016)) in the latent space of the model. Figure 3 provides an overview of the three main components of the MuZero algorithm: represen- tation (encoder) ha, a dynamics (transition) function ge, and prediction head fo- Given an observed state St, EfficientZero projects the state to a latent representation z, = (s;), and predicts future latent states 1 and instantaneous rewards ft using an action-conditional la- tent dynamies function 1,+1, Ft go(2,, For each latent state, a prediction network fe estimates a probability distribution e over (valid) actions
Obearostion
Soarch with
Learned Model
elpolict
brelas)
(mwwrd)
Figure MuiZero/EfficienZZerr combines MCTS with a learned representation network (h), latent dynamics function (g), and predic- tion head (f).
A. well as the expected state value of the given state, i.., é. P: (2x). Intuitively, hg and so allow EfficientZero to search for actions entirely in its latent space before executing actions in the real environment, and fo predicts quantities that help guide the search towards high-return action sequences. Concretely, é provides a return estimate for nodes at the lookahead horizon (as opposed to truncating the cumulative sum of expected rewards) and ý provides an action distribution prior that helps guide the search. We describe EfficientZero's learning objective between model prediction (p. 0, f) and quantity targets (T, 2, u) in Appendix A. EfficientZero improves the sample-efficiena of MuZero by introducing additional auxiliary losses during training. We adopt EfficientZero as our backbone model and learning algorithm, but emphasize that our framework is applicable most model- algorithms, including those for continuous action spaces (Hafner et al., 2019a: 9a; Hansen et al., 2022a).
MODEL-BASER CROSS- -TASK TRANSFER
We propose Model- -ased Cross- Task Transfer (XTRA), a. a two-stage framework for offline multi-task pretraining and cross-task transfer of learned world models by finetuning with online RL. Specifically, we first pretrain a world model on offline data from a set of diverse pretraining tasks, and then iteratively finetune the pretrained model on data from target task collected by online interaction. In the following, we introduce each of the two stages pretraining and finetuning in detail.
3.1 OFFLINE MULTI-TASK PRETRAINING
In this stage, we aim to learn a single world model with general perceptive and dynamics priors across a diverse set of offline tasks. We emphasize that the goal of pretraining is not to obtain a truly generalist agent, but rather to learn good initialization for finetuning to unseen tasks. Learning single RL agent for a diverse set of tasks is however difficult in practice, which is only exacerbated by extrapolation errors due the offline RL setting (Kumar et al., 2020). To address such a challenge,
3

---

Published as a conference paper at ICLR 2023
onhi Tak
Offime bi
Onlin Finetuning
Office
Moddl
Tonk Weiets Uodite
Targ:: Tal
Task Similrity Measurement with Gradient Cosine Simiarity
from Online Terget Task and Offine Tasks
Re-weigh: Offine Tasks Periodically
Figure Illustration of our proposed Concurrent Cross-Task Learning strategy, where we selectively include a subset of the available pretraining tasks while finetuning on a target task.
we propose to pretrain the model following a student-teacher training setup in the same spirit to DQN multi-task policy distillation in Rusu et al. (2016) and Actor- in Parisotto al. (2016), where teacher models are trained separately by offline RL for each task, and then distilled into a single multi-task model using novel instantiation of the MuZero Reanalyze (Schrittwieser al., 2021).
For each pretraining task assume access to fixed dataset {Di |1 S m) that consists of trajectories from an unknown (and potentially sub-optimal) behavior policy. Importantly, we do not make any assumptions about the quality the source of trajectories in the dataset, i.e., do not assume datasets to consist of expert trajectories. We first train individual EfficientZero teacher models on each dataset for a fixed number of iterations in single-task (offline) RL setting, resulting in m teacher models { [1 S i S m}. After training, we store the model predictions, (p, 0). from each teacher model # together with environment reward as the student's 's quantity targets (IT, 2, u)
respectively for given game M (see Appendix A for the definition of each quantity). Next, we learn multi-task student model Te by distilling the task-specific teachers into single model via these quantity targets. Specifically, we optimize the student policy by sampling data uniformly from all pretraining tasks, and use value/policy targets generated by their respective teacher models rather than bootstrapping from student predictions commonly done in the (single-task) MuZero Reanalyze algorithm. This step can be seen as learning multiple tasks simultaneously with direct supervision by distilling predictions from multiple teachers into single model. Empirically, we find this to be a key component in scaling up the number of pretraining tasks. Although teacher models may not be optimal depending on the provided offline datasets, we find that they provide stable (due to fixed parameters during distillation) targets of sufficiently good quality. The simpler alternative - training multi-task model all pretraining tasks simultaneously using RL is found to not scale beyond couple of tasks in practice, as we will demonstrate our experiments in Appendix C. After distilling multi-task student model, we now have single of pretrained parameters that can be used for finetuning to variety of tasks via online interaction, which we introduce in the following section.
3.2 ONLINE FINETUNING ON A TARGET TASK
In this stage, we iteratively interact with a target task (environment) to collect interaction data, and finetune the pretrained model data from the target task. However, empirically observe that directly finetuming a pretrained model often leads to poor performance on the target task due to catastrophical forgetting. Specifically, the initial sub-optimal data collected from the target task can cause a large perturbation in the original pretrained model parameters, ultimately erasing inductive priors learned during pretraining. To overcome this challenge, we propose a concurrent cross-task learning strategy: we retain offline data from the pretraining stage, and concurrently finetune the model on both data from the target task, as well as data from the pretraining tasks. While this procedure addresses catastrophical forgetting, interference between the target task and certain pretraining tasks can be harmful for the sample-efficiency in online RL. As a solution, gradient contributions from offline tasks are periodically re- weighted in an unsupervised manner based on their similarity to the target task. Figure 4 shows the specific concurrent cross-task learning procedure for target task finetuning in our framework.
At each training step t, we jointly optimize the target online task M and m offline (auxiliary) tasks {M JM + M,, < S m) that were used during the offine multi-task pretraining stage. Our online finetuning objective is defined as:
padapt (0) LT (M) +
(1)
1=1
4

---

Published as a conference paper at ICLR 2023
Assault
Camiial
Centipede
Phoenix
Demon Attack
0.
0.4
0.2-
0.0
310
Accsul
Contipede
Alian
Amidar
Bank
Wizard Of Wor
Ms Pacman
1.0
Camtiol
Vesaemmn
0.0
100
Awiderr WOReECORe
(a) Cross-Tass Task Transfer from 4 offline games (left) to target game (right).
(b) Task weights.
Figure 5. Visualization of Concurrent Cross-Tast Learning. (left) the model adants the online target game while concurrently learns offline games. (right) the figure shows the task weights of offline games that are periodically recomputed based on their gradient similarity to the target games (Demon Attack and MsPacman).
where _ is the ordinary (single-task) EfficientZero objective (see Appendix A), and ry' dynami- cally (and independently) updated task weights for each of the m pretraining tasks. The target task loss term maintains a constant task weight of 1. During online finetuning, we distillation targets from teachers obtained from each pretraining game, and use MuZero Reanalyze to compute targets for the target task for which we have no teacher available.
In order to dynamically re-weight task weights 1/' throughout the training process, we break down the total number of environment steps (i.e., 100k in our experiments) into even 7'-ttep cycles (intervals). Within each cycle, we spend first N-steps to compute an updated 17 corresponding to each offline task M The new rf will then be fixed during the remaining T - N steps in the current cycle and the first N steps in the next cycle. We dynamically assign the task weights by measuring the "relevance" between each offline task M' and the (online) target task M. Inspired by the conflicting gradients measurement for multi-task learning in Yu et al. (2020), we compute the cosine similarity between loss gradients 9h from L (M') and Go from (MM) given by
Gh Gr.
Sim(M MM)
(2)
[ ||||
Within the N-step update, we maintain a task-specific counter s' and the new task weights 1y' can be reset by nf Iv at the beginning of each every T-cycle. The procedure for obtaining s is described in Appendix B. Concretely, Sim(M', .M) measures the angle between two task gradients 9 and Gn. Intuitively, we aim to (approximately) prevent gradient contributions from the offline tasks from conflicting with the gradient update direction for the target task by regulating offline tasks objectives with task weights 1). While re-weighting task weights at every gradient update would result in the least amount of conflicting gradients, it prohibitively costly to do so in practice. However, we empirically find the cosine similarity of task gradients to be strongly correlated in time, i.e., the cosine similarity does not change much between consecutive gradient steps. By instead updating task weights every N steps, our proposed technique mitigates gradient conflicts at negligible computational cost in contrast to the compute-intensive intensive gradient modification method proposed in Yu et al. (2020). Figure 5 shows adjustments task weights during finetuning for each of two distinct sets of pretraining and target tasks.
EXPERIMENTS
We evaluate our method and basclines on 14 tasks from the limited-interaction Ataril00k benchmark (Kaiser al., 2020) where only 100k environment steps permitted. We provide implementation of our method at https: //ppcckkasaaasssee
gi githin thub. /x.tr We seek to answer:
How does our proposed framework compare to alternative pretraining and online RL approaches with limited online interaction from the target task?
5

---

Published as a conference paper at ICLR 2023
Table 1. Atari100k benchmark results (similar pretraining tasks). Methods are evaluated at 100k environ- ment steps. For cach game, XTRA is first pretrained on all 4 other games from the same category. Our main result is highlighted We also include three ablations that remove (i) cross-task optimization in finetuning (only online RL). (ii) the pretraining stage (random initialization), and (ii)) task re- weightingg (constant weights of 1). We also include aero-shot performance of our method for target tasks in comparison to behavioral cloning.
Mean of 5 seeds and 32 evaluation episodes.
Ablations (XTRA)
Zero-Shot
Category
Game
BC (Bnstunat)
Efficiett
Effiion
XTRA
W.A.
V.B.
Wa Bank
BC
XTRA
Zum
Zenel
(Oues)
croseddad
pestraining
Weight
(Oura)
Shooter
Assault
838.4
1027.1
1011.6
1294,6
1216.1
1257.5
1164.2
0.0
92.8
Carnival
1952.4
3022.1
2784.3
3860,9
3544.4
2370.0
3071.6
93.75
719.3
Centipede
1814.1
3322.7
275077
5681.4
3833.2
6322.7
5484.1
162.2
1206.8
Demon Attack
825.5
11523.0
4691.0
14140,9
6381.5
9486.8
51045.9
73.8
113.6
Phoonix
427.6
10954.9
3071.0
14579.8
10797.3
9010.6
22873.9
0.0
8073,4
Mcan Improvenent
1,02
2,06
0.02
0.29
Median Improvement
0.55
1.00
0.83
1.28
1.15
0.82
1.65
0.01
0.24
Maze
Alien
152.9
695.0
641.5
954.8
722.8
703.6
633,6
108.1
294,1
Amiddar
121.8
0.0
5.2
Bank Heist
178.8
246.1
244.5
304,9
280.1
225.1
261.4
0.0
7.3
Ms. Pacman
$50.0
1281.4
1172.8
1459.7
1011.1
1122.6
809.2
147.6
448.9
Wizard Wor
928.8
1246.1
263.5
100.0
9.4
Mean Improvement
1,06
0.70
0.17
Median Improvement
0.23
1.00
0.92
1.14
1.11
0.88
0.64
0.10
0.05
Overall
Mean Improvement
0.39
1.00
0.79
1.23
1.04
0.96
1,38
0.05
0.23
Median Improvement
1.12
1.04
0.02
0.16
How do the individual components of our framework influence its success?
When can we empirically expect finetuning to be successful?
Experimental setup. We base our architecture and backbone learning algorithm on EfficientZero (Ye et al., 2021) and focus our efforts the pretraining and finetuning aspects of our problem setting. We consider EfficientZero with two different network sizes to better position our results: (i) the same network architecture as in the original EfficientZero implementation which we simply refer to as EfticientZero, and (ii) a larger variant with 4 times more parameters in the representation network (denoted EfficientZZeo--L) We use the EfficientZero-L variant as the default network for our framework through our experiments, unless stated otherwise. However, we find that our EfficientZero baseline generally does not benefit from a larger architecture, and we thus include both variants for a fair comparison. We experiment with cross-task transfer on three subsets of tasks: tasks that share similar game mechanics (for which we consider two Shooter and Maze categories), and tasks that have no discernible properties in common (referred to Diverse). We measure performance individual Atari games by absolute scores, and also provide aggregate results as measured by mean and median scores across games, normalized by human performance EfficientZero performance at 100k environment steps. All of our results are averaged across 5 random seeds (see Appendix D for more details). We provide details on our pretraining dataset in Appendix F.
Baselines. We compare our method against prior methods for online RL that represent the state-of- the-art the Atari100k benchmark (including EfficientZeroo,, multi-task behavior cloning policy pretrained on the same offline data as our method does for zero-shot performance on the target task and the performance after finetuning on the target task (see Appendix O for details), and a direct comparison to CURL (Srinivas et al., 2020), strong model-free RL baseline, under an offline pretraining online finetuning setting. We also include set of ablations that include EfficientZero with several different model sizes and wverriniig/finetuning schemes. The former baselines serve to position our results with respect to the tfate-of-the-art, and the latter baselines and ablations serve to shed light on the key ingredients for successful multi-task pretraining and finetuning.
4.1 RESULTS& & DISCUSSION
We introduce our results in the context of each of our three questions, and discuss our main findings.
How does our proposed framework compare to alternative pretraining and online RL approaches with limited online interaction from the target task?
Tasks with similar game mechanics. We first investigate the feasibility of finetuning models that are pretrained on games with similar mechanics. We select shooter games and 5 maze games for this experiment. Results for our method, baselines, and set of ablations the Atari 00k benchmark

---

Published as a conference paper at ICLR 2023
Table 2. Atari100k benchmark results (diverse pretraining tasks). XTRA results use the same of pretrained model parameters obtained by offline pretraining on diverse games. Mean 5 seeds each with 32 evaluation episodes. Our result is highlighted, All other results are adopted from EfficientZero (Ye al.,
2021). We also report human- numaan-ormmlizzed mean and median scores.
Game
XTRA (Ours)
EticienIZero
Random
Human
SimPLe
OTRainbow
DrQ
SPR
MuZero
CURL
Assault
1742.2
1263.1
222.4
742.0
527.2
351.9
452.4
571.0
500.1
600.6
HarteZoone
14631.3 3
13871.2
2360.0
37187.5
5184.4
406416
12954.0
16651.0
7687.5
14870.0
10631.8
30826.4
2656.6
3736.3
6279.3
Krull
7735.8
5663.3
1598.0
2665.5
4539.9
3277.9
4018.1
3688.9
4890.8
4229.6
Seaquest
749,5
1100.2
68.4
42051.7
683,3
286.9
301.2
583.1
208.0
384.5
Normed Mean
187
1.29
0.00
1,00
0.70
0.41
0.62
0.65
0.77
0.75
Noroed Median
0.35
0.33
0.00
1,00
0.08
0.18
0,30
0.41
0.15
0.36
2.4
1.4
1.0
1.2
Norma Noraatza Score
1.2
100
2.5
0.8-
7.0-
0.6
0
2.5
1.0
2.0
20
40
50
100
0.0
60
800
100
0
70
40
02
100
0.5
75
70
60
20
100
Envrorment stops (x10
Envtronment stops (x10')
Envirorment stopa (x107)
Enxfronment stops 10
Stester
XTTR Jouns
(Durt)
froden)
Figure 6. (a) Effectiveness of Task Relavance, (b) Frozen Representation, (c) Model Size, and (d)
Environment Steps. We visualize model performance on aggregated scores (5 seeds) from 5 shooter games.
shown in Table 1. For completeness, we also provide learning curves in Figure 11 well as an aggregate curve across 5 seeds of all 10 games normalized by EfficientZero's 's performance at 100K environment steps Figure 1. We find that pretraining improves sample-effiiinncy substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25% respectively, overall. Interestingly. XTRA also had a notable zero-shot ability compared to multi- game Behavioral Cloning baseline that is trained on the same offline dataset. We also consider three ablations: (1) XTRA without cross-task: a variant of our method that naively finetunes the pretrained model without any additional offline data from pretraining tasks during finetuning, (2) XTRA without pretraining: a variant that uses our concurrent tross-task learning (i.e., leverages offfine data during finetuning) but is initialized with random parameters (no pretraining), and finally (3) XTRA without task weights: a variant that uses constant weights of for all task loss terms during finetuning. We find that XTRA achieves extremely high performance 2 games (Demon and Phoenix) without dynamic task weights, improving over Efficien(Zero by as much as 343% on DemonAttack. However, its median performance is overall low compared our default variant that uses dynamic weights. We conjecture that this is because some (combinations of) games are more susceptible to gradient conflicts than others.
Tasks with diverse game mechanics. We now consider a more diverse set of pretraining and target games that have no discernible properties in common. Specifically, we use the following tasks for pretraining: Carnival, Centipede, Phoenix, Pooyan, Riverraid, VideoPinball, WizardOFWor, and YarsRevenge, and evaluate our method on tasks from Atari Results are shown in Table 2. We find that XTRA advances the state-of-the-art in a majority of tasks on the Atari 100k benchmark, and achieve a mean human- nuraa-normalized score of 187% vs. 129% for the previous SOTA, EfficientZero. We perform the same set of the ablations and the results are shown in Table from Appendix H. Additionally, we include an ablation that examines the effect of the number of pretrained tasks on later finetuning performance. Details and results for this ablation are shown in Table 7 from Appendix I.
do for tasks with similar game mechanics with XTRA,
Model-free comparisons. For both settings (e.g., tasks with similar & diverse game mechanics), also compare our framework with a strong model-free baseline, CURL (Srinivas et al., 2020), where CURL is pretrained on the same pretraining tasks XTRA is, and later finetuned each of the target tasks. We find that pretraining does not improve the performance of this model-free baseline as consistently as for our model-based framework, XTRA, under both settings. More details and results on this comparison can be found in Table 4 and 5 from Appendix G.
2. How do the individual components of our framework influence its success?
deeper look at task relevance. While our previous experiments established that XTRA benefits from pretraining even when games are markedly different, we now seek to better quantify the
7

---

Published as a conference paper at ICLR 2023
importance of task relevance. We compare the online finetuning performance of XTRA in two different settings: (1) pretraining on 4 shooter games and finetuning to 5 new shooter games, and (2) pretraining on 4 maze games and finetuning to the same 5 shooter games. Aggregate results across all 5 target tasks are shown in Figure 6 (a). Unsurprisingly. we observe that offline pretraining and concurrently learning from other shooter games significantly henefit the online target shooter games through training, with particularly large gains early in training. On the contrary, pretraining on maze games and finetuning to shooter games show similar performance compared to EfficientZero trained from scratch. This result indicates that (1) selecting pretraining tasks that are relevant to the target task is key benefit from pretraining, and (2) in the extreme case where there are no pretraining tasks relevant to the target task, finetuning with XTRA generally does not harm the online RL performance since it can automatically assign small weights to the pretraining tasks.
Which components transfer in model-based RL? Next, we investigate which model component(s) are important to the success of cross-task transfer. We answer this question by only transferring a subset of the different model components representation h. dynamics function g. and prediction head to to the online finetuning stage and simply using a random initialization for the remaining components, Results are shown in Figure 7. Interestingly, we find that only transferring the pretrained representation h to the online RL stage only improves slightly over learning from scratch, especially in the early stages of online RL. In comparison, loading both the pretrained representation and dynamics function accounts for the majority of the gains in XTRA, whereas loading the prediction (but matters for zero-shot performance). We
heads has no significant impact on sample-efficiene conjecture that this is because learning a good dynamics function is relatively more difficult from few samples than learning a task-specific visual representation, and that the prediction head accounts for only a small amount of the overall parameters in the model. Finally, we hypothesize that the visual representation learned during pretraining will be susceptible to distribution shifts as it transferred to an unseen target task. To verify this hypothesis, we consider an additional experiment where we transfer all components to new tasks, but freeze the representation h during finetuning, i.e., it remains fixed. Results for this experiment are shown in Figure 6 (b). We find that, although this variant of our framework improves over training from scratch in the early stages of training, the frozen representation eventually hinders the model from converging to a good model, which is consistent with observations made in (supervised) imitation learning (Parisi et al., 2022).
Scaling model size. In this experiment, we investigate whether XTRA benefits from larger model sizes. Since dynamics network g and prediction network / are used in MCTS search, increasing the parameter counts for these two networks would increase infer- ence/training time complexity significantly. However, increasing the size of the representation network / has relatively small impact on overall inference/training time (see Figure for reference). We com- pare the performance of our method and EfficientZero trained from scratch with each of our two model sizes, the original EfficientZero architecture and a larger variant (denoted Efficien(Zero-L), results are shown in Figure (c). We find that our default, larger variant of XTRA (denoted XTRA-L in the figure) is slightly better than the smaller model size. In comparison, EfficientZero-L, performs significantly worse than the smaller variant of EfficientZeeo..
1.4
Model Components
1.2
1.0
Normaizee
0.8
0.6
0.4
0.2
0.0
0
80 Environment stops (x10')
20
40
60
100
Load None Load h
Load h
Load heg
Figure 7. Effectiveness of model components. The ag- gregated scores from 5 shooter games by loading parameters of different pretrained model compo- nents. Mean of 5 seeds; shaded area indicates 95% Cls.
Relative improvement vs. environment steps,.. Finally, we visualize the average improvement over EfficientZero throughout training in Figure 6 (d). Results show that XTRA is particularly useful in the early stages of training, i.e., in an extremely limited data setting. We therefore envision that cross-tas task pretraining could benefit many real- world applications of where environment interactions are typically constrained due to physical constraints.
When can we empirically expect finetuning to be successful?
Based on Table and 2, we conclude that cross-task transfer with model-based RL is feasible. Further, Figure 6 (a) shows that our XTRA framework benefits from online finetuning when pretraining tasks are relevant, and both representation and dynamics networks contribute to its success (Figure 7).
8

---

Published as a conference paper at ICLR 2023
5 RELATED WORK
Pretrained representations are widely used to improve downstream performance in learning tasks with limited data or resources available, and have been adopted across a multitude of areas such as computer vision (Girshick et al., 2014; Doersch et al., 2015; He et al., 2020), natural language processing (Devlin et al., 2019; Brown et al., 2020), and audio (van den Oord et al., 2018). By first learning a good representation
large dataset, the representation can quickly be finetuned with, e.g., supervised learning on a small labelled dataset to solve a given task (Pan & Yang, 2010). For example, He et al. (2020) show that contrastive pretraining features for ImageNet classification, and Brown et al. (2020) show that a generative model trained on large-scale natural language data can be used solve unseen tasks given only a few examples. While this is a common paradigm for problems that can be cast as (self-)supervised learning problems, it has seen comparably little adoption in RL literature. This discrepancy is, in part, attributed to optimization challenges in RL (Bodnar et al., 2020; Hansen et al., 202 lb; Xiao et al., 2022; Wang et al., 2022), as well as lack of large-scale datasets that capture both the visual, temporal, and control-relevant (actions, rewards) properties of RI. (Hansen et al., 2022b). In this work, we show that - despite these challenges modern model-based RL algorithms can still benefit substantially from pretraining on multi-task datasets, but require more careful treatment during finetuning.
large, unlabelled dataset learns good
Sample-efficient RL. Improving the sample-efficiency of visual RL algorithms is a long-standing problem and has been approached from many largely orthogonal perspectives, including represen- tation learning (Kulkarni et al., 2019; Yarats et al., 2019; Srinivas et al., 2020; Schwarzer et al., 2021), I), data augmentation (Laskin et al., 2020; Kostrikov et al., 2021; Hansen et al., 2021b), bootstrapping from demonstrations (Zhan et al., 2020) or offline datasets (Wang et al., 2022: Zheng et al., 2022; Baker et al., 2022). using pretrained visual representations for model-free RL (Shah & Kumar, 2021; Xiao et al., 2022; Ze et al., 2022), and model-based RL (Ha & Schmidhuber, 2018 10 Finn & Levine, 2017; Nair et al., 2018; Hafner et al., 2019b; Kaiser et al., 2020; Schrittwieser al., 2020; Hafner et al., 2021; Ye et al., 2021; Hansen et al., 2022a; Seo et al., 2022; Hansen et al., 2023). We choose to focus our efforts on sumple-effiiennyy from the perspective of pretraining in a model-based context, i.e., jointly learning perception and dynamics. Several prior works consider these problems independently from each other: Xiao et al. (2022) shows that model-free policies can be trained with a frozen pretrained visual backbone, and Seo et al. (2022) shows that learning world model on top of features from a visual backbone pretrained with video prediction can improve model learning. Our work differs from prior work in that we show it is possible to pretrain and finetune both the representation and the dynamics using model-based RL.
Finetuning in RL. Gradient- based finetuning is well-studied technique for adaptation in (predomi- nantly model-free) RL, and has been used to adapt to either changes in visuals or dynamics (Mishra al., 2017; Yen-Chen al., 2019; Duan et al., 2016; Julian al., 2020; Hansen et al., 2021a; Bodnar et al., 2020; Wang et al., 2022; 7e et al., 2022), or task specification (Xie & Finn, 2021; Walke et al., 2022). For example, Julian et al. (2020) shows that a model-free policy for robotic manipulation lation can adapt to changes in lighting and object shape by finetuning via rewards on a mixture of data from the new and old environment, and recover original performance in less than 800 trials. Similarly, Hansen et al. (2021a) shows that model-free policies can (to some extent) also adapt to small domain shifts by self-supervised finetuning within a single episode. Other works show that pretraining with offline RL on a dataset from a specific task improve sample-efficiency during online finetuning on the same task (Zheng et al., 2022; Wang et al., 2022). Finally, Lee et al. (2022) shows that offline multi-task RL pretraining via sequence modelling can improve offline finetuning on data from unseen tasks. Our approach is most similar to Julian et al. (2020) in that we finetune via rewards on a mixture of datasets. However, our problem setting is fundamentally different: we investigate whether multi-tast pretraining can improve online RL on an unseen task across multiple axes of variation.
CONCLUSION
In this paper, we propose Model-Based Cross-Task Transfer (XTRA), a framework for sample- efficient online RL with scalable pretraining and finetuning of learned world models using extra, auxiliary data from other tasks. We find that XTRA improves sample-efficienge substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25%, respectively, overall. As a feasibility study, we hope that our empirical analysis and findings on cross-task transfer with model-basee based RL will inspire further research in this direction.

---

Published as a conference paper at ICLR 2023
Reproducibility Statement. We base all experiments the publicly available Atari 100k bench- mark, and provide extensive implementation details in the appendices, including detailed network architecture and https: n ck aasaanssnnn.. aasaansn.gggthh.. o xtra.
Code for reproducing our experiments is made available at
Acknowledgements. This work is supported by NSF Award IIS-2127544, NSF TILOS AI Institute, and gifts from Qualcomm AL We would like to thank Weirui Ye, Wenlong Huang, and Weijian Xu for helpful discussions.
REFERENCES
Bruce D. Abramson. The Expected-Outtoome Model Two- TeTTwwo-Pll Two-Player Games. PhD thesis, Columbia University, 1987. AA18827528.
Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lene, Arthur Mensch, Katie Millican, Malcolm Reynolds, al. Flamingo: a visual language model for few-shot learning. arXiv preprint arXiv:2204. 14198, 2022.
Adrià Puigdomënech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. ArXiv, abs/2003. 2020.
Bowen Baker, lgge Akkaya, Peter Zhokhov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. ArXiv, abs/2206. 11792 2022.
Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents (extended abstract). In IJCAI, 2013.
Richard Bellman. A markovian decision process. Journal of Mathematics and Mechanics, 6(5):
679-684, 1957.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Christopher Hesse, Rafal Józefowicz, Scott Gray. Catherine Olsson, Jakub W. Pachocki, Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, llya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota 2 with large scale deep reinforcement learning. ArXiv, abs/1912.06680. 2019.
Cristian Bodnar, Karol Hausman, Gabriel Dulac-Amold, and Rico Jonschkowski. A geometric perspective self-supervised policy adaptation. ArXiv, abs/2011. 107318, 2020.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. ArXiv, abs/2005. 14165, 2020.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.023111 2022.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. ArXiv, 0bs/1912.0115 2020.
Rémi Coulom. Efficient selectivity and backup operators monte-carlo tree search. International conference on computers and games, pp. 72 83. Springer, 2006.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. ArXiv, abs/1810.04805, 2019.
10

---

Published as a conference paper at ICLR 2023
Carl Doersch, Abhinav Kumar Gupta, and Alexei A. Efros. Unsupervised visual representation learning by context prediction. 2015 IEEE International Conference on Computer Vision (ICCV), pp. 1422-1430, 2015.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and P. Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. ArXiv,
2016.
Rachit Dubey, Pulkit Agrawal, Deepak Pathak, Thomas L. Griffiths, and Alexei Efros. Investigating human priors for playing video games. In ICML, 2018.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. 2017 IEEE International Conference on Robotics and Automation (ICRA), pp. 2786-2793, 2017.
Ross B. Girshick, Jeff Donahue, Trevor Darrell, and Jitendra Malik. Rich feature hierarchies for accurate object detection and semantic segmentation. 2014 IEEE Conference Computer Vision and Pattern Recognition, pp. 580-587, 2014.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems, 31, 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01600 2019a.
Danijar Hafner, Timothy Lillicrap, Ian S. Fischer, Ruben Villegas, David R Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. ArXiv, abs/1811.00 2019b.
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. ArXiv, abs/2010.02193, 2021.
Nicklas Hansen, Rishabh Jangir. Yu Sun, Guillem Alenyå, Picter Abbeel, Alexei A. Efros, Lerrel Pinto, and Xiaolong Wang. Self-supervised policy adaptation during deployment. International Conference Learning Representations (ICLR), 2021a.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. In NeurIPS, 2021b.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In ICML, 2022a.
Nicklas Hansen, Zhecheng Yuan, Yanjie Ze, Tongzhou Mu, Aravind Rajeswaran, Hao Su, Huazhe Xu, and Xiaolong Wang. On pre-training for visuo-motor control: Revisiting a Seaming-ffrom-soe scratch baseline, 2022b.
Nicklas Hansen, Yixin Lin, Hao Su. Xiaolong Wang, Vikash Kumar, and Aravind Rajeswaran. Modem: Accelerating visual model-based reinforcement learning with demonstrations. 2023.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. 2020 IEEE/CVF Conference on Computer Vision and 2020.
Pattern Recognition (CVPR), pp.
Ryan C. Julian, Benjamin Swanson, Gaurav S. Sukhatme, Sergey Levine, Chelsea Finn, and Karol Hausman. Efficient adaptation for end-to-end vision-based robotic manipulation. ArXiv, Abs/2004.10199 2020.
Leslie Pack Kaelbling, Michael L. Littman, and Anthony R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial Intelligence, 1998.
Lukasz Kaiser. Mohammad Babaeizadch, Piotr Milos, Blazej Osinski, Roy H. Campbell, K. Czechowski, D. Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Ryan Sepassi, G. Tucker, and Henryk Michalewski. Model-based reinforcement learning for atari. ArXiv, abs/1903.00374, 2020.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. ArXiv, abs/2004. 13649, 2021.
11

---

Published as a conference paper at ICLR 2023
Tejas D. Kulkarni, Ankush Gupta, Catalin Ionescu, Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for perception and control. ArXiv, abs/1906. 1883, 2019.
Aviral Kumar, Aurick Zhou, G. Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. ArXiv, (es//0006.0479,, 2020.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, P. Abbeel, and A. Srinivas. Reinforcement learning with augmented data. ArXiv, (bbbs/2004.14990, 2020.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang. L. Y. Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian S. Fischer, Eric Jang, Henryk Michalewski, and Igor Mordatch. Multi-game decision transformers. ArXiv, abs/2205. 2022.
Chunyuan Li, Haotian Liu, Liunian Harold Li, Pengchuan Zhang. Jyoti Aneja, Jianwei Yang, Ping Jin, Yong Jae Lee, Houdong Hu, Zicheng Liu, et al. Elevater: A benchmark and toolkit for evaluating language-augmented visual models. arXiv preprint arXiv:2204.08790 2022.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Abbeel. Meta-learning with temporal convolu- tions. ArXiv, abs/1707.031 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5609 2013.
Ashvin Nair, Vitchyr H. Pong, Murtaza Dalal, Shikhar Bahl. Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. In NeurlPS, 2018.
Sinno Jialin Pan and Qiang Yang. A survey on transfer learning, IEEE Transactions on Knowledge and Data Engineering. 02:1345-1359, 2010.
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Kumar Gupta. The unsurprising effectiveness of pre-trained vision models for control. In ICML, 2022.
Emilio Parisotto, Lei Jimmy Ba, and Ruslan Salakhurdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. In ICLR (Poster), 2016.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. International Conference on Machine Learning, pp. 8748 $748-8763, 2021.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg. Tom Eccles, Jake Bruce, Ali Razavi, Ashley D. Edwards, Nicolas Manfred Otto Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. ArXiv, abs/2205.06175. 2022.
Andrei Rusu, Sergio Gomez Colmenarejo, Çaglar Gülçchre, Guillaume Desjardins, James Kirk- patrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy distillation. In ICLR (Poster), 2016.
Julian Schritwieserr, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, L. Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with learned model. Nature, 2020.
Julian Schritwwieser, Thomas K Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with a learned model. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL ht https: //wp///eeeeeeeeeeeeeeeeeeeoe.eee net Forum? ddHKttsG 1aepn-Nbbb 1Nb..
12

---

Published as a conference paper at ICLR 2023
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient reinforcement learning with self sueficreeittvve representations. ICLR, 2021.
Younggyo Seo, Kimin Lee, Stephen James, and Abbeel. Reinforcement learning with action-free pre-training from videos. In ICML, 2022.
Rutav Shah and Vikash Kumar. Rrl: Resnet as representation for reinforcement learning. ArXiv, abs/2107.08 03300, 2021.
David Silver, Aja Huang, Christopher Maddison, Arthur Guez, Laurent Sifre, George Driessche, Julian Schritwieser,, loannis Antonoglou, Veda Panneershelvam,, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529:484-489, 01 2016. doi: 10. 10.1000ntttree 1696.
A. Srinivas, Michael Laskin, and P. Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. /CML, 2020.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. ArXiv, abss/1077 wwtp//307.3778.... 2018.
Homer Walke, Jonathan Yang, Albert Yu, Aviral Kumar. Jedrzej Orbik, Avi Singh, and Sergey Levine. Don't start from scratch: Leveraging prior data to automate robotic reinforcement learning. ArXiv, abs/2207.04700 2022.
Che Wang, Xufang Luo, Keith W. Ross, and Dongsheng Li. Vr13: data-driven framework for visual deep reinforcement learning. ArXiv, abs/2202.10322 10324, 2022.
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. ArXiv, abs/2203.06173, 2022.
Annie Xie and Chelsea Finn. Lifelong robotic reinforcement learning by retaining experiences. ArXiv, 099100099900,,,... 2021.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. arXiv, 2019.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, P. Abbeel, and Yang Gao. Mastering atari games with limited data. In NewrIPS, 2021.
Lin Yen- Ven-Chen, Maria Bauza, and Phillip Isola. Experience-embedded visual foresight. In Conference on Robot Learning. 2019.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33: 5824-5836, 2020.
Yanjic Ze, Nicklas Hansen, Yinbo Chen, Mohit Jain, and Xiaolong Wang. Visual reinforcement learning with self-supervised 3d representations. arXiv preprint arXiv:2210.07241,000 2022.
Albert Zhan, Philip Zhao, Lerrel Pinto, Abbeel, and Michael Laskin. A framework for efficient robotic manipulation. ArXiv, abs/2012.079755, 2020.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In ICML, 2022.
13

---

Published as a conference paper at ICLR 2023
A XTRRAEEFICIENTZERR OBJECTIVES
XTRA uses the same learning objective as Efficien(Zero during both offline pretraining and online finetuning, except the quantity targets are predicted by the teacher model during distillation and concurrent learning for pretrained games instead of Muzero Reanalysis procedure.
Here, we explain the objectives of EfficientZero (Ye et al., 2021) and its predecessor MuZero (Schrittwieser et al., 2020). To warrant the latent dynamics that can mirror the true states of the environment, MuZero is trained to predict three necessary quantities directly relevant for planning: (1) the policy target obtained from visit count distribution of the MCTS (2) immediate reward from environment (3) bootstrapped value target where Ei= "²u² + + On top of MuZero, EfficientZero adds a self-supervised consistency loss term, and predicts sum of environment rewards from next steps, E/=u y'u', instead of single-step reward. We refer reader to the original manuscripts for implementation details. We present the learning objective for EfficientZero time step & with & unroll steps:
6 (0)
11 (me R + A tnc (alls" Asl ( clllll
reward
policy
consistency
B TASK WEIGHTS COMPUTATION
Within the N stepss update, we maintain a task-specific counter s and update the counter by As, at each step n as follows:
As 1, if Sim(M MM > 0.1
0, otherwise
s² As
At every N steps, the new task weights ný are updated by 1f N: and used in subsequent finetuning objectives according to Equation 1. In practice, we start task weight updates at 10k steps to ensure enough data from the online target task has been collected for meaningful similarity measure. All task weights are initialized as 1 for the first 10k steps.
Figure shows how task weights are adaptively adjusted by the model during 100k environment steps during online finetuning stage for 10 games reported in Figure 11. Figure 9 shows the adjustments of task weights for games reported in Figure 12.
1.0
Assaut
Cartral
Conticedo
DemrrAttacc
Phoenix
p.0-
0.6
Wizardo/wior
0.2
0.0
0
20
B6
150
20
40
60
1000
20
80
1000
40
60
BC
1000
20
40
B0
300
EIspgemenn stago Stago:14300 (4:09)
(K324)
bnetteenan Steps (x10")
Envecommnnt JACzs 201)
Envennment Stggo (4:11P)
Carrival
Contipddi
Phoonix
Figure Periodic task re-weighting, We visualize task weights as a function of environment steps for each of our 10 tasks from Table 1. First row corresponds to shooter games and the hattom row corresponds to maze games. We evaluate task weights on all tasks from the same category except for the target task itself.
C DISTILLATION VS. MULTI-GAME OFFLINE RL
Our method learns a multi game world model from offline data via distillation of task specific world models trained with offline RI. An alternative way obtain such multi-game model would be
14

---

Published as a conference paper at ICLR 2023
Se5g
0.6
0.4
0.0
20
00
80
1000
40
60
100
80
100
65
100
20
40
100
YosRovvngg Vidoofinaal
Wizardonwior
Figure Periodic task re-weighting. We visualize task weights as a function environment steps for each of our 5 tasks from Table 2. We evaluate task weights on all 8 tasks used during pretraining.
directly train a single world model on the multi game dataset with offline RL. However, we find that learning such a model is difficult. A comparison between the two approaches on four different pretraining games is shown in Figure 10. We observe that multi- game offline RL (green) achieves low scores in all four pretraining games, whereas our wold model obtained by distillation (orange) performs comparably to single-task world models (blue) in 3 out of tasks.
Allen
200
Amidar
BankHHeist
3000
MsPacman
M
1000-
150
200
750-
100
500-
100
1000
0
90
10
50
90
10
50
90
70
11.
steps
steps
steps Ox10P)
steps
Offline Single Game
Offiine Multigamee
Distill Multigame
Figure 10. Distillation multi-game offline RL. Results are shown four pretraining tasks. Single-game models (blue) are trained via offline RI. on each individual task, whereas results for multi-game offline RI. (green) and our proposed distillation (orange) are obtained by evaluating a single set of pretraining parameters. Distillation nearly matches single-task performance.
D SCORES FOR INDIVIDUAL SEEDS
Our results in Table are aggregated across seeds. In Table 3, we report game scores for each individual seed, as well as the mean, median, and standard deviation of game scores for each game. We also list random and human scores obtained from Badia et al. (2020), and calculate the Human Normalized Score based on the formula: (Scoreagent scOrerandom) // /scoreenmnnn score in prior work. To the best of our knowledge, there are no human performance results for the Carnival game, and we therefore exclude this game from the aggregate Human Normalized Mean and Median Scores computed in Table 3.
Table 3. Scores for individual seeds. Per-seed game scores for our method, random behavior baseline, and human performance. We report both unnormalized and normalized scores (Human Normalized Scores), well their aggregate results. Random and human scores are obtained from Badia (2020). We evaluate each random seed on 32 evaluation episodes at 100k steps.
Gane Score per Seed
Aggregated Metries
Referenees
Seed
Mean
Median
Std
Random
Homan
Normed
Assoult
1450. 19
1356. 53
1236.19
1215.12
1214-72
1294, 55
1236. 19
105.06
222.40
742.00
2.06
Camival
3865.31
4867.50
4155.62
2601.58
3814.38
3860.94
3865.31
819.67
Centipede
7596. 36
6179.25
5380.41
5300.03
3950
5681
5380-41
1336 58
2090.90
12017.00
0.36
DemonnAmack
10420.78
8051.25
27574.06
8117.81
16490.47
14140.88
10570.78
8258.
152.10
1921.00
10521.86
14709.06
7242.60
Allen
807.50
954.75
814.06
7127.70
Amidar
93.00
104.34
76.47
97.58
79.59
90.15
93.00
11.54
5,30
1719.50
0.05
BanksHeast
303.12 12
316.56
270.62
270.00
364.05
304.88
303.12
38.83
14.20
753.10
0.39
McPacman
1109.66
1960.00
1865.94
1228 44
1134 38
1459.69
1228.44
417.48
3077 30
6951.60
0.17
WizardOfWar
213.62
563.50
4756.50
Human Nomred Mediam
15

---

Published as a conference paper at ICLR 2023
Assaut
Carrinal
Certipedo
DerrenAitack
Phoerte
6000
15000
1000
15000
10000
5.00
2000
5000
3000
0
Allen
Amidrr
Banketis
NGPooman
1000
300
1500
100
202
1000
1000
50
500
20
40
50
80
60
80
20
40
300
20
40 G3
60
80
100
Envfromment steos
Enwironment EX10P1
Enwironment
Enviromment stegi
steps
(x110)
RC
--- XTRA tours remm-shot)
F FfficienPer
Fffrieen7een
XTRA (nurs)
Figure 11. Atari100k benchmark results (similar pretraining tasks). We report unnormalized scores,
aggregated across seeds per game. Shaded area indicates 95% confidence intervals.
E ADDITIONAL EVALUATION CURVES OF XTRA ON ATARI1OOK BENCHMARK
For completeness, Figure 11 and 12 include evaluation curves of XTRA on the games for which we report final performance in Table and 2.
5000
8000
30000
6000
500
7500
1000-
5000
4000
400-
300-
5000
2000
200
66
20
(x20")
(x207)
(x122)
Training (x12P)
(X1PP)
(ours)
End Enuironment
Figure 12. Atari100k benchmark results (diverse pretraining tasks). report unnormalized scores. aggregated across seeds per game. Shaded area indicates 95% confidence intervals.
F OFFLINE DATA PREPARATION
To train the model in offline multi-task pretraining stage, we use trajectories collected by EfficientZero (Ye et al., 2021) on the Atari 100k benchmark. For each pretraining game, assume we can access model checkpoints obtained every 10k steps from 120k training steps (the environment step is capped 100k), resulting 12 model checkpoints. For each checkpoint, we evaluate the model performance on the Atari environment following the same procedure from EfficientZero and collect 64 trajectories. This translates to an average of 1M transitions per game, varies depending on episode length for example, this only results in 636k transitions for the game of Assault. Since trajectories are collected from model checkpoints both at the early and late training stage within the 120k training steps, the collected data does not necessarily come from expert agent. Thus, we show that XTRA is effective even when pretraining data is suboptimal, allowing us to learn from very diverse data sources.
G PRETRAINING FINETUNING IN MODEL-FREE RL
We compare the effectiveness of our framework, XTRA, with strong model-free baseline, CURL, that we also implement following a similar pretraining and finetuning scheme. This in contrast to the original formulation of CURL that does not leverage pretraining, We implement our training scheme for CURL with the following setup: (1) pretrain a multi-task CURL model on the same pretraining tasks our framework uses (using offline data generated from training individual CURL agents), and (2) directly finetune the pretrained model on the target task with online RL for 100k environment steps.
16

---

Published as a conference paper at ICLR 2023
Table 4. Model-based XTRA comparison with the model-free pretraining + finetuning scheme for tasks that share diverse game mechanies. The model is pretrained with Camival, Centipede, Phoenix, Pooyan, Riverraid, VideoPinball. WizardOFWor, and YarsRevenge. Results for EfficientZero, CURL, Random, and Human adopted from EfficientZero et al., 2021). other results are based on the average 5 runs.
Model-Based
Model-Free
Game
Efficient
XTRA
CURL
CURL"
Random
Human
Zero
(Ours)
Assault
1263.1
1742.2
600.6
588.6
222.4
742.0
BattleZone
13871.2
14631.3
14870.0
16450.0
2360.0
37187.5
Hero
9315,9
10631.8
6279.3
6294.5
1027.0
30826.4
Krull
5663.3
7735.8
4229.6
3472.8
1598.0
2665.5
Seaquest
1100.2
749.5
384.5
385.5
68.4
42054.7
Normed Mean
1.29
1.87
0.75
0.60
0.00
1.00
Normed Median
0.33
0.35
0.36
0,40
0.00
1.00
Table Model-based XTRA comparison with model-free pretraining + finctuning scheme tasks
that share similar game mechanics. For each target game for finetuning, the model is first pretrained on all other 4 games from the same category. All results are based on the average of 5 runs.
Model-Based
Model-Free
Category
Game
Efficient
XTRA
CURL
CURL"
Zero
(Ours)
Shooter
Assault
1027.1
1294.6
590.2
461.2
Carnival
3022.1
3860.9
591.6
714.8
Centipede
3322.7
5681.4
4137.7
3731.0
DemonAttack
11523.0
14140.9
908.3
638.9
Phoenix
10954.9
14579.8
901.2
1168.4
Mean Improvement
1.00
1.36
0.44
0.39
Median Improvement
1.00
1.28
0.20
0.24
Maze
Alien
695.0
954.8
905.2
782.6
Amidar
109.7
90.2
109.7
169.9
BankHeist
246.1
304,9
151.8
86.2
MsPacman
1281.4
1459.7
1421.6
1234.
WizardOfWor
1033.1
985
1262.0
1244.4
Mean Improvement
1.00
1.11
1.05
1.04
Median Improvement
1.00
1.14
1.11
1.13
Overall
Mean Improvement
1.00
1.23
0.74
0.72
Median Improvement
1.00
1.25
0.81
0.71
H XTRA ABLATIONS FOR TASKS WITH DIVERSE GAME MECHANICS
We perform the same set of ablations for tasks that share diverse game mechanics as we do for tasks that share similar game mechanies for XTRA. Results are shown in Table 6. Details of each ablation can be found in Section 4.
Table XTRA ablation tasks that share diverse game mechanics. Results EfficientZero, Random, and Human are adopted from EfficientZero (Ye et al., 2021). All other results are based on the average of runs.
Ablations (XTRA)
Game
W.b. pretraining
(Ours)
weights
Assault
1263.1 1
1742.2
1716.11
1183.58
1605.07
222.4
742.0
BattleZone
13871.2
14631.3
12918.8
8718.8
10087.5
2360.0
37187.5
Hero
9315.9
10631.8
8275.3
8672.9
7755.4
1027.0
30826.4
Krull
5663.3
7735.8
5910.7
6767.3
7104.7
1598.0
2665.5
Seaquest
1100.2
749.5
811.4
540.6
493.1
68.4
42054.7
Normed Mean
1.29
1.87
1.50
1.43
0.00
1.00
Normed Median
0.33
0.35
0.30
0.26
0.23
0.00
1.00
17

---

Published as a conference paper at ICLR 2023
Table XTRA ablation (number tasks pretraining & cross-task finetuning) for tasks that share diverse game mechanics. Results for EfficientZero are adopted from EfficientZero (Ye et al., 2021). All other results are based on the average of 5 runs.
XTRA
Ablations (XTRA)
EfficientZero
Game
8 Games
4 Games
2 Games
0 Games
0 Games
Assault
1742.2
1676.7
1463.8
1255.9
1263.1
BattleZone
14631.3
9581.3
9550.0
10125.0
13871.2
Hero
10631.8
9654.9
8506.5
6815.1
9315.9
Krull
7735.8
7375.6
7348.9
5590.6
5663.3
Seaquest
749.5
656.4
627.5
770.8
1100.2
Normed Mean
1.87
1.74
1.65
1,23
1.29
Normed Median
0.35
0.29
0.25
0.22
0.33
I EFFECTS NUMBER OF TASKS PRETRAINING AND CROSS-TASKS IN FINETUNING
We perform an additional ablation for tasks that share diverse game mechanics whether changing the number of tasks during pretraining would help hurt the performance in later cross-task finetuning. By reducing the number of tasks in pretraining. the model is exposed to (1) less diverse game mechanics and (2) less offline training data in pretraining, and fewer cross-tasks in finetuning. In this ablation, we gradually reduce the number of pretrained tasks from (Carnival, Centipede, Phoenix, Pooyan, Riverraid, VideoPinball, WizardOfWor, and YarsRevenge), 4 (Phoenix, WizardOfWor, VidcoPinball, YarsRevenge), and to (Phoenix, VideoPinball). XTRA is reduced to EfficientZero-L when the number of pretrained tasks and cross-tasks during finetuning is set O. We find that increasing the number of tasks during pretraining (and later cross-task finetuning) mostly consistently improves XTRA performance.
ARCHITECTURAL DETAILS
We adopt the architecture of EfficientZero (Ye et al., 2021). For EfficientZero-L and XTRA, increase the number of residual blocks from (default) to 4 (ours) in the representation network, which we find to improve pretraining slightly for XTRA. However, we find that our baseline EfficientZero (without pretraining) performs significantly worse with larger representation network. Therefore, we use the default EfficientZero as the main point of comparison throughout this work and only include Efficient-L. for completeness.
The architecture the representation networks as follows:
convolution with stride and 32 output planes, output resolution 48x48. (BN + ReLU)
residual block with 32 planes.
residual downsample block with stride and 64 output planes, output resolution 24x24.
residual block with planes.
Average pooling with stride 2, output resolution 12x12. (BN ReLU)
residual block with 64 planes.
Average pooling with stride 2, output resolution 6x6. (BN ReLU)
residual block with 64 planes.
where the kernel size is 3 x 3 for all operations.
The architecture of the dynamics networks is as follows:
Concatenate the input states and input actions into 65 planes.
convolution with stride and 64 output planes. (BN)
A residual link: add up the output and the input states. (ReLU)
residual block with 64 planes.
The architecture the reward prediction network is as follows:
llconvolution and 16 output planes. (BN + ReLU)
Flatten.
18

---

Published as a conference paper at ICLR 2023
LSTM with 512 hidden size. (BN ReLU)
fully connected layers and 32 output dimensions. (BN ReLU)
fully connected layers and 601 output dimensions.
The architecture of the value and policy prediction networks is as follows:
residual block with 64 planes.
convollutio and 16 output planes. (BN + ReLU)
Flatten.
fully connected layers and 32 output dimensions. (BN ReLU)
fully connected layers and output dimensions.
where D = 601 in the value prediction network and D = |A| in the policy prediction network.
K HYPER-PARAMETERS
We adopt our hyper-parameters from EfficientZero (Ye et al., 2021) with minimal modification. Because XTRA uses data from offline tasks to perform cross-task transfer during online finetuning. we have an additional syper-parameter for mini-bated size for offline tasks, which is set to 256 (default). We list all hyper-parameters Table 8 for completeness. Lastly, we note that EfficientZero performs an additional 20k gradient steps at 100k environment steps, with a 10x smaller learning rate. We follow this procedure when comparing to previous state-of-the-art methods (Table 2), but for simplicity we omit these additional gradient steps in the remainder of our experiments for both XTRA and baselines.
Table 8. (2021) with minimal modification but included here for completeness.
We list all relevant paramm-rrsettrsss below. Values are adopted from Ye et al.
Parameter
Setting
Observation down-sampling
96 96
Frames stacked
Frames skip
Reward clipping
True
Terminal on loss of life
True
Max frames per episode
108K
Discount factor
0.997
Minibatch size (offline tasks)
256
Minibatch size (larget task)
256
Optimizer
SGD
Optimizer: learning rate
0.2
Optimizer: momentum
0,9
Optimizer: weight decay (c)
0.0001
Learning rate schedule
0.02
Max gradient norm Priority exponent (a)
0.6
Priority correction (3))
0.4 +
Training steps
100K/120K
Evaluation episodes
Min replay size for sampling
2000
Self-play network updating inerval
100
Target network updating interval
200
Unroll steps (lunrot)
steps (k)
Policy loss coefficient (A1)
Value loss coefficient (Qz)
0.25
Self Self-supervised consistency loss coefficient (Ag)
LSTM horizontal length (()
Dirichlet noise ratio (E)
0.3
Number of simulations in MCTS (Nam)
50
Reanalyzed policy ratio
1.0
19

---

Published as a conference paper at ICLR 2023
L EFFECT OF MINI-BATCH SIZE
1.4
Batch Size
1.2
During online finetuning (stage 2), we finetune both with data
1.0
from the target task, and data from the pretraining tasks. We maintain the same batch size (256) for the target task data as the non-preet pontrrrrain baselines, but add additional data from the pretraining tasks with a 1:1 ratio. Thus, our effective batch
Normalize Mimalized
0.8
0.6
0.4
0.2
size that of the baselines. To verify that performance improvements stem from our pretraining (stage I) and not the larger batch size, we compare our method to a variant of our EfficientZero-L baseline that uses 2x larger batch size (512). Results are shown in Figure 13. We do not observe any significant change in performance by doubling the batch size for the baseline. Thus, we conclude that a larger (effective) batch size not the source of our performance gains, but rather our pretraining and inclusion of pretraining tasks during the online finetuning.
20
40
60
80
100
Environment steps (x10²)
EfficientZero (BZ=256)
Efficienterrl L (8Z-256)
EfficientZero- (8Z-512)
(Ours. 62'=256, 82"
Figure 13. Effeet of mini-batch size (BZ). For XTRA, we denote the batch size of the Target task and Offline tasks BZ" and BZ respectively.
M PER-GAME IMPROVEMENT OVER EFFICIENTZERO
We visualize XTRA's mean improvement over EfficientZero and Efficien(Zero-L under similar and diverse task settings on per-game basis. The improvement is calculated based on the formula: (scoreXtRA/ score 1. The Baseline can be either EfficientZero or EfficientZero-L, depending on the visualization setting.
(a) Mean Improvement
(b) Mean Improvement
Improvement over Effi efficientZero
0.4 10
(c)
Mcan Improvement
0.6
0.3
0.2-
0.4
0.14
2.0
2015
0.2
1.5
Improvement
-0.1
-0.2
-0.3-
Conppede
Allen
Phoenta
Corniaa
Aussutt
Bankzoor
CermorAtacc
MyPacmann
Wiz Wicasroorooo
Amidar
Photeit
Crenorstecck
Contiedde
Allens
Camvvai
Baokitese
Melacmar
Assaut
Amnder
WirddWWer
Assaul
Krull
Hero
BattiZZan
Saldwest
Mr
Game
Game
Game
Figure 14. (a) Mean improvement over EfficientZeeoo from Table (similar games), (b) Mean improvement
over EfficientZero-L. from Table 1 (similar games), and (c) Mean improvement over EfficientZero from Table 2 (diverse games). Each result is aggregated across 5 seeds.
N GAME INFORMATION
In this section, we aim to provide additional context about the games that we consider during both pretraining and finetuning. Table lists core properties for each game. The Similar Task column marks all games used in Table (similar games), and the two Diverse Task columns mark all games used in Table (diverse games) for pretraining and finetuning, respectively. We further categorize games into five categories based on game mechanics: Maze, Shooter, Tank, Adventure, and Ball Tracking, and also report whether the scene is static or dynamic, as well the (valid) action space for each task. The maximum dimensionality of the action space is 18 for Atari games.
O BEHAVIORAL CLONING BASELINE
We the representation prediction network in XTRA for the behavioral cloning (BC) study. The BC (finetune) from Table follows an offline pretraining + offline finetuning paradigm. The model is finetuned on offline data for the target task (also generated by the EfficientZero baseline). We find
20

---

Published as a conference paper at ICLR 2023
Table Game information. We consider a variety of games our experiments. Here, we provide more context to our selection of games. In our similar experiments (Table 1), we finetune model to each game after pretraining it on the other games within the same category. In diverse experiments (Table 2), we finetune single model pretrained all eight games to each of the target games.
Games
Similar Task
Diverse Task
Diverse Task
Category
Scene Continuity
Action Space
Fine
Alien
Maze
18
Amidar
Maze
10
Assault
Shooter
7 18
Bank Hcist
Maze
Carnival
Shooter
6 18
Centipede
Shooter
DemonAttack
Shooter
6 9
MsPacman
Maze
Phoenix
Shooter
8
WizardOFWor
Maze
10
BattleZone
Tank
18
Hero
Adventure
18
Krull
Adventure
18
Seaquest
Shooter
18
Pooyan
Shooter
6
Riverruid
Shooter
18
VidcoPinball
Ball Tracking
9
YarsRevenge
Shooter
BC (finetune) underperforms the EfficientZero baseline. We also report zero-shot performance the pretrained BC on their designated target tasks directly.
21

---

Published as a conference paper at ICLR 2023
P GAME VISUALIZATIONS
We here provide sample trajectories for each game. These visualizations aim to highlight that (1) reasonable to expect some degree of cross-task transfer in similar game transfer (maze maze, shooter + shooter) due to similarity in visuals and game mechanics, and (2) that games from our diverse category indeed are diverse, both in terms of visuals and mechanies. Figure 15 shows trajectories for five Maze games, Figure 16 shows trajectories for five Shooter games, and Figure 17 shows trajectories for eight games from our diverse game category.
time
Alien
Amidar
BankHeist
MsPacman
WizardOIWon
Figure 15. Visualization of trajectories from Maze game category. We visualize key frames in sample trajectories for five tasks. Actual trajectory lengths vary greatly between games.
22

---

Published as a conference paper at ICLR 2023
time
Assault
Carnival
Centipede
110
188
DemonAAtack
SOBR
6000
6050
Phoenix
Figure 16. Visualization of trajectories from the Shooter game category. We visualize key frames sample trajectorics for five tasks. Actual trajectory lengths vary greatly between games.
23

---

Published as a conference paper at ICLR 2023
time
BattleZone
Hero
Krull
120
118
125
Pooyan
Riverraid
Seaquest
VideoPinbar
1
M
1
YarsRevenge
Figure 17. Visualization of trajectories from the diverse game category. We visualize key frames in sample trajectorics for eight tasks. Actual trajectory lengths vary greatly between games.
24


Reviews:
[{"review_id": "wetr3DPEun", "reviewer": "Reviewer_k2wV", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper studies the problem of cross-task transferability in model-based reinforcement learning by introducing a framework with offline pretraining of world models and online finetuning (XTRA).\nXTRA builds on top of EfficientZero and introduces a novel weighted loss formulation to tackle two common problems: (1) catastrophic forgetting and (2) non-aligned downstream task with the training tasks. The experimental section compares the resulting approach with the original EfficientZero in several ablation studies, which highlight under which conditions the proposed solution shows positive results. \nThe authors conduct experiments on 14 visual tasks from the ALE benchmark and explicitly consider a reduced-budget setting.", "summary_of_the_review": "Considering the lack of empirical support for the gradient-weighting mechanism, and the partial novelty in the overall framework, I vote for a weak rejection. By addressing these criticisms, the paper would be solid and mature for publication. I would recommend the authors to advance possible theses of the observed performance and validate them with additional experiments. For example, I would like to understand (1) to what extent the gradient weighting captures the task similarities and if the problem should be addressed by other mechanisms; and (2) if there is any model bias dependent on visual features that is playing a role in this performance gap.", "strength_and_weaknesses": "The paper is well-written, and the problem under study is clearly defined. The proposed contribution in the loss formulation looks technically sound and the authors fully convey its rationale. Moreover, the experimental section shows an exemplary structure with several ablations that analyse different aspects of the proposed solution. \n\nThe main criticism is on the empirical evidence of the effect of gradient weighting. In the first experiment, the authors evaluate XTRA in tasks with similar game mechanics. While the performances are generally promising, one of the baselines uses constant weights and shows much higher performance in two games. While I acknowledge we cannot expect to have a solution for everything, I think the large performance gap disconfirms one of the main contributions of this work without properly commenting on it. \n\nMoreover, even if focusing on model-based reinforcement learning, I would like to see the performance compared with model-free approaches. Even if the training budget is not comparable, reporting their asymptotic performance would give more context to the readers.\n\nMinor errors:\n- Missing reference to the original paper on World Models by Ha et al. (NeurIPS 2018).\n- In Section 3.1, it is unclear what 'u' refers to when saying about \"task specific quantities (pi, u, z)\". From the appendix, 'u' is used for immediate reward, but I could not find its definition in the main text.\n- Typo in 3.1: \"Learning a single RL agent for a diverse set of tasks is however a difficult in practice\" instead of \"Learning a single RL agent for a diverse set of tasks is however difficult in practice\".\n- End of paragraph \"Scaling model size\": the sentence \"XTRA is robust to model size and can be further benefit from scaling up the model size in the future\" sounds a bit overselling and, in my opinion, it is not supported by sufficient experiments to claim something like that.\n- Typo in Related work: \"or example, He et al (2020)\" instead of \"For example, He et al (2020)\"", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "eMAMeQy0Od", "reviewer": "Reviewer_8xCu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper studies the question of whether pre-trained model-based RL can be effectively fine-tuned to solve downstream tasks. To answer this question, during pre-training, this paper proposed to deploy a teacher-student framework: learn several sing-task teacher models via offline RL and distill these single-task policies into a multi-task student policy. Then, the multi-task student model is fine-tuned with the online target task. Offline data and gradient re-weighting are used to prevent catastrophic forgetting and performance degradation during fine-tuning.", "summary_of_the_review": "This paper studies the question of whether pre-trained model-based RL can be effectively fine-tuned to solve downstream tasks.  However, the experiments lack important analysis of the effect of the offline data set.", "strength_and_weaknesses": "### Strength\n- The proposed question is quite interesting.\n- The method is easy to follow.\n- EfficientZero is a reasonable choice to study this question with offline pre-training.\n\n### Weakness\n- Major\n    - Could the authors please provide more details about the offline data set used during the pre-training stage? How do you obtain the offline data set? What is the size of this offline data set? How are the quality and size affect the downstream performance?\n    - Since the proposed model has an offline pre-training stage, the improvement in sample efficiency seems not significant compared with EfficientZero, as shown in Figure 6. \n    - The BC baseline is too weak, there is no policy improvement stage such as the fine-tuning used for the proposed model.\n- Minor\n    - Figure 1 and figure 2 seemed not being referred anywhere in the paper.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "bQ6_z4eP43X", "reviewer": "Reviewer_MLsj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes XTRA, a framework for finetuning a model-based RL agent that is pretrained on offline datasets to a target task through online interactions. Unlike previous work that uses the same task for offline pretraining and online finetuning, this paper tackles the setting where there are multiple pretraining tasks that are different from the target task. For multi-task offline pretraining, the paper shows that simply learning a single model for all tasks using RL cannot work. Hence, it proposes to first learn a teacher model for each task, and then distill all teacher models into a single student model. For online finetuning, to overcome catastrophic forgetting, the paper proposes to use offline data together with online data. To minimize interference from irrelevant tasks in the offline data, the paper further proposes a re-weighting technique based on task similarity. The experiments focus on the Atari 100k benchmark, with EfficientZero being the backbone and main baseline. There are two settings: (1) pretraining on games that are similar to the target game, and (2) pretraining on diverse games. XTRA outperforms EfficientZero in both settings. Further analysis shows that it is crucial to finetune both the pretrained representation and the dynamics function, and that XTRA has a large advantage over EfficientZero in the early stage of training.", "summary_of_the_review": "I am leaning toward acceptance. The paper makes a first attempt at multi-task offline pretraining and online finetuning using a world model. While additional comparisons to single-task / model-free settings are nice to have, the paper already provides sufficient new knowledge to the community.", "strength_and_weaknesses": "- Strengths\n    - The paper tackles an important new problem that is of interest to the community. Pretraining on multiple tasks and then finetuning on a different target task is more realistic than pretraining and finetuning on the same task as done in previous work, and has the potential for better data efficiency.\n    - The paper is well written and easy to follow. The method is well motivated.\n    - The ablation and analysis are quite informative, answering interesting questions such as which components are beneficial to finetune and which stage of training does the proposed model bring the most improvement. \n- Main Weaknesses\n    - The baselines are mainly online RL methods which do not have access to offline data. The paper can be strengthened by comparing to methods that also do offline pretraining and online finetuning. For example, it is mentioned in related work that Online Decision Transformer (ODT) pretrains and finetunes on the same task. The paper can be much stronger if it shows that when pretrained on the target task, XTRA acheives similar performance as ODT, and when including other pretraining tasks, it outperforms ODT.\n    - Although model-based RL enjoys high sample efficiency in online single-task setting, it seems unclear whether it makes sense or is beneficial to use model-based RL in this paper's setting (as opposed to model-free RL). I would expect different tasks to have very different dynamics, so that it is not only difficult to learn a single model for all pretraining tasks, but it is also difficult to transfer the learned model to the target task. A comparison to model-free online finetuning would help clarify.\n- Questions and Minor Issues\n    - Can you elaborate on what is catastrophic forgetting in online finetuning? This can make the paper more self-contained and help the reader better understand why you want to retain the offline data.\n    - There is not much ablation in the diverse task setting. It would be interesting to see if pretraining and/or retaining offline data helps, and if the task weights are close to zero.\n    - I am curious why you use distillation during pretraining but switch back to the EfficientZero loss during finetuning (Eq 1). Is it possible to still use the distillation loss in finetuing?\n    - It seems the improvement over EfficientZero is even larger in the diverse task setting than in the similar task setting. Can you comment on that? Also, why is the maze-to-shooter result much worse than the diverse task setting?\n    - Figure 5 seems not mentioned in the main text.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "TSsOswnwG1", "reviewer": "Reviewer_bd3b", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper introduces XTRA, a model-based algorithm for reinforcement learning transfer. XTRA extends EfficientZero (SOTA model-based RL on Atari games) to the RL transfer setting in 2 steps:\n\n1. Offline multi-task pretraining: distilling N EfficientZero teachers (pretrained on N tasks via offline RL) into a single EfficientZero student.\n2. Online finetuning: where the student keeps optimizing for the mixture of pretraining tasks while also learning the downstream task via online RL.\n\nBesides this pretraining-finetuning setup for model-based RL, the authors propose a scheme to reweight pretraining tasks according to their affinity with the downstream task — this aims to mimic the effects of gradient-surgery (Yu et al., 2020).\n\nThey evaluate XTRA on the Atari 100k benchmark (learning with only 100k frames), both when transferring between tasks that are similar (e.g., pretraining and downstream are instances of shooter tasks) and when they are dissimilar (i.e., not from any particular game family).", "summary_of_the_review": "- Strengths:\n    - The authors tackle a hard RL transfer problem, and show promising results.\n    - Some ablation studies and analyses — I wish there were more.\n- Weaknesses:\n    - The reported results are not 100% compelling. I’ve asked the authors to clarify.\n    - The paper could include better baselines, some of which need to be discussed as prior work (Actor-Mimic / Policy Distillation).", "strength_and_weaknesses": "- Strengths:\n    - The authors tackle a notoriously hard problem: transfer across different tasks in RL, from images. They also chose one of the most competitive benchmarks to evaluate their method, namely Atari 100k, and include the most competitive baselines. In that sense, the results are impressive and XTRA is indeed successful at leveraging pretraining tasks to more efficiently solve the downstream one (my concerns on this point below).\n    - To their credit, the authors also include some ablations studies and analyses (see 2nd part of 4.1) where they discuss task relevance, model size, and, more interestingly, which component of the model-based architecture is more transferrable and when. (Transferring the dynamics yields in the largest boost in the early finetuning phases.) I only wished those questions were explored in more depth, to help guide the community understand the mechanics of transfer in model-based RL on Atari.\n- Weaknesses:\n    - Since most of the experiments focus on benchmarking, the merits of the paper rest on XTRA’s performance. While it shows promising results, I also wonder if the results are as positive as described in the paper. First, how much more data than EfficientZero does XTRA require when taking pretraining into account (I couldn’t find those numbers in the main text)? Second, the confidence intervals are quite large when transferring to similar tasks (see Fig. 6: Carnival, DemonAttack, Alien, Amidar, Bankheist, MsPacman, Wizard of Wor) — are the reported results significant? It’s also surprising that XTRA does better relative to EfficientZero when downstream tasks are dissimilar from pretraining (see Table 2, confidence intervals) — could the authors comment on this point, and describe how the confidence intervals are computed in this setting?\n        \n        Depending on the answer to those questions, I wonder if this paper isn’t reporting negative results for cross-task RL transfer.\n        \n    - A more minor point: I wish the paper included better baselines — not just ablations of the proposed method. For example, how does XTRA perform against on-policy alternatives such as Actor-Mimic or Policy Distillation (e.g., on top of DrQ-v2)? In a similar vein, the authors claim their task-weighting scheme yields the same benefits as gradient surgery (see 3.2) but when is that empirically demonstrated?\n    - One question re “Most improvements happen within the first  stages of training.”: Do the authors have any insight on why that’s the case? From Fig. 8 it seems that the dynamics are the reason for this boost (but we’re missing “h+f” on this figure). Is it because the dynamics already have a good initialization (i.e., don’t move too much during finetuning) or is it because the initialization lets them quickly move through parameter space (i.e., large movement at first, but not much after)?\n    - Details:\n        - What is the definition (formula) of $s^i$? I’m assuming it’s Eq. 2 but please clarify.\n        - 3.1, p. 3: “tasks is however a difficult” → remove a.\n        - 3.1, p. 4: you could cite Parisi et al.’s “Actor-Mimic”, and Rusu et al.’s “Policy Distillation”.\n        - 3.1, p. 4: $(\\hat{\\pi}, \\hat{u}, \\hat{z})$ should be $(\\hat{\\pi}, \\hat{v}, \\hat{z})$?\n        - 3.2, eq. 1: The sum should be $\\sum_i^{m}$, right?\n        - 4.1, p. 7: “our proposed frames” → “our proposed framework”?\n        - Table 2, p. 7: \\cite{} → \\citep{} for Ye et al., 2021.\n        - 4.1, p. 8: “This results indicates that (i) selecting […] and (2) in the […]” → enumeration consistency.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster