O N THE F EASIBILITY OF C ROSS -T ASK T RANSFER WITH M ODEL -B ASED R EINFORCEMENT L EARNING
A BSTRACT
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 CrossTask
Tra nsfer ( 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 on the Atari100k benchmark 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 .
I NTRODUCTION
[Figure]
shaded area Reinforcement Learning (RL) has achieved great feats across a 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-efficiency and require millions (or even billions) of environment interactions to solve tasks – especially 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 et al., 2018).
In an effort to reliably benchmark and improve the sample-efficiency of image-based RL across a variety of problems, the Arcade
Learning Environment (ALE; (Bellemare et al., 2013)) has become a long-standing challenge for RL.
This task suite has given rise to numerous successful and increasingly sampleefficient algorithms (Mnih et al., 2013; Badia et al., 2020;
Kaiser et 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 a model of the environment.
Most recently, EfficientZero Ye et al.
(2021) – a model-based RL algorithm – has demonstrated impressive sample-efficiency, surpassing human-level performance with as little as 2 hours of realtime game play in select Atari 2600 games from the ALE.
This achievement is 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, [Figure] [Figure] 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 tabula rasa (from scratch) for each new task.
On the contrary, humans rely heavily on prior knowledge and visual cues when learning new skills.
For example, a study found that human players easily pick up on visual cues about game mechanics and objectives when exposed to a video game for the very first time, 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 as computer vision and natural language processing, large-scale unsupervised/self-supervised/supervised pretraining on large-scale datasets (Devlin et al., 2019;
Brown et al., 2020; Li et al., 2022; Radford et al., 2021; Chowdhery et al., 2022) has emerged as a powerful framework for solving numerous downstream tasks with few samples (Alayrac et al., 2022).
This pretraining paradigm has recently been extended to visuo-motor control in various forms, e.g. , by leveraging frozen (no finetuning) pretrained representations (Xiao et al., 2022; Parisi et al., 2022) or by finetuning in 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 a 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 an online RL setting, and across markedly distinct tasks.
Specifically, we seek to 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 to 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;
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 Tra nsfer ( XTRA ), a framework for sample-efficient online RL with scalable pretraining and finetuning of learned world models using extra, auxiliary data from other tasks.
Concretely, our framework consists of two stages: (i) offline multi-task pretraining of a world model on an offline dataset from m diverse tasks, a (ii) finetuning stage where the world model is jointly finetuned on a target task in addition to the 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, we adaptively re-weight gradient contributions in an unsupervised manner based on similarity to the target task.
We evaluate our method and a set of strong baselines extensively across 14 Atari 2600 games from the
Atari100k benchmark (Kaiser et al., 2020) that require algorithms to be extremely sample-efficient.
We find that XTRA improves sample-efficiency substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25% , respectively.
B ACKGROUND
Problem setting.
We model image-based agent-environment interaction as an episodic Partially
Observable Markov Decision Process (POMDP; Kaelbling et al.
(1998)) defined by the tuple M = ⟨O , A , P , ρ, r, γ ⟩ , where O is the observation space (pixels), A is the action space, P : O×A
(cid:55)→ O is a transition function, ρ is the initial state distribution, r : O × A
(cid:55)→ R is a scalar reward function, and γ ∈ [0 , 1) is a discount factor.
As is standard practice in ALE (Bellemare et al., 2013), we convert M to a fully observable Markov Decision Process (MDP; Bellman (1957)) by approximating state s t ∈ S at time t as a stack of frames s t .
= { o t , o t − 1 , o t − 2 , . . .
} where o ∈ O (Mnih et al., 2013), and redefine
P , ρ, r to be functions of s .
Our goal is then to find a (neural) policy π θ (
a | s ) parameterized by θ that maximizes discounted return E π θ [ (cid:80) tt
=1 γ t r ( s t , a t )] where a t ∼ π θ ( a | s ) , s t ∼ P ( s t , a t ) , s 0 ∼ ρ , and T 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 a "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, we seek to first obtain a 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-efficiency of online RL.
We do so by first pretraining the model on an offline (fixed) dataset that consists of transitions ( s , a , r, s ′ ) collected by unknown behavior policies in m environments { ˆ M i
| ˆ M i ̸
= M , 1 ≤ i ≤ m } , and then finetune the model by online interaction on the target task. Search with Learned Model
Observation p (policy) v (value) p (policy) v (value) r (reward) Zh f g f
Z EfficientZero (Ye et al., 2021) is a model-based RL algorithm based on MuZero (Schrittwieser et al., 2020) that learns a discrete-action latent dynamics model from environment 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: a representation (encoder)
h θ , a dynamics (transition) function g θ , and a prediction head f θ .
Given an observed state s t ,
EfficientZero projects the state to a latent representation z t
= h θ ( s t ) , and predicts future latent states z t +1 and instantaneous rewards r t using an action-conditional latent dynamics function z t +1 , r t = g θ ( z t , a t ) .
For each latent state, a prediction network f θ estimates a probability distribution ˆ p over (valid) actions a ∈ A , as well as the expected state value ˆ v of the given state, i.e. , ˆ v t , ˆ p t = f θ ( z t ) .
Intuitively, h θ and g θ allow EfficientZero to search for actions entirely in its latent space before executing actions in the real environment, and f θ predicts quantities that help guide the search towards high-return action sequences.
Concretely, ˆ v provides a return estimate for nodes at the lookahead horizon (as opposed to truncating the cumulative sum of expected rewards) and ˆ p provides an action distribution prior that helps guide the search.
EfficientZero improves the sample-efficiency 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 to most model-based algorithms, including those for continuous action spaces (Hafner et al., 2019a; Hansen et al., 2022).
M ODEL -B ASED C ROSS -T ASK T RANSFER
We propose Model-Based Cross -Task Tra nsfer ( XTRA ), 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 a target task collected by online interaction.
In the following, we introduce each of the two stages – pretraining and finetuning – in detail. [Figure]
3. O FFLINE M ULTI -T ASK P RETRAINING
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, however, that the goal of pretraining is not to obtain a truly generalist agent, but rather to learn a good initialization for finetuning to unseen tasks.
Learning a single RL agent for a diverse set of tasks is however a difficult in practice, which is only exacerbated by extrapolation errors due to the offline RL setting (Kumar et al., 2020).
To address the challenge of multi-task learning, we propose to pretrain the model following a student-teacher training setup, where teacher models are trained separately by offline RL for each task, and then distilled into a single multi-task model using a novel instantiation of the MuZero Reanalyze (Schrittwieser et al., 2021) algorithm.
For each pretraining task we assume access to a fixed dataset { ˆ D i | 1 ≤ i ≤ m } that consists of trajectories from an unknown (and potentially sub-optimal) behavior policy.
Importantly, we do not make any assumptions about the quality or the source of trajectories in the dataset, i.e. , we 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 a single-task (offline) RL setting, resulting in m teacher models { ˆ π iψ | 1 ≤ i ≤ m } .
After training, each teacher model ˆ π iψ has learned to produce task-specific quantities ( ˆ π , ˆ u , ˆ z ) for a given game ˆ M i .
Next, we learn a multi-task student model ˆ π θ by distilling the task-specific teachers into a single model.
Specifically, we optimize the student policy by sampling data uniformly from all pretraining tasks, and generate value/policy targets using the respective teacher models rather than bootstrapping from student predictions as 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 a 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 a multi-task model on all m pretraining tasks simultaneously using
RL is found to not scale beyond a couple of tasks in practice, as we will demonstrate our experiments in Appendix C. After distilling teacher models into the multi-task student model, we now have a single set of pretrained parameters that can be used for finetuning to a variety of tasks via online interaction, which we introduce in the following section.
3. O NLINE F INETUNING ON A T ARGET T ASK
In this stage, we iteratively interact with a target task (environment) to collect interaction data, and finetune the pretrained model on data from the target task.
However, we empirically observe that directly finetuning the pretrained model often leads to catastrophical forgetting, and consequently poor performance on the target task.
To overcome this challenge, 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 during 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]
[Figure] At each training step t , we jointly optimize the target online task M and m offline (auxiliary) tasks { ˆ M i
| ˆ M i ̸
= M , 1 ≤ i ≤ m } that were used during the offline multi-task pretraining stage.
Our online finetuning objective is defined as:
[Equation]
(1) where L ez is the ordinary (single-task) EfficientZero objective (see Appendix A), and η i are dynamically (and independently) updated task weights for each of the m pretraining tasks.
The target task loss term maintains a constant task weight of 1 .
In order to dynamically re-weight task weights η i throughout the training process, we break down the total number of environment steps ( i.e. , 100k in our experiments) into even T -step cycles (intervals).
Within each cycle, we spend first N -steps to compute an updated η i corresponding to each offline task ˆ M i .
The new η i 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 i 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 ˆ G in from L ez n ( ˆ M i ) and G n from L ez n ( M ) given by
[Equation]
Within the N -step update, we maintain a task-specific counter s i and the new task weights η i can be reset by η i
= s i N at the beginning of each every T -cycle following the procedure stated in Appendix
B. Concretely, Sim( ˆ M i , M ) measures the angle between two task gradients ˆ G in and G n .
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 η .
While re-weighting task weights at every gradient update would result in the least amount of conflicting gradients, it is 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 a negligible computational cost in contrast to the compute-intensive gradient modification method proposed in (Yu et al., 2020).
E XPERIMENTS
We evaluate our method and baselines on 14 games (tasks) from the limited-interaction Atari100k benchmark (Kaiser et al., 2020) where only 100k environment steps are permitted.
We provide an im-
[Table] plementation of our method at https://github.com/mlpc-ucsd/XTRA . Through experimental evaluation, we seek to answer the following three questions:
• How does our proposed framework compare to alternative pretraining and online RL approaches with limited online interaction from the target task?
• 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 on 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 EfficientZero , and (ii) a larger variant with 4 times more parameters in the representation network (denoted EfficientZero-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 as Diverse ).
We measure performance on individual Atari games by absolute scores, and also provide aggregate results as measured by mean and median scores across games, normalized by either human performance or EfficientZero performance at 100k environment steps.
All of our results are averaged across 5 random seeds to ensure reliability.
See Appendix E for more details.
Baselines.
We compare our method against 7 prior methods for online RL that represent the stateof-the-art on the Atari100k benchmark (including EfficientZero), as well as a multi-task behavior cloning policy trained on the full pretraining dataset, and a set of ablations that include EfficientZero with several different model sizes and pretraining/finetuning schemes.
The former baselines serve to position our results with respect to the state-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. R ESULTS & D ISCUSSION
We introduce our results as they relate to 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?
[Figure] [Figure] [Table]
Tasks that share similar game mechanics.
We first investigate the feasibility of finetuning models that are pretrained on games with similar mechanics.
We select 5 shooter games and 5 maze games for this experiment.
Results for our method, baselines, and a set of ablations on the Atari100k benchmark are shown in Table 1.
For completeness, we also provide learning curves in Figure 6.
We find that pretraining improves sample-efficiency 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 a multi-game behavior 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 cross-task learning ( i.e., leverages offline 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 1 for all task loss terms during finetuning.
We find that XTRA achieves extremely high performance on 2 games (DemonAttack and Phoenix) without dynamic task weights, improving over EfficientZero by as much as 343 % on DemonAttack.
However, its median performance is overall low compared to 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 5 tasks from Atari100k.
Results are shown in Table 2.
We find that XTRA advances the state-of-the-art in a majority of tasks on the Atari100k benchmark, and achieve a mean human-normalized score of 187 % vs. 129 % for the previous SOTA, EfficientZero.
This suggests that, while task similarity may play a role in the success of XTRA, the algorithmic advances of our proposed frames are a bigger factor in the strong empirical performance.
2. How do the individual components of our framework influence its success?
A 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 importance of task relevance.
We compare the online finetuning performance of XTRA in two [Figure] different settings: (1) pretraining on shooter games and finetuning and to
Aggregate results across all tasks are shown in Figure 7 (a).
Unsurprisingly, we observe that offline pretraining and concurrently learning from other shooter games significantly benefit 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 (i) selecting pretraining tasks that are relevant to the target task is key to 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. [Figure]
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 f – to the online finetuning stage and simply using a random initialization for the remaining components.
Results are shown in Figure 8.
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 heads has no significant impact on sample-efficiency (but matters for zero-shot performance).
We 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 is 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 are shown in Figure 7 (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 model-free RL (Wang et al., 2022; Ze 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 f are used in MCTS search, increasing the parameter counts for these two networks would increase inference/training time complexity significantly.
However, increasing the size of the representation network h has a relatively small impact on overall inference/training time (see Figure 3 for reference).
We compare 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 EfficientZero -L ); results are shown in Figure 7 (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 EfficientZero.
This indicates that, unlike EfficientZero trained from scratch, XTRA is robust to model size and can potentially benefit from further scaling up the model size (along with more pretraining tasks) in the future.
Relative improvement vs. environment steps.
Finally, we visualize the average improvement over EfficientZero throughout training in Figure 7 (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-task pretraining could benefit many real-world applications of RL, where environment interactions are typically constrained due to physical constraints.
When can we empirically expect finetuning to be successful?
We have demonstrated that model-based RL can benefit from cross-task transfer, both when the tasks are similar to the target task (Table 1), and when tasks are diverse ( i.e. , dissimilar to the target task; Table 2).
While transfer is more successful overall when tasks are similar ( e.g. , shooter games that feature similar game mechanics), our proposed approach to finetuning is found to successfully reduce (potentially harmful) gradient contributions of dissimilar tasks as measured by gradient cosine similarity to the target game. Interestingly, we find that both the representation and dynamics networks benefit from pretraining despite differences in game mechanics (Figure 8), but that leveraging a frozen pretrained representation reduces sample-efficiency compared to finetuning. This finding confirms that (1) pretraining a model-based algorithm requires special considerations vs. pretraining a model-free method, and (2) the observation that a frozen representation generally is inferior to finetuning in model-free RL (Wang et al., 2022; Ze et al., 2022) also applies to model-based RL.
R ELATED W ORK
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 on a 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 on a large, unlabelled dataset learns good features for ImageNet classification, and Brown et al. (2020) show that a generative model trained on large-scale natural language data can be used to 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., 2021b; Xiao et al., 2022; Wang et al., 2022), as well as a lack of large-scale datasets that capture both the visual, temporal, and control-relevant (actions, rewards) properties of RL.
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 a more careful treatment during finetuning.
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 representation learning (Kulkarni et al., 2019;
Yarats et al., 2019; Srinivas et al., 2020; Schwarzer et al., 2021), 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 (Finn & Levine, 2017; Nair et al., 2018;
Hafner et al., 2019b; Kaiser et al., 2020; Schrittwieser et al., 2020; Hafner et al., 2021; Ye et al., 2021;
Hansen et al., 2022; Seo et al., 2022).
We choose to focus our efforts on sample-efficiency 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 a 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 a well-studied technique for adaptation in (predominantly model-free) RL, and has been used to adapt to either changes in visuals or dynamics (Mishra et al., 2017; Yen-Chen et al., 2019; Duan et al., 2016; Julian et al., 2020; Hansen et al., 2021a; Bodnar et al., 2020; Wang et al., 2022; Ze 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 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-task pretraining can improve online RL on an unseen task across multiple axes of variation.
C ONCLUSION
In this paper, we propose Model-Based Cross -Task Tra nsfer ( XTRA ), a framework for sampleefficient online RL with scalable pretraining and finetuning of learned world models using extra, auxiliary data from other tasks.
We find that XTRA improves sample-efficiency 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-based RL will inspire further research in this direction. Acknowledgements. This work is supported by NSF Award IIS-2127544 and gifts from Qualcomm AI.
We thank Weirui Ye and Wenlong Huang for their helpful discussions
A E FFICIENT Z ERO O BJECTIVE
MuZero (Schrittwieser et al., 2020) learns to predict three quantities used in {+MCTS: (1) policy target π obtained from the visit count distribution of the MCTS, (2) instantaneous reward u from environment transitions, and (3) bootstrapped value target z = (cid:80) k − 1 i =
 γ i u i + γ k v t + k .+}
EfficientZero (Ye et al., 2021) extends MuZero by adding a self-supervised consistency loss term, and predicts the cumulative sum of rewards instead of individual single-step rewards.
Concretely, the EfficientZero objective for time step t and with a total of K unroll steps is given by
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where λ 1:3 and c are constant coefficients that balance the five terms.
However, we note that our approach is agnostic to the specific choice of algorithm; we only include this additional discussion for completeness.
We refer the reader to Ye et al.
(2021) for more algorithmic details.
B P ERIODIC T ASK R E -W EIGHTING
Between each N -step update, we maintain a task-specific counter s i for the i th task, and update the counter by ∆ s in at each step n as follows:
[Equation]
if Sim( ˆ M i , M ) > 0 .0 , otherwise s
[Equation] = [Equation] in .
At every N steps
[Figure]
the new task weights η i are updated by η i =
s i N , and used in subsequent finetuning objectives according to Equation 1.
In practice
[Figure]
we start task weight updates at 10 k steps to ensure enough data from the online target task has been collected for a meaningful similarity measure.
All task weights are initialized as 1 for the first 10 k steps.
Figure 5 shows task weights as a function of environment steps for two target games; additional results are provided in Figure
[Figure]
C D ISTILLATION VS . M ULTI -G AME O FFLINE RL
Our method learns a multi-game world model from offline data via distillation of task-specific world models trained with offline RL.
An alternative way to obtain such a multi-game model would be to 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 4 tasks.
[Figure]
D E FFECT OF M INI -B ATCH S IZE
[Figure]
During online finetuning (stage 2), we finetune both with data from the target task
[Figure]
and data from the pretraining tasks.
We maintain the same batch size ( 256 ) for the target task data as the non-pretraining baselines
[Figure]
but add additional data from the pretraining tasks with a 1:1 ratio.
Thus, our effective batch size is 2 × that of the baselines.
To verify that performance improvements stem from our pretraining (stage 1) and not the larger batch size, we compare our method to a variant of our EfficientZero-L baseline that uses a 2 × larger batch size ( 512 ).
Results are shown in Figure 11.
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 is not the source of our performance gains, but rather our pretraining and inclusion of pretraining tasks during the online finetuning.
E S CORES FOR I NDIVIDUAL S EEDS
Our results in Table 1 are aggregated across 5 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: ( score agent − score random ) /
( score human − score random ) as 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]
[Figure]
F A DDITIONAL E VALUATION CURVES OF XTRA ON A TARI 100 K BENCHMARK
For completeness, Figure 12 includes evaluation curves of XTRA on the five diverse games for which we report final performance in Table
[Figure]
G P ER -G AME I MPROVEMENT OVER E FFICIENT Z ERO
We visualize XTRA’s mean improvement over EfficientZero and EfficientZero-L under similar and diverse task settings on a per-game basis.
The improvement is calculated based on the formula: ( score XTRA / score Baseline ) − 1 .
The Baseline can be either EfficientZero or EfficientZero-L
[Figure]
depending on the visualization setting.
[Figure]
d e A li e n
P h
o e n i x C a r n i v a l A ss a u l t
[Figure]
H e i s t
[Figure]
e m o n A tt a c k
M s P a c m a n
W
[Figure]
a r d O f
[Figure]
14% -5% -18% (a) Mean Improvement
P h
o e n i x
D e m o n A tt a c k
C e n t
i p
e
d e A li e n
C a r n i v a l
B a n k H
e
i s t
M s P a c m a n
A ss a u l t
A m i d a r
W i z a r d
O f W o r
Game 0.
I m p r
o v e m e n t o v e r E
ff i
c i e n t Z e r
o − L 375% 201% 107% 49% 39% 25% 24% 24% 7% 6% (b) Mean Improvement
A ss a u l t K r u ll H e r o
B a tt l e Z o n e
S e a q u e s t
Game −0.
−0.
−0.
I m p r
o v e m e n t o v e r E
ff i
c i e n t Z e r o 38% 37% 14% 5% -32%
(c) Mean Improvement
H A RCHITECTURAL D ETAILS
We adopt the architecture of EfficientZero (Ye et al., 2021).
For EfficientZero-L and XTRA, we increase the number of residual blocks from 1 (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 a 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 of the representation network is as follows:
• 1 convolution with stride 2 and 32 output planes, output resolution 48x48.
(BN + ReLU) • 1 residual block with 32 planes.
• 1 residual downsample block with stride 2 and 64 output planes, output resolution 24x24.
• 1 residual block with 64 planes.
• Average pooling with stride 2, output resolution 12x12.
(BN + ReLU) • 1 residual block with 64 planes.
• Average pooling with stride 2, output resolution 6x6.
(BN + ReLU) • 1 residual block with 64 planes.
where the kernel size is 3 × 3 for all operations.
The architecture of the dynamics network is as follows:
• Concatenate the input states and input actions into 65 planes.
• 1 convolution with stride 2 and 64 output planes.
(BN) •
A residual link: add up the output and the input states.
(ReLU) • 1 residual block with 64 planes.
The architecture of the reward prediction network is as follows:
• 1 1x1convolution and 16 output planes.
(BN + ReLU)
• Flatten.
• LSTM with 512 hidden size.
(BN + ReLU) • 1 fully connected layers and 32 output dimensions.
(BN + ReLU) • 1 fully connected layers and 601 output dimensions.
The architecture of the value and policy prediction networks is as follows:
• 1 residual block with 64 planes.
• 1 1x1convolution and 16 output planes.
(BN + ReLU)
• Flatten.
• 1 fully connected layers and 32 output dimensions.
(BN + ReLU) • 1 fully connected layers and D output dimensions.
where D = 601 in the value prediction network and D = |A| in the policy prediction network.
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 hyper-parameter for mini-batch size for offline tasks, which is set to(default). We list all hyper-parameters in Table 4 for completeness.
Lastly, we note that Ye et al.
(2021) performs an additional 20k gradient steps at 100k environment steps, with a 10 × 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.
J G AME I NFORMATION
In this section, we aim to provide additional context about the games that we consider during both pretraining and finetuning.
Table 5 lists core properties for each game.
The Similar Task column marks all games used in Table 1 ( similar games), and the two Diverse Task columns mark all games used in Table 2 ( 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 action space for Atari games is 18.
K G AME V ISUALIZATIONS
We here provide sample trajectories for each game.
These visualizations aim to highlight that (1) it is reasonable to expect some degree of cross-task transfer in similar game transfer (maze →
[Table]
[Table]
[Table]
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 mechanics.
Figure 14 shows trajectories for five Maze games, Figure 15 shows trajectories for five Shooter games, and Figureshows trajectories for eight games from our diverse game category.
[Figure]
[Figure]
A ss a u lt
C a r n i v a l C e n ti p e d e D e m on A tt ac k P ho e n i x time −→
B a ttl e Z on e
H e r o
K r u ll
P ooy a n
R i v e rr a i d
S ea qu e s t
V i
d e o P
i nb a ll
Y a r s R e v e ng e