Original: T RANSFORMER - BASED W ORLD M ODELS A RE H APPY W ITH 100 K I NTERACTIONS
Revised: T RANSFORMER - BASED W ORLD M ODELS A RE H APPY W ITH 100 K I NTERACTIONS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Deep neural networks have been successful in many reinforcement learning settings.
Revised: Deep neural networks have been successful in many reinforcement learning settings.

Original: However, compared to human learners they are overly data hungry.
Revised: However, compared to human learners they are overly data hungry.

Original: To build a sample-efficient world model, we apply a transformer to real-world episodes in an autoregressive manner: not only the compact latent states and the taken actions but also the experienced or predicted rewards are fed into the transformer, so that it can attend flexibly to all three modalities at different time steps.
Revised: To build a sample-efficient world model, we apply a transformer to real-world episodes in an autoregressive manner: not only the compact latent states and the taken actions but also the experienced or predicted rewards are fed into the transformer, so that it can attend flexibly to all three modalities at different time steps.

Original: The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state.
Revised: The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state.

Original: By utilizing the Transformer-XL architecture, it is able to learn long-term dependencies while staying computationally efficient.
Revised: By utilizing the Transformer-XL architecture, it is able to learn long-term dependencies while staying computationally efficient.

Original: Our transformer-based world model (TWM) generates meaningful, new experience, which is used to train a policy that outperforms previous model-free and model-based reinforcement learning algorithms on the Atari 100k benchmark.
Revised: Our transformer-based world model (TWM) generates meaningful, new experience, which is used to train a policy that outperforms previous model-free and model-based reinforcement learning algorithms on the Atari 100k benchmark.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep reinforcement learning methods have shown great success on many challenging decision making problems.
Revised: Deep reinforcement learning methods have shown great success on many challenging decision making problems.

Original: Notable methods include DQN (Mnih et al., 2015), PPO (Schulman et al., 2017), and MuZero (Schrittwieser et al., 2019).
Revised: Notable methods include DQN (Mnih et al., 2015), PPO (Schulman et al., 2017), and MuZero (Schrittwieser et al., 2019).

Original: However, most algorithms require hundreds of millions of interactions with the environment, whereas humans often can achieve similar results with less than 1% of these interactions, i.e., they are more sample-efficient.
Revised: However, most algorithms require hundreds of millions of interactions with the environment, whereas humans often can achieve similar results with less than 1% of these interactions, i.e., they are more sample-efficient.

Original: The large amount of data that is necessary renders a lot of potential real world applications of reinforcement learning impossible.
Revised: The large amount of data that is necessary renders a lot of potential real world applications of reinforcement learning impossible.

Original: Recent works have made a lot of progress in advancing the sample efficiency of RL algorithms: model-free methods have been improved with auxiliary objectives (Laskin et al., 2020b), data augmentation (Yarats et al., 2021, Laskin et al., 2020a), or both (Schwarzer et al., 2021).
Revised: Recent works have made a lot of progress in advancing the sample efficiency of RL algorithms: model-free methods have been improved with auxiliary objectives (Laskin et al., 2020b), data augmentation (Yarats et al., 2021, Laskin et al., 2020a), or both (Schwarzer et al., 2021).

Original: Model-based methods have been successfully applied to complex image-based environments and have either been used for planning, such as EfficientZero (Ye et al., 2021), or for learning behaviors in imagination, such as SimPLe (Kaiser et al., 2020).
Revised: Model-based methods have been successfully applied to complex image-based environments and have either been used for planning, such as EfficientZero (Ye et al., 2021), or for learning behaviors in imagination, such as SimPLe (Kaiser et al., 2020).

Original: o
Revised: o

Original: [Equation]
Revised: [Equation]

Original: z t − ℓ ˆ o t − ℓ a t − ℓ h t − ℓ r t − ℓ · · · · · · o t −z t −ˆ o t −a t −h t −r t −o t z t ˆ o t a t h t ˆ r t ˆ γ t ˆ z t +
Revised: z t − ℓ ˆ o t − ℓ a t − ℓ h t − ℓ r t − ℓ · · · · · · o t −z t −ˆ o t −a t −h t −r t −o t z t ˆ o t a t h t ˆ r t ˆ γ t ˆ z t +

Original: Transformer
Revised: Transformer

Original: A promising model-based concept is learning in imagination (Ha & Schmidhuber, 2018; Kaiser et al., 2020; Hafner et al., 2020; Hafner et al., 2021): instead of learning behaviors from the collected experience directly, a generative model of the environment dynamics is learned in a (self-)supervised manner.
Revised: A promising model-based concept is learning in imagination (Ha & Schmidhuber, 2018; Kaiser et al., 2020; Hafner et al., 2020; Hafner et al., 2021): instead of learning behaviors from the collected experience directly, a generative model of the environment dynamics is learned in a (self-)supervised manner.

Original: Such a so-called world model can create new trajectories by iteratively predicting the next state and reward.
Revised: Such a so-called world model can create new trajectories by iteratively predicting the next state and reward.

Original: This allows for potentially indefinite training data for the reinforcement learning algorithm without further interaction with the real environment.
Revised: This allows for potentially indefinite training data for the reinforcement learning algorithm without further interaction with the real environment.

Original: A world model might be able to generalize to new, unseen situations, because of the nature of deep neural networks, which has the potential to drastically increase the sample efficiency.
Revised: A world model might be able to generalize to new, unseen situations, because of the nature of deep neural networks, which has the potential to drastically increase the sample efficiency.

Original: This can be illustrated by a simple example: in the game of Pong, the paddles and the ball move independently.
Revised: This can be illustrated by a simple example: in the game of Pong, the paddles and the ball move independently.

Original: In the best case, a successfully trained world model would imagine trajectories with paddle and ball configurations that have never been observed before, which enables learning of improved behaviors.
Revised: In the best case, a successfully trained world model would imagine trajectories with paddle and ball configurations that have never been observed before, which enables learning of improved behaviors.

Original: In this paper, we propose to model the world with transformers (Vaswani et al., 2017), which have significantly advanced the field of natural language processing and have been successfully applied to computer vision tasks (Dosovitskiy et al., 2021).
Revised: In this paper, we propose to model the world with transformers (Vaswani et al., 2017), which have significantly advanced the field of natural language processing and have been successfully applied to computer vision tasks (Dosovitskiy et al., 2021).

Original: A transformer is a sequence model consisting of multiple self-attention layers with residual connections.
Revised: A transformer is a sequence model consisting of multiple self-attention layers with residual connections.

Original: In each self-attention layer the inputs are mapped to keys, queries, and values.
Revised: In each self-attention layer the inputs are mapped to keys, queries, and values.

Original: The outputs are computed by weighting the values by the similarity of keys and queries.
Revised: The outputs are computed by weighting the values by the similarity of keys and queries.

Original: Combined with causal masking, which prevents the self-attention layers from accessing future time steps in the training sequence, transformers can be used as autoregressive generative models.
Revised: Combined with causal masking, which prevents the self-attention layers from accessing future time steps in the training sequence, transformers can be used as autoregressive generative models.

Original: The Transformer-XL architecture (Dai et al., 2019) is much more computationally efficient than vanilla transformers at inference time and introduces relative positional encodings, which remove the dependence on absolute time steps.
Revised: The Transformer-XL architecture (Dai et al., 2019) is much more computationally efficient than vanilla transformers at inference time and introduces relative positional encodings, which remove the dependence on absolute time steps.

Original: Our contributions: The contributions of this work can be summarized as follows:
Revised: Our contributions: The contributions of this work can be summarized as follows:

Original: We present a new autoregressive world model based on the Transformer-XL (Dai et al., 2019) architecture and a model-free agent trained in latent imagination.
Revised: We present a new autoregressive world model based on the Transformer-XL (Dai et al., 2019) architecture and a model-free agent trained in latent imagination.

Original: Running our policy is computationally efficient, as the transformer is not needed at inference time.
Revised: Running our policy is computationally efficient, as the transformer is not needed at inference time.

Original: This is in contrast to related works (Hafner et al., 2020; 2021; Chen et al., 2022) that require the full world model during inference.
Revised: This is in contrast to related works (Hafner et al., 2020; 2021; Chen et al., 2022) that require the full world model during inference.

Original: Our world model is provided with information on how much reward has already been emitted by feeding back predicted rewards into the world model.
Revised: Our world model is provided with information on how much reward has already been emitted by feeding back predicted rewards into the world model.

Original: As shown in our ablation study, this improves performance.
Revised: As shown in our ablation study, this improves performance.

Original: We rewrite the balanced KL divergence loss of Hafner et al.
Revised: We rewrite the balanced KL divergence loss of Hafner et al.

Original: (2021) to allow us to fine-tune the relative weight of the involved entropy and cross-entropy terms.
Revised: (2021) to allow us to fine-tune the relative weight of the involved entropy and cross-entropy terms.

Original: We introduce a new thresholded entropy loss that stabilizes the policy’s entropy during training and hereby simplifies the selection of hyperparameters that behave well across different games.
Revised: We introduce a new thresholded entropy loss that stabilizes the policy’s entropy during training and hereby simplifies the selection of hyperparameters that behave well across different games.

Original: We propose a new effective sampling procedure for the growing dataset of experience, which balances the training distribution to shift the focus towards the latest experience.
Revised: We propose a new effective sampling procedure for the growing dataset of experience, which balances the training distribution to shift the focus towards the latest experience.

Original: We demonstrate the efficacy of this procedure with an ablation study.
Revised: We demonstrate the efficacy of this procedure with an ablation study.

Original: We compare our transformer-based world model (TWM) on the Atari 100k benchmark with recent sample-efficient methods and obtain excellent results.
Revised: We compare our transformer-based world model (TWM) on the Atari 100k benchmark with recent sample-efficient methods and obtain excellent results.

Original: Moreover, we report empirical confidence intervals of the aggregate metrics as suggested by Agarwal et al.
Revised: Moreover, we report empirical confidence intervals of the aggregate metrics as suggested by Agarwal et al.

Original:  M ETHOD
Revised:  M ETHOD

Original: We consider a partially observable Markov decision process (POMDP) with discrete time steps t ∈ N , scalar rewards r t ∈ R , high-dimensional image observations o t ∈ R h × w × c , and discrete actions a t ∈ { 1 , . . .
Revised: We consider a partially observable Markov decision process (POMDP) with discrete time steps t ∈ N , scalar rewards r t ∈ R , high-dimensional image observations o t ∈ R h × w × c , and discrete actions a t ∈ { 1 , . . .

Original: , m } , which are generated by some policy a t ∼ π ( a t | o 1: t , a 1: t − 1 ) , where o 1: t and a 1: t − 1 denote the sequences of observations and actions up to time steps t and t − 1 , respectively.
Revised: , m } , which are generated by some policy a t ∼ π ( a t | o 1: t , a 1: t − 1 ) , where o 1: t and a 1: t − 1 denote the sequences of observations and actions up to time steps t and t − 1 , respectively.

Original: Episode ends are indicated by a boolean variable d t ∈ { 0 , 1 } .
Revised: Episode ends are indicated by a boolean variable d t ∈ { 0 , 1 } .

Original: Observations, rewards, and episode ends are jointly generated by the unknown environment dynamics o t , r t , d t ∼ p ( o t , r t , d t | o 1: t − 1 , a 1: t − 1 ) .
Revised: Observations, rewards, and episode ends are jointly generated by the unknown environment dynamics o t , r t , d t ∼ p ( o t , r t , d t | o 1: t − 1 , a 1: t − 1 ) .

Original: The goal is to find a policy π that maximizes the expected sum of discounted rewards E π (cid:2)(cid:80) ∞ t =1 γ t − 1 r t (cid:3) , where γ ∈ [0 , 1) is the discount factor.
Revised: The goal is to find a policy π that maximizes the expected sum of discounted rewards E π (cid:2)(cid:80) ∞ t =1 γ t − 1 r t (cid:3) , where γ ∈ [0 , 1) is the discount factor.

Original: Learning in imagination consists of three steps that are repeated iteratively: learning the dynamics, learning a policy, and interacting in the real environment.
Revised: Learning in imagination consists of three steps that are repeated iteratively: learning the dynamics, learning a policy, and interacting in the real environment.

Original: In this section, we describe our world model and policy, concluding with the training procedure.
Revised: In this section, we describe our world model and policy, concluding with the training procedure.

Original: 2. W ORLD M ODEL
Revised: 2. W ORLD M ODEL

Original: Our world model consists of an observation model and a dynamics model, which do not share parameters.
Revised: Our world model consists of an observation model and a dynamics model, which do not share parameters.

Original: Figure 1 illustrates our combined world model architecture.
Revised: Figure 1 illustrates our combined world model architecture.

Original: Observation Model: The observation model is a variational autoencoder (Kingma & Welling, 2014), which encodes observations o t into compact, stochastic latent states z t and reconstructs the observations with a decoder, which in our case is only required to obtain a learning signal for z t :
Revised: Observation Model: The observation model is a variational autoencoder (Kingma & Welling, 2014), which encodes observations o t into compact, stochastic latent states z t and reconstructs the observations with a decoder, which in our case is only required to obtain a learning signal for z t :

Original: Observation encoder:
Revised: Observation encoder:

Original: [Equation]
Revised: [Equation]

Original: Observation
Revised: Observation

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

Original: We adopt the neural network architecture of DreamerV2 (Hafner et al., 2021) with slight modifications for our observation model.
Revised: We adopt the neural network architecture of DreamerV2 (Hafner et al., 2021) with slight modifications for our observation model.

Original: Thus, a latent state z t is discrete and consists of a vector ofcategorical variables with 32 categories.
Revised: Thus, a latent state z t is discrete and consists of a vector ofcategorical variables with 32 categories.

Original: The observation decoder reconstructs the observation and predicts the means of independent standard normal distributions for all pixels.
Revised: The observation decoder reconstructs the observation and predicts the means of independent standard normal distributions for all pixels.

Original: The role of the observation model is to capture only non-temporal information about the current time step, which is different from Hafner et al.
Revised: The role of the observation model is to capture only non-temporal information about the current time step, which is different from Hafner et al.

Original: However, we include short-time temporal information, since a single observation o t consists of four frames (aka frame stacking, see also Section 2.2).
Revised: However, we include short-time temporal information, since a single observation o t consists of four frames (aka frame stacking, see also Section 2.2).

Original: Autoregressive Dynamics Model: The dynamics model predicts the next time step conditioned on the history of its past predictions.
Revised: Autoregressive Dynamics Model: The dynamics model predicts the next time step conditioned on the history of its past predictions.

Original: The backbone is a deterministic aggregation model f ψ which computes a deterministic hidden state h t based on the history of the ℓ previously generated latent states, actions, and rewards.
Revised: The backbone is a deterministic aggregation model f ψ which computes a deterministic hidden state h t based on the history of the ℓ previously generated latent states, actions, and rewards.

Original: Predictors for the reward, discount, and next latent state are conditioned on the hidden state.
Revised: Predictors for the reward, discount, and next latent state are conditioned on the hidden state.

Original: The dynamics model consists of these components:
Revised: The dynamics model consists of these components:

Original: [Equation]
Revised: [Equation]

Original: Reward predictor
Revised: Reward predictor

Original: [Equation]
Revised: [Equation]

Original: Latent state predictor:
Revised: Latent state predictor:

Original: [Equation]
Revised: [Equation]

Original: The aggregation model is implemented as a causally masked Transformer-XL (Dai et al., 2019), which enhances vanilla transformers (Vaswani et al., 2017) with a recurrence mechanism and relative positional encodings.
Revised: The aggregation model is implemented as a causally masked Transformer-XL (Dai et al., 2019), which enhances vanilla transformers (Vaswani et al., 2017) with a recurrence mechanism and relative positional encodings.

Original: With these encodings, our world model learns the dynamics independent of absolute time steps.
Revised: With these encodings, our world model learns the dynamics independent of absolute time steps.

Original: Following Chen et al.
Revised: Following Chen et al.

Original: (2021), the latent states, actions, and rewards are sent into modality-specific linear embeddings before being passed to the transformer.
Revised: (2021), the latent states, actions, and rewards are sent into modality-specific linear embeddings before being passed to the transformer.

Original: The number of input tokens is 3 ℓ − 1 , because of the three modalities (latent states, actions, rewards) and the last reward not being part of the input.
Revised: The number of input tokens is 3 ℓ − 1 , because of the three modalities (latent states, actions, rewards) and the last reward not being part of the input.

Original: We consider the outputs of the action modality as the hidden states and disregard the outputs of the other two modalities (see Figure 1; orange boxes vs. gray boxes).
Revised: We consider the outputs of the action modality as the hidden states and disregard the outputs of the other two modalities (see Figure 1; orange boxes vs. gray boxes).

Original: The latent state, reward, and discount predictors are implemented as multilayer perceptrons (MLPs) and compute the parameters of a vector of independent categorical distributions, a normal distribution, and a Bernoulli distribution, respectively, conditioned on the deterministic hidden state.
Revised: The latent state, reward, and discount predictors are implemented as multilayer perceptrons (MLPs) and compute the parameters of a vector of independent categorical distributions, a normal distribution, and a Bernoulli distribution, respectively, conditioned on the deterministic hidden state.

Original: The next state is determined by sampling from p ψ (ˆ z t +1 | h t ) .
Revised: The next state is determined by sampling from p ψ (ˆ z t +1 | h t ) .

Original: The reward and discount are determined by the mean of p ψ (ˆ r t | h t ) and p ψ (ˆ γ t | h t ) , respectively.
Revised: The reward and discount are determined by the mean of p ψ (ˆ r t | h t ) and p ψ (ˆ γ t | h t ) , respectively.

Original: As a consequence of these design choices, our world model has the following beneficial properties:
Revised: As a consequence of these design choices, our world model has the following beneficial properties:

Original: The dynamics model is autoregressive and has direct access to its previous outputs.
Revised: The dynamics model is autoregressive and has direct access to its previous outputs.

Original: Training is efficient since sequences are processed in parallel (compared with RNNs).
Revised: Training is efficient since sequences are processed in parallel (compared with RNNs).

Original: Inference is efficient because outputs are cached (compared with vanilla Transformers).
Revised: Inference is efficient because outputs are cached (compared with vanilla Transformers).

Original: Long-term dependencies can be captured by the recurrence mechanism.
Revised: Long-term dependencies can be captured by the recurrence mechanism.

Original: We want to provide an intuition on why a fully autoregressive dynamics model is favorable: First, the direct access to previous latent states enables to model more complex dependencies between them, compared with RNNs, which only see them indirectly through a compressed recurrent state.
Revised: We want to provide an intuition on why a fully autoregressive dynamics model is favorable: First, the direct access to previous latent states enables to model more complex dependencies between them, compared with RNNs, which only see them indirectly through a compressed recurrent state.

Original: This also has the potential to make inference more robust, since degenerate predictions can be ignored more easily.
Revised: This also has the potential to make inference more robust, since degenerate predictions can be ignored more easily.

Original: Second, because the model sees which rewards it has produced previously, it can react to its own predictions.
Revised: Second, because the model sees which rewards it has produced previously, it can react to its own predictions.

Original: This is even more significant when the rewards are sampled from a probability distribution, since the introduced noise cannot be observed without autoregression.
Revised: This is even more significant when the rewards are sampled from a probability distribution, since the introduced noise cannot be observed without autoregression.

Original: Loss Functions: The observation model can be interpreted as a variational autoencoder with a temporal prior, which is provided by the latent state predictor.
Revised: Loss Functions: The observation model can be interpreted as a variational autoencoder with a temporal prior, which is provided by the latent state predictor.

Original: The goal is to keep the distributions of the encoder and the latent state predictor close to each other, while slowly adapting to new observations and dynamics.
Revised: The goal is to keep the distributions of the encoder and the latent state predictor close to each other, while slowly adapting to new observations and dynamics.

Original: Hafner et al.
Revised: Hafner et al.

Original: (2021) apply a balanced KL divergence loss, which lets them control which of the two distributions should be penalized more.
Revised: (2021) apply a balanced KL divergence loss, which lets them control which of the two distributions should be penalized more.

Original: To control the influences of its subterms more precisely, we disentangle this loss and obtain a balanced cross-entropy loss that computes the cross-entropy H ( p ϕ ( z t +1 | o t +1 ) , p ψ (ˆ z t +1 | h t )) and the entropy H ( p ϕ ( z t | o t )) explicitly.
Revised: To control the influences of its subterms more precisely, we disentangle this loss and obtain a balanced cross-entropy loss that computes the cross-entropy H ( p ϕ ( z t +1 | o t +1 ) , p ψ (ˆ z t +1 | h t )) and the entropy H ( p ϕ ( z t | o t )) explicitly.

Original: Our derivation can be found in Appendix A.2.
Revised: Our derivation can be found in Appendix A.2.

Original: We call the cross-entropy term for the observation model the consistency loss , as its purpose is to prevent the encoder from diverging from the dynamics model.
Revised: We call the cross-entropy term for the observation model the consistency loss , as its purpose is to prevent the encoder from diverging from the dynamics model.

Original: The entropy regularizes the latent states and prevents them from collapsing to one-hot distributions.
Revised: The entropy regularizes the latent states and prevents them from collapsing to one-hot distributions.

Original: The observation decoder is optimized via negative log-likelihood, which provides a rich learning signal for the latent states.
Revised: The observation decoder is optimized via negative log-likelihood, which provides a rich learning signal for the latent states.

Original: In summary, we optimize a self-supervised loss function for the observation model that is the expected sum over the decoder loss, the entropy regularizer and the consistency loss
Revised: In summary, we optimize a self-supervised loss function for the observation model that is the expected sum over the decoder loss, the entropy regularizer and the consistency loss

Original: [Equation]
Revised: [Equation]

Original: where the hyperparameters α 1 , α 2 ≥ 0 control the relative weights of the terms.
Revised: where the hyperparameters α 1 , α 2 ≥ 0 control the relative weights of the terms.

Original: For the balanced cross-entropy loss, we also minimize the cross-entropy in the loss of the dynamics model, which is how we train the latent state predictor.
Revised: For the balanced cross-entropy loss, we also minimize the cross-entropy in the loss of the dynamics model, which is how we train the latent state predictor.

Original: The reward and discount predictors are optimized via negative log-likelihood.
Revised: The reward and discount predictors are optimized via negative log-likelihood.

Original: This leads to a self-supervised loss for the dynamics model
Revised: This leads to a self-supervised loss for the dynamics model

Original: L Dyn.
Revised: L Dyn.

Original: ψ
Revised: ψ

Original: [Equation]
Revised: [Equation]

Original: with coefficients β 1 , β 2 ≥ 0 and where γ t = 0 for episode ends ( d t = 1 ) and γ t = γ otherwise.
Revised: with coefficients β 1 , β 2 ≥ 0 and where γ t = 0 for episode ends ( d t = 1 ) and γ t = γ otherwise.

Original: 2. P OLICY
Revised: 2. P OLICY

Original: Our policy π θ ( a t | ˆ z t ) is trained on imagined trajectories using a mainly standard advantage actorcritic (Mnih et al., 2016) approach.
Revised: Our policy π θ ( a t | ˆ z t ) is trained on imagined trajectories using a mainly standard advantage actorcritic (Mnih et al., 2016) approach.

Original: We train two separate networks: an actor a t ∼ π θ ( a t | ˆ z t ) with parameters θ and a critic v ξ (ˆ z t ) with parameters ξ .
Revised: We train two separate networks: an actor a t ∼ π θ ( a t | ˆ z t ) with parameters θ and a critic v ξ (ˆ z t ) with parameters ξ .

Original: We compute the advantages via Generalized
Revised: We compute the advantages via Generalized

Original: Advantage Estimation (Schulman et al., 2016) while using the discount factors predicted by the world model ˆ γ t instead of a fixed discount factor for all time steps.
Revised: Advantage Estimation (Schulman et al., 2016) while using the discount factors predicted by the world model ˆ γ t instead of a fixed discount factor for all time steps.

Original: As in DreamerV2 (Hafner et al., 2021), we weight the losses of the actor and the critic by the cumulative product of the discount factors, in order to softly account for episode ends.
Revised: As in DreamerV2 (Hafner et al., 2021), we weight the losses of the actor and the critic by the cumulative product of the discount factors, in order to softly account for episode ends.

Original: Thresholded Entropy Loss: We penalize the objective of the actor with a slightly modified version of the usual entropy regularization term (Mnih et al., 2016).
Revised: Thresholded Entropy Loss: We penalize the objective of the actor with a slightly modified version of the usual entropy regularization term (Mnih et al., 2016).

Original: Our penalty normalizes the entropy and only takes effect when the entropy falls below a certain threshold
Revised: Our penalty normalizes the entropy and only takes effect when the entropy falls below a certain threshold

Original: [Equation]
Revised: [Equation]

Original: where 0 ≤ Γ ≤ 1 is the threshold hyperparameter, H ( π θ ) is the entropy of the policy, m is the number of discrete actions, and ln( m ) is the maximum possible entropy of the categorical action distribution.
Revised: where 0 ≤ Γ ≤ 1 is the threshold hyperparameter, H ( π θ ) is the entropy of the policy, m is the number of discrete actions, and ln( m ) is the maximum possible entropy of the categorical action distribution.

Original: By doing this, we explicitly control the percentage of entropy that should be preserved across all games independent of the number of actions.
Revised: By doing this, we explicitly control the percentage of entropy that should be preserved across all games independent of the number of actions.

Original: This ensures exploration in the real environment and in imagination without the need for ϵ -greedy action selection or changing the temperature of the action distribution.
Revised: This ensures exploration in the real environment and in imagination without the need for ϵ -greedy action selection or changing the temperature of the action distribution.

Original: We also use the same stochastic policy when evaluating our agent in the experiments.
Revised: We also use the same stochastic policy when evaluating our agent in the experiments.

Original: The idea of applying a hinge loss to the entropy was first introduced by Pereyra et al.
Revised: The idea of applying a hinge loss to the entropy was first introduced by Pereyra et al.

Original: (2017) in the context of supervised learning.
Revised: (2017) in the context of supervised learning.

Original: In Appendix A.1 we show the effect of this loss.
Revised: In Appendix A.1 we show the effect of this loss.

Original: Choice of Policy Input: The policy computes an action distribution π θ ( a t | x t ) given some view x t of the state.
Revised: Choice of Policy Input: The policy computes an action distribution π θ ( a t | x t ) given some view x t of the state.

Original: For instance, x t could be o t , z t , or [ z t , h t ] at inference time, i.e., when applied to the real environment, or the respective predictions of the world model ˆ o t , ˆ z t , or [ˆ z t , h t ] at training time.
Revised: For instance, x t could be o t , z t , or [ z t , h t ] at inference time, i.e., when applied to the real environment, or the respective predictions of the world model ˆ o t , ˆ z t , or [ˆ z t , h t ] at training time.

Original: This view has to be chosen carefully, since it can have a significant impact on the performance
Revised: This view has to be chosen carefully, since it can have a significant impact on the performance

Original: [Figure]
Revised: [Figure]

Original: of the policy
Revised: of the policy

Original: and it affects the design choices for the world model.
Revised: and it affects the design choices for the world model.

Original: Using x t =
Revised: Using x t =

Original: o t (or ˆ o t ) is relatively stable even with imperfect reconstructions ˆ o t , as the underlying distribution of observations p ( o ) does not change during training.
Revised: o t (or ˆ o t ) is relatively stable even with imperfect reconstructions ˆ o t , as the underlying distribution of observations p ( o ) does not change during training.

Original: However, it is also less computationally efficient, since it requires reconstructing the observations during imagination and additional convolutional layers for the policy.
Revised: However, it is also less computationally efficient, since it requires reconstructing the observations during imagination and additional convolutional layers for the policy.

Original: Using x t =
Revised: Using x t =

Original: z t (or ˆ z t ) is slightly less stable, as the policy has to adopt to the changes of the distributions p ϕ
Revised: z t (or ˆ z t ) is slightly less stable, as the policy has to adopt to the changes of the distributions p ϕ

Original: ( z t | o t ) and p ψ (ˆ z t +1 | h t ) during training.
Revised: ( z t | o t ) and p ψ (ˆ z t +1 | h t ) during training.

Original: Nevertheless, the entropy regularizer and consistency loss in Equation (3) stabilize these distributions.
Revised: Nevertheless, the entropy regularizer and consistency loss in Equation (3) stabilize these distributions.

Original: Using x t =
Revised: Using x t =

Original: [ z t , h t ] (or [ˆ z t , h t ] ) provides the agent with a summary of the history of experience, but it also adds the burden of running the transformer at inference time.
Revised: [ z t , h t ] (or [ˆ z t , h t ] ) provides the agent with a summary of the history of experience, but it also adds the burden of running the transformer at inference time.

Original: Model-free agents already perform well on most
Revised: Model-free agents already perform well on most

Original: Atari games when using a stack of the most recent frames (e.g., Mnih et al. 2015; Schulman et al.
Revised: Atari games when using a stack of the most recent frames (e.g., Mnih et al. 2015; Schulman et al.

Original: Therefore, we choose x t = z t and apply frame stacking at inference time in order to incorporate short-time information directly into the latent states.
Revised: Therefore, we choose x t = z t and apply frame stacking at inference time in order to incorporate short-time information directly into the latent states.

Original: At training time we use x t = ˆ z t , i.e., the predicted latent states, meaning no frame stacking is applied.
Revised: At training time we use x t = ˆ z t , i.e., the predicted latent states, meaning no frame stacking is applied.

Original: As a consequence, our policy is computationally efficient at training time (no reconstructions during imagination) and at inference time (no transformer when running in the real environment).
Revised: As a consequence, our policy is computationally efficient at training time (no reconstructions during imagination) and at inference time (no transformer when running in the real environment).

Original: 2. T RAINING
Revised: 2. T RAINING

Original: As is usual for learning with world models, we repeatedly (i) collect experience in the real environment with the current policy, (ii) improve the world model using the past experience, (iii) improve the policy using new experience generated by the world model.
Revised: As is usual for learning with world models, we repeatedly (i) collect experience in the real environment with the current policy, (ii) improve the world model using the past experience, (iii) improve the policy using new experience generated by the world model.

Original: During training we build a dataset D =
Revised: During training we build a dataset D =

Original: [( o 1 , a 1 , r 1 , d 1 ) , . .
Revised: [( o 1 , a 1 , r 1 , d 1 ) , . .

Original: , ( o T , a T , r T , d T )] of the collected experience.
Revised: , ( o T , a T , r T , d T )] of the collected experience.

Original: After collecting new experience with the current policy, we improve the world model by sampling N sequences of length ℓ from D and optimizing the loss functions in Equations (3) and (4) using stochastic gradient descent.
Revised: After collecting new experience with the current policy, we improve the world model by sampling N sequences of length ℓ from D and optimizing the loss functions in Equations (3) and (4) using stochastic gradient descent.

Original: After performing a world model update, we select M of the N × ℓ observations and encode them into latent states to serve as initial states for new trajectories.
Revised: After performing a world model update, we select M of the N × ℓ observations and encode them into latent states to serve as initial states for new trajectories.

Original: The dynamics model iteratively generates these M trajectories of length H based on actions provided by the policy.
Revised: The dynamics model iteratively generates these M trajectories of length H based on actions provided by the policy.

Original: Subsequently, the policy is improved with standard model-free objectives, as described in Section 2.2.
Revised: Subsequently, the policy is improved with standard model-free objectives, as described in Section 2.2.

Original: In Algorithm 1 we present pseudocode for training the world model and the policy.
Revised: In Algorithm 1 we present pseudocode for training the world model and the policy.

Original: Balanced Dataset Sampling: Since the dataset grows slowly during training, uniform sampling of trajectories focuses too heavily on early experience, which can lead to overfitting especially in the low data regime.
Revised: Balanced Dataset Sampling: Since the dataset grows slowly during training, uniform sampling of trajectories focuses too heavily on early experience, which can lead to overfitting especially in the low data regime.

Original: Therefore, we keep visitation counts v 1 , . .
Revised: Therefore, we keep visitation counts v 1 , . .

Original: . , v T , which are incremented every time an entry is sampled as start of a sequence.
Revised: . , v T , which are incremented every time an entry is sampled as start of a sequence.

Original: These counts are converted to probabilities using the softmax function ( p 1 , . .
Revised: These counts are converted to probabilities using the softmax function ( p 1 , . .

Original: . , p T )
Revised: . , p T )

Original: [Equation]
Revised: [Equation]

Original: 6) where τ > 0 is a temperature hyperparameter.
Revised: 6) where τ > 0 is a temperature hyperparameter.

Original: With our sampling procedure, new entries in the dataset are oversampled and are selected more often than old ones.
Revised: With our sampling procedure, new entries in the dataset are oversampled and are selected more often than old ones.

Original: Setting τ = ∞ restores uniform sampling as a special case, whereas reducing τ increases the amount of oversampling.
Revised: Setting τ = ∞ restores uniform sampling as a special case, whereas reducing τ increases the amount of oversampling.

Original: See Figurefor a comparison.
Revised: See Figurefor a comparison.

Original: We empirically show the effectiveness in Section 3.3.
Revised: We empirically show the effectiveness in Section 3.3.

Original: [Figure]
Revised: [Figure]

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: To compare data-efficient reinforcement learning algorithms, Kaiser et al.
Revised: To compare data-efficient reinforcement learning algorithms, Kaiser et al.

Original: (2020) proposed the Atari 100 k benchmark, which uses a subset of 26 Atari games from the Arcade Learning Environment (Bellemare et al., 2013) and limits the number of interactions per game to 100 k. This corresponds to 400 k frames (because of frame skipping) or roughly 2 hours of gameplay, which is 500 times less than the usual 200 million frames (e.g., Mnih et al. 2015; Schulman et al. 2017; Hafner et al. 2021).
Revised: (2020) proposed the Atari 100k benchmark, which uses a subset of 26 Atari games from the Arcade Learning Environment (Bellemare et al., 2013) and limits the number of interactions per game to 100 K. This corresponds to 400 K frames (because of frame skipping) or roughly 2 hours of gameplay, which is 500 times less than the usual 200 million frames (e.g., Mnih et al. 2015; Schulman et al. 2017; Hafner et al. 2021).

Original: We compare our method with five strong competitors on the Atari 100k benchmark: (i) SimPLe (Kaiser et al., 2020) implements a world model as an action-conditional video prediction model and trains a policy with PPO (Schulman et al., 2017), (ii) DER (van Hasselt et al., 2019) is a variant of Rainbow (Hessel et al., 2018) fine-tuned for sample efficiency,
Revised: We compare our method with five strong competitors on the Atari 100k benchmark: (i) SimPLe (Kaiser et al., 2020) implements a world model as an action-conditional video prediction model and trains a policy with PPO (Schulman et al., 2017), (ii) DER (van Hasselt et al., 2019) is a variant of Rainbow (Hessel et al., 2018) fine-tuned for sample efficiency,

Original: (iii) CURL (Laskin et al., 2020b) improves representations using contrastive learning as an auxiliary task and is combined with DER, (iv) DrQ (Yarats et al., 2021) improves DQN by averaging Q-value estimates over multiple data augmentations of observations, and (v) SPR (Schwarzer et al., 2021) forces representations to be consistent across multiple time steps and data augmentations by extending Rainbow with a selfsupervised consistency loss.
Revised: (iii) CURL (Laskin et al., 2020b) improves representations using contrastive learning as an auxiliary task and is combined with DER, (iv) DrQ (Yarats et al., 2021) improves DQN by averaging Q-value estimates over multiple data augmentations of observations, and (v) SPR (Schwarzer et al., 2021) forces representations to be consistent across multiple time steps and data augmentations by extending Rainbow with a selfsupervised consistency loss.

Original: 3. R ESULTS
Revised: 3. R ESULTS

Original: [Figure]
Revised: [Figure]

Original: We follow the advice of Agarwal et al.
Revised: We follow the advice of Agarwal et al.

Original: (2021) who found significant discrepancies between reported point estimates of mean (and median) scores and a thorough statistical analysis that includes statistical uncertainty.
Revised: (2021) who found significant discrepancies between reported point estimates of mean (and median) scores and a thorough statistical analysis that includes statistical uncertainty.

Original: Thus, we report confidence interval estimates of the aggregate metrics median, interquartile mean (IQM), mean, and optimality gap in Figure 3 and performance profiles in Figure 4, which we created using the toolbox provided by Agarwal et al.
Revised: Thus, we report confidence interval estimates of the aggregate metrics median, interquartile mean (IQM), mean, and optimality gap in Figure 3 and performance profiles in Figure 4, which we created using the toolbox provided by Agarwal et al.

Original: The metrics are computed on human normalized scores, which are calculated as (score agent - score random)/ (score human - score random) .
Revised: The metrics are computed on human normalized scores, which are calculated as (score agent - score random)/ (score human - score random) .

Original: We report the unnormalized scores per game in Table 1.
Revised: We report the unnormalized scores per game in Table 1.

Original: We compare with new scores for DER, CURL, DrQ, and
Revised: We compare with new scores for DER, CURL, DrQ, and

Original: SPR that were evaluated on 100 runs and provided by Agarwal et al.
Revised: SPR that were evaluated on 100 runs and provided by Agarwal et al.

Original: They report scores for the improved DrQ( ε ), which is DrQ evaluated with standard ε -greedy parameters.
Revised: They report scores for the improved DrQ( ε ), which is DrQ evaluated with standard ε -greedy parameters.

Original: We perform 5 runs per game and compute the average score overepisodes at the end of training for each run.
Revised: We perform 5 runs per game and compute the average score overepisodes at the end of training for each run.

Original: TWM shows a significant improvement over previous approaches in all four aggregate metrics and brings the optimality gap closer to zero.
Revised: TWM shows a significant improvement over previous approaches in all four aggregate metrics and brings the optimality gap closer to zero.

Original: [Table]
Revised: [Table]

Original: reward.

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: 3. A NALYSIS
Revised: 3. A NALYSIS

Original: In Figure 5 we show imagined trajectories of our world model.
Revised: In Figure 5 we show imagined trajectories of our world model.

Original: In Figure 6 we visualize an attention map of the transformer for an imagined sequence.
Revised: In Figure 6 we visualize an attention map of the transformer for an imagined sequence.

Original: In this example, a lot of weight is put on the current action and the last three states.
Revised: In this example a lot of weight is put on the current action and the last three states.

Original: However, the transformer also attends to states and rewards of the past, with only past actions being mostly ignored.
Revised: However, the transformer also attends to states and rewards in the past, with only past actions being mostly ignored.

Original: The two high positive rewards also get high attention, which confirms that the rewards in the input sequence are used by the world model.
Revised: The two high positive rewards also get high attention, which confirms that the rewards in the input sequence are used by the world model.

Original: We hypothesize that these rewards correspond to some events that happened in the environment and this information can be useful for prediction.
Revised: We hypothesize that these rewards correspond to some events that happened in the environment and this information can be useful for prediction.

Original: An extended analysis can be found in Appendix A.1, including more imagined trajectories and attention maps (and a description of the generation of the plots).
Revised: An extended analysis can be found in Appendix A.1, including more imagined trajectories and attention maps (and a description of the generation of the plots), sample efficiency, stochasticity of the world model, long sequence imagination, and frame stacking.

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

Original:  is higher with balanced sampling, demonstrating its importance.
Revised: The score is higher with balanced sampling, demonstrating its importance.

Original: 3. A BLATION S TUDIES
Revised: 3. A BLATION S TUDIES

Original: Uniform Sampling:
Revised: Uniform Sampling:

Original: To show the effectiveness of the sampling procedure described in Section 2.3, we evaluate three games with uniform dataset sampling, which is equivalent to setting τ = ∞ in Equation (6).
Revised: To show the effectiveness of the sampling procedure described in Section 2.3, we evaluate three games with uniform dataset sampling, which is equivalent to setting τ = ∞ in Equation (6).

Original: In Figure 17 we show that balanced dataset sampling significantly improves the performance in these games.
Revised: In Figure 18 we show that balanced dataset sampling significantly improves the performance in these games.

Original: At the end of training, the dynamics loss from Equation (4) is lower when applying balanced sampling.
Revised: At the end of training, the dynamics loss from Equation (4) is lower when applying balanced sampling.

Original: One reason might be that the world model overfits on early training data and performs bad in later stages of training.
Revised: One reason might be that the world model overfits on early training data and performs bad in later stages of training.

Original: No Rewards: As described in Section 2.1, the predicted rewards are fed back into the transformer.
Revised: No Rewards: As described in Section 2.1, the predicted rewards are fed back into the transformer.

Original: In Figure 8 we show on three games that this can significantly increase the performance.
Revised: In Figure 8 we show on three games that this can significantly increase the performance.

Original: In some games the performance is equivalent, probably because the world model can make correct predictions solely based on the latent states and actions.
Revised: In some games the performance is equivalent, probably because the world model can make correct predictions solely based on the latent states and actions.

Original: In Appendix A.1 we perform additional ablation studies, including the thresholded entropy loss, a shorter history length,and conditioning the policy on [ z, h ] .
Revised: In Appendix A.1 we perform additional ablation studies, including the thresholded entropy loss, a shorter history length, conditioning the policy on [ z, h ] , and increasing the sample efficiency.

Original: [Figure]
Revised: [Figure]

Original: Steps
Revised: Steps

Original: H u m a n N o r m a li z e d S c
Revised: H u m a n N o r m a li z e d S c

Original: o r e
Revised: o r e

Original: RewardsNo Rewards
Revised: RewardsNo Rewards

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

Original: The Dyna architecture (Sutton, 1991) introduced the idea of training a model of the environment and using it to further improve the value function or the policy.
Revised: The Dyna architecture (Sutton, 1991) introduced the idea of training a model of the environment and using it to further improve the value function or the policy.

Original: Ha & Schmidhuber (2018) introduced the notion of a world model , which tries to completely imitate the environment and is used to generate experience to train a model-free agent.
Revised: Ha & Schmidhuber (2018) introduced the notion of a world model , which tries to completely imitate the environment and is used to generate experience to train a model-free agent.

Original: They implement a world model as a VAE (Kingma & Welling, 2014) and an RNN and learn a policy in latent space with an evolution strategy.
Revised: They implement a world model as a VAE (Kingma & Welling, 2014) and an RNN and learn a policy in latent space with an evolution strategy.

Original: With
Revised: With

Original: SimPLe, Kaiser et al.
Revised: SimPLe, Kaiser et al.

Original: (2020) propose an iterative training procedure that alternates between training the world model and the policy.
Revised: (2020) propose an iterative training procedure that alternates between training the world model and the policy.

Original: Their policy operates on pixel-level and is trained using PPO (Schulman et al., 2017).
Revised: Their policy operates on pixel-level and is trained using PPO (Schulman et al., 2017).

Original: Hafner et al.
Revised: Hafner et al.

Original: (2020) present Dreamer and implement a world model as a stochastic RNN that splits the latent state in a stochastic part and a deterministic part; this idea was first introduced by Hafner et al., 2019.
Revised: (2020) present Dreamer and implement a world model as a stochastic RNN that splits the latent state in a stochastic part and a deterministic part; this idea was first introduced by Hafner et al., 2019.

Original: This allows their world model to capture the stochasticity of the environment and simultaneously facilitates remembering information over multiple time steps.
Revised: This allows their world model to capture the stochasticity of the environment and simultaneously facilitates remembering information over multiple time steps.

Original: Robine et al.
Revised: Robine et al.

Original: (2020) use a VQ-VAE to construct a world model with drastically lower number of parameters.
Revised: (2020) use a VQ-VAE to construct a world model with drastically lower number of parameters.

Original: DreamerV2 (Hafner et al., 2021) achieves great performance on the Atari 50 M benchmark after making some changes to Dreamer, the most important ones being categorical latent variables and an improved objective.
Revised: DreamerV2 (Hafner et al., 2021) achieves great performance on the Atari 50 M benchmark after making some changes to Dreamer, the most important ones being categorical latent variables and an improved objective.

Original: Another direction of model-based reinforcement learning is planning, where the model is used at inference time to improve the action selection by looking ahead several time steps into the future.
Revised: Another direction of model-based reinforcement learning is planning, where the model is used at inference time to improve the action selection by looking ahead several time steps into the future.

Original: The most prominent work is MuZero (Schrittwieser et al., 2019), where a learned sequence model of rewards and values is combined with Monte-Carlo Tree Search (Coulom, 2006) without learning explicit representations of the observations.
Revised: The most prominent work is MuZero (Schrittwieser et al., 2019), where a learned sequence model of rewards and values is combined with Monte-Carlo Tree Search (Coulom, 2006) without learning explicit representations of the observations.

Original: MuZero achieves impressive performance on the Atari 50 M benchmark, but it is also computationally expensive and requires significant engineering effort.
Revised: MuZero achieves impressive performance on the Atari 50 M benchmark, but it is also computationally expensive and requires significant engineering effort.

Original: EfficientZero (Ye et al., 2021) improves MuZero and achieves great performance on the Atari 100k benchmark.
Revised: EfficientZero (Ye et al., 2021) improves MuZero and achieves great performance on the Atari 100k benchmark.

Original: Transformers (Vaswani et al., 2017) advanced the effectiveness of sequence models in multiple domains, such as natural language processing and computer vision (Dosovitskiy et al., 2021).
Revised: Transformers (Vaswani et al., 2017) advanced the effectiveness of sequence models in multiple domains, such as natural language processing and computer vision (Dosovitskiy et al., 2021).

Original: Recently, they have also been applied to reinforcement learning tasks.
Revised: Recently, they have also been applied to reinforcement learning tasks.

Original: The Decision Transformer (Chen et al., 2021) and the Trajectory Transformer (Janner et al., 2021) are trained on an offline dataset of trajectories.
Revised: The Decision Transformer (Chen et al., 2021) and the Trajectory Transformer (Janner et al., 2021) are trained on an offline dataset of trajectories.

Original: The Decision Transformer is conditioned on states, actions, and returns, and outputs optimal actions.
Revised: The Decision Transformer is conditioned on states, actions, and returns, and outputs optimal actions.

Original: The Trajectory Transformer trains a sequence model of states, actions, and rewards, and is used for planning.
Revised: The Trajectory Transformer trains a sequence model of states, actions, and rewards, and is used for planning.

Original: Chen et al.
Revised: Chen et al.

Original: (2022) replace the RNN of Dreamer with a transformer and outperform Dreamer on Hidden Order Discovery tasks.
Revised: (2022) replace the RNN of Dreamer with a transformer and outperform Dreamer on Hidden Order Discovery tasks.

Original: However, their transformer has no access to previous rewards and they do not evaluate their method on the Atari 100k benchmark.
Revised: However, their transformer has no access to previous rewards and they do not evaluate their method on the Atari 100k benchmark.

Original: Moreover, their policy depends on the outputs of the transformer, leading to higher computational costs during inference time.
Revised: Moreover, their policy depends on the outputs of the transformer, leading to higher computational costs during inference time.

Original: Concurrent to and independent from our work, Micheli et al.
Revised: Concurrent to and independent from our work, Micheli et al.

Original: (2022) apply a transformer to sequences of frame tokens and actions and achieve state-of-the-art results on the Atari 100k benchmark.
Revised: (2022) apply a transformer to sequences of frame tokens and actions and achieve state-of-the-art results on the Atari 100k benchmark.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we discuss a reinforcement learning approach using transformer-based world models.
Revised: In this work, we discuss a reinforcement learning approach using transformer-based world models.

Original: Our method (TWM) outperforms previous model-free and model-based methods in terms of human normalized score on the 26 games of the Atari 100k benchmark.
Revised: Our method (TWM) outperforms previous model-free and model-based methods in terms of human normalized score on the 26 games of the Atari 100k benchmark.

Original: By using the transformer only during training, we were able to keep the computational costs low during inference, i.e., when running the learned policy in a real environment.
Revised: By using the transformer only during training, we were able to keep the computational costs low during inference, i.e., when running the learned policy in a real environment.

Original: We show how feeding back the predicted rewards into the transformer is beneficial for learning the world model.
Revised: We show how feeding back the predicted rewards into the transformer is beneficial for learning the world model.

Original: Furthermore, we introduce the balanced cross-entropy loss for finer control over the trade-off between the entropy and cross-entropy terms in the loss functions of the world model.
Revised: Furthermore, we introduce the balanced cross-entropy loss for finer control over the trade-off between the entropy and cross-entropy terms in the loss functions of the world model.

Original: A new thresholded entropy loss effectively stabilizes the entropy of the policy.
Revised: A new thresholded entropy loss effectively stabilizes the entropy of the policy.

Original: Finally, our novel balanced sampling procedure corrects issues of naive uniform sampling of past experience.
Revised: Finally, our novel balanced sampling procedure corrects issues of naive uniform sampling of past experience.

Original:  PPENDIX 
Revised: A A PPENDIX A. E XTENDED E XPERIMENTS

Original: [Figure]
Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Original: A. E XTENDED E XPERIMENTS Additional Analysis:
Revised: [Figure] [Figure] [Figure] Additional Analysis:

Original: We provide more example trajectories in Figure 9. 2. 
Revised: We provide more example trajectories in Figure 9. 2. We

Original: We present more attention plots in Figures10 and11.
Revised: present more attention plots in Figures and

Original: All attention maps are generated using the attention rollout method by Abnar & Zuidema (2020).
Revised: All attention maps are generated using the attention rollout method by Abnar & Zuidema (2020)

Original: Note that we had to modify the method slightly, in orderto take the causal masks into account.
Revised: Note that we had to modify the method slightly, in order

Revised: to

Original: Stochasticity : The stochastic prediction of the next state allows the world model to sample a variety of trajectories, even from the same starting state, as can be seen in Figure 12. 4.
Revised: take the causal masks into account. Sample Efficiency : We provide the scores of our main experiments after different amounts of interactions with the environment in Table 2.

Revised: After 50 K interactions, our method already has a higher mean normalized score than previous sample-efficient methods

Revised: Our mean normalized score is higher than DER, CURL, and SimPLe after 25 K interactions.

Revised: This demonstrates the high sample efficiency of our approach.

Revised: Stochasticity : The stochastic prediction of the next state allows the world model to sample a variety of trajectories, even from the same starting state, as can be seen in Figure 12.

Original: Imagining Long Sequences : The world model is trained using sequences of length ℓ = 16 , however, it generalizes well to very long trajectories, as shown in 
Revised: Long Sequence Imagination : The world model is trained using sequences of length ℓ = 16 , however, it generalizes well to very long trajectories, as shown in Figure 13. 6.

Original:  Figure 13.
Revised: Frame Stacking : In Figure 14 we visualize the learned stacks of frames. This shows that the world model encodes and predicts the motion of objects.

Original: Additional Ablation Studies:
Revised: Additional Ablation Studies:

Original: Thresholded Entropy Loss :
Revised: Thresholded Entropy Loss :

Original: In Figure 14 we compare (i) our thresholded entropy loss for the policy (see Section 2.2) with (ii) the usual entropy penalty.
Revised: In Figure 15 we compare (i) our thresholded entropy loss for the policy (see Section 2.2) with (ii) the usual entropy penalty.

Original: For (i) we use the same hyperparameters as in our main experiments, i.e., η = 0 .
Revised: For (i) we use the same hyperparameters as in our main experiments, i.e., η = 0 .

Original: 01 and Γ = 0 .
Revised: 01 and Γ = 0 .

Original: For (ii) we set η = 0 .
Revised: For (ii) we set η = 0 .

Original: 001 and Γ = 1 . 0 ,  effectively disabling the threshold.
Revised: 001 and Γ = 1 . 0 , which effectively disables the threshold.

Original: Without a threshold, the entropy is more likely to either collapse or diverge.
Revised: Without a threshold, the entropy is more likely to either collapse or diverge.

Revised: When the threshold is used, the score is

Revised: higher as well, probably because the entropy is in a more sensible range for the explorationexploitation trade-off.

Original: This cannot be solved by adjusting the penalty coefficient η alone, since it would increase or decrease the entropy in all games.
Revised: This cannot be solved by adjusting the penalty coefficient η alone, since it would increase or decrease the entropy in all games.

Original: History Length : We trained our world model with a shorter history and set ℓ = 4 instead of ℓ = 16 .
Revised: History Length : We trained our world model with a shorter history and set ℓ = 4 instead of ℓ = 16 .

Original: This has a negative impact on the score, as can be seen in Figure 15, demonstrating that more time steps into the past are important.
Revised: This has a negative impact on the score, as can be seen in Figure 16, demonstrating that more time steps into the past are important.

Original: Choice of Policy Input :
Revised: Choice of Policy Input :

Original: In Section 2.2 we explained why the input to the policy is only the latent state, i.e., x = z .
Revised: In Section 2.2 we explained why the input to the policy is only the latent state, i [Table] e., x = z .

Original: In Figure 16 we show that using x = [ z, h ] can result in lower final scores.
Revised: In Figure 17 we show that using x = [ z, h ] can result in lower final scores.

Original: We hypothesize that the policy network has a hard time “keeping up” with the changes of the space of h during training and is not able to ignore this additional information.
Revised: We hypothesize that the policy network has a hard time keeping up with the changes of the space of h during training and cannot ignore this additional information.

Original: s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a  10 [-97 64 
-]
Revised: Increasing the Sample Efficiency : To find out whether we can further increase the sample efficiency shown in Table 2, we train a random subset of games again on 10 [-97 64 
-]K, 25 K, and 50 K interactions with the full training budget that we used for the 100 K interactions. In Figure 19 we see that this can lead to significant improvements in some cases, which could mean that the policy benefits from more training on imagined trajectories, but can even lead to worse performance in other cases, which could possibly be caused by overfitting of the world model. When the performance stays the same even with longer training, this could mean that better exploration in the real environment is required to get further improvements.

Original: Krull h 0 .0 .0 .0 .1 .
Revised: [Table] 121.8 Assault 222.

Original: [-s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
-]
Revised: [-s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
-]682.6 Asterix 210. 1116.6 BankHeist 14. 466.7 BattleZone 2360.0 37187.

Original: Hero h 0 .0 .0 .0 .1 .
Revised: 5068.0 Boxing 0.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: [-s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
-]
Revised: [-s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
-][Table]

Original: BankHeist h 0 .0 .0 .0 .1 .
Revised: [Table] [Table] 77.5 Breakout 1.

Original: [-s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
-]
Revised: [-s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
-]

Original: [Table]
Revised: [Table]

Original: h 0 .0 .0 .0 .1 .
Revised: 20.0 ChopperCommand

Original: [Table]
Revised: [Table]

Revised: 1697.4 CrazyClimber 10780.5 35829.4 62583.6 36700.5 7425.0 14773.2 39456.8 45916.0 67766.2 71820.4 DemonAttack 152.

Original:  Freeway h 0 .0 .0 .0 .1 .
Revised: 350.2 Freeway 0. 24.3 Frostbite 65.

Original: [Figure]
Revised: 1475.6 Gopher 257. 1674.8 Hero 1027.0 30826. 7254.0 Jamesbond 29. 362.4 Kangaroo 52.

Original: s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 3 0.
Revised: 1240.0 Krull 1598. 6349.2 KungFuMaster 258.5 22736.3 14862.5 14783. 574.0 16828.0 16368.0 17946.0 22936.0 24554.6 MsPacman 307.

Original: Freeway h 0 .0 .0 .0 .1 .
Revised: 1588.4 Pong -20. 18.8 PrivateEye 24.9 69571. 86.6 Qbert 163.9 13455. [Table]

Original: Wall-Clock Times
Revised: Wall-Clock Times

Original: : For each run, we give the agent a total training and evaluation budget of roughly 10 hours on a single NVIDIA A100 GPU.
Revised: : For each run, we give the agent a total training and evaluation budget of roughly 10 hours on a single NVIDIA A100 GPU.

Original: The time can vary slightly, since the budget is based on the number of updates.
Revised: The time can vary slightly, since the budget is based on the number of updates.

Original: An NVIDIA GeForce RTX 3090 requires 12-13 hours for the same amount of training and evaluation.
Revised: An NVIDIA GeForce RTX 3090 requires 12-13 hours for the same amount of training and evaluation.

Revised: When using a vanilla transformer, which does not use the memory mechanism of the Transformer-XL architecture (Dai et al., 2019), the runtime is roughly 15 .

Revised: 5 hours on an NVIDIA A100 GPU, i.e., 1 . 5 times higher.

Original: To find out how time consuming the individual parts are, we investigate the throughput of the models, with the batch sizes we use in  our main experiments.
Revised: We compare the runtime of our method with previous methods in Table 3. Our method is more than 20 times faster than SimPLe, but slower than model-free methods. However, our method should be as fast as other model-free methods during inference. In Table 2 we have shown that our method

Revised: [Table]

Revised: achieves a higher human normalized score than previous sample-efficient methods after 50K interactions.

Revised: This suggests that our method could potentially outperform previous methods with shorter training, which would take less than 23 .

Revised: 3 hours.

Revised: To determine how time-consuming the individual parts of our method are, we investigate the throughput of the models, with the batch sizes of our main experiments.

Revised: The Transformer-XL version is almost twice as fast, which again shows the importance of this design choice.

Original: The throughputs were measured on a NVIDIA A100 GPU and are given in (approximate) samples per second (higher is better):
Revised: The throughputs were measured on an NVIDIA A100 GPU and are given in (approximate) samples per second:

Original: • World model training: 16,800 samples/s • World model imagination: 39,000 samples/s •  Policy training: 700,000 samples/s
Revised: • World model training: 16,800 samples/s • World model imagination (Transformer-XL): 39,000 samples/s • World model imagination (vanilla): 19,900 samples/s • Policy training: 700,000 samples/s

Original: We also examine how fast the policy can run in an Atari game.
Revised: We also examine how fast the policy can run in an Atari game.

Original: We measured the (approximate) frames per second on a CPU (since the batch size is 1 ).
Revised: We measured the (approximate) frames per second on a CPU (since the batch size is 1 ).

Original: Conditioning the policy on [ z, h ] is abouttimes slower than z , since the transformer is required:
Revised: Conditioning the policy on [ z, h ] is abouttimes slower than z , since the transformer is required:

Original: • Policy conditioned on z : 653 frames/s • Policy conditioned on [ z, h ] : 213 frames/s
Revised: • Policy conditioned on z : 653 frames/s • Policy conditioned on [ z, h ] : 213 frames/s

Original: A. D ERIVATION OF B ALANCED C ROSS -E NTROPY L OSS
Revised: A. D ERIVATION OF B ALANCED C ROSS -E NTROPY L OSS

Original: Hafner et al.
Revised: Hafner et al.

Original: (2021) propose to use a balanced KL divergence loss to jointly optimize the observation encoder q θ and state predictor p θ with shared parameters θ , i.e.,
Revised: (2021) propose to use a balanced KL divergence loss to jointly optimize the observation encoder q θ and state predictor p θ with shared parameters θ , i.e.,

Original: [Equation]
Revised: [Equation]

Original: where sg ( · ) denotes the stop-gradient operation and λ ∈ [0 , 1] controls how much the state predictor adapts to the observation encoder and vice versa.
Revised: where sg ( · ) denotes the stop-gradient operation and λ ∈ [0 , 1] controls how much the state predictor adapts to the observation encoder and vice versa.

Original: We use the identity D KL ( q ∥ p )
Revised: We use the identity D KL ( q ∥ p )

Original: =
Revised: =

Original: H ( q, p ) − H ( q ) , where H ( q, p ) is the cross-entropy of distribution p relative to distribution q , and show that our loss functions lead to the same gradients as the balanced KL objective, but with finer control over the individual components:
Revised: H ( q, p ) − H ( q ) , where H ( q, p ) is the cross-entropy of distribution p relative to distribution q , and show that our loss functions lead to the same gradients as the balanced KL objective, but with finer control over the individual components:

Original: [Equation]
Revised: [Equation]

Original: since ∇ θ H ( sg ( q θ ))
Revised: since ∇ θ H ( sg ( q θ ))

Original: = 0 and by defining λ 1 = λ and λ 2 = 1 − λ and λ 3 = 1 − λ .
Revised: = 0 and by defining λ 1 = λ and λ 2 = 1 − λ and λ 3 = 1 − λ .

Original: In this form, we have control over the cross-entropy of the state predictor relative to the observation encoder and vice versa.
Revised: In this form, we have control over the cross-entropy of the state predictor relative to the observation encoder and vice versa.

Original: Moreover, we explicitly penalize the entropy of the observation encoder, instead of being entangled inside of the KL divergence.
Revised: Moreover, we explicitly penalize the entropy of the observation encoder, instead of being entangled inside of the KL divergence.

Original: [Figure]{+As common in the literature, we define the loss function by omitting the gradient in Equation (10), so that automatic differentiation computes this gradient. For our world model, we split the objective into two loss functions, as the observation encoder and state predictor have separate parameters, yielding Equations (3) and (4). s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
 Krull h 0 .0 .0 .0 .1 .+}
Revised: {+As common in the literature, we define the loss function by omitting the gradient in Equation (10), so that automatic differentiation computes this gradient. For our world model, we split the objective into two loss functions, as the observation encoder and state predictor have separate parameters, yielding Equations (3) and (4). s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
 Krull h 0 .0 .0 .0 .1 .+}

Original: {+s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
+}
Revised: {+s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
+}

Revised: Hero h 0 .0 .0 .0 .1 .

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

Original: {+s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
+}
Revised: {+s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
+}

Revised: BankHeist h 0 .0 .0 .0 .1 .

Original: {+s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
+}
Revised: {+s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 
+}

Original: [Figure]
Revised: [Table]

Original: [Figure] [Figure]
Revised: h 0 .0 .0 .0 .1 . [Table] Freeway h 0 .0 .0 .0 .1 .

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure] [Figure]
Revised: s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a r s a 10 97 64 3 0. Freeway h 0 .0 .0 .0 .1 .

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

Original: n

Original: [Figure]
Revised: [Figure]

Original:  o r m a li z e d S c
Revised: H u m a n N o r m a li z e d S c

Original: o r e
Revised: o r e

Original: [Figure]
Revised: [Figure]

Original: z x

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

Original: As common in the literature, we define the loss function by omitting the gradient in Equation (10), so that automatic differentiation computes this gradient. For our world model, we split the objective into two loss functions, as the observation encoder and state predictor have separate parameters, yielding Equations (3) and (4).
Revised: [Figure]

Original: A. A DDITIONAL T RAINING D ETAILS
Revised: A. A DDITIONAL T RAINING D ETAILS

Original: In Algorithm 1 we present pseudocode for training the world model and the actor-critic agent.
Revised: In Algorithm 1 we present pseudocode for training the world model and the actor-critic agent.

Original: We use the SiLU activation function (Elfwing et al., 2018) for all models.
Revised: We use the SiLU activation function (Elfwing et al., 2018) for all models.

Original: In Table 2 we summarize all hyperparameters that we used in our experiments.
Revised: In Table 4 we summarize all hyperparameters that we used in our experiments.

Original: In Table 3 we provide the number of parameters of our models.
Revised: In Table 5 we provide the number of parameters of our models.

Original: Pretraining for Better Initialization: During training we need to correctly balance the amount of world model training and policy training, since the policy has to “keep up” with the distributional shift of the latent space.
Revised: Pretraining for Better Initialization: During training we need to correctly balance the amount of world model training and policy training, since the policy has to keep up with the distributional shift of the latent space.

Original: However, we can spend some extra training time on the world model with pre-collected data (included in the 100k interactions) at the beginning of training in order to obtain a reasonable initialization for the latent states.
Revised: However, we can spend some extra training time on the world model with pre-collected data (included in the 100 K interactions) at the beginning of training in order to obtain a reasonable initialization for the latent states.

Original: Algorithm 1 Training the world model and the actor-critic agent.
Revised: Algorithm 1 Training the world model and the actor-critic agent.

Original: function train world model( ) // sample sequences of observations, // rewards, actions and discounts o,a,r,d = sample from dataset() z = encode(o) o hat = decode(z) h
Revised: function train world model( ) // sample sequences of observations, // rewards, actions and discounts o,a,r,d = sample from dataset() z = encode(o) o hat = decode(z) h

Original: = transformer(z,a,r) r hat,d hat,z hat =
Revised: = transformer(z,a,r) r hat,d hat,z hat =

Original: predict(h) //
Revised: predict(h) //

Original: optimize world model via // self-supervised learning optim observation(o,z,o hat,z hat) optim dynamics(r,d,z,r hat,d hat,z hat) //
Revised: optimize world model via // self-supervised learning optim observation(o,z,o hat,z hat) optim dynamics(r,d,z,r hat,d hat,z hat) //

Original: z will be used for imagination return z function train actor critic(z) //
Revised: z will be used for imagination return z function train actor critic(z) //

Original: imagine trajectories of states, // rewards, actions and discounts; // use z as starting point imag =
Revised: imagine trajectories of states, // rewards, actions and discounts; // use z as starting point imag =

Original: [z] for t = 0 until H do a = actor(z) imag.append(a)
Revised: [z] for t = 0 until H do a = actor(z) imag.append(a)

Original: h =transformer(imag)r,d,z=predict(h) imag.extend([r,d,z]) //
Revised: h=transformer(imag) r,d,z = predict(h) imag.extend([r,d,z]) //

Original: optimize actor-critic via // reinforcement learning optim actor critic(imag)
Revised: optimize actor-critic via // reinforcement learning optim actor critic(imag)

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