Transformer-based World Models Are Happy With 100k Interactions
Authors: Jan Robine
Marc Höftmann
Tobias Uelwer
Stefan Harmeling
Venue: ICLR 2023 poster (2023)
Keywords: Model-based Reinforcement Learning
World Models
Transfomers
Atari 100k benchmark

Abstract: Deep neural networks have been successful in many reinforcement learning settings. However, compared to human learners they are overly data hungry. 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. The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state. By utilizing the Transformer-XL architecture, it is able to learn long-term dependencies while staying computationally efficient. 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. Our code is available at https://github.com/jrobine/twm.

Published as a conference paper at ICLR 2023
TRANSFORMER-BASER WORLD MODELS ARE HAPPY WITH 100K INTERACTIONS
Jan Robine, Mare Höftmann, Tobias Uelwer, Stefan Harmeling Department of Computer Science, Technical University of Dortmund, Germany
ABSTRACT
Deep neural networks have been successful in many reinforcement learning set- tings. However, compared to human learners they are overly data hungry. To build a sample-efficient world model, we apply a transformer to real-world episodes 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. The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state. By utilizing the Transformer-Xi archi- tecture, it is able to learn long-term dependencies while staying computationally efficient. Our transformer-bae based world model (TWM) generates meaningful, new experience, which is used to train policy that outperforms previous model-free and model-based reinforcement learning algorithms on the Atari 100k benchmark. Our code is available at https:://// gigghu..b...... com/ jrobi ne
INTRODUCTION
Deep reinforcement learning methods have shown great success many challenging decision mak- ing problems. Notable methods include DQN (Mnih et al.] 2015). PPO (Schulman et al., (017)). and MuZero (Schrittwieser et al. 2019). 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. The large amount of data that is necessary renders a lot of potential real world applications of reinforcement learning impossible.
Recent works have made 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 aug- mentation (aarats al. 202 Laskin et al.) (020a). or both (Schwarzer et al. [2021). Model-based methods have been successfully applied to complex image-based environments and have either been used for planning, such as EfficientZero Yee et al., 2021), or for learning behaviors in imagination,
such as SimPLe (Kaiser et al. 2020).
ht-6
h
h
Transformer
Zt
a
Ti-E
1
at-1
Zt
Pt
ive
2 +++1
Ot-F
Ot-1
Ot
Ôt
Figure Our world model architecture. Observations oc-it are encoded using a CNN. Linear embeddings of stochastic, discrete latent states 2t-Eits actions 0g-Lit. and rewards Tt-é: are fed into a transformer, which computes a deterministic hidden state h at each time step. Predictions of the reward rt, discount factor Ys and next latent state 21+1 computed based on hy using MLPs.

---

Published as a conference paper at ICLR 2023
A promising model-based concept is learning in imagination (Ha & Schmidhuber 2018 Kaiser ct al. Hainer et al. |020 Hafner et al. (021): instead of learning behaviors from the collected experience directly, a generative model of the environment dynamics is learned in (self- Jsupervised manner. Such a so-called world model can create new trajectories by iteratively predicting the next state and reward. This allows for potentially indefinite training data for the reinforcement learning algorithm without further interaction with the real environment. 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. This can be illustrated by simple example: in the game of Pong, the paddles and the ball move independently. 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.
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 A transformer is a sequence model consisting of multiple self-attention layers with residual connections. each self-attention layer the inputs are mapped to keys, queries, and values. The outputs computed by weighting the values by the sim- ilarity of keys and queries. Combined with causal masking, which prevents the self-ateention layers from accessing future time steps in the training sequence, transformers can be used as autoregressive generative models. The Transformer-Xi architecture (Dai et al. 2019) is much more computation- ally efficient than vanilla transformers at inference time and introduces relative positional encodings, which remove the dependence on absolute time steps.
Our contributions: The contributions of this work can be summarized follows:
1. We present new autoregressive world model based on the Transformer-XI. Dai et al. 2019) architecture and a model-free agent trained in latent imagination. Running our policy is computationally efficient, as the transformer is not needed at inference time. This is in contrast to related works (Hafner et al. 2020 2021 world model during inference.
Chen et al., 2022) that require the full
2. Our world model is provided with information on how much reward has already been emit- ted by feeding back predicted rewards into the world model. As shown in our ablation study, this improves performance.
3. We rewrite the balanced KL divergence loss Hafner et al. (202)) to allow us to fine-tune the relative weight of the involved entropy and cross-entropy terms.
4. 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.
5. We propose new effective sampling procedure for the growing dataset of experience, which balances the training distribution to shift the focus towards the latest experience. We demonstrate the efficacy of this procedure with an ablation study.
6. We compare our transformer-basee world model (TWM) on the Atari 100k benchmark with recent sample-efficient methods and obtain excellent results. Moreover, we report empirical confidence intervals of the aggregate metrics as suggested by Agarwal et al. (202)
2 METHOD
We consider a partially observable Markov decision process (POMDP) with discrete time steps I E N. scalar rewards E IR. high-dimensional image observations o, E Rhxxwe, and discrete actions 0, E {l...., m). which are generated by some policy aj ~ T (a 01. 01:-1). where 0n. and G1--1 denote the sequences of observations and actions up time steps t and t- 1, respectively. Episode ends are indicated by a boolean variable d, E 10. Observa- tions, rewards, and episode ends are jointly generated by the unknown environment dynamics ot, Tt. d, ~ p(0t,, d, 01 (1:--)) The goal is to find a policy л that maximizes the expected sum of discounted rewards E E 'r where y E [0, 1) is the discount factor. Learning in imagination consists of three steps that repeated iteratively: learning the dynamics, learning policy, and interacting in the real environment. In this section, we describe our world model and policy, concluding with the training procedure.
2

---

Published as a conference paper at ICLR 2023
2.1 WORLD MODEL
Our world model consists of an observation model and a dynamics model, which do not share parameters. Figure Tigureeiillsee Iilluttrate combined world model architecture.
Observation Model: The observation model is a variational autoencoder (Kingma & Welling 2014), which encodes observations or into compact, stochastic latent states 2 and reconstructs the observations with decoder, which in our case is only required to obtain a learning signal for 2y:
Observation encoder:
Po(Zz lxx
(1)
Observation decoder: ô₃ ~ Pp(0 21).
We adopt the neural network architecture of Dreamer V2 (Hafner et al. 2021 with slight modifi-
cations for our observation model. Thus, latent state 2, is discrete and consists of a vector of 32 categorical variables with 32 categories. The observation decoder reconstructs the observation and predicts the means of independent standard normal distributions for all pixels. The role of the ob- servation model is capture only non-temporal information about the current time step, which is different from Hafner et al. (202I). However, we include short-time temporal information, since single observation or consists of four frames (aka frame stacking, see also Section 2.2).
Autoregressive Dynamics Model: The dynamics model predicts the next time step conditioned on the history of its past predictions. The backbone is a deterministic aggregation model fe which computes deterministic hidden state he based on the history of the l previously generated latent states, actions, and rewards. Predictors for the reward, discount, and next latent state are conditioned on the hidden state. The dynamics model consists of these components:
Aggregation model:
h4 f At-Ed, 11-81-1)
Reward predictor:
P. (f he)
Discount predictor:
"e ~ polfi he)
Latent state predictor:
4+ ~ pr h2)-
The aggregation model is implemented causally masked Transformer-Xi (Dai et al. 2019), which enhances vanilla transformers (Vaswani et al. (217)) with a recurrence mechanism and relative positional encodings. With these encodings, our world model learns the dynamics independent of absolute time steps. Following Chen et al. (202T), the latent states, actions, and rewards are sent into modality-specific linear embeddings before being passed to the transformer. The number of input tokens 3 1, because of the three modalities (latent states, actions, rewards) and the last reward not being part of the input. We consider the outputs of the action modality as the hidden states and disregard the outputs the other two modalities (see Figure 1 orange boxes vs. gray boxes).
The latent state, reward, and discount predictors implemented multilayer perceptrons (MLPs) and compute the parameters of a vector of independent categorical distributions, a normal distribu- tion, and Bernoulli distribution, respectively, conditioned on the deterministic hidden state. The next state is determined by sampling from Py (f hx). The reward and discount are determined by the mean of (ft hτ) and p he), respectively.
As consequence of these design choices, our world model has the following beneficial properties:
1. The dynamics model is autoregressive and has direct access to its previous outputs.
2. Training is efficient since sequences are processed in parallel (compared with RNNs).
3. Inference is efficient because outputs are cached (compared with vanilla Transformers).
4. Long-term dependencies can be captured by the recurrence mechanism.
We want to provide an intuition 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. This also has the potential to make inference more robust, since degenerate predictions can be ignored more easily. Second, because the model to its own predictions. This is even more significant when the rewards are sampled from a probability distribution, since the introduced noise cannot be observed without autoregression.
which rewards it has produced previously, it can react
3

---

Published as a conference paper at ICLR 2023
Loss Functions: The observation model can be interpreted variational autoencoder with temporal prior, which is provided by the latent state predictor. The goal is to keep the distribu- tions of the encoder and the latent state predictor close each other, while slowly adapting to new observations and dynamics. Hafner et al. (2021) apply a balanced KI. divergence loss, which lets them control which of the distributions should be penalized more. 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 [[ explicitly. Our derivation can be found in Appendix A.2 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. The entropy regularizes the latent states and prevents them from collapsing to one-hot distributions. The observation decoder is optimized via negative log-likelihood which provides rich learning signal for the latent states. 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
P.p()) he)) and the entropy H(po(z) o.)
2
E
an Hp (ze oo)) a2g((pp(z Po h(-1)) In m))
(3)
decoder
eatropy regularizer
consistency
where the hyperparameters 0n, 02 2 0 control the relative weights of the terms.
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. The reward and discount predictors are optimized via negative log-likelihng noellikelhoood. This leads to self-supervised loss for the dynamics model
T Σ
Dyn.
E
H(po(ze11 Pp I(a))
:3 In po hw) - G2 In pro pre (he))
(4)
latent state predictor
reward predictor
discount predictor
with coefficients B1, 32 2 and where 7/t for episode ends (de = 1) and 7/2 = 7 otherwise.
2.2 POLICY
Our policy <a (2)) is trained imagined trajectories using mainly standard advantage actor- critic (Mnih et al. 2016) approach. We train two separate networks: an actor aę ~> Ty (o 2) with parameters and critic (&t) with parameters Ç. We compute the advantages via Generalized Advantage Estimation (Schulman et al. 2016) while using the discount factors predicted by the world model ft instead of a fixed discount factor for all time steps. As in DreamerV2 (Hafner 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.
Thresholded Entropy Loss:
We penalize the objective of the actor with a slightly modified ver- sion of the usual entropy regularization term (Mnih et al. (016). Our penalty normalizes the entropy and only takes effect when the entropy falls below certain threshold
( Eat max 0, 1- In(m)
C
(5)
where 0 S r S. 1 is the threshold hyperparameter, H(T)) is the entropy of the policy, is the number of discrete actions, and In(m) is the maximum possible entropy of the categorical action distribution. By doing this, we explicitly control the percentage of entropy that should be preserved across all games independent of the number of actions. This ensures exploration in the real environ- ment and in imagination without the need for e-greedy action selection changing the temperature of the action distribution. We also use the same stochastic policy when evaluating our agent in the experiments. The idea of applying hinge loss the entropy was first introduced by Pereyra et al. (2017) in the context of supervised learning. In Appendix A.T we show the effect of this loss.
The policy computes an action distribution (a xx) given some view I, of the state. For instance, 2; could be ot, 2t, [&,, hy] at inference time, i.e., when applied to the real environment, or the respective predictions of the world model ô, 2, or [&,, hy] at training time. This view has to be chosen carefully, since it can have significant impact on the performance
Choice of Policy Input:
4

---

Published as a conference paper at ICLR 2023
8000
1.0
6000
Unitorm (r x]
0.8
(cm-a)
* Sampled
0.6
4000
0.4
0.2
0.0
0K
20K
40K
60K
BOK
100K
0K 20K
40K
60K
BOK
100K
Dataset
Dataset
Figure 2: Comparing our balanced dataset sampling procedure (see Equation 6 ) for different values of r with uniform sampling (T oo). The x-axes correspond to the entries in dataset D in the order they are experienced. The left plot shows the number of times an entry has been selected for training the world model. The right plot shows the relative amount of training time that has been spent on the data up to that entry. E.g., with uniform sampling, 50% of the training time is used for the first 19K entries, whereas for 20 approximately the same time is spend on both halves of the dataset.
of the policy and it affects the design choices for the world model. Using 2y Or (or o,) is rela- tively stable even with imperfect reconstructions o,, as the underlying distribution of observations p(o) does change during training. However, it is also less computationally efficient, since requires reconstructing the observations during imagination and additional convolutional layers for the policy. Using 2₁ = 2t (or & is slightly less stable, as the policy has to adopt to the changes of the distributions (2 o;) and pu(( he) during training. Nevertheless, the entropy reg- ularizer and consistency loss Equation (3) stabilize these distributions. Using 2, [2t, he] (or [&,, he)) provides the agent with a summary of the history of experience, but it also adds the bur- den of running the transformer inference time. Model-free agents already perform well on most Atari games when using stack of the most recent frames (e.g.. Mnih et al. 2015 Schulman et al.) 2017]. Therefore, we choose H, = 2, and apply frame stacking at inference time in order to incor- porate short-time information directly into the latent states. At training time use I₁ &, i.e., the predicted latent states, meaning no frame stacking is applied. As a consequence, our policy is computationally efficient at training time (no reconstructions during imagination) and inference time (no transformer when running in the real environment).
2.3 TRAINING
As is usual for learning with world models, we repeatedly (i) collect experience in the real environ- ment 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.
During training build dataset D [(01, 01, (r,d1,d)))) ri,dd),..... (or, ar. dr-)] of the collected perience. 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 33) and (4) using stochastic gradient descent. After performing a world model update, we select M of the N x observations and encode them into latent states to serve as initial states for new trajectories. The dynamics model iteratively generates these M trajectories of length H based on actions provided by the policy. Subsequently, the policy is improved with standard model-free objectives, as described in Section 2.2 In Algorithm we present pseudocode for training the world model and the policy.
Balanced Dataset Sampling: Since the dataset grows slowly during training, uniform sampling of trajectories focuses too heavily on early experience, which lead overfitting especially in the low data regime. Therefore, we keep visitation counts U1,.. ...,, DT, which are incremented every time an entry is sampled as start of sequence. These counts are converted to probabilities using the softmax function
(6) where 7 > is a temperature hyperparameter. With our sampling procedure, new entries in the dataset are oversampled and selected more often than old ones. Setting F restores uniform sampling as a special case, whereas reducing 7 increases the amount of oversampling. See Figure 2 for comparison. We empirically show the effectiveness in Section 3.3
(p......py) = softmax (-
5

---

Published as a conference paper at ICLR 2023
Median
IQM
Mean
Optimality Gap
TWM (ours) SPR DrQ(e)
CURL
DER
SimPLe
0.15 0.30 0.45
0.60
0.15
0.30
0.45
0.50
0.75 1.00
0.56 0.64 0.72
Human Normalized Score
Figure 3: Aggregate metrics on the Atari 100k benchmark with 95% stratified bootstrap confidence intervals (Agarwal et al., 202T). Higher median, interquartile mean (IQM), and mean, but lower optimality gap indicate better performance. Scores for previous methods are from Agarwal et al. (202T) with 100 runs per game (except SimPLe with 5 runs). We evaluate runs per game, leading to wider confidence intervals.
3 EXPERIMENTS
To compare data-efficient reinforcement learning algorithms, Kaiser et al. (2020) proposed the Atari 100k benchmark, which uses subset of 26 Atari games from the Arcade Learning Environment [Bellemare et al. (2013) and limits the number of interactions per game to 100K. This corresponds 400K frames (because of frame skipping) or roughly hours of gameplay, which 500 times less than the usual 200 million frames (eigh
Mnih et al. 2015 Schulman et al. 2017 Hafner et al. 2021).
We compare our method with five strong competitors on the Atari 100k benchmark: (i) SimPLe (Kaiser et al. 2020) implements world model as an conditionl video prediction model and trains 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, (ini) CURL. (Laskin et al. 2020b)
improves representations using contrastive learning as an auxiliary task and is combined with DER,
(iv) DrQ (Yarats al. 2021) improves DQN by averaging Q-value estimates over multiple data augmentations of observations, and (v) SPR (Schwarzer et al. (021)) forces representations to be consistent across multiple time steps and data augmentations by extending Rainbow with self- supervised consistency loss.
3.1 RESULTS
We follow the advice of Aggaraal et al. (2021) who found significant discrepancies between reported point estimates of mean (and median) scores and thorough statistical analysis that includes statis-
Fractoon of runs with score 7
1.00
TWM (ours)
Dro
0.75
DURL
tical uncertainty. Thus, we report confidence in- terval estimates of the aggregate metrics median,
0.50
interquartile mean (IQM), mean, and optimality gap in Figure 3 and performance profiles in Fig- ure 4, which we created using the toolbox provided
0.25
by Agarwal al.] (2021). The metrics puted on human normalized scores, which are cal- culated as (score _agetttaaat____ score score_rradoml random) _arore_human__
4
8
Human Normalized Score (t)
Figure 4: Performance profiles on the Atari 100k benchmark based on score distributions (Agarwal et al.) 2021). It shows the fraction of runs across all games (y-axis) above hu- man normalized score (x-axis). Shaded re- gions show pointwise 95% confidence bands.
(score human score randeor_a_om...___ We report the unnormalized scores per game in Table T We compare with new scores for DER, CURL, DrQ, and SPR that were evaluated on 100 runs and provided by Agarwal al. (2021). They report scores for the improved DrQ()), which is DrQ evaluated with stan- dard e-greedy parameters. We perform 5 runs per game and compute the average over 100 episodes at the end training for each run. TWM shows a significant improvement over previous approaches in all four aggregate metrics and brings the optimality gap closer to

---

Published as a conference paper at ICLR 2023
Table 1: Mean scores on the Atari 100k benchmark per game as well the aggregated human normalized mean and median. We perform 5 runs per game and compute the average over 100 episodes the end of training for each run. Bold numbers indicate the best scores.
Model-free
Imagination
Game
Random
Human
DER
CURL
DiQ(e)
SPR
SimPLe
TWM (ours)
Alien
227.8
7127.7
802.3
711.0
865.2
841.9
616.9
674.6
Amidar
5.8
1719.5
125.9
113.7
137.8
179.7
74.3
121.8
527.2
682.6
Asterix
210.0
8503.3
535.4
567.2
763.6
962.5
1128.3
1116.6
BeikH
185.5
65.3
466.7
BatileZone
2360.0
37187.5
8977.0
8997.8
10165.3
14834,1
4031.2
5068.0
Boxing
0.1
12.1
-0.3
0.9
9.0
35.7
7.8
77.5
Breakout
1.7
30.5
19.8
19.6
16.4
20.0
ChopperCommand
811.0
7387.8
925,9
783.5
844.6
946.3
979.4
1697.4
Crimerr
36700,5
71820.4
ArtooAttack
152.1
1971.0
627.6
646.5
1321.5
517.6
208.1
350.2
Freeway
0.0
29.6
20.9
28.3
20.3
19.3
16.7
Froshite
65.2
4334.7
871.0
1226.5
1014.2
1170.7
236.9
1475.6
Gopher
257.6
2412.5
467.0
400.9
621.6
660.6
$96.8
1674.8
Hero
1027.0
30826.4
6226.0
4987.7
4167.9
5858.6
7254.0
Jamesbond
29.0
302.8
275.7
33LD
349.1
366.5
100.5
362.4
Kangaroo
Krull
1598.0
2665.5
3256.9
3049.2
4402.1
3681.6
2204.8
63499.2
KuMstera
258.5
22733.3
6580.1
8155.6
11467.4
14783.2
14862.5
24554.6
MisPacman
307.3
6951.6
1187.4
1064.0
1218.1
1318.4
1480.0
1588.4
Pong
-20.7
14,6
-9.7
-18.5
-9.1
-5.4
12.8
18.8
PrivateEye
81.9
35.0
Qbert
163.9
13455.0
1773.5
727.0
1810.7
866.3
1288.8
3330.8
ReadRunner
Seaquest
68.4
42054,7
304.6
315.2
352.3
558.1
683.3
774,4
UpNDown
533.4
11693.2
3075.0
2646.4
4324.5
10859.2
3350.3
15981.7
Normalized Mcan
0.000
1,000
0.350
0.261
0.465
0.616
0.332
0.956
Nomalized Median
0.505
0.0
0.0
0.0
0,0
0.0
A
0.0
0.8
0.0
0.0
0.0
E
3
(a) Boxing. The player (white) presses fire, hits the opponent, and gets reward.
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
(b) Freeway. The player moves and bumps into car. The world model correctly pushes the player down, although up is still pressed. The movement of the cars is modeled correctly.
Trajectories imagined by our world model. Above each frame show the performed action and the produced reward.
Figure 5:
3.2 ANALYSIS
In Figure we show imagined trajectories of our world model. In Figure & we visualize an attention map of the transformer for imagined sequence. In this example lot of weight is put the current action and the last three states. However, the transformer also attends states and rewards in the past, with only past actions being mostly ignored. The two high positive rewards also get high attention, which confirms that the rewards in the input sequence are used by the world model. We hypothesize that these rewards correspond some events that happened in the environment and this information can be useful for prediction.
An extended analysis attention maps (and description of the generation the plots), sample efficiency, stochasticity of the world model, long sequence imagination, and frame stacking.
be found in Appendix A.I including more imagined trajectories and

---

Published as a conference paper at ICLR 2023
Figure 6: Attention map the learned transformer for the current hidden state h, computed on an imagined trajectory for the game Assault. The x-axis corresponds to the input sequence with the three modalities (states, actions, rewards), where the two rightmost columns are the current state and action. The y-axis corresponds to the layer of the transformer.
0.330
Go Luunced -01
1200
1.20
0:20
E600
6.400
3.500
5.152
6.200
3.200
0.300
Figure 7: Comparison of the proposed balanced sampling procedure with uniform sampling random subset of games. We show the development of the human normalized score in the course of training. The score is higher with balanced sampling, demonstrating its importance.
3.3 ABLATION STUDIES
Uniform Sampling: To show the effectiveness of the sampling procedure described in Section 2.3 evaluate three games with uniform dataset sampling, which is equivalent setting in Equation (6. In Figure 7 we show that balanced dataset sampling significantly improves the performance in these games. At the end of training, the dynamics loss from Equation (4) is lower when applying balanced sampling. One reason might be that the world model overfits on early training data and performs bad in later stages of training.
oo
No Rewards: As described in Section 2.1 the predicted rewards are fed back into the transformer. In Figure g we show on three games that this can significantly increase the performance. In some games the performance is equivalent, probably because the world model can make correct predic- tions solely based on the latent states and actions.
In Appendix I.I perform additional ablation studies, including the thresholded entropy loss, shorter history length, conditioning the policy on [2, h). and increasing the sample efficiency.
Barkslent
Dresiont
CowryCi
Rawards
0200
1.000
0.500
3,000
0.300
Poro
Obert
2.500
Fawards
1200
3.200
31500
6300
6200
0120
2170
0.030
6.000
2008
14B
Figure 8: Effect of removing rewards from the input. We show the human normalized score dur- ing training of a random subset of games. Conditioning on rewards can significantly increase the performance. Some games do not benefit from the rewards and the score stays roughly the same.
8

---

Published as a conference paper at ICLR 2023
4 RELATED WORK
The Dyna architecture (Sutton 1991) introduced the idea of training a model of the environment and using it to further improve the value function the policy. [Ha & Schmidhuber (2018) intro- duced the notion of a world model, which tries to completely imitate the environment and is used to generate experience to train model-free agent. They implement world model VAE (Kingma & Welling, 2014) and an RNN and learn policy in latent space with an evolution strategy. With SimPLe, Kaiser et al. (220)) propose an iterative training procedure that alternates between train- ing the world model and the policy. Their policy operates on pixel-level and is trained using PPO (Schulman et al., 2017). |Hafner et al. (2020) present Dreamer and implement a world model as a stochastic RNN that splits the latent state in stochastic part and deterministic part; this idea was first introduced by Hafner et al. 2019 This allows their world model to capture the stochasticity of the environment and simultancously facilitates remembering information over multiple time steps. Robine et al. 2020) use a VQ- VAE to construct a world model with drastically lower number of pa- rameters. DreamerV2 (Hafner et al. (202T) achieves great performance on the Atari 50M benchmark after making some changes to Dreamer, the most important ones being categorical latent variables and improved objective.
Another direction of model-based reinforcement learning is planning, where the model is used at inference time improve the action selection by looking ahead several time steps into the future. 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. MuZero achieves impressive performance on the Atari 50M benchmark, but is also computationally expensive and requires significant engineering effort. EfficientZero (Ye et al. 2021) improves MuZero and achieves great performance on the Atari 100k benchmark.
Transformers (Vaswani et al. 2017) advanced the effectiveness of sequence models in multiple do- mains, such as natural language processing and computer vision (Dosovitskiy et al. (02I). Recently, they have also been applied to reinforcement learning tasks. The Decision Transformer (Chen al. 2021) and the Trajectory Transformer (Janner et al., 2021) are trained on an offline dataset of trajec- tories. The Decision Transformer is conditioned
states, actions, and returns, and outputs optimal actions. The Trajectory Transformer trains a sequence model of states, actions, and rewards, and is used for planning. Chen et al. (2022) replace the RNN of Dreamer with transformer and out- perform Dreamer on Hidden Order Discovery tasks. However, their transformer has no access to previous rewards and they do not evaluate their method on the Atari 100k benchmark. Moreover, their policy depends on the outputs of the transformer, leading to higher computational costs during inference time. Concurrent to and independent from our work, Micheli et al. (2022) apply trans- former to sequences of frame tokens and actions and achieve state-of-thed art results on the Atari 100k benchmark.
CONCLUSION
In this work, we discuss reinforcement learning approach using transformer- based world models. 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. By using the transformer only dur- ing training, we were able to keep the computational costs low during inference, i.e., when running the learned policy in real environment. We show how feeding back the predicted rewards into the transformer is beneficial for learning the world model. Furthermore, 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. A new thresholded entropy loss effectively stabilizes the en- tropy of the policy. Finally, our novel balanced sampling procedure corrects issues of naive uniform sampling of past experience.
introduce the balanced
REFERENCES
Samira Abnar and Willem H. Zuidema. Quantifying attention flow in transformers. In Dan Jurafsky, Joyce Chai, Natalie Schluter. and Joel R. Tetreault (eds.). Proceedings of the 58th Annual Meet- ing of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pp-

---

Published as a conference paper at ICLR 2023
4190- Association for Computational Linguistics, 2020. doi: 10. 10.18633//1/2020..r main. 385. [ttt[ttts////o.cccmm orq///00......//////////// wt....////....333.......
11w.116633// (1//////////200../////// ac1-2a000...mmman 385
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C. Courville, and Marc G. Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. In Marc' Au- relio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.). Advances Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021. December 6-14, 2021, virtual, pp. 29320,, 2021. URL htttps::/proceeeiiiie ngs nwwrips....... hash/ F514cec15 cb 4tw.44ccccc11115151544141411 (ewe:4755577724d5---- 42eeeeeee-eeetaeat.... ht.ll
cc/ neguriiippp poccpapprr22221
Mare G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning envi- ronment: An evaluation platform for general agents. J. Artif. Intell. Res., 47:253-279, 2013, doi:
10.16133jjir.3911. URL attps:://oii wrt//0///........ 1ttpr//...11111.......1
16l116113//1 jair. 3912
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learnin
with transformer world models. CoRR, abs/2202.09481 2022. URL
abs/2202. 09481
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Rein- forcement learning via sequence modeling. In Marc' Aurelio Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Infor- mation Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. URL
150977 2021.
http://prrorddding....ee....ppp. nwwrips. ec/ pttp://ppppe///////// htmll
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In H. Jaap van den Herik, Paolo Ciancarini, and H. H. L. M. Donkers (eds.), Computers and Games, 5th In- ternational Conference, CG 2006, Turin, Italy, May 29-31, 2006. Revised Papers, volume 4630 of Lecture Notes in Computer Science, pp. 72 83. Springer, 2006. doi: 10. 0001009777----400-00 0. 100 1ww.11100000---1 0-000777776-------5-----81
_7. URL bttps:://oi..ooo
Zihang Dai, Zhilin Yang. Yiming Yang, Jaime G. Carbonell, Quoc Viet Lc, and Ruslan Salakhut- dinov. Transformer-xl. Attentive language models beyond a fixed-length context. In Anna Ko- rhonen, David R. Traum, and Lluís Mårquez (eds.). Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pp. 2978-2988, Association for Computational Linguistics, 2019. doi: 10. 10.18653/w1/p19-1225 URL //ooi.///oo/i/ org/ 10. www.1166///11111110
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference Learning Representutiinn, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. id-YichEe dNTTy
net forum?
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018. doi: 10. URL 1011
10.00.0006600 10w.//0...e..e.......mmmmmm nww.ee.......... 2017.12.
012
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Ro- man Gamett (eds.), Advances in Neural Information Processing Systems 31: Annual Confer- ence on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 2455-2467, 2018. URL https: / http::///oceeeeeiig.ssem neurips. cc/ Abst .html
Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri
10

---

Published as a conference paper at ICLR 2023
and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2555-25655 PMLR, 2019. URL:t:t::/prroeeiiiii
ml ww...pressvv99. /ha fnerl 1 9a. html
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In 8th International Conference on Learning Repre- sentations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL
new ew net/ fo
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL [tLLLtttt////peeeee//// met//fouue www.ouummmdddoaaaaaaaeee
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Gheshlaghi Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Sheila A. Mcllraith and Kilian Q. Weinberger (eds.). Proceedings of the Thirty-Secone AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp. 3215-3222. AAAI Press, 2018. URL https: /www www..aa............a............ .org/
ocs/ ocs/index.om . php AAAIAAAIAAAAAAAA
wwwppaaaipiiaaaa 17204
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 1273-1286, 2021. URL http:: htttt::::/ppoceeeiiioo Saper/2021.
www.riiiiiinnnnnnns..eee.o neurips. cc/ html
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL attps://cpenreseee net//
Forum? S1dd JHE
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Bansf, AB, Canada, April 14-16, 2014, Conference Track Proceedings. 2014. URL htttt::/tttt...o/oo/
abs//312.61J1 6114
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Had- sell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Process- ing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 2020, virtual, 2020a. URL httllats://rrccceiinnneeo www.:eeeeeddingg...n...s.eese..e neurips...ppppp.p co/ haper/2020/hasse
.tml
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: contrastive unsupervised representa- tions for reinforcement learning. In Proceedings of the 37th International Conference on Ma- chine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 5639-5650, PMLR, 2020b. URL http:: mlr. wwwppppwsss....peee 11t119911aann0001 html
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are efficient world models. CoRR, abs/2209.0055e2 2022. doi: 10448850/arxii.22200000 00588. URL https://doi. org/ 10. 40.488500arxiv.. 220. 00588
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen,
11

---

Published as a conference paper at ICLR 2023
Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforeement learning. Nat., 518(75400:522-533 2015. doi: 10.1038/maturee44236. URL https:://///... /tt:://dooi.ooo org/10.
1038/naturre 4236
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicraa, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Maria-Florina Balcan and Kilian Q. Weiinberger (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19- 24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 1928-1937, JMLR.org, 2016.
.ml.engsmmmmr.....mmmmmmm.m press/ / v48 1148/mnihalle html
Gabriel Pereyra, George Tucker, Jan Chorowski, Lukasz Kaiser, and Geoffrey E. Hinton. Regu- larizing neural networks by penalizing confident output distributions. In 5th International Con- ference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Workshop Track Proceedings. OpenReview.net, 2017. URL
HyhbYrGYe
Jan Robine, Tobias Uelwer, and Stefan Harmeling, Smaller world models for reinforcement learning. CoRR, abs/2010.05767, 2020. URL
05767
Julian Schritwwessr, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. CoRR, abs/191 1.08265, 2019. [RLphtpp: / /tatt::///xxii.....// www.ariii....ooooooooooooo abs /991. 08265
John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. Yoshua Bengio and Yann LeCun (eds.), 4th International Conference Learning Representations, ICLR 2016, San Juan, Puerto Rico May 2-4, 2016, Conference Track Proceedings, 2016. URL http:
//arxiv.org abs 1566. 02438
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/1707. 066447777777 2017. URL http:://rxiv..rrt abs/
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
www.nee//ae fonttfoormm?id
uCQfPZwRaUK
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bull,, 164) 1991. doi: 10.1145/122344112237, URL https://doi.orgg
1145/ 01255/12344.... 122377
Hado van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in reinforcement learning? In Hanna M. Wallach, Hugo Larochelle, Alina Beygelz- imer, Florence d' A'Alhh--Bue Emily B. Fox, and Roman Garnett (eds.). Advances in Neu- ral Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 14322-14333, 2019. URL https::/proceeddineee ngs nwwrips....... httpcccccpp//11////// Abstract.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Gameett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998 0008,600800000000 2017. URL o hee. http://wwppppppppppp///////aa/aa
12

---

Published as a conference paper at ICLR 2023
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In 9th International Conference on Learning Repre- sentations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https:://peenevie ew. net / forum? 1d= 68 TvGaf!
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Perey Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34; Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 13476-254488 2021. URL http:tttss://ppoceeeeess
ngs neur ips..pps... cc/
228.://20haaaa11aaa0000000222223 -bbataaccbaa html
13

---

Published as a conference paper at ICLR 2023
A
APPENDIX
A.1
EXTENDED EXPERIMENTS
1
(a) Boxing. The player (white) presses fire, misses the opponent, gets reward, and retreats.
0.0
(b) Freeway.
The player moves up, steps back because of approaching car, and continues move up.
(c) Jamesbond. The player presses fire fire missile, and gets reward when it hits.
(d) Gopher. The player (farmer) closes created by the opponent (gopher). The gopher keeps moving independent from the selected actions, indicating that the world model has correctly learned the conrelations between the player and the actions. Note that the gopher briefly disappears and reappcars on the other side.
Figure 9: Additional example trajectories generated by our world model.
Additional Analysis:
1. We provide more example trajectories in Figure 9
2. We present more attention plots in Figures 10 and II All attention maps generated using the attention rollout method by Abnar Zuidema 2020). Note that we had to modify the method slightly, in order take the causal masks
account.
3. Sample Efficiency: We provide the scores of our main experiments after different amounts of interactions with the environment in Table After 50K interactions, method already has a higher mean normalized score than previous sample-efficient methods. Our mean normalized score higher than DER, CURL, and SimPLe after 25K interactions. This demonstrates the high sample efficiency of our approach.
4. 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 [2]
5. Long Sequence Imagination: The world model is trained using sequences of length 16, however, it generalizes well to very long trajectories, as shown Figure
6. Frame Stacking: In Figure 14 we visualize the learned stacks of frames. This shows that the world model encodes and predicts the motion objects.
Additional Ablation Studies:
1. Thresholded Entropy Loss: In Figure 15 compare (i) our thresholded entropy loss for the policy (see Section 2.2 with (ii) the usual entropy penalty. For (i) we use the same hyperparameters in main experiments, i.c., 7 0.01 and I = 0.1. For (ii) we set 77 = 0.001 and I = 1.0, which effectively disables the threshold. Without a threshold, the entropy is more likely to either collapse diverge. When the threshold is used, the score is
14

---

Published as a conference paper at ICLR 2023
higher well, probably because the entropy is in more sensible range for the exploration- exploitation trade-off. This cannot be solved by adjusting the penalty coefficient 7) alone, since it would increase decrease the entropy in all games.
2. History Length: We trained our world model with a shorter history and set & = 4 instead of
8 8= 16. This has a negative impact on the score, as can be seen in Figure T6 demonstrating
that more time steps into the past are important.
3. Choice of Policy Input: In Section 2.2 we explained why the input to the policy is only the
latent state, ie.. I z. In Figure 17 we show that using æ [z, h] can result in lower final scores, We hypothesize that the policy network has hard time keeping up with the changes of the space of h during training and cannot ignore this additional information.
4. Increasing the Sample Efficiency: To find out whether we can further increase the sample efficiency shown in Table [] we train a random subset of games again on 10K, 25K, and 50K interactions with the full training budget that used for the 100K interactions. In Figure 18] 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.
Table 2:
Performance of our method at different stages of training compared with final scores of previous methods. We show individual game scores and mean human normalized scores. The normalized mean of our method higher than SimPLe after only 25K interactions, and higher than previous methods after 50K interactions.
TWM (ours)
Game
Random
Human
SimPI.e
5K
10K
25K
50K
75K
100K
Alien
227,8
7127.7
616.9
841,9
202.8
383.2
463.6
$32.0
776.6
674.6
Amidar
5.8
1719.5
74.3
179.7
3.8
35.4
54.9
101.3
103.0
121.8
Assault
222.4
742.0
527.2
565.6
241.5
315.4
418.7
466.8
627.8
682.6
Asterix
210.0
8503.3
1128.3
962.5
277.0
297.0
536.0
886.0
1116.6
BankHeist
14,2
7
34.2
345,4
17.6
4.4
17.4
125.2
288.4
466.7
BattleZone
2360.0
37187.5
4031.2
14834.1
2640.0
3120.0
2700.0
3740.0
5260.0
5068.0
Boxing
0.1
12.1
7.8
35.7
0.8
3.4
28.5
60,1
67.1
77.5
Breakout
1.7
30.5
16.4
19.6
1.0
5.9
6.9
12.5
15.0
20.0
ChopperCommand
811.0
7387.8
979.4
946.3
928.0
1044.0
1358.0
1306.0
1438.0
1697.4
CruzClimber
10780.5
35829.4 4
62583.6
36700.5
7425.0
14773.2
39456.8
45916.0
67766.2
71820.4
DemonArtack
152.1
1971.0
208.1
517.6
174.7
184.4
216.8
391.4
350.2
Frecway
0.0
29.6
16.7
19.3
0.0
4.6
20.8
23.7
23.9
24.3
Frostbite
65.2
4334.7
236.9
1170.7
66.2
204.6
297.8
247.6
1165.4
1475.6
Gopher
257.6
2412.5
596.8
660.6
345.2
414.0
593.2
1213.2
1549.2
1674.8
Hero
1027.0
30826.4
2656.6
5858.6
448.9
1552.6
4790.9
6302.7
9403.8
7254.0
Jamesbond
29.0
302.8
100.5
366.5
117.0
172.0
215.0
322.0
362.4
Kangaroo
$2.0
3035.0
51.2
3617.4
28.0
92.0
476.0
724.0
876.0
1240.0
Krull
1598.0
2665.5
2204.8
3681.6
1763.6
2552.8
4234.0
4699.2
5848.0
6349.22
KungPuMasser
258.5
22736.3
14862.5
14783,2
574.0
16828.0
16368.0
17946.0
22936.0
24554.6
MsPacman
307.3
6951.6
1480.0
1318.4
245.9
$35.1
1077.5
1224.3
1287.6
1588.4
Pong
-20.7
14.6
12.8
-5.4
-20.4
-19.8
-7.7
8.D
19.9
18.8
PrivateFye
24.9
69571.3
35.0
86.0
61.0
80.0
80.0
3.2
88.8
86.6
Qbert
163.9
13455.0
1288.8
866.3
151.0
298.5
703.5
1046.5
1788.5
3330.8
RoadRunner
11.5
7845.0
5640.6
12213.1
24.0
1120.0
5178.0
7436.0
8034.0
9109.0
Seaquest
68.4
42054.7
683.3
558.1 1
76.8
221.2
428.4
572.0
704.0
774.4
UpNDown
533.4
11693.2
3350.3
10839.2
385.8
1963.0
2905.6
4922.8
10478.6
15981.7
Nomalized Mean
0.000
1.000
0.332
0.616
0.007
0.133
0.408
0.624
0.832
0.956
For each run, we give the agent a total training and evaluation budget of roughly 10 hours on a single NVIDIA A A100 GPU. The time can vary slightly, since the budget is based on the number of updates. An NVIDIA GeForce RTX 3090 requires 12-13 hours for the same amount of training and evaluation. 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.5 hours on an NVIDIA A100 GPU, i.e., 1.5 times higher.
Wall-Clock Times:
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 method
15

---

Published as a conference paper at ICLR 2023
Table 3: Approximate runtime (i.e., training and evaluation time for single run) of our method compared with previous methods that also evaluate on the Atari 100k benchmark. Runtimes of previous methods are taken from [chwarzer et al. (202T). They used an improved version of DER (van Hasselt al., 2019). which is roughly equivalent to DrQ Yarats et al., [02T], so the specified runtime might differ from the original DER implementation. There are data augmented versions for SPR and DER. All runtimes are measured on a single NVIDIA PI00 GPU.
Method
Model-basd based
Runtime in hours
SimPLe
500
TWM (ours)
23,3
SPR (with aug.)
4.6
SPR (w/o aug.)
3.0
DER/DrQ (with aug.)
2.1
DER/DrQ (w/o aug.)
1.4
achieves higher human normalized score than previous sample-efficient methods after 50K inter- actions. This suggests that our method could potentially outperform previous methods with shorter training, which would take less than 23.3 hours.
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. The Transformer-XL ver- sion is almost twice as fast, which again shows the importance of this design choice. The throughputs were measured on an NVIDIA A100 GPU and are given in (approximate) samples per second:
World model training: 16,800 samples/s
World model imagination (Transformer-Xi XL)): 39,000 samples/s
World model imagination (vanilla): 19,900 samples/s
Policy training: 700,000 samples/s
We also examine how fast the policy can run an Atari game. We measured the (approximate)
frames per second on CPU (since the batch size is 1). Conditioning the policy on [2, h. is about times slower than 2, since the transformer is required:
Policy conditioned on 2: 653 frames/s
Policy conditioned on [2, h]: 213 frames/s
A.2 DERIVATION OF BALANCED CROSS-ENTROPY LOSS
Hafner al. (2021) propose to use a balanced KL divergence loss to jointly optimize the observation encoder qe and state predictor pe with shared parameters 0, i.e.,
DRL (sg(g)) || pe) + (1 A) DKL. (ge sg(p))):
(7)
where sg(.) denotes the stop- gop-gradient operation and À E [0, 1] controls how much the state predictor adapts to the observation encoder and vice versa. We use the identity D (g || p) H(q. p) H
where H(q. p) is the cross-entropy of distribution relative to distribution g, and show that our loss
functions lead to the same gradients the balanced KL objective, but with finer control the individual components:
Ve [A (sg(ge) || pe) + (1- A) Dxl || sg(pe))]
(8)
Ve (H((() po) H(sg(g)) + (1- A) (H(go. sg(pe)) - H(go))]
(9)
V. |A H(sg(go) po) >₂ H(go, sg(po)) As H(go)),
(10)
and by defining A X and A = X and X3 = 1 X. In this form, have control over the cross-entropy of the state predictor relative to the observation encoder and vice versa. Moreover, we explicitly penalize the entropy of the observation encoder, instead of being entangled inside of the KL divergence.
since V. H(s((go)
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 world model, we split the objective into two loss functions, as the observation encoder and state predictor have separate parameters, yielding Equations 3 and
16

---

Published as a conference paper at ICLR 2023
Kndl
(a) This world model focuses on previous states.
Her
(b) This world model focuses previous actions, indicating that the effect of ac- tions can last longer than single time step.
Ronk
(c) This world model attends all three modalities in the recent past.
Bottiaoos
(d) This world model attends states at steps. probably because complexity of this 3D game.
(c) This world model mainly focuses on four states at specific time steps.
Figure 10: Average attention maps of the transformer, computed over many time steps. They show how different games require a different focus on modalities and time steps.
Frezwey
Figure 11: Attention map for Freeway for a single time step. At this point the player hits a car and gets pushed back (see also Figure 5b) and the world model puts more attention to past states and rewards, compared with the average attention at other time steps, as shown in Figure T0c The world model has learned to handle this situation separately.
17

---

Published as a conference paper at ICLR 2023
11
11B
Mi
0.0
N
0.0
0.0
0.0 0.0
Figure 12: Three trajectories for the game KungFuMaster generated by our world model, using the starting state. Because of its stochastic nature, the world model is able to generate three different situations (one opponent, two opponents, one other type of opponent). Note that we only show every third frame cover more time steps.
0.0
A
0.0
0.0
0.0
0.0
0.0
A
1.0
0.0
2.0
A
5,0
5.0
8
1.0
A
0.0
0.0
0.0
0,0
0,0
0.0
0.0
0.0
0.0
0.0
A
0.0
A
0.0
6
0.0
0.0
0,0
Figure 13: A long trajectory imagined by our world model for game Hero. The player traverses five different rooms and the world model is able correctly predict the state and reward dynamics. Note that we only show every fifth frame to cover more time steps (the rewards lying in- between are summed up). The total number of time steps is 230.
(a) Assault
(b) Hero
(c) KungFuMaster
(d) Pong
(c) Pong
Figure 14: Visualization of frame stacks reconstructed from predicted states &. Each frame in the stack is visualized by different color. The world model is able to encode and predict movements.
18

---

Published as a conference paper at ICLR 2023
Briakout
Prosia
HuNaomin
Pong
120
6.630
1.40
6600
0.400
0.200
0.00
2:40
0.300
0.000
0.2
0.2
34
0.2
0.4
0.6
0.4
0.1
0.
35
Frest
0.20
Pong
0.20
0.20
Figure 15: Effect of disabling the proposed thresholded entropy loss (by setting 1) on the performance and the entropy in random subset of games. The thresholded version stabilizes the entropy and leads to a better score in Breakout and Pong, while the entropy behaves unfavorably without threshold.
Acnas?
Holhrmnn
Hero
6259
6:400
0.300
5.00
6.900
2755
0.200
0.00
0.2
4.2
6.6
2.
Figure 16: Comparison of the history length 16 used in our main experiments with 4 on random subset of games. We observe a lower human normalized score for E = 4.
KargniMazzer
Uptrooon
3.250
3.59
5.000
(c.N)
4.000
3.500
2.
Figure 17: Conditioning the policy on [z, h) compared with the usual z. In some cases the perfor- mance can be better during training, but the final score is lower or equal.
CrazyClmber
Seaauest
UpNDuum
6423
5.000
5.331
4.000
2.3.0
3.000
2.300
0.200
J3C
18c
165
Figure 18: Scores on random subset of games when we train with lower number of interactions but the same training budget. This only leads to significant improvement for UpNDown, where the final score is higher with only 50K interactions.
19

---

Published as a conference paper at ICLR 2023
6450
3.300
0355
5.400
0.300
0.32
62:00
0.30
0.45
0.50
145
365
Figure 19: Comparison of the proposed balanced sampling procedure with uniform sampling. shows the development of the dynamics loss from Equation (4), which is lower at the end of training all cases.
A.3 ADDITIONAL TRAINING DETAILS
In Algorithm D we present pseudocode for training the world model and the actor-critic agent. We the SiLU activation function (Elfwing al. 2018) for all models. In Table we summarize all hyperparameters that we used in our experiments. In Table 5 we provide the number of parameters of our models.
Pretraining for Better Initialization: During training we need to correctly balance the amount of world model training and policy training, since the policy has keep up with the distributional shift of the latent space. However, we can spend some extra training time on the world model with pre-collected data (included the 100K interactions) at the beginning of training in order to obtain reasonable initialization for the latent states.
Algorithm Training the world model and the actor-critic agent.
function train_ train_world_ (
function
actor_critic(z)
// sample sequences of observations,
// imagine trajectories states,
// rewards, actions and discounts
// rewards, actions and discounts;
e,arr,d = fro_le__som__dttsstttt (n)ode
// use as starting point
imag [z]
o hat = decode
for until H do actor
transformer(z,a,> (z,a,)
r. r_hat,d_hat,z,
predict(h)
imag (ppend
h transformert (imag)
optimize world model via
r, r,d,z predict (h) emag. ([r,d,z])
// self-supervised learming optin_observatiinno...... b___tt,___att))
@ytim_dynamics(C ,, 2,r_hatt _aat 2a__tt,___at,__hatt
// optimize actor-critic via
// will be used for imagination
// reinforcement learning
optim_actor cratiicrriiiiiiee (imag)
return
20

---

Published as a conference paper at ICLR 2023
Table Hyperparameters used in our experiments.
Description
Symbol
Value
Dataset sampling temperature
20
Discount factor
0.99
GAE parameter
0.95
World model batch size
N
100
History length
E
16
Imagination batch size Imagination horizon
M
400
H
Encoder entropy coefficient
011
5.0
Consistency loss coefficient
0.01
Reward coefficient
31
10.0
Discount coefficient
32
50.0
Actor entropy coefficient
7
0.01
Actor entropy threshold
F
0.1
Environment steps
100K
Frame skip
4
Frame down-sampling
64 x 64
Frame gray-scaling
Yes
Frame stack
Terminate on live loss Max frames per episode
Yes
108K
Max no-ops
30
Observation learning rate Dynamics learning rate
0.0001
0.0001
Actor learning
0.0001
Critic learning rate
0.00001
Transformer emhedding size
256
Transformer layers
Transformer heads
x 64
Transformer feedforward size Latent state predictor units
1024
4 512
Reward predictor units
4
Discount predictor units
4 X 256
Actor units
4 X 512
Critic units
4 x 512
Activation function
SiLU
Table 5: Number of parameters of models.
Model
Symbol
# Parameters
Observation model
8.2M
Dynamics model
15
10.SM
Actor
0
1.3M
Critic
S
1.3M
World model
19M
Actor-critic
2.6M
Total
21.6M
Encoder actor (at inference time)
4.4M
21


Reviews:
[{"review_id": "j57zjapQNu5", "reviewer": "Reviewer_qSfN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed a transformer-based world model. Utilizing the Transformer-XL, the agent is trained in imagination rollouts in a computationally efficient way. Though inspired by prior work, such as the Dreamer frameworks, several modifications were made and ablation studies showed benefits to the final performance. Experiments on the Atari 100K benchmark showed an overall improvement over a set of the prior model-free and model-based methods.", "summary_of_the_review": "This paper proposed a transformer-based world model. Experiments on the Atari 100K benchmark showed an overall improvement over a set of the prior model-free and model-based methods. However, the motivation is not clear and some key statements need further explanations.", "strength_and_weaknesses": "### Strength\n- The paper is well-written generally, providing clear reasoning throughout the whole manuscript.\n- Experiments on the Atari 100K benchmark showed an overall improvement over several model-free and model-based methods.\n- Detailed ablation studies have been provided.\n\n### Weakness\n\n- Major:\n    - Is the input for policy learning $z_t$ or the Transformer predicted $\\hat{z}_t$ during policy learning? The description in the `Choice of Policy Input` section looks confusing. If it's $\\hat{z}_t$, did you apply the frame stacking?\n    - It seems this work has been inspired by the Dreamer framework a lot, why didn't the author choose to compare it with Dreamer-V2?\n    - The results on the effect of taking rewards as the input seems too weak, with benefits only shown on one task. Could the authors please explain how are the tasks selected in the ablation studies?\n    - Since Atari games do not require long-term memory in general, how does the Transformer help on this Atari 100K benchmark?\n    \n- Minor:\n    - The objectives seem complex, could the authors analyze the sensitivity of the set of hyper-parameters used in the objectives?\n    - How are the M trajectories generated given N sequences used to train the world model?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "gAeCKRxeyA", "reviewer": "Reviewer_swpJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper investigates the Transformer-based world model for model-based reinforcement learning. \n\nInspired by the successful applications of Transformers to supervised learning or generative learning, applying Transformers to reinforcement learning has been one of the most interesting topics among RL community members [1, 2, 3, 4, 5]. However, one of the pain points of Transformer applications is the much larger computation overheads when interacting with a real environment or learning in imagination [1, 2, 3]. The world model design in this paper reduces the overhead while it limits the agent to attend to the past at inference time.\n\nIn addition, they suggest reasonable techniques to train the agent on top of this world model to be more stable and better and show the ablation studies with/without the techniques too. Their model outperforms baselines, including the state-of-the-art model-free and model-based methods.\n\n", "summary_of_the_review": "This paper does not consider some points about why Transformer is good for the world model and how their model can handle the long-term dependent knowledge. However, they showed impressive performance by applying Transformer to world model with their techniques, which is good enough to be shared in our community. Especially by densely analyzing their techniques, they made the reader easily understand why the techniques are required. Hopefully, the authors will revise this paper by addressing my concerns. \n\n\n[1] Parisotto, Emilio, et al. \"Stabilizing transformers for reinforcement learning.\" International conference on machine learning. PMLR, 2020.\n\n[2] Parisotto, Emilio, and Ruslan Salakhutdinov. \"Efficient transformers in reinforcement learning using actor-learner distillation.\" arXiv preprint arXiv:2104.01655 (2021).\n\n[3] Chen, Chang, et al. \"Transdreamer: Reinforcement learning with transformer world models.\" arXiv preprint arXiv:2202.09481 (2022).\n\n[4] Esslinger, Kevin, Robert Platt, and Christopher Amato. \"Deep Transformer Q-Networks for Partially Observable Reinforcement Learning.\" arXiv preprint arXiv:2206.01078 (2022).\n\n[5] Micheli, Vincent, Eloi Alonso, and François Fleuret. \"Transformers are sample efficient world models.\" arXiv preprint arXiv:2209.00588 (2022).", "strength_and_weaknesses": "Strength\n- The paper is well-written and easy to follow.\n- They also support ablation studies for suggested techniques, balanced dataset sampling, threshold entropy loss, and feeding back the predicted reward. Thus, the reader can understand why they applied those theoretically and empirically.\n- Using $z_t$ as the input of the policy, not the $h_t$. This choice can interact with real environment more quickly, which is different from previous works [1, 3].\n\nWeakness\n- The motivation why the Transformer-based World Model will work well is not considered (e.g., in [1,2,3,4], they address the long-term dependency issue). Maybe better transition modeling could be one of the motivations. Then it should make this paper more concrete if there were empirical results about the world model qualities when using Transformer and recurrent modules.\n- As pointed out above, their modeling cleverly uses $z_t$ as an input of policy, while it is limited for attending the history. It could be optional, like for the task requiring long-term dependency, using $h_t$ and usually using $z_t$, but when using $h_t$, the inference could be slower.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RgKHFomwWt7", "reviewer": "Reviewer_6X3a", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper introduces a model-based RL approach that learns a transformer-based world model and additional techniques to further improve the sample-efficiency. Specifically, VAE is trained to learn visual information from raw pixels, and autoregressive dynamics model based on Transformer-XL architecture is trained upon a sequence consisting of states and rewards. The method is evaluated on a set of Atari tasks within 100k steps to measure sample-efficiency of the proposed method.", "summary_of_the_review": "While the paper proposed an intuitive approach that achieves strong sample-efficiency, the main weakness of the paper is in that it lacks analysis and experimental results to support the claims made in the paper, which makes me difficult to see that the current draft is strong enough to be accepted as a ICLR paper (I'm trying to avoid borderline score following the suggestion from the conference this year). Hence I'm recommending the paper to be rejected at the current status but willing to update the score per the discussion within the rebuttal period.", "strength_and_weaknesses": "## Strengths\n\n- Strong performance with simple and intuitive method\n- Paper reads well\n\n## Weaknesses\n\n- Claims within the paper are not justified enough except that the proposed method performs well.\n    - It is not clear whether transformer-based model is indeed better than previous architecture with GRU of DreamerV2. From the current results, we cannot know whether the gain comes from additional techniques or transformer-based world models or additional computes. How does it compare to DreamerV2 trained with the proposed techniques, similar computes (not parameters), and hyperparameters to be more sample-efficient?\n    - The paper says training is more efficient as it's parallel, but does not provide any measurement or investigation to explicitly show this compared to previous architecture based on a recurrent architecture. Gviven that the method utilizes a quite large model consisting of 10-layers Transformers, this should be further thoroughly supported.\n    - It is written that leveraging Transformer-XL [Dai et al., 2019] can be good for long-term inference, but in the paper it's not investigated whether this is indeed the case. Additional results that compared with vanilla Transformer architecture could be included to support this.\n    - It is claimed that long-term dependencies can be captured with Transformers but final policy is trained on top of the representation from observation consisting of 4 stcaked frames. So it's difficult to know it's really good for capturing long-term dependencies in terms of either prediction for generating training samples or inference. It should be evaluated on more challenging tasks that require memories?\n    - For thresholded entropy loss, I can understand that it can make normalized entropy stable, but is this necessarily better than changing entropy throughout the training? How does the performance is related to the change in the entropy here?\n    - For sampling scheme, it would be nice to further describe and formulate how this cheme would affect the optimization objectives compared to using uniform sampling scheme. Moreover, ablation to see how this affects performance would be also nice.\n- Investigation into the effect of conditioning the policy on $z_t$ is not complete. It's okay to say that this would incur more costs, but it seems an important difference to prior approaches, so there should be an investigation into this with additional results and also corresponding wall times.\n- Comparison with prior works can be much more informative if it includes some metrics in terms of training costs -- as it seems like the proposed method would require much more resources than other approaches (of course it's my guess because related information is not included anywhere, correct if i'm wrong), and it also does some pre-training on initial exploration samples.\n- Missing comparison with TransDreamer, and not clear why this is only included in Related Work? \n\n[Dai et al., 2019] Dai, Zihang, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. \"Transformer-xl: Attentive language models beyond a fixed-length context.\" arXiv preprint arXiv:1901.02860 (2019).\n\n[Chen et al., 2022] Chen, Chang, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. \"Transdreamer: Reinforcement learning with transformer world models.\" arXiv preprint arXiv:2202.09481 (2022).\n\n---\n\nUpdating the score to 6 after the discussion with AC and other reviewers. I would like to strongly encourage the authors to include results to compare with RNN-based world models or tone-down on claims on the benefit of transformer over RNN-based world models.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OFg2uE2Wut", "reviewer": "Reviewer_gDHw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors investigate the effectiveness of transformers as world models. Their main innovation is to input not just states and actions (as was done in earlier works like Ha & Schmidhuber, 2018), but also to add rewards (predicted or real). This additional constraint on the input interacts well with the transformer architecture to yield sample efficient and performant agents on the Atari 100k benchmark.\n\nUPDATE\n\nraising my score due to the reviewers responses and work over the rebuttal period", "summary_of_the_review": "I was really excited about this paper for the possibility of understanding why transformers might make sample efficient world models. Instead, this paper is a standard benchmark paper describing a recipe for performant transformer-based world models. I am far less interested in the latter, but I can appreciate that the contribution required a lot of work and is nontrivial. So I'm borderline on this paper. It would make a fantastic workshop paper. I'm less sure if ICLR is the right venue. Looking forward to seeing what the other reviewers think.", "strength_and_weaknesses": "Strengths\n- Their approach outperforms baselines across the board on Atari 100K.\n- They train a lighter weight policy using the transformer world model, which makes inference faster than it would be otherwise.\n- The method is well explained.\n- The authors included code as their SI.\n\nWeaknesses\n- Missing a citation and discussion of Micheli et al., Transformers are Sample Efficient World Models. I do not believe this is published in a conference yet so we can consider it concurrent work but you should include it in your paper.\n- There's no analysis of the world model! Tell me more about what the transformer has learned: is it a veridical representation of the games, are certain elements overemphasized over others? What's going on with the world model? How does its representations change when you add reward as an input?\n- Results are only reported for one challenge (Atari 100K) and there's clearly a large amount of tweaking that went into this work in terms of loss and data engineering. That's understandable to an extent, but as it is paired with the limited insights into what is happening in the world model, I was left underwhelmed by this work.\n- The title is flashy, but it also raises an important question which I don't believe was properly addressed in this paper: what is the sample efficiency of the transformer world model when you train on different amounts of data? For instance, try log-spaced intervals from 10 to the 100K reported in the main text. Understanding what failures occur at which dataset sizes would go a long way in illustrating the benefits (and persistent failure modes) of transformers as world models, and how introducing reward changes that sample efficiency curve.\n\nQuestions\n- In your intro you write: \"This can be illustrated by a simple example: in the game of Pong, the paddles and the ball move independently. A successfully trained world model is able to imagine trajectories with paddle and ball configurations that have never been observed before, which enables learning of improved behaviors.\" Is there evidence for this? Or are you positing that world models could help an agent solve this out-of-domain challenge?\n- \"In contrast to Hafner et al. (2021), the role of the observation model is to capture only non-temporal information about the current time step. However, we include short-time temporal information, since a single observation ot consists of four frames (aka frame stacking, see also Section 2.2).\" I think the first sentence here is funky. Did you mean \"in contrast to Hafner, *where* the role...\"\n- \"Due to these design choices, our world model combines the benefits of transformers and recurrent neural networks...\" I'm not sure any of the four bullets are areas where RNNs are better than transformers. My thought is the key benefit of RNNs over transformers is in the amount of memory needed to run them (i.e., RNNs operate on a single token at a time whereas transformers look at many at once). ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster