Original: O N THE F EASIBILITY OF C ROSS -T ASK T RANSFER WITH M ODEL -B ASED R EINFORCEMENT L EARNING
Revised: O N THE F EASIBILITY OF C ROSS -T ASK T RANSFER WITH M ODEL -B ASED R EINFORCEMENT L EARNING

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Reinforcement Learning (RL) algorithms can solve challenging control problems directly from image observations, but they often require millions of environment interactions to do so.
Revised: Reinforcement Learning (RL) algorithms can solve challenging control problems directly from image observations, but they often require millions of environment interactions to do so.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We propose Model-Based CrossTask
Revised: We propose Model-Based Cross -Task

Original: Tra nsfer ( XTRA ), a framework for sample-efficient online RL with scalable pretraining and finetuning of learned world models.
Revised: Tra nsfer ( XTRA ), a framework for sample-efficient online RL with scalable pretraining and finetuning of learned world models.

Original: By proper pretraining andconcurrent cross-task online fine-tuning, 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: [Figure]
Revised: [Figure]

Original: Reinforcement Learning (RL) has achieved great feats across a wide range of areas, most notably game-playing
Revised: ) has achieved great feats across a wide range of areas, most notably game-playing

Original: (Mnih et al., 2013; Silver et al., 2016; Berner et al., 2019; Cobbe et al., 2020).
Revised: (Mnih et al., 2013; Silver et al., 2016; Berner et al., 2019; Cobbe et al., 2020).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: In an effort to reliably benchmark and improve the sample-efficiency of image-based RL across a variety of problems, the Arcade
Revised: In an effort to reliably benchmark and improve the sample-efficiency of image-based RL across a variety of problems, the Arcade

Original: Learning Environment (ALE; (Bellemare et al., 2013)) has become a long-standing challenge for RL.
Revised: Learning Environment (ALE; (Bellemare et al., 2013)) has become a long-standing challenge for RL.

Original: This task suite has given rise to numerous successful and increasingly sampleefficient algorithms (Mnih et al., 2013; Badia et al., 2020;
Revised: This task suite has given rise to numerous successful and increasingly sampleefficient algorithms (Mnih et al., 2013; Badia et al., 2020;

Original: Kaiser et al., 2020; Schrittwieser et al., 2020; Kostrikov et al., 2021; Hafner et al., 2021;
Revised: Kaiser et al., 2020; Schrittwieser et al., 2020; Kostrikov et al., 2021; Hafner et al., 2021;

Original: Ye et al., 2021), notably most of which are model-based, i.e. , they learn a model of the environment (Ha & Schmidhuber, 2018).
Revised: Ye et al., 2021), notably most of which are model-based, i.e. , they learn a model of the environment (Ha & Schmidhuber, 2018).

Original: Most recently, EfficientZero Ye et al.
Revised: Most recently, EfficientZero Ye et al.

Original: (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.
Revised: (2021), a model-based RL algorithm, has demonstrated impressive sample-efficiency, surpassing human-level performance with as little as 2 hours of real-time game play in select Atari 2600 games from the ALE.

Original: 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, thus reducing reliance on real environment interactions for skill acquisition.
Revised: 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, thus reducing reliance on real environment interactions for skill acquisition.

Original: However, current RL algorithms – including EfficientZero – are still predominantly assumed to learn both perception, model, and skillstabula 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
Revised: However, current RL algorithms, including EfficientZero, are still predominantly assumed to learn both perception, model, and skills

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure] 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).
Revised: tabula rasa (from scratch) for each new task. Conversely, humans rely heavily on 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 a new game, and that human performance is severely degraded if such cues are removed or conflict with prior experiences (Dubey et al., 2018).

Original: 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;
Revised: 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;

Original: 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).
Revised: 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).

Original: 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).
Revised: 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).

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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?
Revised: 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?

Original: 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.
Revised: 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.

Original: 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;
Revised: 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;

Original: He et al., 2020;
Revised: He et al., 2020;

Original: Reed et al., 2022; Lee et al., 2022) is found to be unsuccessful or outright harmful in an RL context.
Revised: Reed et al., 2022; Lee et al., 2022) is found to be unsuccessful or outright harmful in an RL context.

Original: 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 (see Figure 2).
Revised: 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 (see Figure 2).

Original: 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 m offline tasks.
Revised: 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 m offline tasks.

Original: By leveraging offline data both in pretraining and finetuning, XTRA overcomes the challenges of catastrophical forgetting.
Revised: By leveraging offline data both in pretraining and finetuning, XTRA overcomes the challenges of catastrophical forgetting.

Original: To prevent harmful interference from certain offline tasks, we adaptively re-weight gradient contributions in unsupervised manner based on similarity to target task.
Revised: To prevent harmful interference from certain offline tasks, we adaptively re-weight gradient contributions in unsupervised manner based on similarity to target task.

Original: We evaluate our method and a set of strong baselines extensively across 14 Atari 2600 games from the
Revised: We evaluate our method and a set of strong baselines extensively across 14 Atari 2600 games from the

Original: Atari100k benchmark (Kaiser et al., 2020) that require algorithms to be extremely sample-efficient.
Revised: Atari100k benchmark (Kaiser et al., 2020) that require algorithms to be extremely sample-efficient.

Original: From Figure 1 and Table 1, We observe  XTRA improves sample-efficiency substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25% , respectively.
Revised: From Figure 1 and Table 1, we observe that XTRA improves sample-efficiency substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25% , respectively.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Problem setting.
Revised: Problem setting.

Original: We model image-based agent-environment interaction as an episodic Partially
Revised: We model image-based agent-environment interaction as an episodic Partially

Original: Observable Markov Decision Process (POMDP; Kaelbling et al.
Revised: Observable Markov Decision Process (POMDP; Kaelbling et al.

Original: (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
Revised: (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

Original: (cid:55)→ R is a scalar reward function, and γ ∈ [0 , 1) is a discount factor.
Revised: (cid:55)→ R is a scalar reward function, and γ ∈ [0 , 1) is a discount factor.

Original: 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 .
Revised: 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 .

Original: = { o t , o t − 1 , o t − 2 , . . .
Revised: = { o t , o t − 1 , o t − 2 , . . .

Original: } where o ∈ O (Mnih et al., 2013), and redefine
Revised: } where o ∈ O (Mnih et al., 2013), and redefine

Original: P , ρ, r to be functions of s .
Revised: P , ρ, r to be functions of s .

Original: Our goal is then to find a (neural) policy π θ (
Revised: Our goal is then to find a (neural) policy π θ (

Original: a | s ) parameterized by θ that maximizes discounted return E π θ [ (cid:80) tt
Revised: a | s ) parameterized by θ that maximizes discounted return E π θ [ (cid:80) tt

Original: =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.
Revised: =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.

Original: For clarity, we denote all parameterization by θ throughout this work.
Revised: For clarity, we denote all parameterization by θ throughout this work.

Original: 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.
Revised: 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.

Original: Define M as the target task that we aim to solve.
Revised: Define M as the target task that we aim to solve.

Original: 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.
Revised: 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.

Original: 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 | ˜
Revised: 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 | ˜

Original: M i ̸
Revised: M i ̸

Original: = M , 1 ≤ i ≤ m } , and then finetune the model by online interaction on the target task. 
Revised: = M , 1 ≤ i ≤ m } , and then finetune the model by online interaction on the target task. [Table] [Figure] [Table]

Original: [Figure]
Revised: [Figure]

Original: network ( h ), latent dynamics function ( g ), and prediction head ( f ).
Revised: [Table] [Figure] [Table] [Figure]

Original:  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;
Revised: (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;

Original: Coulom, 2006; Silver et al., 2016)) in the latent space of the model.
Revised: Coulom, 2006; Silver et al., 2016)) in the latent space of the model.

Original: Figure 3 provides an overview of the three main components of the MuZero algorithm: a representation (encoder)
Revised: Figure 3 provides an overview of the three main components of the MuZero algorithm: a representation (encoder)

Original: h θ , a dynamics (transition) function g θ , and a prediction head f θ .
Revised: h θ , a dynamics (transition) function g θ , and a prediction head f θ .

Original: Given an observed state s t ,
Revised: Given an observed state s t ,

Original: EfficientZero projects the state to a latent representation z t
Revised: EfficientZero projects the state to a latent representation z t

Original: = 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 =
Revised: = 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 =

Original: g θ
Revised: g θ

Original: ( z t , a t ) .
Revised: ( z t , a t ) .

Original: 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 ) .
Revised: 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 ) .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We describe EfficientZero’s learning objective between model prediction ( ˆ p, ˆ v, ˆ r ) and quantity targets ( π, z, u ) in Appendix A EfficientZero improves the sample-efficiency of MuZero by introducing additional auxiliary losses during training.
Revised: We describe EfficientZero’s learning objective between model prediction ( ˆ p, ˆ v, ˆ r ) and quantity targets ( π, z, u ) in Appendix A. EfficientZero improves the sample-efficiency of MuZero by introducing additional auxiliary losses during training.

Original: We adopt EfficientZero as our backbone model and learning algorithm, but emphasize that our framework is applicable to most model-based algorithms, including  continuous action spaces (Hafner et al., 2019a; Hansen et al., 2022).
Revised: 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., 2022a).

Original:  M ODEL -B ASED C ROSS -T ASK T RANSFER
Revised:  M ODEL -B ASED C ROSS -T ASK T RANSFER

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In the following, we introduce each of the two stages – pretraining and finetuning – in detail.
Revised: In the following, we introduce each of the two stages – pretraining and finetuning – in detail.

Original: 3. O FFLINE M ULTI -T ASK P RETRAINING
Revised: 3. O FFLINE M ULTI -T ASK P RETRAINING

Original: In this stage, we aim to learn a single world model with general perceptive and dynamics priors across a diverse set of offline tasks.
Revised: In this stage, we aim to learn a single world model with general perceptive and dynamics priors across a diverse set of offline tasks.

Original: 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.
Revised: We emphasize 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.

Original: Learning a single RL agent for a diverse set of tasks is however difficult in practice, which is only exacerbated by extrapolation errors due to the offline RL setting (Kumar et al., 2020).
Revised: Learning a single RL agent for a diverse set of tasks is however difficult in practice, which is only exacerbated by extrapolation errors due to the offline RL setting (Kumar et al., 2020).

Original: To address the challenge of multi-task learning, we propose to pretrain the model following a student-teacher training setup
Revised: To address such a challenge,

Original: [Figure]
Revised: [Figure]

Original:  in the same spirit to DQN multi-task policy distillation in Rusu et al.
Revised: 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.

Original: (2016) and Actor-Mimic in Parisotto et al.
Revised: (2016) and Actor-Mimic in Parisotto et al.

Original: (2016) , 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.
Revised: (2016), 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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 } .
Revised: 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 } .

Original: After training, we store the model predictions, ( ˆ p, ˆ v ), from each teacher model ˜ π iψ together withthe environment reward u as  student’s quantity targets ( π, z, u ) respectively for a given game ˜ M i (see Appendix A for the definition of each quantity).
Revised: After training, we store the model predictions, ( ˆ p, ˆ v ), from each teacher model ˜ π iψ together with environment reward u as the student’s quantity targets ( π, z, u ) respectively for a given game ˜ M i (see Appendix A for the definition of each quantity).

Original: Next, we learn a multi-task student model π θ by distilling the task-specific teachers into a single model via these quantity targets.
Revised: Next, we learn a multi-task student model π θ by distilling the task-specific teachers into a single model via these quantity targets.

Original: 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 as commonly done in the (single-task)
Revised: 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 as commonly done in the (single-task)

Original: MuZero Reanalyze algorithm.
Revised: MuZero Reanalyze algorithm.

Original: This step can be seen as learning multiple tasks simultaneously with direct supervision by distilling predictions from multiple teachers into a single model.
Revised: This step can be seen as learning multiple tasks simultaneously with direct supervision by distilling predictions from multiple teachers into a single model.

Original: Empirically, we find this to be a key component in scaling up the number of pretraining tasks.
Revised: Empirically, we find this to be a key component in scaling up the number of pretraining tasks.

Original: 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.
Revised: 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.

Original: 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 a 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.
Revised: 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 a 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.

Original: 3. O NLINE F INETUNING ON A T ARGET T ASK
Revised: 3. O NLINE F INETUNING ON A T ARGET T ASK

Original: 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.
Revised: 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.

Original: However, we empirically observe that directly finetuning the pretrained model often leads to poor performance on the target task due to catastrophical forgetting .
Revised: However, we empirically observe that directly finetuning a pretrained model often leads to poor performance on the target task due to catastrophical forgetting .

Original: Specifically, the initial sub-optimal data collected from online target task will cause a large perturbation in the original pretrained model parameters, which often erases useful inductive priors learned during pretraining.
Revised: 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.

Original: To overcome this challenge, we propose  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.
Revised: 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.

Original: 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.
Revised: 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.

Original: As a solution, gradient contributions from offline tasks are periodically re-weighted in an unsupervised manner based on their similarity to the target task.
Revised: As a solution, gradient contributions from offline tasks are periodically re-weighted in an unsupervised manner based on their similarity to the target task.

Original: Figure 4 shows the specific concurrent cross-task learning procedure for target task finetuning in our framework.
Revised: Figure 4 shows the specific concurrent cross-task learning procedure for target task finetuning in our framework.

Original: At each training step t , we jointly optimize the target online task M and m offline (auxiliary) tasks { ˜ M i | ˜
Revised: At each training step t , we jointly optimize the target online task M and m offline (auxiliary) tasks { ˜ M i | ˜

Original: M i ̸
Revised: M i ̸

Original: = M , 1 ≤ i ≤ m } that were used during the offline multi-task pretraining stage.
Revised: = M , 1 ≤ i ≤ m } that were used during the offline multi-task pretraining stage.

Original: Our online finetuning objective is defined as:
Revised: Our online finetuning objective is defined as:

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: 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.
Revised: 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.

Original: The target task loss term maintains a constant task weight of 1 .
Revised: The target task loss term maintains a constant task weight of 1 .

Original: During online finetuning, we use distillation targets from teachers for offline pretrained games and  MuZero Reanalysis to compute targets for online target task where we have no teacher.
Revised: During online finetuning, we use 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.

Original: 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).
Revised: 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).

Original: Within each cycle, we spend first N -steps to compute an updated η i corresponding to each offline task
Revised: Within each cycle, we spend first N -steps to compute an updated η i corresponding to each offline task

Original: ˜
Revised: ˜

Original: M i .
Revised: M i .

Original: 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.
Revised: 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.

Original: We dynamically assign the task weights by measuring the “relevance” between each offline task ˜
Revised: We dynamically assign the task weights by measuring the “relevance” between each offline task ˜

Original: M i and the (online) target task M .
Revised: M i and the (online) target task M .

Original: Inspired by the conflicting gradients measurement for multi-task learning in (Yu et al., 2020), we compute the cosine similarity between loss gradients ˜
Revised: Inspired by the conflicting gradients measurement for multi-task learning in Yu et al. (2020), we compute the cosine similarity between loss gradients ˜

Original: G in from L ez n ( ˜ M i ) and G n from L ez n ( M ) given by
Revised: G in from L ez n ( ˜ M i ) and G n from L ez n ( M ) given by

Original: [Equation]
Revised: [Equation]

Original: Within the N -step update, we maintain a task-specific counter s i and the new task weights η i can be reset by η i
Revised: Within the N -step update, we maintain a task-specific counter s i and the new task weights η i can be reset by η i

Original: = s i N at the beginning of each every T -cycle.
Revised: = s i N at the beginning of each every T -cycle.

Original: Procedure to obtain s i is stated in Appendix B. Concretely, Sim( ˜ M i , M ) measures the angle between two task gradients ˜
Revised: The procedure for obtaining s i is described in Appendix B. Concretely, Sim( ˜ M i , M ) measures the angle between two task gradients ˜

Original: G in and G n .
Revised: G in and G n .

Original: 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 η .
Revised: 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 η .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: By instead updating task weights every
Revised: By instead updating task weights every

Original: 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).
Revised: 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.

Original: Figure 5 demonstrates adjustments of task weights during finetuning with visualizations from different pretraining and finetuning tasks.
Revised: Figure 5 shows adjustments to task weights during finetuning for each of two distinct sets of pretraining and target tasks.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate our method and baselines on 14 tasks from the limited-interaction Atari100k benchmark (Kaiser et al., 2020) where only 100k environment steps are permitted.
Revised: We evaluate our method and baselines on 14 tasks from the limited-interaction Atari100k benchmark (Kaiser et al., 2020) where only 100k environment steps are permitted.

Revised: We provide an implementation of our method at https://github.com/mlpc-ucsd/XTRA .

Original: We seek to answer:
Revised: We seek to answer:

Original: • When can we empirically expect finetuning to be successful?
Revised: • How does our proposed framework compare to alternative pretraining and online RL approaches with limited online interaction from the target task?

Original: [Table]
Revised: [Table]

Original: • How does our proposed framework compare to alternative pretraining and online RL approaches with limited online interaction from the target task?

Original: • How do the individual components of our framework influence its success?
Revised: • How do the individual components of our framework influence its success?

Revised: • When can we empirically expect finetuning to be successful?

Original: Experimental setup.
Revised: Experimental setup.

Original: 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.
Revised: 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.

Original: 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 ).
Revised: 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 ).

Original: We use the EfficientZero-L variant as the default network for our framework through our experiments, unless stated otherwise.
Revised: We use the EfficientZero-L variant as the default network for our framework through our experiments, unless stated otherwise.

Original: 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.
Revised: 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.

Original: 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 ).
Revised: 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 ).

Original: 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 human performance or EfficientZero performance at 100k environment steps.
Revised: 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 human performance or EfficientZero performance at 100k environment steps.

Original: All of our results are averaged across 5 random seeds (See Appendix D for more details).
Revised: All of our results are averaged across 5 random seeds (see Appendix D for more details).

Original: We also include offline data preparation procedure in Appendix F.
Revised: We provide details on our pretraining dataset in Appendix F.

Original: Baselines.
Revised: Baselines.

Original: 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.
Revised: We compare our method against 7 prior methods for online RL that represent the state-ofthe-art on the Atari100k benchmark (including EfficientZero), a 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), a strong model-free RL baseline, under an offline pretraining + online finetuning setting. We also include a set of ablations that include EfficientZero with several different model sizes and pretraining/finetuning schemes.

Original: 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.
Revised: 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.

Original: 4. R ESULTS & D ISCUSSION
Revised: 4. R ESULTS & D ISCUSSION

Original: We introduce our results in the context of each of our three questions, and discuss our main findings.
Revised: We introduce our results in the context of each of our three questions, and discuss our main findings.

Original: How does our proposed framework compare to alternative pretraining and online RL approaches with limited online interaction from the target task?
Revised: How does our proposed framework compare to alternative pretraining and online RL approaches with limited online interaction from the target task?

Original: Tasks that share similar game mechanics.
Revised: Tasks with similar game mechanics

Original: We first investigate the feasibility of finetuning models that are pretrained on games with similar mechanics.
Revised: We first investigate the feasibility of finetuning models that are pretrained on games with similar mechanics.

Original: We select 5 shooter games and 5 maze games for this experiment.
Revised: We select 5 shooter games and 5 maze games for this experiment

Original: Results for our method, baselines, and a set of ablations on the Atari100k benchmark  are shown in Table 1.
Revised: Results for our method, baselines and a set of ablations on the Atari100k benchmark [Table] [Figure] are shown in Table 1.

Original: For completeness, we also provide learning curves in Figure 11 as well as an aggregate curve across 5 seeds of all 10 games normalized by EfficientZero’s performance at 100K environment steps in Figure 1.
Revised: For completeness, we also provide learning curves in Figure 11 as well as an aggregate curve across 5 seeds of all 10 games normalized by EfficientZero’s performance at 100K environment steps in Figure 1.

Original: We find that pretraining improves sample-efficiency substantially[Table] [Figure] across most tasks, improving mean and median performance of EfficientZero by 23 % and 25 % , respectively, overall.
Revised: We find that pretraining improves sample-efficiency substantially across most tasks, improving mean and median performance of EfficientZero by 23 % and 25 % , respectively, overall.

Original: Interestingly, XTRA also had a notable zero-shot ability compared to a multigame Behavioral Cloning baseline that is trained on the same offline dataset.
Revised: Interestingly, XTRA also had a notable zero-shot ability compared to a multigame Behavioral Cloning baseline that is trained on the same offline dataset.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: However, its median performance is overall low compared to our default variant that uses dynamic weights.
Revised: However, its median performance is overall low compared to our default variant that uses dynamic weights.

Original: We conjecture that this is because some (combinations of) games are more susceptible to gradient conflicts than others.
Revised: We conjecture that this is because some (combinations of) games are more susceptible to gradient conflicts than others.

Original: We also include offline pretraining + offline finetuning Behavioral

Original: Cloning baseline to compare with XTRA’s performance in Table 1.

Original: See Appendix J for details.

Original: Tasks with diverse game mechanics.
Revised: Tasks with diverse game mechanics.

Original: We now consider a more diverse set of pretraining and target games that have no discernible properties in common.
Revised: We now consider a more diverse set of pretraining and target games that have no discernible properties in common.

Original: Specifically, we use the following tasks for pretraining: Carnival, Centipede, Phoenix, Pooyan, Riverraid, VideoPinball, WizardOfWor, and
Revised: Specifically, we use the following tasks for pretraining: Carnival, Centipede, Phoenix, Pooyan, Riverraid, VideoPinball, WizardOfWor, and

Original: YarsRevenge, and evaluate our method on 5 tasks from Atari100k.
Revised: YarsRevenge, and evaluate our method on 5 tasks from Atari100k.

Original: Results are shown in Table 2.
Revised: Results are shown in Table 2.

Original: 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.
Revised: 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.

Original: We also include a direct comparison to CURL Srinivas et al.
Revised: We perform the same set of the ablations as we do for tasks with similar game mechanics with XTRA, and the results are shown in Table 6 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.

Original: (2020), a strong model-freeRL baseline,under the same offline pretraining + online finetuning setting as XTRA. We implement CURL w/ pretraining + finetuning by following the experimental setup as for our model-based method: (1) we first pretrain on the same 8 offline tasks as our method, and then (2) finetune to  the target task with online RL for 100k environment steps.
Revised: Model-free comparisons. For both settings (e.g., tasks with similar & diverse game mechanics), we also compare our framework with a strong model-free baseline, CURL (Srinivas et al., 2020), where CURL is pretrained on the same pretraining tasks as XTRA is, and later finetuned to each of the target tasks.

Original: We find that pretraining does not improve  performance of the model-free method CURL as consistently as for our model-based framework, XTRA.
Revised: 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.

Revised: More details and results on this comparison can be found in Table 4 and 5 from Appendix G.

Original:  How do the individual components of our framework influence its success?
Revised: 2. How do the individual components of our framework influence its success?

Original: A deeper look at task relevance.
Revised: A deeper look at task relevance.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Aggregate results across all 5 target tasks are shown in Figure 6 (a).
Revised: Aggregate results across all 5 target tasks are shown in Figure 6 (a) .

Original: 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.
Revised: 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.

Original: On the contrary, pretraining on maze games and finetuning to shooter games show similar performance compared to EfficientZero trained from scratch.
Revised: On the contrary, pretraining on maze games and finetuning to shooter games show similar performance compared to EfficientZero trained from scratch.

Original: This result indicates that (1) 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.
Revised: This result indicates that (1) 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.

Original: Which components transfer in model-based RL?
Revised: Which components transfer in model-based RL?

Original: Next, we investigate which model component(s) are important to the success of cross-task transfer.
Revised: Next, we investigate which model component(s) are important to the success of cross-task transfer.

Original: 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.
Revised: 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.

Original: Results are shown in Figure 7.
Revised: Results are shown in Figure 7.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Results for this experiment are shown in Figure 6 (b).
Revised: Results for this experiment are shown in Figure 6 (b) .

Original: 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).
Revised: 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).

Original: [Figure]
Revised: [Figure]

Original: Scaling model size.
Revised: Scaling model size.

Original: In this experiment, we investigate whether
Revised: In this experiment, we investigate whether

Original: XTRA benefits from larger model sizes.
Revised: XTRA benefits from larger model sizes.

Original: 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.
Revised: 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.

Original: However, increasing the size of the representation network h has a relatively small impact on overall inference/training time (see Figure 3 for reference).
Revised: However, increasing the size of the representation network h has a relatively small impact on overall inference/training time (see Figure 3 for reference).

Original: 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 6 (c).
Revised: 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 6 (c) .

Original: We find that our default, larger variant of XTRA (denoted XTRA-L in the figure) is slightly better than the smaller model size.
Revised: We find that our default, larger variant of XTRA (denoted XTRA-L in the figure) is slightly better than the smaller model size.

Original: In comparison, EfficientZero-L, performs significantly worse than the smaller variant of EfficientZero.
Revised: In comparison, EfficientZero-L, performs significantly worse than the smaller variant of EfficientZero.

Original: Relative improvement vs. environment steps.
Revised: Relative improvement vs. environment steps.

Original: Finally, we visualize the average improvement over EfficientZero throughout training in
Revised: Finally, we visualize the average improvement over EfficientZero throughout training in

Original: Figure 6 (d).
Revised: Figure 6 (d) .

Original: The curve indicates that XTRA is particularly useful in the early stages of training, i.e. , in an extremely limited data setting.
Revised: Results show that XTRA is particularly useful in the early stages of training, i.e. , in an extremely limited data setting.

Original: 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.
Revised: 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.

Original: When can we empirically expect finetuning to be successful?
Revised: When can we empirically expect finetuning to be successful?

Original: From Table 1, 2, we expect cross-task transfer with model-based RL is feasible, and Figure 6 (a) shows  our XTRA framework benefits  online finetuning when it is relevant to the pretrained tasks. We observe both representation, dynamics networks contribute to the successful transfer from Figure 7.
Revised: Based on Table 1 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).

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: 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;
Revised: 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;

Original: 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).
Revised: 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).

Original: 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).
Revised: 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).

Original: For example, He et al.
Revised: For example, He et al.

Original: (2020) show that contrastive pretraining on a large, unlabelled dataset learns good features for ImageNet classification, and Brown et al.
Revised: (2020) show that contrastive pretraining on a large, unlabelled dataset learns good features for ImageNet classification, and Brown et al.

Original: (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.
Revised: (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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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 (Hansen et al., 2022b).

Original: 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.
Revised: 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.

Original: Sample-efficient RL.
Revised: Sample-efficient RL.

Original: 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;
Revised: 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;

Original: Yarats et al., 2019; Srinivas et al., 2020; Schwarzer et al., 2021), data augmentation (Laskin et al., 2020;
Revised: Yarats et al., 2019; Srinivas et al., 2020; Schwarzer et al., 2021), data augmentation (Laskin et al., 2020;

Original: 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;
Revised: 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;

Original: Baker et al., 2022), using pretrained visual representations for model-free RL (Shah & Kumar, 2021; 
Revised: 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; Finn &

Original: Xiao et al., 2022), and model-based RL (Ha & Schmidhuber, 2018; 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;
Revised: 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;

Original: Hansen et al., 2022; Seo et al., 2022).
Revised: Hansen et al., 2022a; Seo et al., 2022; Hansen et al., 2023).

Original: 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.
Revised: 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.

Original: Several prior works consider these problems independently from each other: Xiao et al.
Revised: Several prior works consider these problems independently from each other: Xiao et al.

Original: (2022) shows that model-free policies can be trained with a frozen pretrained visual backbone, and Seo et al.
Revised: (2022) shows that model-free policies can be trained with a frozen pretrained visual backbone, and Seo et al.

Original: (2022) shows that learning a world model on top of features from a visual backbone pretrained with video prediction can improve model learning.
Revised: (2022) shows that learning a world model on top of features from a visual backbone pretrained with video prediction can improve model learning.

Original: 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.
Revised: 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.

Original: Finetuning in RL.
Revised: Finetuning in RL.

Original: 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), or task specification (Xie & Finn, 2021; Walke et al., 2022).
Revised: 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).

Original: For example, Julian et al.
Revised: For example, Julian et al.

Original: (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.
Revised: (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.

Original: Similarly, Hansen et al.
Revised: Similarly, Hansen et al.

Original: (2021a) shows that model-free policies can (to some extent) also adapt to small domain shifts by self-supervised finetuning within a single episode.
Revised: (2021a) shows that model-free policies can (to some extent) also adapt to small domain shifts by self-supervised finetuning within a single episode.

Original: 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).
Revised: 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).

Original: Finally, Lee et al.
Revised: Finally, Lee et al.

Original: (2022) shows that offline multi-task RL pretraining via sequence modelling can improve offline finetuning on data from unseen tasks.
Revised: (2022) shows that offline multi-task RL pretraining via sequence modelling can improve offline finetuning on data from unseen tasks.

Original: Our approach is most similar to Julian et al.
Revised: Our approach is most similar to Julian et al.

Original: (2020) in that we finetune via rewards on a mixture of datasets.
Revised: (2020) in that we finetune via rewards on a mixture of datasets.

Original: 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.
Revised: 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.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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.
Revised: 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.

Original: We find that XTRA improves sample-efficiency substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25% , respectively, overall.
Revised: We find that XTRA improves sample-efficiency substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25% , respectively, overall.

Original: We hope  our  analysis and findings onthe feasibility of cross-task transfer with model-based RL can greatly benefit future studies in this direction.
Revised: 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.

Original: Reproducibility Statement.
Revised: Reproducibility Statement.

Original: We base all experiments on the publicly available Atari100k benchmark, and provide extensive implementation details in the appendices, including detailed network architecture and hyper-parameters.
Revised: We base all experiments on the publicly available Atari100k benchmark, and provide extensive implementation details in the appendices, including detailed network architecture and hyper-parameters.

Original: We are committed to releasing our code publicly.
Revised: Code for reproducing our experiments is made available at https://nicklashansen.github.io/xtra . Acknowledgements. This work is supported by NSF Award IIS-2127544 and gifts from Qualcomm AI.

Revised: We thank Weirui Ye, Wenlong Huang and Weijian Xu for helpful discussions

Original: AA PPENDIX : XTRA/E FFICIENT Z ERO O BJECTIVES
Revised: A XTRA/E FFICIENT Z ERO O BJECTIVES

Original: XTRA uses the same learning objective as EfficientZero 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.
Revised: XTRA uses the same learning objective as EfficientZero 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.

Original: Here, we explain the objectives of EfficientZero (Ye et al., 2021) and its predecessor MuZero (Schrittwieser et al., 2020).
Revised: Here, we explain the objectives of EfficientZero (Ye et al., 2021) and its predecessor MuZero (Schrittwieser et al., 2020).

Original: 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 u from environment (3) bootstrapped value target z
Revised: 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 u from environment (3) bootstrapped value target z

Original: where z = (cid:80)
Revised: where z = (cid:80)

Original: k − 1 i =

Revised: k − 1 i =

Original: γ i u i + γ k v t
Revised: γ i u i + γ k v t

Original: + k .
Revised: + k .

Original: On top of
Revised: On top of

Original: MuZero, EfficientZero adds a self-supervised consistency loss term, and predicts sum of environment rewards from next k steps, (cid:80) k − 1 i =

Revised: MuZero, EfficientZero adds a self-supervised consistency loss term, and predicts sum of environment rewards from next k steps, (cid:80) k − 1 i =

Original: γ i u i , instead of single-step reward.
Revised: γ i u i , instead of single-step reward.

Original: We refer reader to the original manuscript for implementation details (Ye et al., 2021).
Revised: We refer reader to the original manuscripts for implementation details.

Original: We present the learning objective for EfficientZero at time step t with k unroll steps:
Revised: We present the learning objective for EfficientZero at time step t with k unroll steps:

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: BA PPENDIX : T ASK W EIGHTS C OMPUTATION
Revised: B T ASK W EIGHTS C OMPUTATION

Original: Within the N -steps update, we maintain a task-specific counter s i and update the counter by ∆ s in at each step n as follows:
Revised: Within the N -steps update, we maintain a task-specific counter s i and update the counter by ∆ s in at each step n as follows:

Original: [Equation]
Revised: [Equation]

Original: After N steps, the new task weights η i can be reset by η i =
Revised: At every N steps [Figure] the new task weights η i are updated by η i =

Original: s i N , and be adopted in our online adaptation objective in  Equation 1.
Revised: s i N , and used in subsequent finetuning objectives according to Equation 1.

Original: In practice, we start task weights update at 10k steps to ensure enough data from the online target task is collected.
Revised: 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.

Original: All task weights are initialized as 1 for the first 10k steps, and reset at the beginning of each every T -cycle.
Revised: All task weights are initialized as 1 for the first 10 k steps.

Original: Figure 8 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.
Revised: Figure 8 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.

Original: Figure 9 shows the adjustments of task weights for 5 games reported in Figure 12.
Revised: Figure 9 shows the adjustments of task weights for 5 games reported in Figure 12.

Original: [Figure]
Revised: [Figure]

Original: C A PPENDIX : C OMPARISON BETWEEN M ULTI -G AME O FFLINE T RAINING AND D ISTILLATION
Revised: C D ISTILLATION VS . M ULTI -G AME O FFLINE RL

Original: In Figure 10, we present results frommulti- game offline training and multi-game distillationforgames. The multi-game results are evaluated from a single set of model parameters trained with 
Revised: 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

Original: [Figure]
Revised: [Figure]

Original: all 4 games. While the model can easily learn each game individually (offline single game)
Revised: to directly train a single world model on the multi-game dataset with offline RL.

Original: [Figure]
Revised: However, we find that learning such a model is difficult.

Original: it fails to learn 4 games simultaneously in the offline setting. On the opposite [Figure]
Revised: A comparison between the two approaches on four different pretraining games is shown in Figure 10.

Original: our multi-game  distillation successfully learns 4 games from teacher targets in multi-task fashion
Revised: 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.

Original: [Figure]
Revised: [Figure]

Original: DA PPENDIX : XTRA G AME S CORES FOR T ABLE 1.
Revised: D S CORES FOR I NDIVIDUAL S EEDS

Original: We report raw game scores from our XTRA stated in Table 1 for all 5 seeds we have run. We attach the mean, median, and standard deviation of game scores from 5 seeds of eachindividual game.
Revised: 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.

Original: We also listthe random and human score from Badia et al.
Revised: We also list random and human scores obtained from Badia et al.

Original: (2020), and calculate the HumanNormalized
Revised: (2020), and calculate the Human

Original:  Score based on the formula: ( score agent − score random ) / ( score human − score random ) .
Revised: Normalized Score based on the formula: ( score agent − score random ) / ( score human − score random ) as in prior work.

Original: There is norandom or human score available for Carnival, and  therefore  the aggregated Human Normalized Mean and MedianScores are calculated from the other 9 games.
Revised: 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

Original: Details are in Table 3.
Revised: Scores computed in Table 3.

Original: [Table]
Revised: [Table]

Revised: [Figure]

Original: E APPENDIX : A DDITIONAL E VALUATION CURVES OF XTRA ON A TARI 100 K BENCHMARK
Revised: E A DDITIONAL E VALUATION CURVES OF XTRA ON A TARI 100 K BENCHMARK

Original: For completeness, Figure 11 and 12 include evaluation curves of XTRA on the games for which we report final performance in Table 1 and2. [Figure]
Revised: For completeness, Figure 11 and 12 include evaluation curves of XTRA on the games for which we report final performance in Table 1 and

Original: [Figure]
Revised: [Figure]

Original: FA PPENDIX : O FFLINE D ATA P REPARATION
Revised: F O FFLINE D ATA P REPARATION

Original: To train the model in offline multi-task pretraining stage, we usecollected trajectories of EfficientZero Ye et al. (2021) from the Atari100k benchmark.
Revised: To train the model in offline multi-task pretraining stage, we use trajectories collected by EfficientZero (Ye et al., 2021) on the Atari100k benchmark.

Original: For each pretraining game, we assume we can access model checkpoints obtained from 100k environment steps and 120k training steps, resulting inmodel checkpoints.
Revised: For each pretraining game, we assume we can access model checkpoints obtained every 10k steps from 120k training steps (the environment step is capped at 100k), resulting in 12 model checkpoints.

Original: For eachmodel checkpoint, we repeat the self-play procedure in EfficientZero Ye et al. (2021) to collect 64 trajectories.
Revised: For each checkpoint, we evaluate the model performance on the Atari environment following the same procedure from EfficientZero and collect 64 trajectories.

Original: This translates to an average of 1M transitions per game, but varies depending on episode length – for example, this only results in 636k transitions for the game of Assault.
Revised: This translates to an average of 1M transitions per game, but varies depending on episode length – for example, this only results in 636k transitions for the game of Assault.

Revised: 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 an expert agent.

Original: Thus, we show that XTRA is effective even when pretraining data is suboptimal, allowing us to learn from very diverse data sources.
Revised: Thus, we show that XTRA is effective even when pretraining data is suboptimal, allowing us to learn from very diverse data sources.

Revised: G P RETRAINING + F INETUNING IN M ODEL - FREE RL

Revised: We compare the effectiveness of our framework, XTRA, with a strong model-free baseline, CURL, that we also implement following a similar pretraining and finetuning scheme.

Revised: This is in contrast to the original formulation of CURL that does not leverage pretraining.

Revised: We implement our training scheme for CURL with the following setup: (1) pretrain a multi-task CURL model on the same pretraining tasks as 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.

Original: G A PPENDIX : HYPER - PARAMETERS FOR XTRAAND E FFICIENT Z ERO M ODEL A RCHITECTURE
Revised: [Table] [Table] H XTRA A BLATIONS FOR T ASKS WITH D IVERSE G AME M ECHANICS

Original: We follow the EfficientZero’s official architecture implementation to make straight comparison.
Revised: We perform the same set of ablations for tasks that share diverse game mechanics as we do for tasks that share similar game mechanics for XTRA. Results are shown in Table 6. Details of each ablation can be found in Section 4. [Table]

Revised: [Table]

Revised: I E FFECTS OF N UMBER OF T ASKS IN P RETRAINING AND C ROSS -T ASKS IN F INETUNING

Revised: We perform an additional ablation for tasks that share diverse game mechanics – whether changing the number of tasks during pretraining would help or hurt the performance in later cross-task finetuning.

Revised: 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.

Revised: In this ablation

Revised: [Table]

Original: For
Revised: we gradually reduce the number of pretrained tasks from 8 (Carnival, Centipede, Phoenix, Pooyan, Riverraid, VideoPinball, WizardOfWor [Table] and YarsRevenge), to 4 (Phoenix, WizardOfWor,

Original: EfficientZero-L and XTRA, we increase number of convolutions layers from 1 to 4 to linearly scale up the size of representation networks.
Revised: VideoPinball, YarsRevenge), and to 2 (Phoenix, VideoPinball). XTRA is reduced to EfficientZero-L when the number of pretrained tasks and cross-tasks during finetuning is set to 0. We find that increasing the number of tasks during pretraining (and later cross-task finetuning) mostly consistently improves XTRA performance. J A RCHITECTURAL D ETAILS

Revised: We adopt the architecture of EfficientZero (Ye et al., 2021).

Revised: 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.

Revised: However, we find that our baseline EfficientZero (without pretraining) performs significantly worse with a larger representation network.

Revised: Therefore, we use the default EfficientZero as the main point of comparison throughout this work and only include Efficient-L for completeness.

Original: The architecture of the representation networks is as follows:
Revised: The architecture of the representation networks is as follows:

Original: • 1 convolution with stride 2 and 32 output planes, output resolution 48x48.
Revised: • 1 convolution with stride 2 and 32 output planes, output resolution 48x48.

Original: (BN + ReLU) • 1 residual block with 32 planes.
Revised: (BN + ReLU) • 1 residual block with 32 planes.

Original: • 1 residual downsample block with stride 2 and 64 output planes, output resolution 24x24.
Revised: • 1 residual downsample block with stride 2 and 64 output planes, output resolution 24x24.

Original: • 1 residual block with 64 planes.
Revised: • 1 residual block with 64 planes.

Original: • Average pooling with stride 2, output resolution 12x12.
Revised: • Average pooling with stride 2, output resolution 12x12.

Original: (BN + ReLU) • 1 residual block with 64 planes.
Revised: (BN + ReLU) • 1 residual block with 64 planes.

Original: • Average pooling with stride 2, output resolution 6x6.
Revised: • Average pooling with stride 2, output resolution 6x6.

Original: (BN + ReLU) • 1 residual block with 64 planes.
Revised: (BN + ReLU) • 1 residual block with 64 planes.

Original: , where the kernel size is 3 × 3 for all operations.
Revised: , where the kernel size is 3 × 3 for all operations.

Original: The architecture of the dynamics networks is as follows:
Revised: The architecture of the dynamics networks is as follows:

Original: • Concatenate the input states and input actions into 65 planes.
Revised: • Concatenate the input states and input actions into 65 planes.

Original: • 1 convolution with stride 2 and 64 output planes.
Revised: • 1 convolution with stride 2 and 64 output planes.

Original: (BN) •
Revised: (BN) •

Original: A residual link: add up the output and the input states.
Revised: A residual link: add up the output and the input states.

Original: (ReLU) • 1 residual block with 64 planes.
Revised: (ReLU) • 1 residual block with 64 planes.

Original: The architecture of the reward prediction network is as follows:
Revised: The architecture of the reward prediction network is as follows:

Original: • 1 1x1convolution and 16 output planes.
Revised: • 1 1x1convolution and 16 output planes.

Original: (BN + ReLU)
Revised: (BN + ReLU)

Original: • Flatten.
Revised: • Flatten.

Original: • LSTM with 512 hidden size.
Revised: • LSTM with 512 hidden size.

Original: (BN + ReLU) • 1 fully connected layers and 32 output dimensions.
Revised: (BN + ReLU) • 1 fully connected layers and 32 output dimensions.

Original: (BN + ReLU) • 1 fully connected layers and 601 output dimensions.
Revised: (BN + ReLU) • 1 fully connected layers and 601 output dimensions.

Original: The architecture of the value and policy prediction networks is as follows:
Revised: The architecture of the value and policy prediction networks is as follows:

Original: • 1 residual block with 64 planes.
Revised: • 1 residual block with 64 planes.

Original: • 1 1x1convolution and 16 output planes.
Revised: • 1 1x1convolution and 16 output planes.

Original: (BN + ReLU)
Revised: (BN + ReLU)

Original: • Flatten.
Revised: • Flatten.

Original: • 1 fully connected layers and 32 output dimensions.
Revised: • 1 fully connected layers and 32 output dimensions.

Original: (BN + ReLU) • 1 fully connected layers and D output dimensions.
Revised: (BN + ReLU) • 1 fully connected layers and D output dimensions.

Original: , where D = 601 in the value prediction network and D = action space in the policy prediction network.
Revised: where D = 601 in the value prediction network and D = |A| in the policy prediction network.

Original: H A PPENDIX : H YPER - PARAMETERSFOR O NLINE F INETUNING XTRA
Revised: K H YPER - PARAMETERS

Original: In this paper, we adopted all hyper-parameters  for EfficientZero to make straight comparison except training steps and reanalyzed policy ratio.
Revised: 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 8 for completeness.

Original: In the original EfficientZero paper, EfficientZero trains the model for additional 20k  steps after collecting 100k environment steps data and decay learning rate by 0.1.
Revised: Lastly, we note that EfficientZero performs an additional 20k gradient steps at 100k environment steps, with a 10 × smaller learning rate.

Original: We set training steps to 120k as original EfficientZero setting for Table 2 to make direct and fair comparison to the reported number from EfficientZero paper.
Revised: We follow this procedure when comparing to previous state-of-the-art methods (Table 2) [Table]

Original: We set training steps to 100k to match exact 100k environment steps with no additional training for all other results reported in this paper forrunning bothEfficientZero baselines (EfficientZero and EfficientZero-L) and our XTRA models.
Revised: but for simplicity we omit these additional gradient steps in the remainder of our experiments for both XTRA and baselines.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Revised: L E FFECT OF M INI -B ATCH S IZE

Original: [Table]
Revised: [Figure]

Original: [Figure]
Revised: During online finetuning (stage 2), we finetune both with data 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-pretraining baselines, 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 ).

Original: [Table]
Revised: 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 is not the source of our performance gains, but rather our pretraining and inclusion of pretraining tasks during the online finetuning. M P ER -G AME I MPROVEMENT OVER E FFICIENT Z ERO

Original: [Figure]
Revised: 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, depending on the visualization setting. C e n t

Original: [Table]
Revised: i p e d e A li e n P h

Revised: o e n i x C a r n i v a l A ss a u l t

Original: [Figure]
Revised: [Figure]

Revised: a n

Original: [Table]
Revised: [Figure] H e i s t D e [Figure]

Revised: n

Original: [Figure]
Revised: [Figure]

Revised: tt a c

Original: [Table]
Revised: [Figure]

Original: I A PPENDIX : G AME I NFORMATION
Revised: P [Figure] a n [Figure]

Revised: d

Original: We provide information on games appeared in pretraining and finetuning for XTRA.
Revised: [Figure] r A m [Figure] I m p r

Original: The Similar Task column marks all the games used in Table 1 , and the two Diverse Task column mark all the games used in Table 2 for pretraining and finetuning the model.
Revised: o v e m e n t o v e r E ff i

Original: We define all the games in the following categories: Maze, Shooter, Tank, Adventure, and Ball Tracking.
Revised: c i e n t Z e r o

Original: In the Maze category, the actor is required to move horizontally or vertically to find a path to (1) pick up objectives to gain scores or special abilities, and (2) avoid being caught by moving enemies.
Revised: [Figure]

Original: In the
Revised: N G AME I NFORMATION

Original: Shooter category, the actor is required to shoot horizontally or vertically toward objectives, and at the same time avoid being shot by enemies.
Revised: In this section, we aim to provide additional context about the games that we consider during both pretraining and finetuning.

Original: In the Tank category, the actor is required to move a fighting vehicle into any directions on a plane, shoot to the enemies appearing on the map and avoid being hit; it is worth mentioning that in this category, the visual observation appears in the First-Person Perspective (FPP), which is different from the visual observation from the shooter category, where the visual observation appears in Third-Person Perspective (TPP) with absolute positioning (e.g. the background is fixed).
Revised: Table 9 lists core properties for each game.

Original: In the adventure category, the actor is required to perform a diverse set of skills which can differ from game to game; these skills often include shooting, hitting the enemies, solving mazes, and protecting targets, some of which can also be seen from other game categories. In the
Revised: 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

Original: Ball Tracking category, the actor is required to keep track of a moving ball, andhit it to prevent the ball from moving off the gaming area. We also include attributes in terms of scene continuity and action space . A continuous scene represents that within the game, there will not be sudden changes of a visual observation possibly due to actor moving out of the current gaming zone, leveling up, etc. Action spacerefers to the number of actions that the actor is able to use for a game.
Revised: Tracking , and also report whether the scene is static or dynamic, as well the (valid) action space for each task.

Original: The maximum  action space  for Atari games is [Table] [Table]
Revised: The maximum dimensionality of the action space is 18 for Atari games.

Original: J A PPENDIX : B EHAVIORAL C LONING B ASELINE
Revised: O B EHAVIORAL C LONING B ASELINE

Original: We use the representation + prediction network in XTRA for the behavioral cloning (BC) study.
Revised: We use the representation + prediction network in XTRA for the behavioral cloning (BC) study.

Original: The
Revised: The

Original: BC (finetune) from Table 1 follows an offline pretraining + offline finetuning paradigm.
Revised: BC (finetune) from Table 1 follows an offline pretraining + offline finetuning paradigm.

Original: The model is finetuned on offline data for the target task (also generated by the EfficientZero baseline).
Revised: The model is finetuned on offline data for the target task (also generated by the EfficientZero baseline).

Original: We find
Revised: We find

Original:  BC (finetune) underperforms the EfficientZero baseline.
Revised: [Table] BC (finetune) underperforms the EfficientZero baseline.

Original: We also report zero-shot performance of the pretrained BC on their designated target tasks directly.
Revised: We also report zero-shot performance of the pretrained BC on their designated target tasks directly [Table]

Original: K A PPENDIX : G AME V ISUALIZATIONS
Revised: P G AME V ISUALIZATIONS

Original: To better understand the feasibility of cross-task transfer, we show sample trajectory from our experiment games. Figure 13 and 14 include all games involved inTable 1 (tasks sharing similar game mechanics).
Revised: 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 → 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.

Original: Figure 15 include pretraining/finetuning games involved in Table 2 (tasks sharing diverse game mechanics) excluding ones that are shown in Figure 13 and Figure 14.
Revised: Figure 15 shows trajectories for five Maze games, Figure 16 shows trajectories for five Shooter games, and Figureshows trajectories for eight games from our diverse game category.

Original: [Figure]
Revised: [Figure]

Original: −→ A li e n A m i d a r
Revised: i d a r

Original: B a nk
Revised: B a nk

Original: H e i s t
Revised: H e i s t

Original: M s P ac m a n
Revised: M s P ac m a n

Original: W i za r d O f
Revised: W i za r d O f

Original: W o r [Figure] −→
Revised: W o r time −→

Original: A s s a u lt
Revised: A ss a u lt

Original: C a r n i v a l
Revised: C a r n i v a l

Original: C e n ti p e d e
Revised: C e n ti p e d e

Original: D e m on A tt ac k
Revised: D e m on A tt ac k

Original: P ho e n i x time −→
Revised: P ho e n i x time −→

Original: B a ttl e Z on e
Revised: B a ttl e Z on e

Original: H e r o
Revised: H e r o

Original: K r u ll
Revised: K r u ll

Original: P ooy a n
Revised: P ooy a n

Original: R i v e rr a i d
Revised: R i v e rr a i d

Original: S ea qu e s t
Revised: S ea qu e s t

Original: V i
Revised: V i

Original: d e o P
Revised: d e o P

Original: i nb a ll
Revised: i nb a ll

Original: Y a r s R e v e ng e
Revised: Y a r s R e v e ng e