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

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

Original: Recently, model-based RL algorithms have greatly improved sample-efficiency by concurrently learning an internal model of the world, and supplementing real environment interactions with imagined rollouts for policy improvement.
Revised: Recently, model-based RL algorithms have greatly improved sample-efficiency by concurrently learning an internal model of the world, and supplementing real environment interactions with imagined rollouts for policy improvement.

Original: However, learning an effective model of the world from scratch is challenging, and in stark contrast to humans that rely heavily on world understanding and visual cues for learning new skills.
Revised: However, learning an effective model of the world from scratch is challenging, and in stark contrast to humans that rely heavily on world understanding and visual cues for learning new skills.

Original: In this work, we investigate whether internal models learned by modern model-based RL algorithms can be leveraged to solve new, distinctly different tasks faster.
Revised: In this work, we investigate whether internal models learned by modern model-based RL algorithms can be leveraged to solve new, distinctly different tasks faster.

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

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

Original: By offline multi-task pretraining and online cross-task finetuning, we achieve substantial improvements over a baseline trained from scratch; we improve mean performance of model-based algorithm
Revised: By offline multi-task pretraining and online cross-task finetuning, we achieve substantial improvements over a baseline trained from scratch; we improve mean performance of model-based algorithm

Original: EfficientZero by 23%, and by as much as 71% in some instances.
Revised: EfficientZero by 23%, and by as much as 71% in some instances.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: [Figure]
Revised: [Figure]

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

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

Original: However, traditional RL algorithms often suffer from poor sample-efficiency and require millions (or even billions) of environment interactions to solve tasks – especially when learning from high-dimensional observations such as images.
Revised: However, traditional RL algorithms often suffer from poor sample-efficiency and require millions (or even billions) of environment interactions to solve tasks – especially when learning from high-dimensional observations such as images.

Original: This is in stark contrast to humans that have a remarkable ability to quickly learn new skills despite very limited exposure (Dubey et al., 2018).
Revised: This is in stark contrast to humans that have a remarkable ability to quickly learn new skills despite very limited exposure (Dubey et al., 2018).

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

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

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

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

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

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

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

Original: This achievement is attributed, in part, to the algorithm concurrently learning an internal model of the environment from interaction, and using the learned model to imagine (simulate) further interactions for planning and policy improvement, thus reducing reliance on real environment interactions for skill acquisition.
Revised: This achievement is attributed, in part, to the algorithm concurrently learning an internal model of the environment from interaction, and using the learned model to imagine (simulate) further interactions for planning and policy improvement, thus reducing reliance on real environment interactions for skill acquisition.

Original: However, current RL algorithms, including EfficientZero, are still predominantly assumed to learn both perception, model, and skills
Revised: However, current RL algorithms, including EfficientZero, are still predominantly assumed to learn both perception, model, and skills

Original: [Figure]
Revised: [Figure]

Original: tabula rasa (from scratch) for each new task.
Revised: tabula rasa (from scratch) for each new task.

Original: Conversely, humans rely heavily on prior knowledge and visual cues when learning new skills – a study found that human players easily identify visual cues about game mechanics when exposed to a new game, and that human performance is severely degraded if such cues are removed or conflict with prior experiences (Dubey et al., 2018).
Revised: Conversely, humans rely heavily on prior knowledge and visual cues when learning new skills – a study found that human players easily identify visual cues about game mechanics when exposed to a new game, and that human performance is severely degraded if such cues are removed or conflict with prior experiences (Dubey et al., 2018).

Original: In related areas such as computer vision and natural language processing, large-scale unsupervised/self-supervised/supervised pretraining on large-scale datasets (Devlin et al., 2019;
Revised: In related areas such as computer vision and natural language processing, large-scale unsupervised/self-supervised/supervised pretraining on large-scale datasets (Devlin et al., 2019;

Original: Brown et al., 2020; Li et al., 2022; Radford et al., 2021; Chowdhery et al., 2022) has emerged as a powerful framework for solving numerous downstream tasks with few samples (Alayrac et al., 2022).
Revised: Brown et al., 2020; Li et al., 2022; Radford et al., 2021; Chowdhery et al., 2022) has emerged as a powerful framework for solving numerous downstream tasks with few samples (Alayrac et al., 2022).

Original: This pretraining paradigm has recently been extended to visuo-motor control in various forms, e.g. , by leveraging frozen (no finetuning) pretrained representations (Xiao et al., 2022; Parisi et al., 2022) or by finetuning in a supervised setting (Reed et al., 2022; Lee et al., 2022).
Revised: This pretraining paradigm has recently been extended to visuo-motor control in various forms, e.g. , by leveraging frozen (no finetuning) pretrained representations (Xiao et al., 2022; Parisi et al., 2022) or by finetuning in a supervised setting (Reed et al., 2022; Lee et al., 2022).

Original: However, the success of finetuning for online RL has mostly been limited to same-task initialization of model-free policies from offline datasets (Wang et al., 2022; Zheng et al., 2022), or adapting policies to novel instances of a given task (Mishra et al., 2017; Julian et al., 2020; Hansen et al., 2021a), with prior work citing high-variance objectives and catastrophical forgetting as the main obstacles to finetuning representations with RL (Bodnar et al., 2020; Xiao et al., 2022).
Revised: However, the success of finetuning for online RL has mostly been limited to same-task initialization of model-free policies from offline datasets (Wang et al., 2022; Zheng et al., 2022), or adapting policies to novel instances of a given task (Mishra et al., 2017; Julian et al., 2020; Hansen et al., 2021a), with prior work citing high-variance objectives and catastrophical forgetting as the main obstacles to finetuning representations with RL (Bodnar et al., 2020; Xiao et al., 2022).

Original: In this work, we explore whether such positive transfer can be induced with current model-based RL algorithms in an online RL setting, and across markedly distinct tasks.
Revised: In this work, we explore whether such positive transfer can be induced with current model-based RL algorithms in an online RL setting, and across markedly distinct tasks.

Original: Specifically, we seek to answer the following questions: when and how can a model-based RL algorithm such as EfficientZero benefit from pretraining on a diverse set of tasks?
Revised: Specifically, we seek to answer the following questions: when and how can a model-based RL algorithm such as EfficientZero benefit from pretraining on a diverse set of tasks?

Original: We base our experiments on the ALE due to cues that are easily identifiable to humans despite great diversity in tasks, and identify two key ingredients – cross-task finetuning and task alignment – for model-based adaptation that improve sample-efficiency substantially compared to models learned tabula rasa.
Revised: We base our experiments on the ALE due to cues that are easily identifiable to humans despite great diversity in tasks, and identify two key ingredients – cross-task finetuning and task alignment – for model-based adaptation that improve sample-efficiency substantially compared to models learned tabula rasa.

Original: In comparison, we find that a naïve treatment of the finetuning procedure as commonly used in supervised learning (Pan & Yang, 2010; Doersch et al., 2015;
Revised: In comparison, we find that a naïve treatment of the finetuning procedure as commonly used in supervised learning (Pan & Yang, 2010; Doersch et al., 2015;

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

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

Original: Based on our findings, we propose Model-Based Cross -Task Tra nsfer ( XTRA ), a framework for sample-efficient online RL with scalable pretraining and finetuning of learned world models using extra, auxiliary data from other tasks (see Figure 2).
Revised: Based on our findings, we propose Model-Based Cross -Task Tra nsfer ( XTRA ), a framework for sample-efficient online RL with scalable pretraining and finetuning of learned world models using extra, auxiliary data from other tasks (see Figure 2).

Original: Concretely, our framework consists of two stages: (i) offline multi-task pretraining of a world model on an offline dataset from m diverse tasks, a (ii) finetuning stage where the world model is jointly finetuned on a target task in addition to m offline tasks.
Revised: Concretely, our framework consists of two stages: (i) offline multi-task pretraining of a world model on an offline dataset from m diverse tasks, a (ii) finetuning stage where the world model is jointly finetuned on a target task in addition to m offline tasks.

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

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

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

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

Original: From Figure 1 and Table 1, we observe 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 that XTRA improves sample-efficiency substantially across most tasks, improving mean and median performance of EfficientZero by 23% and 25% , respectively.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Problem setting.
Revised: Problem setting.

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

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

Original: (1998)) defined by the tuple M =
Revised: (1998)) defined by the tuple M =

Original: ⟨O , A , P , ρ, r, γ ⟩ , where O is the observation space (pixels), A is the action space, P : O×A (cid:55)→ O is a transition function, ρ is the initial state distribution, r : O × A
Revised: ⟨O , A , P , ρ, r, γ ⟩ , where O is the observation space (pixels), A is the action space, P : O×A (cid:55)→ O is a transition function, ρ is the initial state distribution, r : O × A

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

Original: As is standard practice in ALE (Bellemare et al., 2013), we convert M to a fully observable Markov Decision Process (MDP; Bellman (1957)) by approximating state s t ∈ S at time t as a stack of frames s t .
Revised: As is standard practice in ALE (Bellemare et al., 2013), we convert M to a fully observable Markov Decision Process (MDP; Bellman (1957)) by approximating state s t ∈ S at time t as a stack of frames s t .

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

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

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

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

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

Original: =1 γ t r ( s t , a t )] where a t ∼ π θ ( a | s ) , s t ∼ P ( s t , a t ) , s 0 ∼ ρ , and T is the episode horizon.
Revised: =1 γ t r ( s t , a t )] where a t ∼ π θ ( a | s ) , s t ∼ P ( s t , a t ) , s 0 ∼ ρ , and T is the episode horizon.

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

Original: To obtain a good policy from minimal environment interaction, we learn a "world model" from interaction data and use the learned model for action search.
Revised: To obtain a good policy from minimal environment interaction, we learn a "world model" from interaction data and use the learned model for action search.

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

Original: Then, we seek to first obtain a good parameter initialization θ that allows us to solve task M using fewer interactions (samples) than training from scratch, i.e. , we wish to improve the sample-efficiency of online RL.
Revised: Then, we seek to first obtain a good parameter initialization θ that allows us to solve task M using fewer interactions (samples) than training from scratch, i.e. , we wish to improve the sample-efficiency of online RL.

Original: We do so by first pretraining the model on an offline (fixed) dataset that consists of transitions ( s , a , r, s ′ ) collected by unknown behavior policies in m environments { ˜ M i | ˜
Revised: We do so by first pretraining the model on an offline (fixed) dataset that consists of transitions ( s , a , r, s ′ ) collected by unknown behavior policies in m environments { ˜ M i | ˜

Original: M i ̸
Revised: M i ̸

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

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: (reward)
Revised: (reward)

Original: Zh f g f
Revised: Zh f g f

Original: Z
Revised: Z

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 =
Revised: = h θ ( s t ) , and predicts future latent states z t +1 and instantaneous rewards ˆ r t using an action-conditional latent dynamics function z t +1 , ˆ r t =

Original: g θ
Revised: g θ

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

Original: For each latent state, a prediction network f θ estimates a probability distribution ˆ p over (valid) actions a ∈ A , as well as the expected state value ˆ v of the given state, i.e. , ˆ v t , ˆ p t = f θ ( z t ) .
Revised: For each latent state, a prediction network f θ estimates a probability distribution ˆ p over (valid) actions a ∈ A , as well as the expected state value ˆ v of the given state, i.e. , ˆ v t , ˆ p t = f θ ( z t ) .

Original: Intuitively, h θ and g θ allow EfficientZero to search for actions entirely in its latent space before executing actions in the real environment, and f θ predicts quantities that help guide the search towards high-return action sequences.
Revised: Intuitively, h θ and g θ allow EfficientZero to search for actions entirely in its latent space before executing actions in the real environment, and f θ predicts quantities that help guide the search towards high-return action sequences.

Original: Concretely, ˆ v provides a return estimate for nodes at the lookahead horizon (as opposed to truncating the cumulative sum of expected rewards) and ˆ p provides an action distribution prior that helps guide the search.
Revised: Concretely, ˆ v provides a return estimate for nodes at the lookahead horizon (as opposed to truncating the cumulative sum of expected rewards) and ˆ p provides an action distribution prior that helps guide the search.

Original: We describe EfficientZero’s learning objective between model prediction ( ˆ p, ˆ v, ˆ r ) and quantity targets ( π, z, u ) in Appendix A.
Revised: We describe EfficientZero’s learning objective between model prediction ( ˆ p, ˆ v, ˆ r ) and quantity targets ( π, z, u ) in Appendix A.

Original: EfficientZero improves the sample-efficiency of MuZero by introducing additional auxiliary losses during training.
Revised: 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 those for continuous action spaces (Hafner et al., 2019a; Hansen et al., 2022a).
Revised: We adopt EfficientZero as our backbone model and learning algorithm, but emphasize that our framework is applicable to most model-based algorithms, including those for continuous action spaces (Hafner et al., 2019a; Hansen et al., 2022a).

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

Original: We propose Model-Based Cross -Task Tra nsfer ( XTRA ), a two-stage framework for offline multi-task pretraining and cross-task transfer of learned world models by finetuning with online RL.
Revised: We propose Model-Based Cross -Task Tra nsfer ( XTRA ), a two-stage framework for offline multi-task pretraining and cross-task transfer of learned world models by finetuning with online RL.

Original: Specifically, we first pretrain a world model on offline data from a set of diverse pretraining tasks, and then iteratively finetune the pretrained model on data from a target task collected by online interaction.
Revised: Specifically, we first pretrain a world model on offline data from a set of diverse pretraining tasks, and then iteratively finetune the pretrained model on data from a target task collected by online interaction.

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

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

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

Original: We emphasize that the goal of pretraining is not to obtain a truly generalist agent, but rather to learn a good initialization for finetuning to unseen tasks.
Revised: We emphasize that the goal of pretraining is not to obtain a truly generalist agent, but rather to learn a good initialization for finetuning to unseen tasks.

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

Original: To address such a challenge,
Revised: To address such a challenge,

Original: [Figure]
Revised: [Figure]

Original: we propose to pretrain the model following a student-teacher training setup in the same spirit to DQN multi-task policy distillation in Rusu et al.
Revised: we propose to pretrain the model following a student-teacher training setup in the same spirit to DQN multi-task policy distillation in Rusu et al.

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

Original: (2016), where teacher models are trained separately by offline RL for each task, and then distilled into a single multi-task model using a novel instantiation of the MuZero Reanalyze (Schrittwieser et al., 2021).
Revised: (2016), where teacher models are trained separately by offline RL for each task, and then distilled into a single multi-task model using a novel instantiation of the MuZero Reanalyze (Schrittwieser et al., 2021).

Original: For each pretraining task we assume access to a fixed dataset { ˜ D i | 1 ≤ i ≤ m } that consists of trajectories from an unknown (and potentially sub-optimal) behavior policy.
Revised: For each pretraining task we assume access to a fixed dataset { ˜ D i | 1 ≤ i ≤ m } that consists of trajectories from an unknown (and potentially sub-optimal) behavior policy.

Original: Importantly, we do not make any assumptions about the quality or the source of trajectories in the dataset, i.e. , we do not assume datasets to consist of expert trajectories.
Revised: Importantly, we do not make any assumptions about the quality or the source of trajectories in the dataset, i.e. , we do not assume datasets to consist of expert trajectories.

Original: We first train individual EfficientZero teacher models on each dataset for a fixed number of iterations in a single-task (offline) RL setting, resulting in m teacher models { ˜ π iψ | 1 ≤ i ≤ m } .
Revised: We first train individual EfficientZero teacher models on each dataset for a fixed number of iterations in a single-task (offline) RL setting, resulting in m teacher models { ˜ π iψ | 1 ≤ i ≤ m } .

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

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

Original: Specifically, we optimize the student policy by sampling data uniformly from all pretraining tasks, and use value/policy targets generated by their respective teacher models rather than bootstrapping from student predictions as commonly done in the (single-task)
Revised: Specifically, we optimize the student policy by sampling data uniformly from all pretraining tasks, and use value/policy targets generated by their respective teacher models rather than bootstrapping from student predictions as commonly done in the (single-task)

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

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

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

Original: Although teacher models may not be optimal depending on the provided offline datasets, we find that they provide stable (due to fixed parameters during distillation) targets of sufficiently good quality.
Revised: Although teacher models may not be optimal depending on the provided offline datasets, we find that they provide stable (due to fixed parameters during distillation) targets of sufficiently good quality.

Original: The simpler alternative – training a multi-task model on all m pretraining tasks simultaneously using RL is found to not scale beyond a couple of tasks in practice, as we will demonstrate our experiments in Appendix C. After distilling a multi-task student model, we now have a single set of pretrained parameters that can be used for finetuning to a variety of tasks via online interaction, which we introduce in the following section.
Revised: The simpler alternative – training a multi-task model on all m pretraining tasks simultaneously using RL is found to not scale beyond a couple of tasks in practice, as we will demonstrate our experiments in Appendix C. After distilling a multi-task student model, we now have a single set of pretrained parameters that can be used for finetuning to a variety of tasks via online interaction, which we introduce in the following section.

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

Original: In this stage, we iteratively interact with a target task (environment) to collect interaction data, and finetune the pretrained model on data from the target task.
Revised: In this stage, we iteratively interact with a target task (environment) to collect interaction data, and finetune the pretrained model on data from the target task.

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

Original: Specifically, the initial sub-optimal data collected from the target task can cause a large perturbation in the original pretrained model parameters, ultimately erasing inductive priors learned during pretraining.
Revised: Specifically, the initial sub-optimal data collected from the target task can cause a large perturbation in the original pretrained model parameters, ultimately erasing inductive priors learned during pretraining.

Original: To overcome this challenge, we propose a concurrent cross-task learning strategy: we retain offline data from the pretraining stage, and concurrently finetune the model on both data from the target task, as well as data from the pretraining tasks.
Revised: To overcome this challenge, we propose a concurrent cross-task learning strategy: we retain offline data from the pretraining stage, and concurrently finetune the model on both data from the target task, as well as data from the pretraining tasks.

Original: While this procedure addresses catastrophical forgetting, interference between the target task and certain pretraining tasks can be harmful for the sample-efficiency in online RL.
Revised: While this procedure addresses catastrophical forgetting, interference between the target task and certain pretraining tasks can be harmful for the sample-efficiency in online RL.

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

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

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

Original: M i ̸
Revised: M i ̸

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

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

Original: [Equation]
Revised: [Equation]

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

Original: where L ez is the ordinary (single-task) EfficientZero objective (see Appendix A), and η i are dynamically (and independently) updated task weights for each of the m pretraining tasks.
Revised: where L ez is the ordinary (single-task) EfficientZero objective (see Appendix A), and η i are dynamically (and independently) updated task weights for each of the m pretraining tasks.

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

Original: During online finetuning, we use distillation targets from teachers obtained from each pretraining game, and use MuZero Reanalyze to compute targets for the target task for which we have no teacher available.
Revised: During online finetuning, we use distillation targets from teachers obtained from each pretraining game, and use MuZero Reanalyze to compute targets for the target task for which we have no teacher available.

Original: In order to dynamically re-weight task weights η i throughout the training process, we break down the total number of environment steps ( i.e. , 100k in our experiments) into even T -step cycles (intervals).
Revised: In order to dynamically re-weight task weights η i throughout the training process, we break down the total number of environment steps ( i.e. , 100k in our experiments) into even T -step cycles (intervals).

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

Original: ˜
Revised: ˜

Original: M i .
Revised: M i .

Original: The new η i will then be fixed during the remaining T − N steps in the current cycle and the first N steps in the next cycle.
Revised: The new η i will then be fixed during the remaining T − N steps in the current cycle and the first N steps in the next cycle.

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

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

Original: Inspired by the conflicting gradients measurement for multi-task learning in Yu et al.
Revised: Inspired by the conflicting gradients measurement for multi-task learning in Yu et al.

Original: (2020), we compute the cosine similarity between loss gradients ˜
Revised: (2020), we compute the cosine similarity between loss gradients ˜

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

Original: [Equation]
Revised: [Equation]

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

Original: =
Revised: =

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

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

Original: G in and
Revised: G in and

Original: G n .
Revised: 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.
Revised: N steps, our proposed technique mitigates gradient conflicts at a negligible computational cost in contrast to the compute-intensive gradient modification method proposed in Yu et al.

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

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

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

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

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

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

Original: [Table]
Revised: [Table]

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: • When can we empirically expect finetuning to be successful?
Revised: • When can we empirically expect finetuning to be successful?

Original: Experimental setup.
Revised: Experimental setup.

Original: We base our architecture and backbone learning algorithm on EfficientZero (Ye et al., 2021) and focus our efforts on the pretraining and finetuning aspects of our problem setting.
Revised: We base our architecture and backbone learning algorithm on EfficientZero (Ye et al., 2021) and focus our efforts on the pretraining and finetuning aspects of our problem setting.

Original: We consider EfficientZero with two different network sizes to better position our results: (i) the same network architecture as in the original EfficientZero implementation which we simply refer to as EfficientZero , and (ii) a larger variant with 4 times more parameters in the representation network (denoted EfficientZero-L ).
Revised: We consider EfficientZero with two different network sizes to better position our results: (i) the same network architecture as in the original EfficientZero implementation which we simply refer to as EfficientZero , and (ii) a larger variant with 4 times more parameters in the representation network (denoted EfficientZero-L ).

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

Original: However, we find that our EfficientZero baseline generally does not benefit from a larger architecture, and we thus include both variants for a fair comparison.
Revised: However, we find that our EfficientZero baseline generally does not benefit from a larger architecture, and we thus include both variants for a fair comparison.

Original: We experiment with cross-task transfer on three subsets of tasks: tasks that share similar game mechanics (for which we consider two Shooter and Maze categories), and tasks that have no discernible properties in common (referred to as Diverse ).
Revised: We experiment with cross-task transfer on three subsets of tasks: tasks that share similar game mechanics (for which we consider two Shooter and Maze categories), and tasks that have no discernible properties in common (referred to as Diverse ).

Original: We measure performance on individual Atari games by absolute scores, and also provide aggregate results as measured by mean and median scores across games, normalized by human performance or EfficientZero performance at 100k environment steps.
Revised: We measure performance on individual Atari games by absolute scores, and also provide aggregate results as measured by mean and median scores across games, normalized by human performance or EfficientZero performance at 100k environment steps.

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

Original: We provide details on our pretraining dataset in Appendix F.
Revised: We provide details on our pretraining dataset in Appendix F.

Original: Baselines.
Revised: Baselines.

Original: We compare our method against 7 prior methods for online RL that represent the state-ofthe-art on the Atari100k benchmark (including EfficientZero), a multi-task behavior cloning policy pretrained on the same offline data as our method does for zero-shot performance on the target task and the performance after finetuning on the target task (see Appendix O for details), and a direct comparison to CURL (Srinivas et al., 2020), a strong model-free RL baseline, under an offline pretraining + online finetuning setting.
Revised: We compare our method against 7 prior methods for online RL that represent the state-ofthe-art on the Atari100k benchmark (including EfficientZero), a multi-task behavior cloning policy pretrained on the same offline data as our method does for zero-shot performance on the target task and the performance after finetuning on the target task (see Appendix O for details), and a direct comparison to CURL (Srinivas et al., 2020), a strong model-free RL baseline, under an offline pretraining + online finetuning setting.

Original: We also include a set of ablations that include EfficientZero with several different model sizes and pretraining/finetuning schemes.
Revised: We also include a set of ablations that include EfficientZero with several different model sizes and pretraining/finetuning schemes.

Original: The former baselines serve to position our results with respect to the state-of-the-art, and the latter baselines and ablations serve to shed light on the key ingredients for successful multi-task pretraining and finetuning.
Revised: The former baselines serve to position our results with respect to the state-of-the-art, and the latter baselines and ablations serve to shed light on the key ingredients for successful multi-task pretraining and finetuning.

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

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

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

Original: Tasks
Revised: Tasks

Original: with
Revised: with

Original: similar game mechanics
Revised: similar game mechanics

Original: We first
Revised: We first

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

Original: We select
Revised: We select

Original: 5 shooter games and
Revised: 5 shooter games and

Original: 5 maze games for this experiment
Revised: 5 maze games for this experiment

Original: Results for our method, baselines and a set of ablations on the Atari100k benchmark
Revised: Results for our method, baselines and a set of ablations on the Atari100k benchmark

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: are shown in Table 1.
Revised: are shown in Table 1.

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

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

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

Original: We also consider three ablations: (1) XTRA without cross-task: a variant of our method that naively finetunes the pretrained model without any additional offline data from pretraining tasks during finetuning, (2) XTRA without pretraining: a variant that uses our concurrent cross-task learning ( i.e., leverages offline data during finetuning ) but is initialized with random parameters (no pretraining), and finally (3) XTRA without task weights: a variant that uses constant weights of 1 for all task loss terms during finetuning.
Revised: We also consider three ablations: (1) XTRA without cross-task: a variant of our method that naively finetunes the pretrained model without any additional offline data from pretraining tasks during finetuning, (2) XTRA without pretraining: a variant that uses our concurrent cross-task learning ( i.e., leverages offline data during finetuning ) but is initialized with random parameters (no pretraining), and finally (3) XTRA without task weights: a variant that uses constant weights of 1 for all task loss terms during finetuning.

Original: We find that XTRA achieves extremely high performance on 2 games (DemonAttack and Phoenix) without dynamic task weights, improving over EfficientZero by as much as 343 % on DemonAttack.
Revised: We find that XTRA achieves extremely high performance on 2 games (DemonAttack and Phoenix) without dynamic task weights, improving over EfficientZero by as much as 343 % on DemonAttack.

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

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

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

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

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

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

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

Original: We find that XTRA advances the state-of-the-art in a majority of tasks on the Atari100k benchmark, and achieve a mean human-normalized score of 187 % vs. 129 % for the previous SOTA, EfficientZero.
Revised: We find that XTRA advances the state-of-the-art in a majority of tasks on the Atari100k benchmark, and achieve a mean human-normalized score of 187 % vs. 129 % for the previous SOTA, EfficientZero.

Original: We perform the same set of the ablations as we do for tasks with similar game mechanics with XTRA, and the results are shown in Table 6 from Appendix H. Additionally, we include an ablation that examines the effect of the number of pretrained tasks on later finetuning performance.
Revised: We perform the same set of the ablations as we do for tasks with similar game mechanics with XTRA, and the results are shown in Table 6 from Appendix H. Additionally, we include an ablation that examines the effect of the number of pretrained tasks on later finetuning performance.

Original: Details and results for this ablation are shown in Table 7 from Appendix I.
Revised: Details and results for this ablation are shown in Table 7 from Appendix I.

Original: Model-free comparisons.
Revised: Model-free comparisons.

Original: For both settings (e.g., tasks with similar & diverse game mechanics), we also compare our framework with a strong model-free baseline, CURL (Srinivas et al., 2020), where
Revised: For both settings (e.g., tasks with similar & diverse game mechanics), we also compare our framework with a strong model-free baseline, CURL (Srinivas et al., 2020), where

Original: CURL is pretrained on the same pretraining tasks as XTRA is, and later finetuned to each of the target tasks.
Revised: CURL is pretrained on the same pretraining tasks as XTRA is, and later finetuned to each of the target tasks.

Original: We find that pretraining does not improve the performance of this model-free baseline as consistently as for our model-based framework, XTRA, under both settings.
Revised: We find that pretraining does not improve the performance of this model-free baseline as consistently as for our model-based framework, XTRA, under both settings.

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

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

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

Original: While our previous experiments established that XTRA benefits from pretraining even when games are markedly different, we now seek to better quantify the importance of task relevance.
Revised: While our previous experiments established that XTRA benefits from pretraining even when games are markedly different, we now seek to better quantify the importance of task relevance.

Original: We compare the online finetuning performance of XTRA in two different settings: (1) pretraining on 4 shooter games and finetuning to 5 new shooter games, and (2) pretraining on 4 maze games and finetuning to the same 5 shooter games.
Revised: We compare the online finetuning performance of XTRA in two different settings: (1) pretraining on 4 shooter games and finetuning to 5 new shooter games, and (2) pretraining on 4 maze games and finetuning to the same 5 shooter games.

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

Original: Unsurprisingly, we observe that offline pretraining and concurrently learning from other shooter games significantly benefit the online target shooter games through training, with particularly large gains early in training.
Revised: Unsurprisingly, we observe that offline pretraining and concurrently learning from other shooter games significantly benefit the online target shooter games through training, with particularly large gains early in training.

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

Original: This result indicates that (1) selecting pretraining tasks that are relevant to the target task is key to benefit from pretraining, and (2) in the extreme case where there are no pretraining tasks relevant to the target task, finetuning with XTRA generally does not harm the online RL performance since it can automatically assign small weights to the pretraining tasks.
Revised: This result indicates that (1) selecting pretraining tasks that are relevant to the target task is key to benefit from pretraining, and (2) in the extreme case where there are no pretraining tasks relevant to the target task, finetuning with XTRA generally does not harm the online RL performance since it can automatically assign small weights to the pretraining tasks.

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

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

Original: We answer this question by only transferring a subset of the different model components – representation h , dynamics function g , and prediction head f – to the online finetuning stage and simply using a random initialization for the remaining components.
Revised: We answer this question by only transferring a subset of the different model components – representation h , dynamics function g , and prediction head f – to the online finetuning stage and simply using a random initialization for the remaining components.

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

Original: Interestingly, we find that only transferring the pretrained representation h to the online RL stage only improves slightly over learning from scratch, especially in the early stages of online RL.
Revised: Interestingly, we find that only transferring the pretrained representation h to the online RL stage only improves slightly over learning from scratch, especially in the early stages of online RL.

Original: In comparison, loading both the pretrained representation and dynamics function accounts for the majority of the gains in XTRA, whereas loading the prediction heads has no significant impact on sample-efficiency (but matters for zero-shot performance).
Revised: In comparison, loading both the pretrained representation and dynamics function accounts for the majority of the gains in XTRA, whereas loading the prediction heads has no significant impact on sample-efficiency (but matters for zero-shot performance).

Original: We conjecture that this is because learning a good dynamics function is relatively more difficult from few samples than learning a task-specific visual representation, and that the prediction head accounts for only a small amount of the overall parameters in the model.
Revised: We conjecture that this is because learning a good dynamics function is relatively more difficult from few samples than learning a task-specific visual representation, and that the prediction head accounts for only a small amount of the overall parameters in the model.

Original: Finally, we hypothesize that the visual representation learned during pretraining will be susceptible to distribution shifts as it is transferred to an unseen target task.
Revised: Finally, we hypothesize that the visual representation learned during pretraining will be susceptible to distribution shifts as it is transferred to an unseen target task.

Original: To verify this hypothesis, we consider an additional experiment where we transfer all components to new tasks, but freeze the representation h during finetuning, i.e. , it remains fixed.
Revised: To verify this hypothesis, we consider an additional experiment where we transfer all components to new tasks, but freeze the representation h during finetuning, i.e. , it remains fixed.

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

Original: We find that, although this variant of our framework improves over training from scratch in the early stages of training, the frozen representation eventually hinders the model from converging to a good model, which is consistent with observations made in (supervised) imitation learning (Parisi et al., 2022).
Revised: We find that, although this variant of our framework improves over training from scratch in the early stages of training, the frozen representation eventually hinders the model from converging to a good model, which is consistent with observations made in (supervised) imitation learning (Parisi et al., 2022).

Original: [Figure]
Revised: [Figure]

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

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

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

Original: Since dynamics network g and prediction network f are used in MCTS search, increasing the parameter counts for these two networks would increase inference/training time complexity significantly.
Revised: Since dynamics network g and prediction network f are used in MCTS search, increasing the parameter counts for these two networks would increase inference/training time complexity significantly.

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

Original: We compare the performance of our method and EfficientZero trained from scratch with each of our two model sizes, the original EfficientZero architecture and a larger variant (denoted EfficientZero -L ); results are shown in Figure 6 (c) .
Revised: We compare the performance of our method and EfficientZero trained from scratch with each of our two model sizes, the original EfficientZero architecture and a larger variant (denoted EfficientZero -L ); results are shown in Figure 6 (c) .

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

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

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

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

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

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

Original: We therefore envision that cross-task pretraining could benefit many real-world applications of RL, where environment interactions are typically constrained due to physical constraints.
Revised: We therefore envision that cross-task pretraining could benefit many real-world applications of RL, where environment interactions are typically constrained due to physical constraints.

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

Original: Based on Table 1 and 2, we conclude that cross-task transfer with model-based RL is feasible.
Revised: Based on Table 1 and 2, we conclude that cross-task transfer with model-based RL is feasible.

Original: Further,
Revised: Further,

Original: Figure 6 (a) shows that our XTRA framework benefits from online finetuning when pretraining tasks are relevant, and both representation and dynamics networks contribute to its success (Figure 7).
Revised: Figure 6 (a) shows that our XTRA framework benefits from online finetuning when pretraining tasks are relevant, and both representation and dynamics networks contribute to its success (Figure 7).

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

Original: Pretrained representations are widely used to improve downstream performance in learning tasks with limited data or resources available, and have been adopted across a multitude of areas such as computer vision (Girshick et al., 2014;
Revised: Pretrained representations are widely used to improve downstream performance in learning tasks with limited data or resources available, and have been adopted across a multitude of areas such as computer vision (Girshick et al., 2014;

Original: Doersch et al., 2015; He et al., 2020), natural language processing (Devlin et al., 2019; Brown et al., 2020), and audio (van den Oord et al., 2018).
Revised: Doersch et al., 2015; He et al., 2020), natural language processing (Devlin et al., 2019; Brown et al., 2020), and audio (van den Oord et al., 2018).

Original: By first learning a good representation on a large dataset, the representation can quickly be finetuned with, e.g. , supervised learning on a small labelled dataset to solve a given task (Pan & Yang, 2010).
Revised: By first learning a good representation on a large dataset, the representation can quickly be finetuned with, e.g. , supervised learning on a small labelled dataset to solve a given task (Pan & Yang, 2010).

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

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

Original: (2020) show that a generative model trained on large-scale natural language data can be used to solve unseen tasks given only a few examples.
Revised: (2020) show that a generative model trained on large-scale natural language data can be used to solve unseen tasks given only a few examples.

Original: While this is a common paradigm for problems that can be cast as (self-)supervised learning problems, it has seen comparably little adoption in RL literature.
Revised: While this is a common paradigm for problems that can be cast as (self-)supervised learning problems, it has seen comparably little adoption in RL literature.

Original: This discrepancy is, in part, attributed to optimization challenges in RL (Bodnar et al., 2020; Hansen et al., 2021b; Xiao et al., 2022; Wang et al., 2022), as well as a lack of large-scale datasets that capture both the visual, temporal, and control-relevant (actions, rewards) properties of RL (Hansen et al., 2022b).
Revised: This discrepancy is, in part, attributed to optimization challenges in RL (Bodnar et al., 2020; Hansen et al., 2021b; Xiao et al., 2022; Wang et al., 2022), as well as a lack of large-scale datasets that capture both the visual, temporal, and control-relevant (actions, rewards) properties of RL (Hansen et al., 2022b).

Original: In this work, we show that – despite these challenges – modern model-based RL algorithms can still benefit substantially from pretraining on multi-task datasets, but require a more careful treatment during finetuning.
Revised: In this work, we show that – despite these challenges – modern model-based RL algorithms can still benefit substantially from pretraining on multi-task datasets, but require a more careful treatment during finetuning.

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

Original: Improving the sample-efficiency of visual RL algorithms is a long-standing problem and has been approached from many – largely orthogonal – perspectives, including representation learning (Kulkarni et al., 2019;
Revised: Improving the sample-efficiency of visual RL algorithms is a long-standing problem and has been approached from many – largely orthogonal – perspectives, including representation learning (Kulkarni et al., 2019;

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

Original: Kostrikov et al., 2021; Hansen et al., 2021b), bootstrapping from demonstrations (Zhan et al., 2020) or offline datasets (Wang et al., 2022; Zheng et al., 2022;
Revised: Kostrikov et al., 2021; Hansen et al., 2021b), bootstrapping from demonstrations (Zhan et al., 2020) or offline datasets (Wang et al., 2022; Zheng et al., 2022;

Original: Baker et al., 2022), using pretrained visual representations for model-free RL (Shah & Kumar, 2021; Xiao et al., 2022; Ze et al., 2022), and model-based RL (Ha & Schmidhuber, 2018; Finn &
Revised: Baker et al., 2022), using pretrained visual representations for model-free RL (Shah & Kumar, 2021; Xiao et al., 2022; Ze et al., 2022), and model-based RL (Ha & Schmidhuber, 2018; Finn &

Original: Levine, 2017; Nair et al., 2018; Hafner et al., 2019b; Kaiser et al., 2020; Schrittwieser et al., 2020;
Revised: Levine, 2017; Nair et al., 2018; Hafner et al., 2019b; Kaiser et al., 2020; Schrittwieser et al., 2020;

Original: Hafner et al., 2021;
Revised: Hafner et al., 2021;

Original: Ye et al., 2021;
Revised: Ye et al., 2021;

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

Original: Hansen et al., 2023).
Revised: Hansen et al., 2023).

Original: We choose to focus our efforts on sample-efficiency from the perspective of pretraining in a model-based context, i.e. , jointly learning perception and dynamics.
Revised: We choose to focus our efforts on sample-efficiency from the perspective of pretraining in a model-based context, i.e. , jointly learning perception and dynamics.

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

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

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

Original: Our work differs from prior work in that we show it is possible to pretrain and finetune both the representation and the dynamics using model-based RL.
Revised: Our work differs from prior work in that we show it is possible to pretrain and finetune both the representation and the dynamics using model-based RL.

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

Original: Gradient-based finetuning is a well-studied technique for adaptation in (predominantly model-free) RL, and has been used to adapt to either changes in visuals or dynamics (Mishra et al., 2017; Yen-Chen et al., 2019; Duan et al., 2016; Julian et al., 2020; Hansen et al., 2021a; Bodnar et al., 2020; Wang et al., 2022; Ze et al., 2022), or task specification (Xie & Finn, 2021; Walke et al., 2022).
Revised: Gradient-based finetuning is a well-studied technique for adaptation in (predominantly model-free) RL, and has been used to adapt to either changes in visuals or dynamics (Mishra et al., 2017; Yen-Chen et al., 2019; Duan et al., 2016; Julian et al., 2020; Hansen et al., 2021a; Bodnar et al., 2020; Wang et al., 2022; Ze et al., 2022), or task specification (Xie & Finn, 2021; Walke et al., 2022).

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

Original: (2020) shows that a model-free policy for robotic manipulation can adapt to changes in lighting and object shape by finetuning via rewards on a mixture of data from the new and old environment, and recover original performance in less than 800 trials.
Revised: (2020) shows that a model-free policy for robotic manipulation can adapt to changes in lighting and object shape by finetuning via rewards on a mixture of data from the new and old environment, and recover original performance in less than 800 trials.

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

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

Original: Other works show that pretraining with offline
Revised: Other works show that pretraining with offline

Original: 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: 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
Revised: (2022) shows that offline multi-task

Original: RL pretraining via sequence modelling can improve offline finetuning on data from unseen tasks.
Revised: 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: As a feasibility study, we hope that our empirical analysis and findings on cross-task transfer with model-based RL will inspire further research in this direction.
Revised: As a feasibility study, we hope that our empirical analysis and findings on cross-task transfer with model-based RL will inspire further research in this direction.

Original: Reproducibility Statement.
Revised: Reproducibility Statement.

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

Original: Code for reproducing our experiments is made available at https://nicklashansen.github.io/xtra .
Revised: Code for reproducing our experiments is made available at https://nicklashansen.github.io/xtra .

Original: Acknowledgements.
Revised: Acknowledgements.

Original: This work is supported by NSF Award IIS-2127544 and gifts from Qualcomm AI.
Revised: This work is supported by NSF Award IIS-2127544, NSF TILOS AI Institute, and gifts from Qualcomm AI.

Original: We  thank Weirui Ye, Wenlong Huang and Weijian Xu for helpful discussions
Revised: We would like to thank Weirui Ye, Wenlong Huang and Weijian Xu for helpful discussions

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

Original: XTRA uses the same learning objective as EfficientZero during both offline pretraining and online finetuning, except the quantity targets are predicted by the teacher model during distillation and concurrent learning for pretrained games instead of Muzero Reanalysis procedure.
Revised: XTRA uses the same learning objective as EfficientZero during both offline pretraining and online finetuning, except the quantity targets are predicted by the teacher model during distillation and concurrent learning for pretrained games instead of Muzero Reanalysis procedure.

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

Original: To warrant the latent dynamics that can mirror the true states of the environment, MuZero is trained to predict three necessary quantities directly relevant for planning: (1) the policy target π obtained from visit count distribution of the MCTS (2) immediate reward u from environment (3) bootstrapped value target z
Revised: To warrant the latent dynamics that can mirror the true states of the environment, MuZero is trained to predict three necessary quantities directly relevant for planning: (1) the policy target π obtained from visit count distribution of the MCTS (2) immediate reward u from environment (3) bootstrapped value target z

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

Original: k − 1 i =

Revised: k − 1 i =

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

Original: + k .
Revised: + k .

Original: On top of
Revised: On top of

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

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

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

Original: We refer reader to the original manuscripts for implementation details.
Revised: We refer reader to the original manuscripts for implementation details.

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

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

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

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

Original: [Equation]
Revised: [Equation]

Original: At every N steps
Revised: At every N steps

Original: [Figure]
Revised: [Figure]

Original: the new task weights η i are updated by η i =
Revised: the new task weights η i are updated by η i =

Original: s i N , and used in subsequent finetuning objectives according to Equation 1.
Revised: s i N , and used in subsequent finetuning objectives according to Equation 1.

Original: In practice
Revised: In practice

Original: [Figure]
Revised: [Figure]

Original: we start task weight updates at 10 k steps to ensure enough data from the online target task has been collected for a meaningful similarity measure.
Revised: we start task weight updates at 10 k steps to ensure enough data from the online target task has been collected for a meaningful similarity measure.

Original: All task weights are initialized as 1 for the first 10 k steps.
Revised: All task weights are initialized as 1 for the first 10 k steps.

Original: Figure 8 shows how task weights are adaptively adjusted by the model during 100k environment steps during online finetuning stage for 10 games reported in Figure 11.
Revised: Figure 8 shows how task weights are adaptively adjusted by the model during 100k environment steps during online finetuning stage for 10 games reported in Figure 11.

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

Original: [Figure]
Revised: [Figure]

Original: C D ISTILLATION VS . M ULTI -G AME O FFLINE RL
Revised: C D ISTILLATION VS . M ULTI -G AME O FFLINE RL

Original: Our method learns a multi-game world model from offline data via distillation of task-specific world models trained with offline RL.
Revised: Our method learns a multi-game world model from offline data via distillation of task-specific world models trained with offline RL.

Original: An alternative way to obtain such a multi-game model would be
Revised: An alternative way to obtain such a multi-game model would be

Original: [Figure]
Revised: [Figure]

Original: to directly train a single world model on the multi-game dataset with offline RL.
Revised: to directly train a single world model on the multi-game dataset with offline RL.

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

Original: A comparison between the two approaches on four different pretraining games is shown in Figure 10.
Revised: A comparison between the two approaches on four different pretraining games is shown in Figure 10.

Original: We observe that multi-game offline RL ( green ) achieves low scores in all four pretraining games, whereas our wold model obtained by distillation ( orange ) performs comparably to single-task world models ( blue ) in 3 out of 4 tasks.
Revised: We observe that multi-game offline RL ( green ) achieves low scores in all four pretraining games, whereas our wold model obtained by distillation ( orange ) performs comparably to single-task world models ( blue ) in 3 out of 4 tasks.

Original: [Figure]
Revised: [Figure]

Original: D S CORES FOR I NDIVIDUAL S EEDS
Revised: D S CORES FOR I NDIVIDUAL S EEDS

Original: Our results in Table 1 are aggregated across 5 seeds.
Revised: Our results in Table 1 are aggregated across 5 seeds.

Original: In Table 3, we report game scores for each individual seed, as well as the mean, median, and standard deviation of game scores for each game.
Revised: In Table 3, we report game scores for each individual seed, as well as the mean, median, and standard deviation of game scores for each game.

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

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

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

Original: ( score human − score random ) as in prior work.
Revised: ( score human − score random ) as in prior work.

Original: To the best of our knowledge, there are no human performance results for the Carnival game, and we therefore exclude this game from the aggregate Human Normalized Mean and Median
Revised: To the best of our knowledge, there are no human performance results for the Carnival game, and we therefore exclude this game from the aggregate Human Normalized Mean and Median

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

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

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

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

Original: [Figure]
Revised: [Figure]

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

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

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

Original: For each checkpoint, we evaluate the model performance on the Atari environment following the same procedure from EfficientZero and collect 64 trajectories.
Revised: For each checkpoint, we evaluate the model performance on the Atari environment following the same procedure from EfficientZero and collect 64 trajectories.

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

Original: Since trajectories are collected from model checkpoints both at the early and late training stage within the 120k training steps, the collected data does not necessarily come from an expert agent.
Revised: Since trajectories are collected from model checkpoints both at the early and late training stage within the 120k training steps, the collected data does not necessarily come from an expert agent.

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

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

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

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

Original: We implement our training scheme for CURL with the following setup: (1) pretrain a multi-task CURL model on the same pretraining tasks as our framework uses (using offline data generated from training individual CURL agents), and (2) directly finetune the pretrained model on the target task with online RL for 100k environment steps.
Revised: We implement our training scheme for CURL with the following setup: (1) pretrain a multi-task CURL model on the same pretraining tasks as our framework uses (using offline data generated from training individual CURL agents), and (2) directly finetune the pretrained model on the target task with online RL for 100k environment steps.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: H XTRA A BLATIONS FOR T ASKS WITH D IVERSE G AME M ECHANICS
Revised: H XTRA A BLATIONS FOR T ASKS WITH D IVERSE G AME M ECHANICS

Original: We perform the same set of ablations for tasks that share diverse game mechanics as we do for tasks that share similar game mechanics for XTRA.
Revised: We perform the same set of ablations for tasks that share diverse game mechanics as we do for tasks that share similar game mechanics for XTRA.

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

Original: Details of each ablation can be found in Section 4.
Revised: Details of each ablation can be found in Section 4.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

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

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

Original: By reducing the number of tasks in pretraining, the model is exposed to (1) less diverse game mechanics and (2) less offline training data in pretraining, and fewer cross-tasks in finetuning.
Revised: By reducing the number of tasks in pretraining, the model is exposed to (1) less diverse game mechanics and (2) less offline training data in pretraining, and fewer cross-tasks in finetuning.

Original: In this ablation
Revised: In this ablation

Original: [Table]
Revised: [Table]

Original: we gradually reduce the number of pretrained tasks from 8 (Carnival, Centipede, Phoenix,
Revised: we gradually reduce the number of pretrained tasks from 8 (Carnival, Centipede, Phoenix,

Original: Pooyan, Riverraid, VideoPinball, WizardOfWor
Revised: Pooyan, Riverraid, VideoPinball, WizardOfWor

Original: [Table]
Revised: [Table]

Original: and YarsRevenge), to 4 (Phoenix, WizardOfWor,
Revised: and YarsRevenge), to 4 (Phoenix, WizardOfWor,

Original: VideoPinball, YarsRevenge), and to 2 (Phoenix, VideoPinball).
Revised: VideoPinball, YarsRevenge), and to 2 (Phoenix, VideoPinball).

Original: XTRA is reduced to EfficientZero-L when the number of pretrained tasks and cross-tasks during finetuning is set to 0.
Revised: XTRA is reduced to EfficientZero-L when the number of pretrained tasks and cross-tasks during finetuning is set to 0.

Original: We find that increasing the number of tasks during pretraining (and later cross-task finetuning) mostly consistently improves XTRA performance.
Revised: We find that increasing the number of tasks during pretraining (and later cross-task finetuning) mostly consistently improves XTRA performance.

Original: J A RCHITECTURAL D ETAILS
Revised: J A RCHITECTURAL D ETAILS

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

Original: For EfficientZero-L and XTRA, we increase the number of residual blocks from 1 (default) to 4 (ours) in the representation network, which we find to improve pretraining slightly for XTRA.
Revised: For EfficientZero-L and XTRA, we increase the number of residual blocks from 1 (default) to 4 (ours) in the representation network, which we find to improve pretraining slightly for XTRA.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: • Flatten.
Revised: • Flatten.

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

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

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

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

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

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

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

Original: • Flatten.
Revised: • Flatten.

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

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

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

Original: K H YPER - PARAMETERS
Revised: K H YPER - PARAMETERS

Original: We adopt our hyper-parameters from EfficientZero (Ye et al., 2021) with minimal modification.
Revised: We adopt our hyper-parameters from EfficientZero (Ye et al., 2021) with minimal modification.

Original: Because XTRA uses data from offline tasks to perform cross-task transfer during online finetuning, we have an additional hyper-parameter for mini-batch size for offline tasks, which is set to(default).
Revised: Because XTRA uses data from offline tasks to perform cross-task transfer during online finetuning, we have an additional hyper-parameter for mini-batch size for offline tasks, which is set to(default).

Original: We list all hyper-parameters in Table 8 for completeness.
Revised: We list all hyper-parameters in Table 8 for completeness.

Original: Lastly, we note that EfficientZero performs an additional 20k gradient steps at 100k environment steps, with a 10 × smaller learning rate.
Revised: Lastly, we note that EfficientZero performs an additional 20k gradient steps at 100k environment steps, with a 10 × smaller learning rate.

Original: We follow this procedure when comparing to previous state-of-the-art methods (Table 2)
Revised: We follow this procedure when comparing to previous state-of-the-art methods (Table 2)

Original: [Table]
Revised: [Table]

Original: but for simplicity we omit these additional gradient steps in the remainder of our experiments for both
Revised: but for simplicity we omit these additional gradient steps in the remainder of our experiments for both

Original: XTRA and baselines.
Revised: XTRA and baselines.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

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

Original: [Figure]
Revised: [Figure]

Original: During online finetuning (stage 2), we finetune both with data from the target task, and data from the pretraining tasks.
Revised: During online finetuning (stage 2), we finetune both with data from the target task, and data from the pretraining tasks.

Original: We maintain the same batch size ( 256 ) for the target task data as the non-pretraining baselines, but add additional data from the pretraining tasks with a 1:1 ratio.
Revised: We maintain the same batch size ( 256 ) for the target task data as the non-pretraining baselines, but add additional data from the pretraining tasks with a 1:1 ratio.

Original: Thus, our effective batch size is 2 × that of the baselines.
Revised: Thus, our effective batch size is 2 × that of the baselines.

Original: To verify that performance improvements stem from our pretraining (stage 1) and not the larger batch size, we compare our method to a variant of our EfficientZero-L baseline that uses a 2 × larger batch size ( 512 ).
Revised: To verify that performance improvements stem from our pretraining (stage 1) and not the larger batch size, we compare our method to a variant of our EfficientZero-L baseline that uses a 2 × larger batch size ( 512 ).

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

Original: We do not observe any significant change in performance by doubling the batch size for the baseline.
Revised: We do not observe any significant change in performance by doubling the batch size for the baseline.

Original: Thus, we conclude that a larger (effective) batch size is not the source of our performance gains, but rather our pretraining and inclusion of pretraining tasks during the online finetuning.
Revised: Thus, we conclude that a larger (effective) batch size is not the source of our performance gains, but rather our pretraining and inclusion of pretraining tasks during the online finetuning.

Original: M P ER -G AME I MPROVEMENT OVER E FFICIENT Z ERO
Revised: M P ER -G AME I MPROVEMENT OVER E FFICIENT Z ERO

Original: We visualize XTRA’s mean improvement over EfficientZero and EfficientZero-L under similar and diverse task settings on a per-game basis.
Revised: We visualize XTRA’s mean improvement over EfficientZero and EfficientZero-L under similar and diverse task settings on a per-game basis.

Original: The improvement is calculated based on the formula: ( score XTRA / score Baseline ) − 1 .
Revised: The improvement is calculated based on the formula: ( score XTRA / score Baseline ) − 1 .

Original: The Baseline can be either EfficientZero or EfficientZero-L, depending on the visualization setting.
Revised: The Baseline can be either EfficientZero or EfficientZero-L, depending on the visualization setting.

Original: C e n t
Revised: C e n t

Original: i p
Revised: i p

Original: e
Revised: e

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

Original: P h
Revised: P h

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

Original: [Figure]
Revised: [Figure]

Original: a n
Revised: a n

Original: [Figure]
Revised: [Figure]

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

Original: D e
Revised: D e

Original: [Figure]
Revised: [Figure]

Original: n
Revised: n

Original: [Figure]
Revised: [Figure]

Original: tt a c
Revised: tt a c

Original: [Figure]
Revised: [Figure]

Original: P
Revised: P

Original: [Figure]
Revised: [Figure]

Original: a n
Revised: a n

Original: [Figure]
Revised: [Figure]

Original: d
Revised: d

Original: [Figure]
Revised: [Figure]

Original: r A m
Revised: r A m

Original: [Figure]
Revised: [Figure]

Original: I m p r
Revised: I m p r

Original: o v e m e n t o v e r E
Revised: o v e m e n t o v e r E

Original: ff i
Revised: ff i

Original: c i e n t Z e r o
Revised: c i e n t Z e r o

Original: [Figure]
Revised: [Figure]

Original: N G AME I NFORMATION
Revised: N G AME I NFORMATION

Original: In this section, we aim to provide additional context about the games that we consider during both pretraining and finetuning.
Revised: In this section, we aim to provide additional context about the games that we consider during both pretraining and finetuning.

Original: Table 9 lists core properties for each game.
Revised: Table 9 lists core properties for each game.

Original: The Similar Task column marks all games used in Table 1 ( similar games), and the two Diverse Task columns mark all games used in Table 2 ( diverse games) for pretraining and finetuning, respectively.
Revised: The Similar Task column marks all games used in Table 1 ( similar games), and the two Diverse Task columns mark all games used in Table 2 ( diverse games) for pretraining and finetuning, respectively.

Original: We further categorize games into five categories based on game mechanics: Maze , Shooter , Tank , Adventure , and Ball
Revised: We further categorize games into five categories based on game mechanics: Maze , Shooter , Tank , Adventure , and Ball

Original: Tracking , and also report whether the scene is static or dynamic, as well the (valid) action space for each task.
Revised: Tracking , and also report whether the scene is static or dynamic, as well the (valid) action space for each task.

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

Original: O B EHAVIORAL C LONING B ASELINE
Revised: O B EHAVIORAL C LONING B ASELINE

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

Original: The
Revised: The

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

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

Original: We find
Revised: We find

Original: [Table]
Revised: [Table]

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

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

Original: [Table]
Revised: [Table]

Original: P G AME V ISUALIZATIONS
Revised: P G AME V ISUALIZATIONS

Original: We here provide sample trajectories for each game.
Revised: We here provide sample trajectories for each game.

Original: These visualizations aim to highlight that (1) it is reasonable to expect some degree of cross-task transfer in similar game transfer (maze → maze, shooter → shooter) due to similarity in visuals and game mechanics, and (2) that games from our diverse category indeed are diverse, both in terms of visuals and mechanics.
Revised: These visualizations aim to highlight that (1) it is reasonable to expect some degree of cross-task transfer in similar game transfer (maze → maze, shooter → shooter) due to similarity in visuals and game mechanics, and (2) that games from our diverse category indeed are diverse, both in terms of visuals and mechanics.

Original: Figure 15 shows trajectories for five Maze games, Figure 16 shows trajectories for five Shooter games, and Figureshows trajectories for eight games from our diverse game category.
Revised: Figure 15 shows trajectories for five Maze games, Figure 16 shows trajectories for five Shooter games, and Figureshows trajectories for eight games from our diverse game category.

Original: [Figure]
Revised: [Figure]

Original: i d a r
Revised: i d a r

Original: B a nk
Revised: B a nk

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

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

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

Original: W o r time −→
Revised: W o r time −→

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

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

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

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

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

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

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

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

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

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

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

Original: V i
Revised: V i

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

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

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