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
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: 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
Revised: We propose

Original: Model-Based CrossTask
Revised: Model-Based CrossTask

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 and concurrent 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 proper pretraining and concurrent 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.

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: Reinforcement Learning (RL) 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.
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 realtime 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 skills tabula rasa (from scratch) for each new task.
Revised: However, current RL algorithms – including EfficientZero – are still predominantly assumed to learn both perception, model, and skills tabula rasa (from scratch) for each new task.

Original: On the contrary, humans rely heavily on prior knowledge and visual cues when learning new skills.
Revised: On the contrary, humans rely heavily on prior knowledge and visual cues when learning new skills.

Original: For example, a study found that human players easily pick up on visual cues about game mechanics and objectives when exposed to a
Revised: For example, a study found that human players easily pick up on visual cues about game mechanics and objectives when exposed to a

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

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

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.
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 tothe 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 inan unsupervised manner based on similarity tothe 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:  We find that 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 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
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

Original: (cid:55)→ O is a transition function, ρ is the initial state distribution, r : O × A
Revised: (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.

Original: [Figure]
Revised: [Figure]

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

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: 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 = g θ( z t , a 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 = g θ

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: ( z t , a t ) . For each latent state, a prediction network f θ estimates a probability distribution ˆ p over (valid) actions a ∈ A , as well as the expected state value ˆ v of the given state, i.e. , ˆ v t , ˆ p t = f θ ( z t ) .

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:  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 continuous action spaces (Hafner et al., 2019a; Hansen et al., 2022).

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, 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.

Original: Learning a single RL agent for a diverse set of tasks is howevera 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, where teacher models are trained separately by offline RL for each task, and then distilled into a
Revised: To address the challenge of multi-task learning, we propose to pretrain the model following a student-teacher training setup

Original: [Figure]
Revised: [Figure]

Original:  single multi-task model using a novel instantiation of the MuZero Reanalyze (Schrittwieser et al., 
Revised: in the same spirit to DQN multi-task policy distillation in Rusu et al. (2016) and Actor-Mimic in Parisotto et al. (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.

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,  each teacher model ˆ π iψ has learned to produce task-specific quantities ( ˆ π , ˆ u, ˆ z )  for a given game ˆ M i .
Revised: After training, we store the model predictions, ( ˆ p, ˆ v ), from each teacher model ˜ π iψ together with the 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).

Original: Next, we learn a multi-task student modelˆ π θ by distilling the task-specific teachers into a single model.
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 generate value/policy targets using the 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 teacher models into the multi-task student model, we now have a single set of pretrained parameters that can be used for finetuning to a variety of tasks via online interaction, which we introduce in the following section.
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 tocatastrophical forgetting, and consequently poor performance on the target task.
Revised: However, we empirically observe that directly finetuning the pretrained model often leads to poor performance on the target task due to catastrophical forgetting .

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

Original: To overcome this challenge, we  retain offline data from the pretraining stage, and concurrently finetune the model on both data from the target task, as well as data from the pretraining tasks.
Revised: 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.

Original: While this procedure addresses catastrophical forgetting, interference between the target task and certain pretraining tasks can be harmful for the sample-efficiency during online RL.
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.

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:  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: [Figure] [Figure] 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 .

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

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 
Revised: Within each cycle, we spend first N -steps to compute an updated η i corresponding to each offline task

Original: [Figure] [Figure]
Revised: ˜

Original: task ˆ 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 ˆ M i and the (online) target task M .
Revised: We dynamically assign the task weights by measuring the “relevance” between each offline task ˜ 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 ˆ G in from L ez n ( ˆ M i ) and G n from L ez n ( M ) given by
Revised: Inspired by the conflicting gradients measurement for multi-task learning in (Yu et al., 2020), we compute the cosine similarity between loss gradients ˜ G in from L ez n ( ˜ M i ) and G n from L ez n ( M ) given by

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 following the procedure stated in Appendix
Revised: = s i N at the beginning of each every T -cycle. Procedure to obtain s i is stated in Appendix

Original: B. Concretely, Sim( ˆ M i , M ) measures the angle between two task gradients ˆ G in and G n .
Revised: B. Concretely, Sim( ˜ M i , M ) measures the angle between two task gradients ˜ 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., 2020).

Revised: Figure 5 demonstrates adjustments of task weights during finetuning with visualizations from different pretraining and finetuning 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.

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

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

Original:  • How does our proposed framework compare to alternative pretraining and online RL approaches with limited online interaction from the target task?
Revised: [Table] • 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?

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 [Table]
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: [Figure] network (denoted EfficientZero-L ). 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 byeither 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 to ensure reliability. See Appendix D for more details.
Revised: All of our results are averaged across 5 random seeds (See Appendix D for more details). We also include offline data preparation procedure 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 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.

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?[Table] [Figure]
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 that share 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 are shown in Table 1.

Original: For completeness, we also provide learning curves in Figure 6.
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  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 [Table] [Figure] 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 multi-game behavior 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.

Revised: We also include offline pretraining + offline finetuning Behavioral

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

Revised: 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: This suggests that, while task similarity may play a role in the success of XTRA, the algorithmic advances of our proposed frames are a bigger factor in the strong empirical performance.
Revised: We also include a direct comparison to CURL Srinivas et al. (2020), a strong model-free RL 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. We find that pretraining does not improve performance of the model-free method CURL as consistently as for our model-based framework, XTRA.

Original: How do the individual components of our framework influence its success?
Revised: 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 7 (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 (i) selecting pretraining tasks that are relevant to the target task is key to benefit from pretraining, and (2) in the extreme case where there are no pretraining tasks relevant to the target task, finetuning with XTRA generally does not harm the online RL performance since it can automatically assign small weights to the pretraining tasks.
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 8.
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 7 (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 7 (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: This indicates that, unlike EfficientZero trained from scratch, XTRA is robust to model size and can potentially benefit from further scaling up the model size in the future.

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 7 (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: The curve indicates 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 7 (a) shows our XTRA framework benefits online finetuning when it is relevant to the pretrained tasks.
Revised: 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.

Original: We observe both representation, dynamics networks contribute to the successful transfer from Figure 8.
Revised: We observe both representation, dynamics networks contribute to the successful transfer from 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: or 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., 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 contrastive pretraining on a large, unlabelled dataset learns good features for ImageNet classification, and Brown et al. (2020) show that a generative model trained on large-scale natural language data can be used to solve unseen tasks given only a few examples.

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.

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;

Original: Xiao et al., 2022), and model-based RL (Finn & Levine, 2017; Nair et al., 2018; Hafner et al., 2019b; Kaiser et al., 2020; Schrittwieser et al., 2020; Hafner et al., 2021; Ye et al., 2021;
Revised: 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;

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

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), 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 on the feasibility of cross-task transfer with model-based
Revised: We hope our analysis and findings on the feasibility of cross-task transfer with model-based

Original: RL can greatly benefit future studies in this direction.
Revised: RL can greatly benefit future studies 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: We are committed to releasing our code publicly.

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

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.

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 MuZero, EfficientZero adds a self-supervised consistency loss term, and predicts sum of  rewards {+from next k steps, (cid:80) k − 1 i =
 γ i u i ,+} instead of single-step reward.
Revised: On top of MuZero, EfficientZero adds a self-supervised consistency loss term, and predicts sum of environment rewards {+from next k steps, (cid:80) k − 1 i =
 γ i u i ,+} instead of single-step reward.

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

Original: Here, 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: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Table][Equation] [Figure] [Equation]
Revised: [Table]

Original: [Table]

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

Original: [Equation] [Table] [Figure] [Table]
Revised: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: B A PPENDIX : T ASK W EIGHTS C OMPUTATION
Revised: B A PPENDIX : 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: After N steps, the new task weights η i can be reset by η i =

Original: s i N , and be adopted in our online adaptation objective in Equation 1.
Revised: s i N , and be adopted in our online adaptation objective in 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, we start task weights update at 10k steps to ensure enough data from the online target task is collected.

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 10k steps, and reset at the beginning of each every T -cycle.

Original: Figure 9 shows how task weights are adaptively adjusted by the model during 100k environment steps during online finetuning stage for 10 games reported in Figure 6.
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. 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 A PPENDIX : C OMPARISON BETWEEN M ULTI -G AME O FFLINE T RAINING AND D ISTILLATION

Original: In Figure 10, we present results from multi-game offline training and multi-game distillation forgames.
Revised: In Figure 10, we present results from multi- game offline training and multi-game distillation forgames.

Original: The multi-game results are evaluated from a single set of model parameters trained with  all 4 games.
Revised: The multi-game results are evaluated from a single set of model parameters trained with [Figure] all 4 games.

Original: While the model can easily learn each game individually (offline single game), it fails to learn 4 games simultaneously in the offline setting.
Revised: While the model can easily learn each game individually (offline single game) [Figure] it fails to learn 4 games simultaneously in the offline setting.

Original: On the opposite, our multi-game distillation successfully learns 4 games from teacher targets in multi-task fashion.
Revised: On the opposite [Figure] our multi-game distillation successfully learns 4 games from teacher targets in multi-task fashion [Figure]

Original: D A PPENDIX : XTRA G AME S CORES FOR T ABLE 1.
Revised: D A PPENDIX : XTRA G AME S CORES FOR T ABLE 1.

Original: We report raw game scores from our XTRA stated in Table 1 for all 5 seeds we have run.
Revised: We report raw game scores from our XTRA stated in Table 1 for all 5 seeds we have run.

Original: We attach the mean, median, and standard deviation of game scores from 5 seeds of each individual game.We [Figure]
Revised: We attach the mean, median, and standard deviation of game scores from 5 seeds of each individual game.

Original:  also list the random and human score from Badia et al.
Revised: We also list the random and human score from Badia et al.

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

Original: Score based on the formula: ( score agent − score random ) / ( score human − score random ) .
Revised: Score based on the formula: ( score agent − score random ) / ( score human − score random ) .

Original: There is no random or human score available for Carnival, and therefore the aggregated Human Normalized Mean and
Revised: There is no random or human score available for Carnival, and therefore the aggregated Human Normalized Mean and

Original: Median Scores are calculated from the other 9 games.
Revised: Median Scores are calculated from the other 9 games.

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

Original: [Table]

Original: [Figure]

Original: [Table]
Revised: [Table]

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

Original: We show additional evaluation curves of XTRA on 5 games shown in Table  2. 
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 2. [Figure] [Figure]

Original: The averaged game score is evaluated by 5 seeds each with 32 evaluation episodes
Revised: To train the model in offline multi-task pretraining stage, we use collected trajectories of EfficientZero

Original: [Table]
Revised: Ye et al. (2021) from the Atari100k benchmark. For each pretraining game, we assume we can access model checkpoints obtained from 100k environment steps and 120k training steps, resulting inmodel checkpoints. For each model checkpoint, we repeat the self-play procedure in EfficientZero

Original: [Figure]
Revised: Ye et al. (2021) to collect 64 trajectories. 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. Thus, we show that XTRA is effective even when pretraining data is suboptimal, allowing us to learn from very diverse data sources.

Original: F A PPENDIX : H YPER - PARAMETERS FOR XTRA AND E FFICIENT Z ERO M ODEL A RCHITECTURE
Revised: G A PPENDIX : H YPER - PARAMETERS FOR XTRA AND E FFICIENT Z ERO M ODEL A RCHITECTURE

Original: We follow the EfficientZero’s official architecture implementation to make straight comparison.
Revised: We follow the EfficientZero’s official architecture implementation to make straight comparison.

Original: For
Revised: For

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: EfficientZero-L and XTRA, we increase number of convolutions layers from 1 to 4 to linearly scale up the size of representation networks.

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 = action space in the policy prediction network.

Original: G A PPENDIX : H YPER - PARAMETERS FOR O NLINE F INETUNING XTRA
Revised: H A PPENDIX : H YPER - PARAMETERS FOR O NLINE F INETUNING XTRA

Original: In this paper, we adopted all hyper-parameters for EfficientZero to make straight comparison except training steps and reanalyzed policy ratio.
Revised: In this paper, we adopted all hyper-parameters for EfficientZero to make straight comparison except training steps and reanalyzed policy ratio.

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

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

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 for running both EfficientZero baselines (EfficientZero and EfficientZero-L) and our XTRA models. 
Revised: We set training steps to 100k to match exact 100k environment steps with no additional training for all other results reported in this paper for running both EfficientZero baselines (EfficientZero and EfficientZero-L) and our XTRA models. [Table] [Figure] [Table]

Revised: [Figure]

Original: [Table]
Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Original: H A PPENDIX : G AME I NFORMATION
Revised: I A PPENDIX : G AME I NFORMATION

Original: We provide information on games appeared in pretraining and finetuning for XTRA.
Revised: We provide information on games appeared in pretraining and finetuning for XTRA.

Original: The Similar
Revised: The Similar

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

Original: We define all the games in the following categories: Maze, Shooter, Tank, Adventure, and Ball Tracking.
Revised: We define all the games in the following categories: Maze, Shooter, Tank, Adventure, and Ball Tracking.

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

Original: In the
Revised: In the

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: Shooter category, the actor is required to shoot horizontally or vertically toward objectives, and at the same time avoid being shot by enemies.

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

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

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

Original: In the
Revised: In the

Original: Ball Tracking category, the actor is required to keep track of a moving ball, and hit it to prevent the ball from moving off the gaming area.
Revised: Ball Tracking category, the actor is required to keep track of a moving ball, and hit it to prevent the ball from moving off the gaming area.

Original: We also include attributes in terms of scene continuity and action space .
Revised: We also include attributes in terms of scene continuity and action space .

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

Original: Action space refers to the number of actions that the actor is able to use for a game.
Revised: Action space refers to the number of actions that the actor is able to use for a game.

Original: The maximum action space for Atari games is
Revised: The maximum action space for Atari games is

Original: [Table] [Table]
Revised: [Table] [Table]

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

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

Revised: The

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

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

Revised: We find

Revised: BC (finetune) underperforms the EfficientZero baseline.

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

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

Original: To better understand the feasibility of cross-task transfer, we show sample trajectory from our experiment games.
Revised: To better understand the feasibility of cross-task transfer, we show sample trajectory from our experiment games.

Original: Figure12 and 13  include all games involved in Table 1 (tasks sharing similar game mechanics).
Revised: Figure 13 and 14 include all games involved in Table 1 (tasks sharing similar game mechanics).

Original: Figure 14 include pretraining/finetuning games involved in Table 2 (tasks sharing diverse game mechanics) excluding ones that are shown in Figure 12 and Figure 13.
Revised: 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.

Original: [Figure]
Revised: [Figure]

Revised: −→

Original: A li e n
Revised: A li e n

Original: A m i d a r
Revised: A m 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
Revised: W o r

Original: [Figure]
Revised: [Figure]

Revised: −→

Original: A s s a u lt
Revised: A s s 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[Figure] 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