Original: Multi-Game Decision Transformers
Revised: Multi-Game Decision Transformers

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: A longstanding goal of the ﬁeld of AI is a strategy for compiling diverse experienceinto a highly capable, generalist agent.
Revised: A longstanding goal of the ﬁeld of AI is a strategy for compiling diverse experienceinto a highly capable, generalist agent.

Original: In the subﬁelds of vision and language, thiswas largely achieved by scaling up transformer-based models and training them onlarge, diverse datasets.
Revised: In the subﬁelds of vision and language, thiswas largely achieved by scaling up transformer-based models and training them onlarge, diverse datasets.

Original: Motivated by this progress, we investigate whether the samestrategy can be used to produce generalist reinforcement learning agents.
Revised: Motivated by this progress, we investigate whether the samestrategy can be used to produce generalist reinforcement learning agents.

Original: Speciﬁcally, we show that a single transformer-based model – with a single set of weights– trained purely ofﬂine can play a suite of up to 46 Atari games simultaneously atclose-to-human performance.
Revised: Speciﬁcally, we show that a single transformer-based model – with a single set of weights– trained purely ofﬂine can play a suite of up to 46 Atari games simultaneously atclose-to-human performance.

Original: When trained and evaluated appropriately, we ﬁndthat the same trends observed in language and vision hold, including scaling ofperformance with model size and rapid adaptation to new games via ﬁne-tuning.
Revised: When trained and evaluated appropriately, we ﬁndthat the same trends observed in language and vision hold, including scaling ofperformance with model size and rapid adaptation to new games via ﬁne-tuning.

Original: We compare several approaches in this multi-game setting, such as online andofﬂine RL methods and behavioral cloning, and ﬁnd that our Multi-Game Decision
Revised: We compare several approaches in this multi-game setting, such as online andofﬂine RL methods and behavioral cloning, and ﬁnd that our Multi-Game Decision

Original: Transformer models offer the best scalability and performance.
Revised: Transformer models offer the best scalability and performance.

Original: We release thepre-trained models and code to encourage further research in this direction.
Revised: We release thepre-trained models and code to encourage further research in this direction.

Original:  Introduction
Revised:  Introduction

Original: Building large-scale generalist models that solve many tasks by training on massive task-agnosticdatasets has emerged as a dominant approach in natural language processing [18, 12], computervision [19, 6] and their intersection [60, 4].
Revised: Building large-scale generalist models that solve many tasks by training on massive task-agnosticdatasets has emerged as a dominant approach in natural language processing [17, 11], computervision [18, 5] and their intersection [55, 3].

Original: These models can adapt to new tasks (such as translation[62, 77]), make use of unrelated data (such as using high-resource language to improve translations oflow-resource languages [17]), or even incorporate new modalities by projecting images into languagespace [46, 74].
Revised: These models can adapt to new tasks (such as translation[57, 70]), make use of unrelated data (such as using high-resource language to improve translations oflow-resource languages [16]), or even incorporate new modalities by projecting images into languagespace [41, 67].

Original: The success of these methods largely derives from a combination of scalable modelarchitectures [76], an abundance of unlabeled task-agnostic data, and continuous improvements inhigh performance computing infrastructure.
Revised: The success of these methods largely derives from a combination of scalable modelarchitectures [69], an abundance of unlabeled task-agnostic data, and continuous improvements inhigh performance computing infrastructure.

Original: Crucially, scaling laws [38, 31] indicate that performancegains due to scale have not yet reached a saturation point.
Revised: Crucially, scaling laws [34, 28] indicate that performancegains due to scale have not yet reached a saturation point.

Original: In this work we argue that a similar progression is possible in the ﬁeld of reinforcement learning, andtake initial steps toward scalable methods that produce highly capable generalist agents.
Revised: In this work we argue that a similar progression is possible in the ﬁeld of reinforcement learning, andtake initial steps toward scalable methods that produce highly capable generalist agents.

Original: In contrastto vision and language domains, reinforcement learning has seen advocacy for the use of smallermodels [16, 49, 8] and is usually either used to solve single tasks, or multiple tasks within the sameenvironment.
Revised: In contrastto vision and language domains, reinforcement learning has seen advocacy for the use of smallermodels [15, 44, 7] and is usually either used to solve single tasks, or multiple tasks within the sameenvironment.

Original: Importantly, training across multiple environments – with very different dynamics,rewards, visuals, and agent embodiments – has been studied less signiﬁcantly.
Revised: Importantly, training across multiple environments – with very different dynamics,rewards, visuals, and agent embodiments – has been studied less signiﬁcantly.

Original: Speciﬁcally, we investigate whether a single model – with a single set of parameters – can be trainedto act in multiple environments from large amounts of expert and non-expert experience.
Revised: Speciﬁcally, we investigate whether a single model – with a single set of parameters – can be trainedto act in multiple environments from large amounts of expert and non-expert experience.

Original: We considertraining on a suite of 41 Atari games [9, 25] for their diversity, informally asking “Can modelslearn something universal from playing many video games?”.
Revised: We considertraining on a suite of 41 Atari games [8, 22] for their diversity, informally asking “Can modelslearn something universal from playing many video games?”.

Original: To train this model, we use only thepreviously-collected trajectories from Agarwal et al.
Revised: To train this model, we use only thepreviously-collected trajectories from Agarwal et al.

Original: [1], but we evaluate our agent interactively.
Revised: [1], but we evaluate our agent interactively.

Original: Weare not striving for mastery or efﬁciency that game-speciﬁc agents can offer, as we believe we arestill in early stages of this research agenda.
Revised: Weare not striving for mastery or efﬁciency that game-speciﬁc agents can offer, as we believe we arestill in early stages of this research agenda.

Original: Rather, we investigate whether the same trends observedin language and vision hold for large-scale generalist reinforcement learning agents.
Revised: Rather, we investigate whether the same trends observedin language and vision hold for large-scale generalist reinforcement learning agents.

Original: [Figure]
Revised: [Table]

Revised: Multi-game models are all trained on a dataset [1] with median human-normalized score of 78%, which Multi-Game DT notably exceeds.

Original: We ﬁnd that we can train a single agent that averages 85% of human-level performance simultaneouslyacross all games after training on ofﬂine expert and non-expert datasets (see Figure 1).
Revised: We ﬁnd that we can train a single agent that averages 85% of human-level performance simultaneouslyacross all games after training on ofﬂine expert and non-expert datasets (see Figure 1).

Original: Furthermore,we see similar trends that mirror those observed in language and vision: rapid ﬁne-tuning to neverbefore-seen games with very little data (Section 4.5), a scaling relationship between performance andmodel size (Section 4.4), and faster training progress for larger models.
Revised: Furthermore,we see similar trends that mirror those observed in language and vision: rapid ﬁne-tuning to neverbefore-seen games with very little data (Section 4.5), a power-law relationship between performanceand model size (Section 4.4), and faster training progress for larger models.

Original: Notably, not all existing approaches to multi-environment training work well.
Revised: Notably, not all existing approaches to multi-environment training work well.

Original: We investigate severalapproaches, including treating the problem as ofﬂine decision transformer-based sequence modeling[14, 35], online RL [53], ofﬂine temporal difference methods [42], contrastive representations [56],and behavior cloning [59].
Revised: We investigate severalapproaches, including treating the problem as ofﬂine decision transformer-based sequence modeling[13, 31], online RL [48], ofﬂine temporal difference methods [38], contrastive representations [51],and behavior cloning [54].

Original: We ﬁnd that decision transformer based models offer the best performanceand scaling properties in the multi-environment regime.
Revised: We ﬁnd that decision transformer based models offer the best performanceand scaling properties in the multi-environment regime.

Original: However, to permit training on both expertand non-expert trajectories, we ﬁnd it is necessary to use a guided generation technique from languagemodeling to generate expert-level actions, which is an important departure from standard decisiontransformers.
Revised: However, to permit training on both expertand non-expert trajectories, we ﬁnd it is necessary to use a guided generation technique from languagemodeling to generate expert-level actions, which is an important departure from standard decisiontransformers.

Original: Our contributions are threefold: First, we show that it is possible to train a single high-performinggeneralist agent to act across multiple environments from ofﬂine data alone.
Revised: Our contributions are threefold: First, we show that it is possible to train a single high-performinggeneralist agent to act across multiple environments from ofﬂine data alone.

Original: Second, we show thatscaling trends observed in language and vision hold.
Revised: Second, we show thatscaling trends observed in language and vision hold.

Original: And third, we compare multiple approaches forachieving this goal, ﬁnding that decision transformers combined with guided generation perform thebest.
Revised: And third, we compare multiple approaches forachieving this goal, ﬁnding that decision transformers combined with guided generation perform thebest.

Original: It is our hope this study can inspire further research in generalist agents.
Revised: It is our hope this study can inspire further research in generalist agents.

Original: To aid this, we plan tomake our pre-trained models and code publicly available.
Revised: To aid this, we plan tomake our pre-trained models and code publicly available

Revised: [Table]

Original: [Figure]
Revised: [Table]

Revised: (cid:211)(cid:342)(cid:300)(cid:211)(cid:303)(cid:315)(cid:713)(cid:300)(cid:261)(cid:176)(cid:343)(cid:713)(cid:275)(cid:268)(cid:713)(cid:176)(cid:261)(cid:261)(cid:713)(cid:315)(cid:303)(cid:176)(cid:243)(cid:268)(cid:243)(cid:268)(cid:233)(cid:713)(cid:233)(cid:176)(cid:267)(cid:211)(cid:307) (cid:354)(cid:268)(cid:211)(cid:693)(cid:315)(cid:320)(cid:268)(cid:211)(cid:713)(cid:315)(cid:275)(cid:713)(cid:268)(cid:211)(cid:341)(cid:713)(cid:233)(cid:176)(cid:267)(cid:211)(cid:307) (cid:664)(cid:664)(cid:664) (cid:664)(cid:664)(cid:664)

Original:  Related Work
Revised:  Related Work

Original: A generalist agent for solving a variety of environments has been a goal for artiﬁcial intelligence(AI) researchers since the inception of AI as a ﬁeld of study [50].
Revised: A generalist agent for solving a variety of environments has been a goal for artiﬁcial intelligence(AI) researchers since the inception of AI as a ﬁeld of study [45].

Original: This same reason motivated theintroduction of the Atari suite (the Arcade Learning Environment, or ALE) as a testbed for learningalgorithms [10]; in their own words, the ALE is for “empirically assessing agents designed forgeneral competency.”
Revised: This same reason motivatedthe introduction of the Atari suite (the Arcade Learning Environment, or ALE) as a testbed forlearning algorithms [9]; in their own words, the ALE is for “empirically assessing agents designedfor general competency.”

Original: While the celebrated deep Q -learning [52] and actor critic [54] agents wereamong the ﬁrst to use a single algorithm for all games, they nevertheless required separate trainingand hyperparameters for each game agent.
Revised: While the celebrated deep Q -learning [47] and actor critic [49] agents wereamong the ﬁrst to use a single algorithm for all games, they nevertheless required separate trainingand hyperparameters for each game agent.

Original: Later works have demonstrated the ability to learn a single neural network agent on multiple Atari games simultaneously, either online [20] or via policydistillation [58, 66].
Revised: Later works have demonstrated the ability to learn a single neural network agent on multiple Atari games simultaneously, either online [19] or via policydistillation [53, 60].

Original: The aim of our work is similar – to learn a single agent for playing multiple
Revised: The aim of our work is similar – to learn a single agent for playing multiple

Original: Atari games – with a focus on ofﬂine learning.
Revised: Atari games – with a focus on ofﬂine learning.

Original: We demonstrate results with human-level competencyon up to 46 games, which is unseen in literature.
Revised: We demonstrate results with human-level competencyon up to 46 games, which is unseen in literature.

Original: A closely related setting is learning to solve multiple tasks within the same or similar environments.
Revised: A closely related setting is learning to solve multiple tasks within the same or similar environments.

Original: For example in the robotics ﬁeld, existing works propose to use language-conditioned tasks [48, 3, 34],while others posit goal-reaching as a way to learn general skills [51], among other proposals [37, 81].
Revised: For example in the robotics ﬁeld, existing works propose to use language-conditioned tasks [43, 2, 30],while others posit goal-reaching as a way to learn general skills [46], among other proposals [33, 74].

Original: In this work, we tackle the problem of learning to act in a large collection of environments withdistinctively different dynamics, rewards, and agent embodiments.
Revised: In this work, we tackle the problem of learning to act in a large collection of environments withdistinctively different dynamics, rewards, and agent embodiments.

Original: This complicated but importantsetting requires a different type of generalization that has been studied signiﬁcantly less.
Revised: This complicated but importantsetting requires a different type of generalization that has been studied signiﬁcantly less.

Original: A concurrent work [64] also aims to train a transformer-based generalist agent based on ofﬂine dataincluding for the ALE.
Revised: A concurrent work [58] also aims to train a transformer-based generalist agent based on ofﬂine dataincluding for the ALE.

Original: This work differs from ours in that the ofﬂine training data is exclusivelynear-optimal and it requires prompting by expert trajectories at inference time.
Revised: This work differs from ours in that the ofﬂine training data is exclusivelynear-optimal and it requires prompting by expert trajectories at inference time.

Original: In contrast, we extenddecision transformers [14] from the Upside-Down RL family [70, 67] to learn from a diverse dataset(expert and non-expert data), predict returns, and pick optimality-conditioned returns.
Revised: In contrast, we extenddecision transformers [13] from the Upside-Down RL family [64, 61] to learn from a diverse dataset(expert and non-expert data), predict returns, and pick optimality-conditioned returns.

Original: Furthermore,we provide comparisons against existing behavioral cloning, online and ofﬂine RL methods, andcontrastive representations [79, 56].
Revised: Furthermore,we provide comparisons against existing behavioral cloning, online and ofﬂine RL methods, andcontrastive representations [72, 51].

Original: Other works that also consider LLM-like sequence modeling fora variety of single control tasks include [65, 83, 35, 23].
Revised: Other works that also consider LLM-like sequence modeling fora variety of single control tasks include [59, 76, 31, 20].

Original:  Method
Revised:  Method

Original: We consider a decision-making agent that at every time t receives an observation of the world o t ,chooses an action a t , and receives a scalar reward r t .
Revised: We consider a decision-making agent that at every time t receives an observation of the world o t ,chooses an action a t , and receives a scalar reward r t .

Original: Our goal is to learn a single optimal policydistribution P ∗ θ ( a t | o ≤ t , a <t , r <t ) with parameters θ that maximizes the agent’s total future return
Revised: Our goal is to learn a single optimal policydistribution P ⇤ ✓ ( a t | o  t , a <t , r <t ) with parameters ✓ that maximizes the agent’s total future return

Original: R t =(cid:80)
Revised: R t =

Original:  k>t r k on all the environments we consider.
Revised: P k>t r k on all the environments we consider.

Original: 3.1 Reinforcement Learning as Sequence Modeling
Revised: 3.1 Reinforcement Learning as Sequence Modeling

Original: Following [14], we pose the problem of ofﬂine reinforcement learning as a sequence modelingproblem where we model the probability of the next sequence token x i conditioned on all tokensprior to it: P θ ( x i | x <i ) , similar to contemporary decoder-only sequence models [12, 15, 61].
Revised: Following [13], we pose the problem of ofﬂine reinforcement learning as a sequence modelingproblem where we model the probability of the next sequence token x i conditioned on all tokensprior to it: P ✓ ( x i | x <i ) , similar to contemporary decoder-only sequence models [11, 14, 56].

Original: Thesequences we consider have the form:
Revised: Thesequences we consider have the form:

Original: [Equation]
Revised: [Equation]

Original: where t represents a time-step, M is the number of image patches per observation (which we furtherdiscuss in Section 3.2), and ˆ R t is the agent’s target return for the rest of the sequence.
Revised: where t represents a time-step, M is the number of image patches per observation (which we furtherdiscuss in Section 3.2), and ˆ R t is the agent’s target return for the rest of the sequence.

Original: Such asequence order respects the causal structure of the environment decision process.
Revised: Such asequence order respects the causal structure of the environment decision process.

Original: Figure 3 presentsan overview of our model architecture.o a a causal transformer decoder t t t return action
Revised: Figure 3 presentsan overview of our model architecture.

Original: [Equation]
Revised: [Equation]

Original: observation
Revised: (cid:3)(cid:194) (cid:3)(cid:106) (cid:3)(cid:106) (cid:3)(cid:106) (cid:303)(cid:211)(cid:315)(cid:320)(cid:303)(cid:268) (cid:176)(cid:201)(cid:315)(cid:243)(cid:275)(cid:268) (cid:125) (cid:3)(cid:106) (cid:125) (cid:664)(cid:664)(cid:664)

Original: [Equation]
Revised: [Figure]

Original: reward ... + Returns, actions, and rewards are tokenized (See Section 3.2 for details), and we train the model topredict the next return, action, and reward discrete token in a sequence via standard cross-entropyloss.
Revised: Returns, actions, and rewards are tokenized (See Section 3.2 for details), and we train the model topredict the next return, action, and reward discrete token in a sequence via standard cross-entropyloss.

Original: The sequence we consider is different from Chen et al.
Revised: The sequence we consider is different from Chen et al.

Original: [14], which has (cid:104) ... , ˆ R t , o t , a t , ... (cid:105) .
Revised: [13], which has h ..., ˆ R t , o t , a t , ... i .

Original: Our design allows predicting the return distribution and sampling from it, instead of relying on a user tomanually select an expert-level return at inference time (See Section 3.4).
Revised: Our design allows predicting the return distribution and sampling from it, instead of relying on a user tomanually select an expert-level return at inference time (See Section 3.4).

Original: Predicting future value and rewards have been shown to be useful objectives for learning betterrepresentations in artiﬁcial reinforcement learning agents [47, 68, 44] and important signals forrepresentation learning in humans [5].
Revised: Predicting future value and rewards have been shown to be useful objectives for learning betterrepresentations in artiﬁcial reinforcement learning agents [42, 62, 40] and important signals forrepresentation learning in humans [4].

Original: Thus, while we may not directly use all of the predictedquantities, the task of predicting them encourages structure and representation learning of ourenvironments.
Revised: Thus, while we may not directly use all of the predictedquantities, the task of predicting them encourages structure and representation learning of ourenvironments.

Original: In this work, we do not attempt to predict future observations due to their non-discretenature and the additional model capacity that would be required to generate images.
Revised: In this work, we do not attempt to predict future observations due to their non-discretenature and the additional model capacity that would be required to generate images.

Original: However,building image-based forward prediction models of the environment has been shown to be a usefulrepresentation objective for RL [28, 27, 29].
Revised: However,building image-based forward prediction models of the environment has been shown to be a usefulrepresentation objective for RL [25, 24, 26].

Original: We leave it for future investigation.
Revised: We leave it for future investigation.

Original: 3.2 Tokenization
Revised: 3.2 Tokenization

Original: To generate returns, actions, and rewards via multinomial distributions similarly to language generation, we convert these quantities to discrete tokens.
Revised: To generate returns, actions, and rewards via multinomial distributions similarly to language generation, we convert these quantities to discrete tokens.

Original: Actions a are already discrete quantities inthe environments we consider.
Revised: Actions a are already discrete quantities inthe environments we consider.

Original: We convert scalar rewards to ternary quantities {− 1 , 0 , +1 } , anduniformly quantize returns into a discrete range shared by all our environments 1 .
Revised: We convert scalar rewards to ternary quantities {  1 , 0 , +1 } , anduniformly quantize returns into a discrete range shared by all our environments 1 .

Original: Inspired by simplicity and effectiveness of transformer architecture for processing images [19], wedivide each observation image into a collection of M patches 2 (see Figure 3).
Revised: Inspired by simplicity and effectiveness of transformer architecture for processing images [18], wedivide each observation image into a collection of M patches 2 (see Figure 3).

Original: Each patch is additivelycombined with a trainable position encoding and linearly projected into the input token embeddingspace.
Revised: Each patch is additivelycombined with a trainable position encoding and linearly projected into the input token embeddingspace.

Original: We experimented with using image tokenizations coming from a convolutional network, butdid not ﬁnd it to have a signiﬁcant beneﬁt and omitted it for simplicity.
Revised: We experimented with using image tokenizations coming from a convolutional network, butdid not ﬁnd it to have a signiﬁcant beneﬁt and omitted it for simplicity.

Original: We chose our tokenization scheme with simplicity in mind, but many other schemes are possible.
Revised: We chose our tokenization scheme with simplicity in mind, but many other schemes are possible.

Original: While all our environments use a shared action space, varying action spaces when controlling differentagent morphologies can still be tokenized using methods of [33, 43, 26].
Revised: While all our environments use a shared action space, varying action spaces when controlling differentagent morphologies can still be tokenized using methods of [29, 39, 23].

Original: And while we used uniformquantization to discretize continuous quantities, more sophisticated methods such as VQ-VAE [75]can be used to learn more effective discretizations.
Revised: And while we used uniformquantization to discretize continuous quantities, more sophisticated methods such as VQ-VAE [68]can be used to learn more effective discretizations.

Original: 3.3 Training Dataset
Revised: 3.3 Training Dataset

Original: To train the model, we use an existing dataset of Atari trajectories (with quantized returns) introducedin [1].
Revised: To train the model, we use an existing dataset of Atari trajectories (with quantized returns) introducedin [1].

Original: The dataset contains trajectories collected from the training progress of a DQN agent [53].
Revised: The dataset contains trajectories collected from the training progress of a DQN agent [48].

Original: Following [25], we select 46 games where DQN performance signiﬁcantly exceeded that of a randomagent.
Revised: Following [22], we select 46 games where DQN performance signiﬁcantly exceeded that of a randomagent.

Original: 41 games are used for training and 5 games are held out for out-of-distribution generalizationexperiments.
Revised: 41 games are used for training and 5 games are held out for out-of-distribution generalizationexperiments.

Original: We chose 5 held-out games representing different game categories including Alien and MsPacman(maze based), Pong (ball tracking), SpaceInvaders (shoot vertically), and StarGunner (shoothorizontally), to ensure out-of-distribution generalization can be evaluated on different types ofgames.
Revised: We chose 5 held-out games representing different game categories including Alien and MsPacman(maze based), Pong (ball tracking), SpaceInvaders (shoot vertically), and StarGunner (shoothorizontally), to ensure out-of-distribution generalization can be evaluated on different types ofgames.

Original: For each of 41 games, we use data from 2 training runs, each containing roll-outs from 50 policycheckpoints, in turn each containing 1 million environment steps.
Revised: For each of 41 games, we use data from 2 training runs, each containing roll-outs from 50 policycheckpoints, in turn each containing 1 million environment steps.

Original: This totals 4.1 billion steps.
Revised: This totals 4.1 billion steps.

Original: Usingthe tokenization scheme in previous sections, the dataset contains almost 160 billion tokens.
Revised: Usingthe tokenization scheme in previous sections, the dataset contains almost 160 billion tokens.

Original: As the dataset contains agent’s behavior at all stages of learning, it contains both expert and nonexpert behaviors.
Revised: As the dataset contains agent’s behavior at all stages of learning, it contains both expert and nonexpert behaviors.

Original: We do not perform any special ﬁltering, curation, or balancing of the dataset.
Revised: We do not perform any special ﬁltering, curation, or balancing of the dataset.

Original: Themotivation to train on such data instead of expert-only behaviors is twofold: Firstly, sub-optimalbehaviors are more diverse than optimal behaviors and may still be useful for learning representationsof the environment and consequences of poor decisions.
Revised: Themotivation to train on such data instead of expert-only behaviors is twofold: Firstly, sub-optimalbehaviors are more diverse than optimal behaviors and may still be useful for learning representationsof the environment and consequences of poor decisions.

Original: Secondly, it may be difﬁcult to create asingle binary criteria for optimality as it is typically a graded quantity.
Revised: Secondly, it may be difﬁcult to create asingle binary criteria for optimality as it is typically a graded quantity.

Original: Thus, instead of assumingonly task-relevant expert behaviors, we train our model on all available behaviors, yet generate expertbehavior at inference time as described in the next section.
Revised: Thus, instead of assumingonly task-relevant expert behaviors, we train our model on all available behaviors, yet generate expertbehavior at inference time as described in the next section.

Original: 3.4 Expert Action Inference
Revised: 3.4 Expert Action Inference

Original: As described above, our training datasets contain a mix of expert and non-expert behaviors, thusdirectly generating actions from the model imitating the data is unlikely to consistently produce expertbehavior (as we conﬁrm in Section 4.7).
Revised: As described above, our training datasets contain a mix of expert and non-expert behaviors, thusdirectly generating actions from the model imitating the data is unlikely to consistently produce expertbehavior (as we conﬁrm in Section 4.7).

Original: Instead, we want to control action generation to consistentlyproduce actions of highly-rewarding behavior.
Revised: Instead, we want to control action generation to consistentlyproduce actions of highly-rewarding behavior.

Original: This mirrors the problem of discriminator-guidedgeneration in language models, for which a variety of methods have been proposed [40, 78, 57].
Revised: This mirrors the problem of discriminator-guidedgeneration in language models, for which a variety of methods have been proposed [36, 71, 52].

Original: We propose an inference-time method inspired by [40] and assume a binary classiﬁer P ( expert t | ... )that identiﬁes whether or not the behavior is expert-level before taking an action at time t .
Revised: We propose an inference-time method inspired by [36] and assume a binary classiﬁer P ( expert t | ... )that identiﬁes whether or not the behavior is expert-level before taking an action at time t .

Original: Following
Revised: Following

Original: Bayes’ rule, the distribution of expert-level returns at time t is then:
Revised: Bayes’ rule, the distribution of expert-level returns at time t is then:

Original: [Equation]
Revised: [Equation]

Original: Similarly to [69, 72, 73, 39], we deﬁne a binary classiﬁer to be proportional to future return withinverse temperature parameter κ :
Revised: Similarly to [63, 65, 66, 35], we deﬁne a binary classiﬁer to be proportional to future return withinverse temperature parameter  :

Original: [Equation]
Revised: [Equation]

Original: where R low is the return lower bound and R high is the return upper bound. This results in a simple auto-regressive procedure where we ﬁrst sample high-but-plausible targetreturns R t according to log-probability log P θ ( R t | ... )
Revised: This results in a simple auto-regressive procedure where we ﬁrst sample high-but-plausible targetreturns R t according to log-probability log P ✓ ( R t | ... )

Original: + κ ( R t− R low ) / ( R high − R low ) , and thensample actions according toP
Revised: +  R t , and then sample actions according to

Original: θ ( a t | R t , ... ) .
Revised: P ✓ ( a t | R t , ... ) .

Original: See Figure 4 for an illustration of this procedure and Appendix for implementation details.
Revised: See Figure 4 for an illustration of this procedure and Appendix for implementationdetails.

Original: It can be seen as a variation of return-conditioned policies[41, 70, 14] that automatically generates expert-level (but likely) returns at every timestep, instead ofmanually ﬁxing them for the duration of the episode. 
Revised: It can be seen as a variation of return-conditioned policies [37, 64, 13] that automaticallygenerates expert-level (but likely) returns at every timestep, instead of manually ﬁxing them for theduration of the episode. [Table] (cid:176)

Original: [Equation]
Revised: [Table]

Revised: (cid:122)(cid:687)(cid:211)(cid:342)(cid:300)(cid:211)(cid:303)(cid:315)(cid:713)(cid:795)(cid:713)(cid:125)(cid:688)

Original: [Figure]
Revised: [Table]

Revised: (cid:122)(cid:687)(cid:713)(cid:176)(cid:713)(cid:795)(cid:713)(cid:125)(cid:713)(cid:668)(cid:713)(cid:688)

Original: Importantly, this formulation only affects the inference procedure of the model – training is entirelyunaffected and can rely on standard next-token prediction frameworks and infrastructure.
Revised: Importantly, this formulation only affects the inference procedure of the model – training is entirelyunaffected and can rely on standard next-token prediction frameworks and infrastructure.

Original: Whilewe chose this formulation for its simplicity, controllable generation is an active area of study andwe expect other more effective methods to be introduced in the future.
Revised: Whilewe chose this formulation for its simplicity, controllable generation is an active area of study andwe expect other more effective methods to be introduced in the future.

Original: As such, our contribution isto point out a connection between problems of controllable generation in language modeling andoptimality conditioning in control.
Revised: As such, our contribution isto point out a connection between problems of controllable generation in language modeling andoptimality conditioning in control.

Original:  Experiments
Revised:  Experiments

Original: We formulate our experiments to answer a number of questions that are addressed in followingsections:
Revised: We formulate our experiments to answer a number of questions that are addressed in followingsections:

Original: • How do different online and ofﬂine methods perform in the multi-game regime?
Revised: • How do different online and ofﬂine methods perform in the multi-game regime?

Original: • How do different methods scale with model size ?
Revised: • How do different methods scale with model size ?

Original: • How effective are different methods at transfer to novel games ?
Revised: • How effective are different methods at transfer to novel games ?

Original: • Does multi-game decision transformer improve upon training data?
Revised: • Does multi-game decision transformer improve upon training data?

Original: • Does expert action inference (Section 3.4) improve upon behavioral cloning?
Revised: • Does expert action inference (Section 3.4) improve upon behavioral cloning?

Original: We additionally qualitatively explore the behavior of these models, such as attention patterns andtoken embedding structure in the Appendix.
Revised: We additionally qualitatively explore the behavior of these models, such as attention patterns andtoken embedding structure in the Appendix.

Original: 4.1 Setup
Revised: 4.1 Setup

Original: Model Variants and Scaling.
Revised: Model Variants and Scaling.

Original: We base our decision transformer (DT) conﬁguration on GPT-2 [12]as summarized in the Appendix.
Revised: We base our decision transformer (DT) conﬁguration on GPT-2 [11]as summarized in the Appendix.

Original: We report DT-200M (a Multi-Game DT with 200M parameters)results if not speciﬁed otherwise.
Revised: We report DT-200M (a Multi-Game DT with 200M parameters)results if not speciﬁed otherwise.

Original: Other smaller variants are DT-40M and DT-10M. We set sequencelength to 4 game frames for all experiments, which results in sequences of 156 tokens.
Revised: Other smaller variants are DT-40M and DT-10M. We set sequencelength to 4 game frames for all experiments, which results in sequences of 156 tokens.

Original: Training and Fine-tuning.
Revised: Training and Fine-tuning.

Original: We train all Multi-Game DT models on TPUv4 hardware and the
Revised: We train all Multi-Game DT models on TPUv4 hardware and the

Original: Jaxline (Babuschkin et al.
Revised: Jaxline (Babuschkin et al.

Original: [7]) framework for 10M steps using the LAMB optimizer [80] with a3 · 10 − 4 learning rate, 4000 steps linear warm-up, no weight decay, gradient clip 1.0, β 1 = 0 .
Revised: [6]) framework for 10M steps using the LAMB optimizer [73] with a3 · 10  4 learning rate, 4000 steps linear warm-up, no weight decay, gradient clip 1.0,  1 = 0 .

Original: 9 andβ 2 = 0 .
Revised: 9 and 2 = 0 .

Original: 999 , and batch size 2048.
Revised: 999 , and batch size 2048.

Original: For ﬁne-tuning on novel games, we train for 100k steps with a10 − 4 learning rate and batch size of 256 instead.
Revised: For ﬁne-tuning on novel games, we train for 100k steps with a10  4 learning rate and batch size of 256 instead.

Original: Metrics.
Revised: Metrics.

Original: We measure performance on individual Atari games by human normalized scores (HNS)[53], i.e. ( score − score random ) / ( score human − score random ) , or DQN-normalized scores, i.e. normalizingby the best DQN scores seen in the training dataset instead of using human scores.
Revised: We measure performance on individual Atari games by human normalized scores (HNS)[48], i.e. ( score  score random ) / ( score human  score random ) , or DQN-normalized scores, i.e. normalizingby the best DQN scores seen in the training dataset instead of using human scores.

Original: To create anaggregate comparison metric across all games, we use median human-normalized score, following
Revised: To create anaggregate comparison metric across all games, we use median human-normalized score, following

Original: Agarwal et al.
Revised: Agarwal et al.

Original: 4.2 Baseline Methods
Revised: 4.2 Baseline Methods

Original: BC Our Decision Transformer (Sec. 3.1) can be reduced to a transformer-based Behavioral Cloning(BC)
Revised: BC Our Decision Transformer (Sec. 3.1) can be reduced to a transformer-based Behavioral Cloning(BC)

Original: [59] agent by removing the target return condition and return token prediction.
Revised: [54] agent by removing the target return condition and return token prediction.

Original: Similar to whatwe do for Decision Transformer, we also learn BC models at different scales (10M, 40M, 200Mparameters) while keeping other conﬁgurations unchanged.
Revised: Similar to whatwe do for Decision Transformer, we also learn BC models at different scales (10M, 40M, 200Mparameters) while keeping other conﬁgurations unchanged.

Original: C51 DQN As a point of comparison for online performance, we use the C51 algorithm [11] whichis a variant of deep Q -learning (DQN) but with a categorical loss for minimizing the temporaldifference (TD) errors.
Revised: C51 DQN As a point of comparison for online performance, we use the C51 algorithm [10] whichis a variant of deep Q -learning (DQN) but with a categorical loss for minimizing the temporaldifference (TD) errors.

Original: Following improvements suggested in Hessel et al.
Revised: Following improvements suggested in Hessel et al.

Original: [30] as well as our ownempirical observations, we use multi-step learning with n = 4 .
Revised: [27] as well as our ownempirical observations, we use multi-step learning with n = 4 .

Original: For the single-game experiments, weuse the standard convolutional neural network (CNN) used in the implementation of C51 [13].
Revised: For the single-game experiments, weuse the standard convolutional neural network (CNN) used in the implementation of C51 [12].

Original: Forthe multi-game experiments, we modify the C51 implementation based on a hyperparameter searchto use an Impala neural network architecture [20] with three blocks using 64 , 128 , and 128 channelsrespectively with a batch size of 128 and update period of 256 .
Revised: Forthe multi-game experiments, we modify the C51 implementation based on a hyperparameter searchto use an Impala neural network architecture [19] with three blocks using 64 , 128 , and 128 channelsrespectively with a batch size of 128 and update period of 256 .

Original: CQL For an ofﬂine TD-based learning algorithm we use conservative Q -learning (CQL) [42].
Revised: CQL For an ofﬂine TD-based learning algorithm we use conservative Q -learning (CQL) [38].

Original: Namely, we augment the categorical loss of C51 with a behavioral cloning loss minimizing− log π Q ( a | s ) , where ( s, a ) is a state-action pair sampled from the ofﬂine dataset and π Q ( ·| s )=
Revised: Namely, we augment the categorical loss of C51 with a behavioral cloning loss minimizing log ⇡ Q ( a | s ) , where ( s, a ) is a state-action pair sampled from the ofﬂine dataset and ⇡ Q ( ·| s )

Original:  softmax( Q ( s, · )) .
Revised: = 210 softmax( Q ( s, · )) .

Original: Following the recommendations in Kumar et al.
Revised: Following the recommendations in Kumar et al.

Original: [42] we weight the contribution of the BC loss by 1 when using 100% of the ofﬂine data and 4 when using 1%.
Revised: [38] we weight the contribution of the BC loss by 1 when using 100% of the ofﬂine data and 4 when using 1%.

Original: As in
Revised: As in

Original: C51, we use a standard CNN network for single-game experiments and an Impala network formulti-game experiments.
Revised: C51, we use a standard CNN network for single-game experiments and an Impala network formulti-game experiments.

Original: For scaling experiments, we vary the number of blocks and channels ineach block of the Impala: the number of blocks and channels is one of (5 blocks , 128 channels) ≈5M params , (10 blocks , 256 channels) ≈ 30M params , (5 blocks , 512 channels) ≈ 60M params ,(10 blocks , 512 channels) 
Revised: For scaling experiments, we vary the number of blocks and channels ineach block of the Impala: the number of blocks and channels is one of (5 blocks , 128 channels) ⇡5M params , (10 blocks , 256 channels) ⇡ 30M params , (5 blocks , 512 channels) ⇡ 60M params ,(10 blocks , 512 channels) ⇡ 120M params .

Original: ≈ 120M params . CPC, BERT, and ACL
Revised: CPC, BERT, and ACL

Original: For rapid adaptation to new games via ﬁne-tuning, we consider representation learning baselines including contrastive predictive coding (CPC) [56], BERT pretraining [18],and attentive contrastive learning (ACL) [79].
Revised: For rapid adaptation to new games via ﬁne-tuning, we consider representation learning baselines including contrastive predictive coding (CPC) [51], BERT pretraining [17],and attentive contrastive learning (ACL) [72].

Original: All state representation networks are implementedas additional multi-layer perceptrons (MLPs) or transformer layers on top of the Impala CNN usedin C51 and CQL baselines.
Revised: All state representation networks are implementedas additional multi-layer perceptrons (MLPs) or transformer layers on top of the Impala CNN usedin C51 and CQL baselines.

Original: CPC uses two additional MLP layers with 512 units each interleavedwith ReLU activation to represent φ ( s ) , which is optimized by maximizing φ ( s ) (cid:62) Wφ ( s (cid:48) ) of truetransitions ( s, s (cid:48) ) and minimizing φ ( s ) (cid:62) Wφ (˜ s ) where ˜ s is a state randomly sampled from the batch(including states from other games).
Revised: CPC uses two additional MLP layers with 512 units each interleavedwith ReLU activation to represent  ( s ) , which is optimized by maximizing  ( s ) > W  ( s 0 ) of truetransitions ( s, s 0 ) and minimizing  ( s ) > W  (˜ s ) where ˜ s is a state randomly sampled from the batch(including states from other games).

Original: For BERT pretraining, we use 2 self-attention layers with 4attention heads of 256 units each and feed-forward dimension 512 , and train φ ( s ) using BERT’smasked self-prediction loss on a trajectory of sequence length 16 .
Revised: For BERT pretraining, we use 2 self-attention layers with 4attention heads of 256 units each and feed-forward dimension 512 , and train  ( s ) using BERT’smasked self-prediction loss on a trajectory of sequence length 16 .

Original: ACL shares the same modelparametrization as BERT, with the inclusion of action prediction in the pretraining objective.
Revised: ACL shares the same modelparametrization as BERT, with the inclusion of action prediction in the pretraining objective.

Original: 4.3 How do different online and ofﬂine methods perform in the multi-game regime?
Revised: 4.3 How do different online and ofﬂine methods perform in the multi-game regime?

Original: We compare different online and ofﬂine algorithms in the multi-game regime and their single-gamecounterparts in Figure 1.
Revised: We compare different online and ofﬂine algorithms in the multi-game regime and their single-gamecounterparts in Figure 1.

Original: We ﬁnd that single-game specialists are still most performant.
Revised: We ﬁnd that single-game specialists are still most performant.

Original: Among multigame generalist models, our Multi-Game Decision Transformer model comes closest to specialistperformance.
Revised: Among multigame generalist models, our Multi-Game Decision Transformer model comes closest to specialistperformance.

Original: Multi-game online RL with non-transformer models comes second, while we struggledto get good performance with ofﬂine non-transformer models.
Revised: Multi-game online RL with non-transformer models comes second, while we struggledto get good performance with ofﬂine non-transformer models.

Original: We note that our multi-game online
Revised: We note that our multi-game online

Original: C51 DQN median score of 68% compares similarly to multi-game median Impala score of 70%,which we calculated from results reported by [20] for our suite of games.
Revised: C51 DQN median score of 68% compares similarly to multi-game median Impala score of 70%,which we calculated from results reported by [19] for our suite of games.

Original: 4.4 How do different methods scale with model size?
Revised: 4.4 How do different methods scale with model size?

Original: In large language and vision models, lowest-achievable training loss typically decreases predictablywith increasing model size.
Revised: In large language and vision models, lowest-achievable training loss typically decreases predictablywith increasing model size.

Original: Kaplan et al.
Revised: Kaplan et al.

Original: [38] demonstrated an empirical scaling relationship betweenthe capacity of a language model (NLP terminology for a next-token autoregressive generative model)and its performance (negative log likelihood on held-out data).
Revised: [34] demonstrated an empirical power law relationshipbetween the capacity of a language model (NLP terminology for a next-token autoregressive generative model) and its performance (negative log likelihood on held-out data).

Original: These trends were veriﬁed over manyorders of magnitude of model size, ranging from few-million parameter models to hundreds of billionparameter models.
Revised: These trends wereveriﬁed over many orders of magnitude of model size, ranging from few-million parameter models tohundreds of billion parameter models.

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

Original: model performance scales with model size, on training set games and novel games. We investigate whether similar trends hold for interactive in-game performance – not just trainingloss – and show a similar  performancescaling trend in Figure 5a.
Revised: We investigate whether similar trends hold for interactive in-game performance – not just training loss– and show a similar power-law performance trend in Figure 5a.

Original: Multi-Game Decision Transformerperformance reliably increases over two orders of magnitude, whereas the other methods eithersaturate, or have much slower performance growth.
Revised: Multi-Game Decision Transformerperformance reliably increases over two orders of magnitude, whereas the other methods eithersaturate, or have much slower performance growth.

Original: We also ﬁnd that larger models train faster, in the sense of reaching higher in-game performance afterobserving the same number of tokens.
Revised: We also ﬁnd that larger models train faster, in the sense of reaching higher in-game performance afterobserving the same number of tokens.

Original: We discuss these resultsin the Appendix.
Revised: We discuss these results the Appendix.

Original: 4.5 How effective are different methods at transfer to novel games?
Revised: 4.5 How effective are different methods at transfer to novel games?

Original: Pretraining for rapid adaptation to new games has not been explored widely on Atari games despitebeing a natural and well-motivated task due to its relevance to how humans transfer knowledge tonew games.
Revised: Pretraining for rapid adaptation to new games has not been explored widely on Atari games despitebeing a natural and well-motivated task due to its relevance to how humans transfer knowledge tonew games.

Original: Nachum and Yang [55] employed pretraining on large ofﬂine data and ﬁne-tunining onsmall expert data for Atari and compared to a set of state representation learning objectives basedon bisimulation [24, 82], but their pretraining and ﬁne-tuning use the same game.
Revised: Nachum and Yang [50] employed pretraining on large ofﬂine data and ﬁne-tunining onsmall expert data for Atari and compared to a set of state representation learning objectives basedon bisimulation [21, 75], but their pretraining and ﬁne-tuning use the same game.

Original: We are insteadinterested in the transfer ability of pretrained agents to new games.
Revised: We are insteadinterested in the transfer ability of pretrained agents to new games.

Original: We hence devise our own evaluation setup by pretraining DT, CQL, CPC, BERT, and ACL on thefull datasets of the 41 training games with 50M steps each, and ﬁne-tuning one model per held-outgame using 1% (500k steps) from each game.
Revised: We hence devise our own evaluation setup by pretraining DT, CQL, CPC, BERT, and ACL on thefull datasets of the 41 training games with 50M steps each, and ﬁne-tuning one model per held-outgame using 1% (500k steps) from each game.

Original: The 1% ﬁne-tuning data is uniformly sampled fromthe 50M step dataset without quality ﬁltering.
Revised: The 1% ﬁne-tuning data is uniformly sampled fromthe 50M step dataset without quality ﬁltering.

Original: DT and CQL use the same objective for pretrainingand ﬁne-tuning, whereas CPC, BERT, and ACL each use their own pretraining objective and are ﬁne-tuned using the BC objective.
Revised: DT and CQL use the same objective for pretrainingand ﬁne-tuning, whereas CPC, BERT, and ACL each use their own pretraining objective and are ﬁne-tuned using the BC objective.

Original: All methods are ﬁne-tuned for 100,000 steps, which is muchshorter than training any agent from scratch.
Revised: All methods are ﬁne-tuned for 100,000 steps, which is muchshorter than training any agent from scratch.

Original: We additionally include training CQL from scratch onthe 1% held-out data to highlight the beneﬁt of rapid ﬁne-tuning.
Revised: We additionally include training CQL from scratch onthe 1% held-out data to highlight the beneﬁt of rapid ﬁne-tuning.

Original: Fine-tuning performance on the held-out games is shown in Figure 6.
Revised: Fine-tuning performance on the held-out games is shown in Figure 6.

Original: Pretraining with the DTobjective performs the best across all games.
Revised: Pretraining with the DTobjective performs the best across all games.

Original: All methods with pretraining outperform training CQLfrom scratch, which veriﬁes our hypothesis that pretraining on other games should indeed help withrapid learning of a new game.
Revised: All methods with pretraining outperform training CQLfrom scratch, which veriﬁes our hypothesis that pretraining on other games should indeed help withrapid learning of a new game.

Original: CPC and BERT underperform DT, suggesting that learning staterepresentations alone is not sufﬁcient for desirable transfer performance.
Revised: CPC and BERT underperform DT, suggesting that learning staterepresentations alone is not sufﬁcient for desirable transfer performance.

Original: While ACL adds an actionprediction auxiliary loss to BERT, it showed little effect, suggesting that modeling the actions inthe right way on the ofﬂine data is important for good transfer performance.
Revised: While ACL adds an actionprediction auxiliary loss to BERT, it showed little effect, suggesting that modeling the actions inthe right way on the ofﬂine data is important for good transfer performance.

Original: Furthermore, we ﬁndthat ﬁne-tuning performance improves as the DT model becomes larger, while CQL ﬁne-tuningperformance is inconsistent with model size (see Figure 5b).
Revised: Furthermore, we ﬁndthat ﬁne-tuning performance improves as the DT model becomes larger, while CQL ﬁne-tuningperformance is inconsistent with model size (see Figure 5b).

Original: [Figure]
Revised: [Table]

Revised: DT performs the best among all methods considered.

Original: 4.6 Does multi-game decision transformer improve upon training data?
Revised: 4.6 Does multi-game decision transformer improve upon training data?

Original: We want to evaluate whether decision transformer with expert action inference is capable of actingbetter than the best demonstrations seen during training.
Revised: We want to evaluate whether decision transformer with expert action inference is capable of actingbetter than the best demonstrations seen during training.

Original: To do this, we look at the top 3 performingdecision transformer model rollouts.
Revised: To do this, we look at the top 3 performingdecision transformer model rollouts.

Original: We use top 3 rollouts instead of the mean across all rolloutsto more fairly compare to the best demonstration, rather than an average expert demonstration.
Revised: We use top 3 rollouts instead of the mean across all rolloutsto more fairly compare to the best demonstration, rather than an average expert demonstration.

Original: Weshow percentage improvement over best demonstration score for individual games in Figure 7.
Revised: Weshow percentage improvement over best demonstration score for individual games in Figure 7.

Original: Wesee signiﬁcant improvement over the training data in a number of games.[Figure] p r o v [Figure]
Revised: Wesee signiﬁcant improvement over the training data in a number of games.

Original: v

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

Original: r

Original: [Figure]

Original: i n

Original: [Figure]

Original: t

Original: [Figure]

Original: s

Original: [Figure]

Original: t

Original: [Figure]

Original: s

Original: [Figure]

Original: t

Original: [Figure]

Original: r n

Original: [Figure] a l C [Figure]
Revised: [Table]

Original: t i pede

Original: [Figure]

Original: hoppe r

Original: [Figure]

Original: o

Original: C r a zy

Original: C li

Original: [Figure]

Original: D e

Original: [Figure]

Original: A tt a c

Original: D

Original: [Figure]

Original: un

Original: [Figure]

Original: ndu

Original: [Figure]

Original: F i s h

Original: i

Original: [Figure]

Original: r

Original: F r

Original: [Figure]

Original: a y

Original: F r o s t b i t e G ophe r G r a v

Original: i t a r H e r o

Original: I c e H

Original: o ck e y

Original: J a m e s bond

Original: K anga r oo K r u ll

Original: K ung F u

Original: M a s

Original: N a m e T h

Original: i s G

Original: P hoen i x

Original: P oo y an Q

Original: [Figure]

Original: r t

Original: R i v e rr a i d

Original: R oad

Original: R unne

Original: R obo t an k

Original: S eaque s t

Original: [Figure]

Original: i

Original: [Figure]

Original: P il o t

Original: U p ND

Original: [Figure]

Original: n

Original: V i deo

Original: [Figure]

Original: i

Original: [Figure]

Original: l [Figure] rather than expert average – demonstration score. 4.7 Does optimal action inference improve upon behavior cloning?
Revised: 4.7 Does optimal action inference improve upon behavior cloning?

Original: In Figure 1 we see that median performance across all games is indeed signiﬁcantly improved bygenerating optimality-conditioned actions.Figure 8 shows the mean and standard deviation of scores A m i d a r A ss a u l t A s t e r i
Revised: In Figure 1 we see that median performance across all games is indeed signiﬁcantly improved bygenerating optimality-conditioned actions.

Original: x

Original: A t l a n t i s

Original: B a n k H

Original: e

Original: i s t

Original: B a tt l e Z o n e

Original: B e a m

Original: R i d e r B o x

Original: i n g

Original: B r e a k o u t

Original: C a r n i v a l

Original: C e n t

Original: i p e d e

Original: C h

Original: o

Original: pp e r C o mm a n d

Original: C r a z y C li m b e r

Original: D e m o n

Original: A tt a c k

Original: D o u b l e

Original: D un k

Original: E n d u r o

Original: F i s h

Original: i n g

Original: D e r b y

Original: F r ee w a y

Original: F r o s t b i t e

Original: G o p h

Original: e r

Original: G r a v i t a r H e r o

Original: I c e H o c k e y

Original: J a m e s b o n d

Original: K a n g a r oo K r u ll

Original: K un g F u

Original: M a s t e r

Original: N a m e T h i s G a m e

Original: P h

Original: o e n i x

Original: P oo y a n Q b e r t

Original: R i v e rr a i d

Original: R o a d R unn e r

Original: R o b o t a n k

Original: S e a q u e s t

Original: T i m e P il o t

Original: U p N D

Original: o w n

Original: V i

Original: d e o P i n b a ll

Original: W i z a r d

Original: O f W o r

Original: Y a r s R e v e n g e

Original: Z a xx o n

Original: D Q N - N o r m a li

Original: z e d

Original: S c o r e [Figure] across all games.
Revised: Figure 8 shows the mean and standard deviation across all games.

Original: While behavior cloning may sometimes produce highly-rewarding episodes, it is [Figure] less likely to do so.
Revised: While behavior cloning may sometimes produce highly-rewarding episodes, it isless likely to do so.

Original: We ﬁnd decision transformer outperforms behavioral cloning in 31 out of 41games.
Revised: We ﬁnd decision transformer outperforms behavioral cloning in 31 out of 41games.

Original:  Conclusion
Revised:  Conclusion

Original: In the quest to develop highly capable and generalist agents, we have made important and measurableprogress.
Revised: In the quest to develop highly capable and generalist agents, we have made important and measurableprogress.

Original: Namely, our results exhibit a clear beneﬁt of using large transformer-based models inmulti-game domains, and the general trends in these results – performance improvements withlarger models and the ability to rapidly ﬁne-tune to new tasks – mirror the successes observed forlarge-scale vision and language models.
Revised: Namely, our results exhibit a clear beneﬁt of using large transformer-based models inmulti-game domains, and the general trends in these results – performance improvements withlarger models and the ability to rapidly ﬁne-tune to new tasks – mirror the successes observed forlarge-scale vision and language models.

Original: Our results also highlight difﬁculties of online RL algorithmsin handling the complexity of multi-game training on Atari.
Revised: Our results also highlight difﬁculties of online RL algorithmsin handling the complexity of multi-game training on Atari.

Original: It is interesting to note that our bestresults are achieved by decision transformers, which essentially learn via supervised learning onsequence data, compared to alternative approaches such as temporal difference learning (more typicalin reinforcement learning), policy gradients, and contrastive representation learning.
Revised: It is interesting to note that our bestresults are achieved by decision transformers, which essentially learn via supervised learning onsequence data, compared to alternative approaches such as temporal difference learning (more typicalin reinforcement learning), policy gradients, and contrastive representation learning.

Original: This begs thequestion of whether online learning algorithms can be modiﬁed to be as “data-absorbent” as DT-likemethods.
Revised: This begs thequestion of whether online learning algorithms can be modiﬁed to be as “data-absorbent” as DT-likemethods.

Original: While even our best generalist agents at times fall short of performance achieved by agentstrained on a single-task, this is broadly consistent with related works that have trained single modelson many tasks [36, 64].
Revised: While even our best generalist agents at times fall short of performance achieved by agentstrained on a single-task, this is broadly consistent with related works that have trained single modelson many tasks [32, 58].

Original: However, our best generalist agents are already capable of outperforming thedata they are trained on.
Revised: However, our best generalist agents are already capable of outperforming thedata they are trained on.

Original: We believe the trends suggest clear paths for future work – that, with largermodels and larger suites of tasks, performance is likely to scale up commensurately.
Revised: We believe the trends suggest clear paths for future work – that, with largermodels and larger suites of tasks, performance is likely to scale up commensurately.

Original: Limitations.
Revised: Limitations.

Original: We acknowledge reasons for caution in over-generalizing our conclusions.
Revised: We acknowledge reasons for caution in over-generalizing our conclusions.

Original: Our resultsare based largely on performance in the Atari suite, where action and observation spaces are alignedacross different games.
Revised: Our resultsare based largely on performance in the Atari suite, where action and observation spaces are alignedacross different games.

Original: It is unclear whether ofﬂine RL datasets such as Atari are of sufﬁcientscale and diversity that we would see similar performance scaling as observed in NLP and visionbenchmarks.
Revised: It is unclear whether ofﬂine RL datasets such as Atari are of sufﬁcientscale and diversity that we would see similar performance scaling as observed in NLP and visionbenchmarks.

Original: Whether we can observe other forms of generalization, such as zero-shot adaptation, aswell as whether our conclusions hold for other settings, remains unclear.
Revised: Whether we can observe other forms of generalization, such as zero-shot adaptation, aswell as whether our conclusions hold for other settings, remains unclear.

Original: Societal Impacts.
Revised: Societal Impacts.

Original: In the current setting, we do not foresee signiﬁcant societal impact as the models 310 are limited to playing simple video games.
Revised: In the current setting, we do not foresee signiﬁcant societal impact as the models 309 are limited to playing simple video games.

Original: We emphasize that our current agents are not intendedto interact with humans or be used outside of self-contained game-playing domains.
Revised: We emphasize that our current agents are not intendedto interact with humans or be used outside of self-contained game-playing domains.

Original: One should 312 exercise increased caution if extending our algorithms and methods to such situations in order toensure any safety and ethical concerns are appropriately addressed.
Revised: One should 311 exercise increased caution if extending our algorithms and methods to such situations in order toensure any safety and ethical concerns are appropriately addressed.

Original: At the same time, the capability 314 of decision making based on reward feedback – rather than purely imitation of the data – has thepotential to be easier to align with human values and goals.Checklist For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
Revised: At the same time, the capability 313 of decision making based on reward feedback – rather than purely imitation of the data – has thepotential to be easier to align with human values and goals.

Original: contributions and scope?

Original: [Yes]

Original: (b) Did you describe the limitations of your work?

Original: [Yes] See Section

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] See

Original: Section

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A]

Original: (b) Did you include complete proofs of all theoretical results?

Original: [N/A]

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main ex-

Original: perimental results (either in the supplemental material or as a URL)?

Original: [No] Code and

Original: models will be released.

Original: (b) Did you specify all the training details (e

Original: g., data splits, hyperparameters,

Original: how they

Original: were chosen)?

Original: [Yes]

Original: (c) Did you report error bars (e

Original: g., with respect to the random seed after running experi-

Original: ments multiple times)?

Original: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type

Original: of GPUs, internal cluster, or cloud provider)?

Original: [Yes]

Original: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]

Original: (b) Did you mention the license of the assets?

Original: [N/A]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: information or offensive content?

Original: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...

Original: text

Original: instructions given

Original: participants and screenshots

Original: participant

Original: Review

Original: [Table]

Original: A Implementation DetailsA.1 Transformer network architecture

Original: The input consists of a sequence of observations, returns, actions and rewards.

Original: Observations areimages in the format B × T × W × H × C .

Original: We use 84 × 84 grayscale images ( i.e. , W = 84 , H =84 , C = 1 ).

Original: Similar to ViT [19], we extract M non-overlapping image patches, perform a linearprojection and then rasterise them into d model -dimensional 1D tokens.

Original: We deﬁne each patch to be14 × 14 pixels ( i.e. , M = 6 × 6 = 36 ).

Original: A learned positional embedding is added to each of the patchtokens

Original: o 1 , ..., o M to retain positional information as in ViT. As described in Section 3.2, returnsare discretized into 120 buckets in {− 20 , ..., 100 } , and rewards are converted to ternary quantities{− 1 , 0 , +1 } .

Original: For the whole sequence (cid:104) ..., o t 1 , ..., o tM , ˆ R t , a t , r t , ... (cid:105) , we learn another positional embedding at eachposition and add to each token embedding.

Original: We experimented with rotary position embedding [71],but did not ﬁnd a signiﬁcant beneﬁt from them in our setting.

Original: On top of the token embeddings, ourtransformer models use a standard transformer decoder architecture.

Original: A standard transformer implementation for sequence modeling would employ a sequential causalattention masking to prevent positions from attending to subsequent positions [76].

Original: However, forthe sequence (cid:104) ..., o t 1 , ..., o tM , ˆ R t , a t , r t , ... (cid:105) that we consider, we do not want to prevent the positioncorresponding to observation token o tm from accessing subsequent observation tokens { o tm (cid:48) : m (cid:48) >m } within the same timestep, since there is no clear sequential causal relation between image patches.

Original: Therefore, we change the sequential causal masking to allow observation tokens within the sametimestep to access each other, but not subsequent positions after o tM , i.e. ˆ R t , a t , r t , o t +11 , ..., o t +1 M , ...

Original: Table 1 summarizes the transformer conﬁgurations we use for each model size.

Original: We train thesemodels on an internal cluster, each with 64 TPUv4.

Original: Due to prohibitively long training times, we onlyevaluated one training seed.

Original: For the 40M DT model, we perform four additional runs with differentseeds to investigate the sensitivity to random seeds.

Original: The mean and standard deviation of the medianscore across 5 random seeds are 0.79 and 0.06; the mean and standard deviation of IQM scores (see

Original: Appendix C for details) are 0.95 and 0.05.

Original: Compared to the differences to the baseline results and theresults of other DT model sizes, we consider the variance to be relatively small.

Original: [Table]

Original: We performed hyperparameter search on DT-40M, and directly applied to DT-200M. With the setof hyperparameters working the best for smaller models, it still manifests a clear scaling trend ofmulti-game and transfer learning performance.

Original: A.2 Fine-tuning protocol for Atari games

Original: In the ﬁne-tuning experiments, we reserved ﬁve games (Alien, MsPacman, Pong, Space Invaders and

Original: Star Gunner) to be used only for ﬁne-tuning.

Original: These games were selected due to their varied gameplaycharacteristics.

Original: Each game was ﬁne-tuned separately to measure the model’s transfer performance fora ﬁxed game.

Original: We use 1% of the original dataset (corresponding to roughly 500 000 transitions) tospeciﬁcally test ﬁne-tuning in low-data regimes.

Original: A.3 Action and return sampling during in-game evaluation

Original: We sample actions from the model with a temperature of 1.

Original: Inspired by Nucleus sampling (Holtzmanet al.

Original: [32]), we only sample from the top 85th percentile action logits for all Decision Transformermodels and Behavioral Cloning models (this parameter was selected to give highest performance forboth models).

Original: While we train the model to predict actions for all timesteps in the sequence, duringin-game evaluation, we execute the last predicted action in the sequence (conditioned on all pastobservations, and past generated actions, rewards, and target returns).

Original: To generate target returns as discussed in Section 3.4, we sample them from the model with thetemperature of 1 and the top 85th percentile logits.

Original: We use κ = 10 in all our experiments.

Original: To avoidstoring the history of previously generated target returns (which may be difﬁcult to incorporate intosome RL frameworks), we experimented with autoregressively regenerating all target returns in thesequence, and found that to work well without requiring any special recurrent state maintenanceoutside of the model.

Original: Algorithm 1 is a pseudocode for our expert return and action inference.

Original: Algorithm

Original: Given an environment E , the current time step t .

Original: κ = 10 , the return upper bound R high = 100 , the return lower bound R low = −

Original: Initialize a context window (cid:104) ..., r t − 2 , o t − 1 1 , ..., o t − 1 M , ˆ R t − 1 , a t − 1 , r t − 1 , o t 1 , ..., o tM (cid:105) , abbreviated as . . .

Original: in the following.

Original: # Autoregressive return and action generation while terminal state not reached yet do

Original: Compute 121 logits ( R t = − 20 , . .

Original: . , 100) for the categorical return distribution P

Original: ( R t | ... )

Original: # Increase logits proportionally to return magnitudes to prefer high magnitude

Original: Deﬁne log P ( R t | expert t , ... )

Original: = log P

Original: ( R t | ... )

Original: + κ

Original: ( R t − R low ) /

Original: ( R high − R low )

Original: # Sample a return

Original: R t ∼ P ( R t | expert t , ... )

Original: Compute logits for the categorical action distribution P ( a t | R t , ... )

Original: # Sample an action a t ∼ P ( a t | R t , ... )

Original: # Interact with the environment o t +11 , ..., o t +1 M , r t ∼ E step ( a t ) t = t + 1 and shift the context window accordingly end while

Original: As an alternative way to generate expert-but-likely returns, we also experimented with simplygenerating N return samples from the model according to log-probability log P θ ( R t | ... ) , and pickingthe highest one.

Original: We then generate the action conditioned on this largest picked return as before.

Original: Thisavoids needing the hyperparameter κ .

Original: In this setting, we found N = 128 , inverse temperature of0.75 for return sampling, no percentile cutoff for return sampling, and sampling from the top 50thpercentile action logits with a temperature of 1 to work similarly well.

Original: A.4 Evaluation protocol and Atari environment details

Original: Our environment is the Atari 2600 Gym environment with pre-processing performed as in Agarwalet al.

Original: Our Atari observations are 84 × 84 grayscale images.

Original: We compress observation imagesto jpeg in the dataset (to keep dataset size small) and during in-game evaluation.

Original: All games use thesame shared set of 18 discrete actions.

Original: For all methods, each game score is calculated by averagingover 16 model rollout episode trials.

Original: To reduce inter-trial variability, we do not use sticky actionsduring evaluation.

Original: A.5 Data augmentation

Original: All models were trained with data augmentations.

Original: We investigate training with the followingaugmentation methods: random cropping, random channel permutation, random pixel permutation,horizontal ﬂip, vertical ﬂip, and random rotations.

Original: We found random cropping and random rotationsto work the best.

Original: (In our random cropping implementation, images of size 84 × 84 are padded oneach side with 4 zero-value pixels, and then randomly cropped to 84 × 84 .)

Original: In general, we aim toexpand the domain of problems solved during training to similar kinds that we hope to generalizeto by encoding useful inductive biases.

Original: We maintain the same random augmentation parameters foreach window sequence.

Original: We apply data augmentation in both pre-training and ﬁne-tuning.

Original: B Baseline Implementation Details

Original: BC Our BC model is effectively the same as our DT model but removing the return token ˆ R t fromthe training sequence:

Original: [Equation]

Original: Instead of predicting a return token (distribution) given observation tokens o t 1 , ..., o tM and theprevious part of the sequence, we directly predict an action token (distribution), which also meansthat we remove return conditioning for the BC model.

Original: During evaluation, we sample actions with atemperature of 1, and sample from the top 85th percentile logits (as discussed in Appendix A.3).

Original: Allother implementation details and conﬁgurations are identical to DT.

Original: C51 DQN For single-game experiments, our implementation and training followed the detailsin [11] except for using multi-step learning with n = 4 .

Original: For multi-game experiments we trained usingthe details provided in the main text; we ran the algorithm for 15M gradient steps ( ≈ 4 B environmentsteps ≈ 16 B Atari frames).

Original: CQL For CQL we use the same optimizer and learning rate as for C51 DQN.

Original: We use a per-replicabatch size of 32 and run for 1M gradient steps on a TPU pod with 32 cores, yielding a global batchsize of 256.

Original: During ﬁnetuning for each game, we copy the entire Q -network trained with CQL, andapply an additional 100k gradient steps of batch size 32 on a single CPU, where each batch is sampledexclusively from the ofﬂine dataset of the ﬁnetuned game.

Original: We also experimented with smallerlearning rates ( 0 . 00003 instead of the default 0 . 00025 ) but found the results largely unchanged.

Original: Wealso tried using ofﬂine C51 and double DQN as opposed to CQL, and found performance to be worse.

Original: CPC For the CPC baseline [56], we apply a contrastive loss between φ ( o t ) , φ ( o t +1 ) using theobjective function−

Original: [Equation]

Original: where W is a trainable matrix and ρ is a non-trainable prior distribution; for mini-batch training weset ρ to be the distribution of states in the mini-batch.

Original: The state representations φ ( o ) is parametrizedby CNNs followed by two MLP layers with 512 units each interleaved with ReLU activation.

Original: Forthe CNN architecture, we used the C51 implementation with an Impala neural network architectureof three blocks using 16, 32, and 32 channels respectively, and trained with a batch size of 256and learning rate of 0.00025 both during pretraining and downstream BC adaptation.

Original: We conductrepresentation learning for a total of 1M gradient steps, and ﬁnetune on 1% data for 100k steps every50k steps of representation learning and report the best ﬁnetuning results.

Original: BERT and ACL Our BERT and ACL baselines are based on the representation learning objectivesdescribed in [79].

Original: For the BERT [18] state representation learning baseline, we (1) take a subtrajectory o t : t + k , a t : t + k , r t : t + k from the dataset (without special tokenization as in DT), (2) randomlymask a subset of these, (3) pass the masked sequence into a transformer, and then (4) for each maskedinput state o t

Original: + i , apply a contrastive loss between its representation φ ( o t + i ) and

Original: the transformeroutput Transformer[ i ] at the corresponding sequence position:

Original: [Equation]

Original: where ρ is the distribution over states in the mini-batch.

Original: For attentive contrastive learning (ACL) [79],we apply an additional action prediction loss to the output of BERT at the sequence positions of theaction inputs.

Original: To parameterize φ , we use the same CNN architecture as in CPC, while the transformer is parameterized by two self-attention layers with 4 attention heads of 256 units each and feed-forward dimension512.

Original: The transformer does not apply any additional directional masking to its inputs.

Original: We usedk = 16 .

Original: Pretraining and ﬁnetuning is analogous to CPC.

Original: Namely, when ﬁnetuning we take the pretrainedrepresentation φ and use a BC objective for learning a neural network (two MLP layers with 512units each) policy on top of this representation.

Original: C Comparisons between methods based on Inter-Quartile Mean Scores

Original: We used median human-normalized scores to aggregate performance over individual games in Figure1 and Figure 5, following [1].

Original: It was discussed in [2] that the median has high variability, and inthe most extreme case, the median is unaffected by zero performance on nearly half of the tasks.

Original: Therefore, we follow the evaluation practice proposed in [2] to also compute Inter-Quartile Mean(IQM) across all games as another aggregate metric.

Original: We present Figure 9 as the IQM version of

Original: Figure

Original: [Figure]

Original: and Figure 10 as the IQM version of Figure 5.

Original: We can observe that our conclusions on thebeneﬁts of Multi-Game Decision Transformers do not change.

Original: Furthermore, in the following sectionsin the Appendix, we report aggregate performance in both median and IQM scores.

Original: [Figure]

Original: [Figure]

Original: ) Scaling of IQM scores for all training games with

Original: [Figure]

Original: D Learning Exclusively from Expert Data

Original: Our DT and baseline models all use the full Atari datasets for training, which include large amountsof sub-optimal, non-expert behavior.

Original: We believe that learning from large, diverse datasets in this wayhelps learning and improves performance.

Original: To verify this hypothesis, we experiment with training our

Original: DT model and baseline BC model on a ﬁltered version of the training data [1], for which we preserveonly the top 10% of episodes from each game according to computed episodic return.

Original: We plot returnhistograms in Figure 11.

Original: We use this expert dataset to train our multi-game decision transformer(DT-40M) and the transformer-based behavioral cloning model (BC-40M).

Original: Figure 12 and Figure 13compare

Original: these models trained on expert data and our DT-40M trained on all data.

Original: We observe that: (1) Training only on expert data improves behavioral cloning; (2) Training on fulldata, including expert and non-expert data, improves Decision Transformer; (3) Decision Transformerwith full data outperforms behavioral cloning trained on expert data.

Original: [Figure]

Original: [Figure]

Original: H u m an - N o r m a li z ed I Q

Original: M S c o r e

Original: [Figure]

Original: E Comparisons between transformers and convolution networks

Original: Decision Transformer is an Upside-Down RL (UDRL) [67, 70] implementation that uses the transformer architecture and considers RL as a sequence modeling problem.

Original: To understand the beneﬁt ofthe transformer architecture, we compare to an UDRL implementation that uses feed-forward, convolutional Impala networks [20].

Original: Speciﬁcally, we use the same return, action, and reward tokenizers asin DT, and only replace the observation (four consecutive Atari frames stacked together) encoding touse the Impala architecture.

Original: Similar to what we do for CQL, we also experiment with different sizesof the Impala architecture by varying the number of blocks and channels in each block of the Impalanetwork: the number of blocks and channels is one of (5 blocks , 128 channels)

Original: ≈ 5M params ,(10 blocks , 256 channels)

Original: ≈ 30M params , (5 blocks , 512 channels)

Original: ≈ 60M params .

Original: We use a(768 , 768) 2 -layer fully-connected head to predict the next return token from observation embedding;another (768 , 768) head to predict the next action token from a concatenation of observation embedding and return token embedding; another (768 , 768) head to predict the next reward token from aconcatenation of observation embedding, return token embedding, and action token embedding.

Original: The input to the model is slightly different from what we have for DT: Instead of considering a

Original: T -timestep sub-trajectory ( T = 4 ) where each timestep contains o t , R t , a t , r t , we stack T imageframes (as common in [53]), and only consider R t , a t , r t from the last timestep.

Original: All other designchoices and evaluation protocols are the same as DT.

Original: [Figure]

Original: [Figure]

Original: Figure 14 shows clear advantages of Decision Transformer over UDRL with the Impala architecture.

Original: In the comparison between UDRL (Impala) and CQL that uses the same Impala network at eachmodel size we evaluated, we observe that UDRL (Impala) outperforms CQL.

Original: The results show thatthe beneﬁts of our method come not only from using network architectures, but also from the UDRL formulation.

Original: While the reasons for the beneﬁts of transformer architecture over other neural networksare still an open question in general (see e.g., [63]), our hypothesis is that transformers allow foreasier discovery of correlations between components of the input and output, due to the fact thattransformers process the input as a ﬂat sequence with attention allowed between any patch, action,or return token.

Original: Although it is not feasible to compare transformer with all possible convolutionalarchitectures due to the broad design space, we believe these empirical results still show a clear trendfavoring both UDRL and transformer architectures.

Original: F Extra Discussion of Online vs. Ofﬂine and Scaling Results

Original: Online vs. ofﬂine results.

Original: (Discussion continued from Section 4.3)

Original: We believe the apparentadvantage of ofﬂine DT compared to online multi-game methods like C51 may be explained inpart through classical differences between online and ofﬂine settings in RL [45].

Original: Online methodsmust balance exploration with the ability to learn and generalize from experience, which could bechallenging in the multi-game setting, whereas ofﬂine DT only needs to learn to distill and generalizefrom the ﬁxed multi-game experience given to it (collected by specialist DQN agents [1]).

Original: Beyondthe difference between online and ofﬂine, one could also argue that C51 suffers from more traininginstability than DT due to the use of a temporal difference (TD) loss, which we discuss in the nextparagraph.

Original: Scaling results.

Original: (Discussion continued from Section 4.4.)

Original: In contrast, in Figures 5a and 5b, weﬁnd that CQL does not improve with increased model size, and actually shows a sharp drop in theperformance of larger models on the ﬁne-tuning tasks.

Original: Temporal Difference (TD) methods suffergreater instability with larger model size in the multi-game setting, leading to this “inverse” scaling.

Original: Indeed, attempts at other objectives closer to pure TD (C51, DQN, DDQN) led to even worse results(which we do not report).

Original: We note that similar conclusions about instability with respect to networksize have been made by other work [22].

Original: G Effect of Model Size on Training Speed

Original: [Figure]

Original: It is believed that large transformer-based language models train faster than smaller models, in thesense that they reach higher performance after observing a similar number of tokens [38, 15].

Original: Weﬁnd this trend to hold in our setting as well.

Original: Figure 15 shows an example of performance on twoexample games as multi-game training progresses.

Original: We see that larger models reach higher scores pernumber of training steps taken (thus tokens observed).

Original: [Figure]

Original: H Qualitative Attention Analysis

Original: [Figure]

Original: We ﬁnd that the Decision Transformer model consistently attends to observation image patches thatcontain meaningful game entities.

Original: Figure 16 visualizes selected attention heads and layers for variousgames.

Original: We ﬁnd heads consistently attend to entities such as player character, player’s free movementspace, non-player objects, and environment features.

Original: [Figure]

Original: I Raw Atari Scores

Original: We report full raw scores of 41 training Atari games for best performing sizes of multi-game modelsin Table

Original: [Table]

Original: [Table]

Original: Frostbite 2,435.

Original: Gopher 9,935.

Original: [Table]

Original: Hero 20,408.

Original: IceHockey -10.

Original: Jamesbond 700.

Original: Kangaroo 12,700.

Original: Krull 8,685.

Original: KungFuMaster 15,562.

Original: NameThisGame 9,056.

Original: Phoenix 5,295.

Original: Pooyan 2,859.

Original: Qbert 13,734.

Original: Riverraid 14,755.

Original: RoadRunner 54,568.

Original: Robotank 63.

Original: Seaquest 5,173.

Original: TimePilot 2,743.

Original: UpNDown 16,291.

Original: [Table]