Transformers are Sample-Efficient World Models
Authors: Vincent Micheli
Eloi Alonso
François Fleuret
Venue: ICLR 2023 notable top 5% (2023)
Keywords: deep learning
reinforcement learning
model-based reinforcement learning
world models
learning in imagination
transformers
discrete autoencoders
generative modeling
sequence modeling

Abstract: Deep reinforcement learning agents are notoriously sample inefficient, which considerably limits their application to real-world problems. Recently, many model-based methods have been designed to address this issue, with learning in the imagination of a world model being one of the most prominent approaches. However, while virtually unlimited interaction with a simulated environment sounds appealing, the world model has to be accurate over extended periods of time. Motivated by the success of Transformers in sequence modeling tasks, we introduce IRIS, a data-efficient agent that learns in a world model composed of a discrete autoencoder and an autoregressive Transformer. With the equivalent of only two hours of gameplay in the Atari 100k benchmark, IRIS achieves a mean human normalized score of 1.046, and outperforms humans on 10 out of 26 games, setting a new state of the art for methods without lookahead search. To foster future research on Transformers and world models for sample-efficient reinforcement learning, we release our code and models at https://github.com/eloialonso/iris.

Published as a conference paper at ICLR 2023
TRANSFORMERS ARE SAMPLE-EFFICIENT WORLD MODELS
Vincent Micheli University of Geneva
Eloi Alonso" University of Geneva
François Fleuret University of Geneva
ABSTRACT
Deep reinforcement learning agents notoriously sample inefficient, which considerably limits their application to real-world problems. Recently, many model-based methods have been designed to address this issue, with learning in the imagination of a world model being one of the most prominent approaches. However, while virtually unlimited interaction with simulated environment sounds appealing, the world model has to be accurate over extended periods of time. Motivated by the success of Transformers in sequence modeling tasks, we introduce IRIS, a data-efficient agent that learns in a world model composed of a discrete autoencoder and an autoregressive Transformer. With the equivalent of only two hours of gameplay in the Atari 100k benchmark, IRIS achieves a mean human normalized score of 1.046, and outperforms humans on 10 out of 26 games, setting a new state of the art for methods without lookahead search. To foster future research on Transformers and world models for sample-efticient reinforcement learning, we release our code and models at htt //git hubb www.w//ww./.......ooo...oo....oo
INTRODUCTION
Deep Reinforcement Learning (RL) has become the dominant paradigm for developing competent agents in challenging environments. Most notably, deep RL algorithms have achieved impressive
performance in multitude of arcade (Mnih et al. 2015 Schrittwieser et 20200; Hafner al.
202T), real-time strategy (Vinyals et al., [219 Berner et al. 2019), board (Silver et al. 2016 2018
Schrittwieser et al., (020) and imperfect information (Schmid al. 2021 Brown al. (020aa games. However, common drawback of these methods is their extremely ly low sample efficiency. Indeed, experience requirements range from months of gameplay for DreamerV2 (Hafner et al., 202T) in
Atari 2600 games (Bellemare et al. (203b)) to thousands of years for OpenA Five in Dota2 (Berner cct al. 2019). While some environments can be sped up for training agents, real-world applications often cannot. Besides, additional cost or safety considerations related to the number of environmental interactions may arise (Yampolskiy 2018). Hence, sample efficiency is a necessary condition to bridge the gap between research and the deployment of deep RL agents in the wild.
Model-based methods (Sutton & Barto 2018) constitute a promising direction towards data efficiency.
Recently, world models were leveraged in several ways: pure representation learning (Schwarzer
et al. 2021), lookahead search (Schrittwieser et al. 2020 Ye et al. 202T), and learning in imagination
(Ha & Schmidhuber 2018 Kaiser et al.) 2020 Hafner al. 2020) 2021). The latter approach is
particularly appealing because training an agent inside a world model frees it from sample efficiency
constraints. Nevertheless, this framework relies heavily on accurate world models since the policy is purely trained in imagination. In a pioneering work, Ha & Schmidhuber (2018) successfully built imagination-base agents in toy environments. SimPLe recently showed promise in the more challenging Atari 100k benchmark (Kaiser et al. 2020). Currently, the best Atari agent learning imagination DreamerV2 (Hafner et al.) 2021). although was developed and evaluated with two hundred million frames available, far from the sample-efficient regime. Therefore, designing new world model architectures, capable of handling visually complex and partially observable environments with few samples, is key to realize their potential as surrogate training grounds.
The Transformer architecture (Vaswani et al. 2017) is now ubiquitous in Natural Language Processing (Devlin al.) 2019; Radford et al. 2019 Brown et al.] 2020b Raffel et al., 2020), and is also
gaining traction in Computer Vision Dosoviiskiy et al. 2021 He et al. 2022), as well as in Offline
"Equal contributions, order determined by a coin flip. Correspondence: (fir last: @un ce

---

Published as a conference paper at ICLR 2023
P-1
E do
di
d
G
G
D
D
8
Figure Unrolling imagination over time. This figure shows the policy т, depicted with purple arrows, taking a sequence of actions in imagination. The green arrows correspond to the encoder E and the decoder D of a discrete autoencoder, whose task is to represent frames in its learnt symbolic language. The backbone G of the world model is GPT-like Transformer, illustrated with blue arrows. For each action that the policy π takes, G simulates the environment dynamics, by autoregressively unfolding new frame tokens that D can decode. G also predicts reward and potential episode termination. More specifically, an initial frame Tn is encoded with E into tokens 20 = (zg. (z,...., 2 E(ro). The decoder D reconstructs an image do = D(z),, from which the policy predicts the action dg- From 20 and ao. G predicts the reward fo. episode termination dg E {0. 1}. and in an autoregressive manner 2 = (é).... 2f ), the tokens for the next frame. A dashed box indicates image tokens for given time step, whereas a solid box represents the input sequence of G, i.e. (zo, (n) at 0, (2o, do, 31, 01) at t = 1, etc. The policy π is purely trained with imagined trajectories, and is only deployed in the real environment to improve the world model (E, G).
Reinforcement Learning (Janner al. 2021 Chen et al., 2021). In particular, the GPT (Radford) et al. 2018, 2019 Brown et al. 2020b) family of models delivered impressive results in language understanding tasks. Similarly to world models, these attention-bant based models are trained with high- dimensional signals and self-supervised learning objective, thus constituting ideal candidates to simulate an environment.
Transformers particularly shine when they operate over sequences of discrete tokens (Devlin et al. 2019; Brown et al. 2020b). For textual data, there are simple ways (Schuster Nakajima 2012 Kudo & Richardson 2018 to build vocabulary, but this conversion is not straightforward with images. A naive approach would consist in treating pixels as image tokens, but standard Transformer architectures scale quadratically with sequence length, making this idea computationally intractable. To address this issue, VQGAN (Esser et al. 2021) and DALL-E (Ramesh et al. 2021) employ a discrete autoencoder (Van Den Oord al. 2017)) as a mapping from raw pixels of image tokens. Combined with an autoregressive Transformer, these methods demonstrate strong unconditional and conditional image generation capabilities. Such results suggest a new approach to design world models.
much smaller amount
In the present work, we introduce IrIs (Imagination with auto- Ritto--egrssion over an Inner Speech), an agent trained in the imagination of a world model composed of a discrete autoencoder and an autoregressive Transformer. IRIS learns behaviors by accurately simulating millions of trajectories. Our approach casts dynamics learning as sequence modeling problem, where an autoencoder builds a language of image tokens and a Transformer composes that language over time. With minimal tuning, IRIS outperforms line of recent methods (Kaiser et al. 2020 Hessel et al.) 2018 Laskin Schwarzer et al., 2022] for sample-efficien RL in the Atari 100k
et al. 2020 Yarats et al. 2021
benchmark (Kiserr al. 2020 After only two hours of real-time experience, it achieves mean human normalized score of 1.046. and reaches superhuman performance on 10 out of 26 games. We
describe IRIS in Section 2 and present our results in Section 3
2

---

Published as a conference paper at ICLR 2023
1ttle
Figure 2: Four imagined trajectories in KungFuMaster We use the same conditioning frame across the four rows, in green, and the world model imagine the rest. As the initial frame only contains the player, there is no information about the enemies that will come next. Consequently, the world model generates different types and numbers of opponents in each simulation. also able reflect an essential game mechanic, highlighted in the blue box, where the first enemy disappears after getting hit by the player.
METHOD
We formulate the problem as a Partially Observable Markov Decision Process (POMDP) with image observations T, R4xwx3 discrete actions az E {1,..... : A), scalar rewards r, R, episode termi- nation d, E {0, 1}, discount factor E (0, 1), initial observation distribution Po. and environment dynamics 2(+1, Ft. de ~ p(Ix+1, (e. de arge, age). The reinforcement learning objective is train policy π that yields actions maximizing the expected sum of rewards E, EE120 rr).
Our method relies on the three standard components to learn in imagination (Sutton & Barto 2018): experience collection, world model learning, and behavior learning. In the vein of Ha Schmidhuber (2018); Kaiser et al. (2000); Hafner et al. (2020) 2021), our agent learns to act exc exclusively within its world model, and we only make use of real experience learn the environment dynamics.
We repeatedly perform the three following steps:
collect_experience gather experience in the real enviromment with the current policy, Shhate_world_mgg improve rewards, episode ends and next observations predictions. aodate__ehhviio in imagination, improve the policy and value functions.
The world model is composed of discrete autoencoder (Van Den Oord et al. (2017). to convert an image to tokens and back, and GPT-like autoregressive Transformer (Vaswani et 2017
Radford et al. 2019 Brown et al. 2020b),, whose task is to capture environment dynamics. Figure illustrates the interplay between the policy and these two components during imagination. We first describe the autoencoder and the Transformer in Sections 2.1 and 2.2 respectively. Section 2.3 then
details the procedure to learn the policy and value functions in imagination. Appendix comprehensive description of model architectures and hyperparameters. Algorithm I summarizes the training protocol.
A provides a
2. FROM IMAGE OBSERVATIONS TOKENS
The discrete autoencoder (E, D) learns symbolic language of its to represent digensiinni images as a small number of tokens. The back and forth between frames and tokens is illustrated with green arrows in Figure

---

Published as a conference paper at ICLR 2023
Figure 3: Pixel perfect predictions in Pong. The top row displays a test trajectory collected in the real environment. The bottom row depicts the reenactment of that trajectory inside the world model. More precisely, we condition the world model with the first two frames of the true sequence, in green. We then sequentially feed it the true actions and let it imagine the subsequent frames. After only 120 games of training, the world model perfectly simulates the ball's trajectory and players' movements. Notably, it also captures the game mechanic of updating the scoreboard after winning an exchange, as shown in the blue box.
More precisely, the encoder F. RAxux3 4 (1...., ,NK converts an input image into K tokens from a vocabulary size N. Let & = [e:]= E RNxxd be the corresponding embedding table of d-dimensional vectors, The input image z; is first passed through Convolutional Neural Network (CNN) (LeCun et al. 1989 producing output Yt E IR K'd We then obtain the output
tokens 24 (2) , 1.....
2,
1..... N) as 2 argmin, Vi ea fi]l22 the index of the closest
embedding vector in Van Den Oord et al. 2017 Esser et al. 2021D. Conversely, the CNN decoder
D: {1..... [I....,N)] NN + IR XXXx3 turns K tokens back into image.
This discrete autoencoder is trained on previously collected frames, with an equally weighted combination of I, reconstruction loss, commitment loss (Van Den Oord al. 2017: Esser et al. 202T), and a perceptual loss (Esser et al. 2021 Johnson et 2016 Larsen et al. 2016). We use
straight-throus through estimator (Bengio et 2013) enable backpropagation training.
2.2 MODELING DYNAMICS
At a high level, the Transformer captures the environment dynamics by modeling the language of the discrete autoencoder time. Its central role of unfolding imagination is highlighted with the blue arrows Figure
Specifically, G operates sequences of interleaved frame and action tokens. An input quence (z)...., =K ag; ........6 2f 24....,,........ 2 a,) is obtained from the raw sequence (ro, ao, 21, 01,...,;4, (4) by encoding the frames with F,, as described in Section 2.1
At each time step f, the Transformer models the three following distributions:
Transition:
2+11 pa (fe+1l 2gr, as)) with at
poo(2tt1 =gg, 1gg, age, 25
(1)
Reward:
ft ~ po Ar 12g0, agt)
(2)
Termination:
de ~ pal de 22g age)
(3)
<k == (z]+1)- k-1 the tokens that Note that the conditioning for the k-th token also includes 2
,+11 were already predicted, i.c. the autoregressive process happens at the token level.
We train G in self-supervised manner on segments of L time steps, sampled from past experience.
We cross-entropy loss for the transition and termination predictors, and mean-squared error loss or cross-entropy loss for the reward predictor, depending on the reward function.
2.3 LEARNING IMAGINATION
Together, the discrete autoencoder (E, and the Transformer G form world model, capable of imagination. The policy n, depicted with purple arrows in Figure [ exclusively learns this imagination MDP.
4

---

Published as a conference paper at ICLR 2023
BOS 15
DO9 4
OBS
B06
Figure 4: Imagining rewards and episode ends in Breakout (top) and Gopher (bottom). Each row depicts an imagined trajectory initialized with single frame from the real environment. Yellow boxes indicate frames where the world model predicts positive reward. Breakout, captures that breaking brick yields rewards, and the brick is correctly removed from the following frames. In Gopher, player has to protect the carrots from rodents. The world model successfully internalizes that plugging hole killing highlighted with red boxes. The world model accurately reflects that missing the ball in Breakout, or letting enemy reach the carrots in Gopher, will result in the end of an episode.
enemy leads to rewards. Predicted episode terminations
At time step f. the policy observes reconstructed image observation ât and samples action a, (a,(#gg)) The world model then predicts the reward fr, the episode end d,, and the next observation -t-1 = 1)) with 2,+1 ~ Pa(2+11 20, ao, & ai,... 2t, at). This imagination procedure is initialized with real observation 26 sampled from past experience, and is rolled out for H steps, the imagination horizon byperparameter. We stop if an episode end is predicted before reaching the horizon. Figure T Iillustrates the imagination procedure.
As we roll out imagination for fixed number of steps, we cannot simply use Monte Carlo estimate for the expected return. Hence, bootstrap the rewards that the agent would get beyond given time step, we have value network V that estimates V(+)) = E,, |E>e >
Many actor-critic methods could be employed to train æ and in imagination (Sutton & Barto 2018 Kaiser al.) 2020 Hafner et al.) 2020). For the sake of simplicity, we opt for the learning objectives and hyperparameters of Dreamer V2 (Hainer et al,, 2021), that delivered strong performance in Atari games. Appendix B gives a detailed breakdown of the reinforcement learning objectives.
EXPERIMENTS
Sample-efficient reinforcement learning is a growing field with multiple benchmarks in complex visual environments (Hafner 2022, Kanervisto et al. 2022). In this work, we focus on the well
established Atari 100k benchmark (Kaiser al.) 2020). We present the benchmark and its baselines in Section 3.I We describe the evaluation protocol and discuss the results in Section 3.2 Qualitative examples of the world model's capabilities given in Section 3.3
5

---

Published as a conference paper at ICLR 2023
Table 1: Returns on the 26 games of Atari 100k after hours of real-time experience, and human- normalized aggregate metrics. Bold numbers indicate the top methods without lookahead search while underlined numbers specify the overall best methods. IRIS outperforms learning-onj methods in terms of number of superhuman games, mean, interquartile mean (IQM), and optimality gap.
Lookahead search
No lookahead seurch
Game
Random
Human
MuZero
EfficientZero
SimPl.e
CURL
DrQ
SPR
IRIS (ours)
Alien
227.8
7127.7
530.0
808.5
616.9
711.0
865.2
841.9
420.0
Amidar
5.8
1719.5
38.8
148.6
74.3
113.7
137.8
179.7
143.0
Assault
222.4
742.0
500.1
1263.1
527.2
500.9
579.6
565.6
1524,4
Asterix
210.0
8503.3
1734.0
25557.8 8
1128.3
567.2
763.6
962.5
853.6
Heist
34.2
232.9
345.4
BatteZoone
2360.0
37187.5
7687.5
13871.2
4031.2
8997.8
10165.3
14834.1
13074.0
Boxing
0.1
12.1
15.1
52.7
7.8
0.9
9.0
35.7
70.1
Breakout
1.7
30.5
48.0
16,4
2.6
83.7
ChopperCommand
811.0
7387.8
1350.0
1117.3
979.4
783,5
844.6
946.3
1565.0
CrazyClimber Clumber
10780.5
35829.4
56937.0
83940.2
62583.6
9154.4
21539.0
36700.5
59324.2
DemonAttack
152.1
1971.0
3527.0
13003.2
208.1
646.5
1321.5
$17.6
2034,4
Freeway
0.0
29.6
21.8
21.8
16.7
28.3
20.3
19.3
31,1
Frostbite
65.2
4334.7
255.0
296.3
236.9
1226.5
1014.2
1170.7
Gopher
257.6
2412.5
1256.0
3260.3
596.8
400,9
621.6
660.6
2236.1
Hero
1027.0
30826.4
3095.0
9315.9
2656.6
4987.7
4167.9
5858.6
7037.4
Jamesbond
29.0
302.8
87.5
517.0
100.5
331.0
349.1
366.5
462.7
Kangaroo
52.0
3035.0
62.5
724.1
51.2
740.2
1088.4
3617.4
838.2
Krull
1598.0
2665.5
4890.8
5663.3
2204.8
3049.2
4402.1
3681.6
6616.4
KungFuMaster
14862.5
11467.4
21759.8
MsPacman
307.3
6951.6
1265.6
1281.2
1480.0
1064.0
1218.1
1318.4
999,1
Pong
-2077
14.6
-6.7
20.1
12.8
-18.5
-9.1
14.6
PrivateEye
69571.3
$6.3
96.7
35.0
3.5
86.0
100.0
Qhert
163.9
13455.0
3952.0
13781.9
1288.8
727.0
1810.7
866.3
745.7
RoadRunner
11.5
7845.0
2500.0
17751_3
5640.6
5006.1
112114
12213.1
9614.6
Seaquest
68.4
42054.7
208
1100.2
683.3
315.2
352.3
558.1
661.3
UpNDown
533.4
11693.2
2896.9
17264.2
3350,3
2646.4
4324.5
10859.2
3546.2
#Superhuman (1)
0
N/A
14
2
3
6
10
Mean (t)
0.000
1.000
0.562
1.943
0.332
0.261
0.465
0.616
1.046
Median (f)
0.000
1.000
0,227
1,090
0.134
0.092
0.313
0.396
0.289
IQM (1)
0,000
1.000
NIA
N/A
0.130
0.113
0,280
0.337
0,501
Optimality Gap (1)
1,000
N/A
0.729
0.768
0.631
0.577
0.512
3.1 BENCHMARK AND BASELINES
Atari 100k consists of 26 Atari games (Bellemare et al. 2013a) with various mechanics, evaluating a wide range of agent capabilities. this benchmark, an agent is only allowed 100k actions in each environment. This constraint is roughly equivalent to 2 hours of human gameplay. By way of comparison, unconstrained Atari agents are usually trained for 50 million steps, a 500 fold increase in experience.
Multiple baselines were compared on the Atari 100k benchmark. SimPLe (Kaiser et al. 2020 trains policy with PPO (Schulman al. 2017) in video generation model. CURL (Laskin et al. 2020) develops off-policy agents from high-level image features obtained with contrastive learning. DrQ (Yarats et al. 2021) augments input images and averages Q-value estimates over several transformations. SPR (Schwarzer et al. [021] enforces consistent representations of input images across augmented views and neighbouring time steps. The aforementioned baselines earry additional techniques to improve performance, such as prioritized experience replay (Schaul et al. 2016). epsilon- greddy scheduling, or data augmentation.
We make distinction between methods with and without lookahead search. Indeed, algorithms relying on search at decision time (Silver et al., 2016)] 2018: Schrittwieser et al. 2020) can vastly improve agent performance, but they come at a premium in computational resources and code complexity. MuZero (Schrittwieser et al., 2020) and EfficientZero (Ye et al. 2021) are the current standard for search-based methods in Atari 100k. MuZero leverages Monte Carlo Tree Search (MCTS) (Kocsis Szepesvári 2006 Coulom 2007) as a policy improvement operator, by unrolling multiple hypothetical trajectories in the latent space of world model. EfficientZero improves upon MuZero by introducing a self-supervised consistency loss, predicting returns over short horizons in shot, and correcting off- off-policy trajectories with its world model.

---

Published as a conference paper at ICLR 2023
Median
Interquartlee
IRIS
SPR
DrQ
CURL
SimPLe
0.75
1.00
0.1
Human Normalized Score
Figure 5: Mean, median, and interquartile mean human normalized scores, computed with stratified bootstrap confidence intervals. runs for IRIS and SimPLe, 100 runs for SPR, CURL, and DrQ (Agarwal al. 2021).
SimPLe
DrQ
IRIS (ours)
1.00
Fraction of runs with score >
CURL
SPR
P(IIIS
0.75
SPR
0.50
Algorithm
DrQ
0.25
CURL
0.00
SimPLe
Human Normalized Score (T)
0.00
0.25
0.50
0.75
1.00
(a) Performance profiles, i.e. fraction of runs above given human normalized score.
(b) Probabilities of improvement, i.e. how likely is for IRIS to outperform baselines any game.
Figure 6: Performance profiles (left) and probabilities of improvement (right) (Agarwal al. 2021
3.2 RESULTS
The human normalized score is the established measure of performance in Atari 100k. It is defined as score agent-------- aopent-scored score_ravdor
Juman sccor_rrraaee score_random" where score reredmmdom comes from random policy, and score haraann is obtained from human players (Wang et al. 2016).
Table displays returns across games and human-normalized aggregate metrics. For MuZero and EfficientZero, we report the averaged results published by Ye et al. (22TT) (3 runs). We use results from the Atari 100k case study conducted by (Agarwll et al. (2021) for the other baselines (100 new runs for CURL, DrQ, SPR, and 5 existing runs for SimPLe). Finally, we evaluate IRIS by computing an average over 100 episodes collected the end of training for each game (5 runs).
Agarwal et al. 2021) discuss the limitations of mean and median scores, and show that substantial discrepancies arise between standard point estimates and interval estimates in RL benchmarks. Following their recommendations, we summarize in Figure the human normalized scores with stratified bootstrap confidence intervals for mean, median, and interquartile mean (IQM). finer comparisons, also provide performance profiles and probabilities of improvement in Figure 6
With the equivalent of only two hours gameplay, IRIS achieves a superhuman mean score of 1.046 (+70%), an IQM of 0.501 (+49%), optimality gap of 0.512 (+11%), and outperforms human players on 10 out of 26 games (+67%), where the relative improvements are computed with respect to SPR (Schwarzer et al. 2021). These results constitute lookahead search in the Atari 100k benchmark. We also note that IRIS outperforms MuZero, although the latter was designed for the sample-efficient regime.
state of the art for methods without

---

Published as a conference paper at ICLR 2023
130
41"
Files
Figure 7: Three consecutive levels in the games Frostbite (left) and Krull (right). In our experiments, the world model struggles to simulate subsequent levels in Frosthite, but not in Krull. Indeed, exiting the first level in Frostbite requires a long and unlikely sequence of actions to first build the igloo, and then go back to it from the bottom of the screen. Such rare events prevent the world model from internalizing new aspects of the game, which will therefore not experienced by the policy in imagination. While Krull features more diverse levels, the world model successfully reflects this variety, and IRIS even sets a new state of the art in this environment. This is likely due to more frequent transitions from stage to the next in Krull, resulting in sufficient coverage of each level.
In addition, performance profiles (Figure (agure/6aaaa reveal that IRTS is on with the strongest baselines for bottom 50% of games, at which point stochastically dominates (Agarwal al.) 2021 Dror al. 2019) the other methods. Similarly, the probability of improvement greater than 0.5 for all baselines (Figure 6b).
In terms of median score, IRIS overlaps with other methods (Figure 5). Interestingly, Schwarzer et al. (202I) note that the median is only influenced by a few decisive games, as evidenced by the width of the confidence intervals for median scores, even with 100 runs for DrQ, CURL and SPR.
We observe that IRIS particularly strong in games that do not suffer from distributional shifts as the training progresses. Examples of such games include Pong, Breakout, and Boxing. On the contrary, the agent struggles when a new level or game mechanic is unlocked through an unlikely event. This sheds light a double exploration problem. IRIS has to first discover its world model to internalize it. Only then may the policy rediscover and exploit it. Figure details this phenomenon in Frostbite and Krull, two games with multiple levels. In summary, long as transitions between levels do not depend on low-probabiling events, the double exploration problem does not hinder performance.
aspect of the game for
Another kind of games difficult to simulate are visually challenging environments where capturing small details is important. As discussed in Appendix increasing the number of tokens to encode frames improves performance, albeit at the cost of increased computation.
3.3 WORLD MODEL ANALYSIS
As IRIS learns behaviors entirely in its imagination, the quality of the world model is the cornerstone of our approach. For instance, it key that the discrete autoencoder correctly reconstructs elements like a ball, a player, or an enemy. Similarly, the potential inability of the Transformer to capture important game mechanics, like reward attribution or episode termination, can severely hamper the agent's 's performance. Hence, matter the amount of imagined trajectories, the agent will learn suboptimal policies the world model is flawed.
provides quantitative evaluation, we aim to complement the analysis with qualitative examples of the abilities of the world model. Figure 2 shows the generation of many plausible futures in the face of uncertainty. Figure deppicts pixel-perfee predictions in Pong. Finally, we illustrate Figure F predictions for rewards and episode terminations, which are crucial to the reinforcement learning objective.
While Section 3.2
8

---

Published as a conference paper at ICLR 2023
4 RELATED WORK
LEARNING IN THE IMAGINATION OF WORLD MODELS
The idea of training policies in a learnt model of the world was first investigated in tabular environ- ments (Sutton & Barto 2018). Ha & Schmidhuber (2018) showed that simple visual environments could be simulated with autoencoders and recurrent networks. SimPLe (Kaiser et al. 2020) demon- strated that PPO policy (Schulman et al., (017)) trained in video prediction model outperformed humans some Atari games. Improving upon Dreamer (Hafner et al., 2020), Dreamer (Hafner et al. 2021 the first agent learning in imagination to achieve human- level performance in the Atari 50M benchmark. Its world model combines a convolutional autoencoder with a recurrent state-ppcc model (RSSM) (Hafner et al. 2019) for latent dynamies Iearning. More recently, Chen et al. 2022) explored a variant of Dreamer V2 where a Transformer replaces the recurrent network in the RSSM and Seo al. (2022) enhance DreamerV2 in the setting where offline dataset of videos is available for pretraining.
REINFORCEMENT LEARNING WITH TRANSFORMERS
Following spectacular advances in natural language processing (Manning Goldie 2022), the reinforcement learning community has recently stepped into the realm of Transformers, Parisotto ctt al. (2020) make the observation that the standard Transformer architecture is difficult to optimize with RL objectives. The authors propose replace residual connections by gating layers to stabilize the learning procedure. Our world model does not require such modifications, which is most likely due to its self-superrvsed learning objective. The Trajectory Transformer (Janner al.) (021) and the Decision Transformer (Chen et al.) (021) represent offline trajectories as a static dataset of sequences, and the Online Decision Transformer (Zheng al., 2022) extends the latter to the online setting. The Trajectory Transformer is trained to predict future returns, states and actions. At inference time, it can thus plan for the optimal action with a reward-driven beam search, yet the approach is limited to low-dimensional states. On the contrary, Decision Transformers can handle image inputs but cannot be easily extended as world models. Ozair et al. (202T) introduce an offline variant of MuZero (Schrittwieser et al., 2020) capable of handling stochastic environments by performing an hybrid search with a Transformer over both actions and trajectory-level discrete latent variables.
VIDEO GENERATION WITH DISCRETE AUTOENCODERS AND TRANSFORMERS
VQGAN EEser et al. 2 and DALL-E (Ramesh al.) 2021 discrete autoencoders compress a frame into a small sequence of tokens, that a transformer can then model autoregressively. Other works extend the approach to video generation. GODIVA (Wu al., 202T) models sequences of frames instead of a single frame for text conditional video generation. VideoGPT (Yan et al.) (022] introduces video-level discrete autoencoders, and Transformers with spatial and temporal attention patterns, for unconditional and action conditional video generation.
5 CONCLUSION
We introduced IRIS, an agent that learns purely in the imagination of a world model composed of discrete autoencoder and an autoregressive Transformer. IRIS sets a new state of the art in the Atari 100k benchmark for methods without lookahead search. We showed that its world model acquires a deep understanding of game mechanics, resulting in pixel perfect predictions in some games. We also illustrated the generative capabilities of the world model, providing a rich gameplay experience when training in imagination. Ultimately, with minimal tuning compared to existing battle-hane hatteeharddend agents, IRIS opens new path towards efficiently solving complex environments.
In the future, IRIS could be scaled up to computationally demanding and challenging tasks that would benefit from the speed of its world model. Besides, its policy currently learns from reconstructed frames, but could probably leverage the internal representations of the world model. Another exciting avenue of research would be combine learning in imagination with MCTS. Indeed, both approaches deliver impressive results, and their contributions to agent performance might be complementary.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
The different components and their training objectives introduced in Section 2 and Appendix B We describe model architectures and list hyperparameters in Appendix A We specify the resources used to produce our results in Appendix G
Algorithm makes explicit the interplay between we provide the source of the reported results for the
components in the training loop. In Section 3.2
baselines, as well as the evaluation protocol.
The code is part of the supplementary materials, and will be open-sourced to ensure reproducible results and foster future research. Minimal dependencies are required to run the codebase and we provide a thorough user guide to get started. Training and evaluation can be launched with simple commands, customization is possible with configuration files, and we include scripts to visualize agents playing and let users interact with the world model.
ETHICS STATEMENT
The development of autonomous agents for real-world environments raises many safety and envi- ronmental concerns. During its training period, an agent may cause serious harm individuals and damage its surroundings. It is our belief that learning in the imagination of world models greatly reduces the risks associated with training new autonomous agents. Indeed, in this work, we propose world model architecture capable of accurately modeling environments with very few samples. However, in a future line of research, one could go one step further and leverage existing data to eliminate the necessity of interacting with the real world.
ACKNOWLEDGMENTS
We would like to thank Maxim Peter, Bálint Máté, Daniele Paliotta, Atul Sinha, and Alexandre Dupuis for insightful discussions and comments. Vincent Micheli was supported by the Swiss National Science Foundation under grant number FNS-187494.
10

---

Published as a conference paper at ICLR 2023
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances neural information processing systems, 144:9304-293200 2021.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013a.
Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling, The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013b.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint
2013.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemystaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06664 2019.
Noam Brown, Anton Bakhtin, Adam Lerer, and Qucheng Gong. Combining deep reinforcement learning and search for systems, 33:1 15:17757-1709, 2020a.
games. Advances in neural information processing
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 20200.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09488, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 2021.
Rémi Coulom. Computing "elo ratings" of move patterns in the game of go. ICGA journal, 30(4): 198-208, 2007.
Jacob Devlin, Ming-War Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x16 words: Transformers for image recognition at scale, International Conference on Learning Representations, 2021.
Rotem Dror, Segev Shlomov, and Roi Reichart. Deep dominance- to properly compare deep neural models. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2773-2785, 2019.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution im- age synthesis. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12873-12883, 2021.
A. Gers, J. Schmidhuber, and Cummins. Learning to forget: Continual prediction with LSTM. Neural Computation, 12(10):451-2471, 2000.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems, 31, 2018.
11

---

Published as a conference paper at ICLR 2023
Danijar Hafner. Benchmarking the spectrum of agent capabilities. In International Conference Learning Representations, 2022.
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 2555-2565. PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 16000-16009, 2022.
Mattco Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Thirty-second AAAI conference artificial intelligence, 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9())
1735-1780, 1997.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. Advances in neural information processing systems, 34, 2021.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In European conference on computer vision, pp. 694-711. Springer, 2016.
Eukasz Kaiser, Mohammad Babaeizadeh, Piotr Mitos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, al. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020.
Anssi Kanervisto, Stephanie Milani, Karolis Ramanauskas, Nicholay Topin, Zichuan Lin, Junyou Li, Jianing Shi, Deheng Ye, Qiang Fu, Wei Yang, Weijun Hong, Zhongyue Huang, Haicheng Chen, Guangjun Zeng. Yue Lin, Vincent Micheli, Eloi Alonso, François Fleuret, Alexander Nikulin, Yury Belousov, Oleg Svidchenko, and Aleksei Shpilman. Minerl diamond 2021 competition: Overview, results, and lessons learned. In Proceedings the NeurlPS 202/ Competitions and Demonstrations Track, Proceedings of Machine Learning Research, 2022. URLbttps: htttttss:/pprcceeinggg mlr ptess/v116//keerriittoot
html
Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations, 2019.
Andrej Karpathy. minGPT: A minimal PyTorch re-implementtatoon of the OpenAl GPT (Generative Pretrained Transformer) training, 2020. WtRlhhtts:///giiii hww.:/uuu.comeceo ka reatty /mitt///mmm
Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In European conference
machine learning, 2006.
Taku Kudo and John Richardson. Sentencepiece: simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71,700 2018.
Anders Boesen Lindbo Larsen, Søren Kaae Sønderby, Hugo Larochelle, and Ole Winther. Autoen- coding beyond pixels using a learned similarity metric. In International conference on machine learning, pp. 15661 PMLR, 2016.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representa- tions for reinforcement learning. International Conference on Machine Learning, pp. 5639-5650. PMLR, 2020.
12

---

Published as a conference paper at ICLR 2023
Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-551, 1989.
Christopher Manning and Anna Goldie. Cs224n natural language processing with deep learning. 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529-533 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1928-19299111 1937. PMLR, 2016.
Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou, Aaron Van Den Oord, and Oriol Vinyals. Vector quantized models for planning. In International Conference on Machine Learning, pp- 0302-8313. PMLR, 2021.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International Conference on Machine Learning, pp. 7487-7498. PMLR, 2020.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21:1-67, 2020.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. International Conference on Machine Learning, pp. 8821-8831. PMLR, 2021.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. In International Conference on Learning Representations, 2016.
Martin Schmid, Matej Moravcik, Neil Burch, Rudolf Kadlec, Josh Davidson, Kevin Waugh, Nolan Bard, Finbarr Timbers, Mare Lanctot, Zach Holland, et al. Player of games. arXiv preprint arXiv:2112.03178, 2021.
Julian Schritwieserr, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, L. Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839): 604-609, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Mike Schuster and Kaisuke Nakajima. Japanese and korean voice search. In 2012 IEEE international conference on acoustics, speech and signal processing (ICASSP), pp. 5149-5152. IEEE, 2012.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self- preficreeietive representations. In International Conference on Learning Representations, 2021.
Younggyo Seo, Kimin Lee, Stephen L James, and Pieter Abbeel. Reinforcement learning with action- free pre-training from videos. In International Conference on Machine Learning, pp. 19561-19579. PMLR, 2022.
13

---

Published as a conference paper at ICLR 2023
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):48--44 489, 2016.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140-1144, 2018.
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. A Bradford Book, Cambridge, MA, USA, 2018.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
Oriol Vinyals, Igor Babuschkin, Wojeiech M Czamecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung. David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, al. Grandmaster level in StarCraft using multi-agent reinforcement learning. Nature, 575(7782):50-334, 2019.
Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International conference on machine learning, pp. 1995-2003, PMLR, 2016.
Chenfei Wu, Lun Huang, Qianxi Zhang, Binyang Li, Lei Ji, Fan Yang, Guillermo Sapiro, and Nan Duan. Godiva: Generating open-domain videos from natural descriptions. arXiv preprint arXiv:2/04. 14806, 2021.
Roman V Yampolskiy. Artificial Intelligence Safery and Security. Chapman & Hall/CRC, 2018.
Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. Videogpt: Video generation using vq-vae and transformers. arXiv preprint arXiv:2104. 1011::2004.0117,, 2021.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021.
Weirui Ye. Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in neural information processing systems, 34, 2021.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In International Conference Machine Learning, pp. 27042-27059. PMLR, 2022.
14

---

Published as a conference paper at ICLR 2023
A MODELS AND HYPERPARAMETERS
A. DISCRETE AUTORNCODER
Our discrete autoencoder is based on the implementation of VQGAN (Esser et al. 202T). We removed the discriminator, essentially turning the VQGAN into a vanilla VQVAE Van Den Oord et al. 2017) with an additional perceptual loss (Johnson et al. 2016 Larsen al. 2016).
The training objective is the following:
CE D.8) = || D - [
-
-(()) pperceptuat (x)
Here, the first term is the reconstruction loss, the next terms constitute the commitment loss (where sg(-) is the stop-gradient operator), and the last term is the perceptual loss.
Table 2: Encoder / Decoder hyperparameters, We list the hyperparameters for the encoder, the same ones apply for the decoder.
Hyperparameter
Value
Frame dimensions (h, w)
64 x 64
Layers
4
Residual blocks per layer
Channels in convolutions
64
Self-attention layers at resolution
8 / 16
Table 3: Embedding table hyperparameters.
Hyperparameter
Value
Vocabulary size (N)
512
Tokens per frame (K)
16
Token embedding dimension (d)
512
Note that during experience collection in the real environment, frames still go through the autoencoder keep the input distribution of the policy unchanged. See Algorithm for details.
A.2 TRANSFORMER
Our autoregressive Transformer is based the implementation of minGPT (Karpathy 2020). It takes as input a sequence L(K + 1) tokens and embeds it into L(K + 1) x tensor using an A D embedding table for actions, and N x D embedding table for frames tokens. This tensor forwarded through M Transformer blocks. We use GPT2-like blocks (Radford et al.) 2019). i.e. each block consists of self-attention module with layer normalization of the input, wrapped with residual connection, followed by per-position multi-layer perceptron with layer normalization of the input, wrapped with another residual connection.
15

---

Published as a conference paper at ICLR 2023
Table 4: Transformer hyperparameters
Hyperparameter
Value
Timesteps (L)
Embedding dimension (D)
256
Layers (M)
Attention heads
Weight decay
0.01
Embedding dropout
0.1
Attention dropout
0.1
Residual dropout
0.1
A.3 ACTOR-CRITEL
The weights of the actor and critic are shared except for layer. The actor-critic takes as input a 64 x 64 x frame, and forwards it through convolutional block followed by LSTM cell (Mnih) ctt al. 2016 Hochreiter Schmidhuber 1997 Gers et al. 2000). The convolutional block consists of the same layer repeated four times: 3x3 convolution with stride 1 and padding I, ReLU activation, and 2x2 max-pooling with stride 2. The dimension the LSTM hidden state is 512. Before starting the imagination procedure from given frame, burn-in (Kapturowski et al. 2019) the 20 previous frames to initialize the hidden state.
Table 5: Training loop & Shared hyperparameters
Hyperparameter
Value
Epochs
600
# Collection epochs
500
Environment steps per epoch
200
Collection
0.01
Eval sampling temperature
0.5
Start autoencoder after epochs
5
Start transformer after epochs
25
Start critic after epochs Autoencoder batch size Transformer batch size Actor-critic batch size
50
256
64
64
Training steps per epoch Learning rate
200
le-4
Optimizer
Adam
Adam 81
0.9
Adam B2
0.999
Max gradient norm
10.0
16

---

Published as a conference paper at ICLR 2023
B ACTOR-CRITIC LEARNING OBJECTIVES
We follow Dreamer (Hafner et al. 2020 2021 in using the generic A-return, that balances bias and variance, as the regression target for the value network. Given an imagined trajectory (0o, ag, Po. do..... -H- aH-1. f dH @H), the X-return can be defined recursively as fol- lows:
A - d - + AAtt1] V((n)
if t< H 1= H
A
The value network V trained to minimize Lv, the expected squared difference with X-returns over imagined trajectories.
H
Lv = E,
(V(+t) - sg( (A.))
(S)
Here, sg(.) denotes the gradient stopping operation, meaning that the target is a constant in the gradient-based optimization, classically established in the literature (Mnih et al. 2015 Hessel
et al. 2018 Halner et al. 2020).
As large amounts of trajectories are generated in the imagination MDP, can straightforward reinforcement learning objective for the policy, such as REINFORCE (Sutton & Barto 2018). To reduce the variance REINFORCE gradients, we use the value V(++) baseline (Sutton & Barto
2018). We also add a weighted entropy maximization objective to maintain sufficient exploration. The actor trained minimize the following REINFORCE objective over imagined trajectories:
C, E
log n (a(l) V(+.)) H((((((((s))
(6)
Table 6: RL training hyperparameters
Hyperparameter
Value
Imagination horizon (H)
20
0.995
X
0.95
0.001
17

---

Published as a conference paper at ICLR 2023
C
OPTIMALITY GAP
Optimality Gap
IRIS (ours)
SPR DrQ CURL SimPLe
0.56
0.64
0.72
Human Normalized Score
Figure 8: Optimality gap. lower is better. human-level score (Agarwal et al. 2021).
The amount by which the algorithm fails to reach
D IRIS ALGORITHM
Algorithm IRIS
Procedure training_loop ():
for epochs do
collect_experieet (steps_collect)
for steps_worlddmmadel do
modee__mmeel ()
for steps behavior do
update_behavior ()
Procedure collect_experience (n))
To + eev.rreeet...... ()
for t 0 to
c - --E(((()) // forward frame through discrete autoencoder
Sample de (
Z141, d, + env. step (a))
if d then
T++1 <- reset ()
D DU (Ir, an Tes 1-
Procedure update_word modell ():
Sample (xt, ar r2, de) + D
Compute 24 = E((x.) and i. = (2x)) for = T...., L
Update E and D
Compute
ft, de žz, Or,,...... 21) for 1 T...., + 1
Update G
Procedure
():
Sample 2g ~
20 E(xo)
to + D(2))
for 0 to H - do
Sample at ~ πI (ar))))
Sample E,,1, Pe, de po oo(
ao..... 20, ao, .... 2,,
Le+1 + +(
Compute V(2;) for t. == O,... H
Update and V
18

---

Published as a conference paper at ICLR 2023
E AUTOENCONDING FRAMES WITH VARYING AMOUNTS OF TOKENS
The sequence length of the Transformer is determined by the number of tokens used to encode a single frame and the number of timesteps in memory. Increasing the number of tokens per frame results in better reconstructions, although it requires more compute and memory.
This tradeoff is particularly important in visually challenging games with high number of possible configurations, where the discrete autoencoder struggles to properly encode frames with only 16 tokens. For instance, Figure Showss shows that, when increasing the number of tokens per frame to 64 in Alien, the discrete autoencoder correctly reconstructs the player, its enemies, and rewards.
Figure 9: Tradcoff between the number of tokens per frame and reconstructions quality in Alien. Each column displays a 64 x 64 frame from the real environment (top), its reconstruction with a discrete encoding of 16 tokens (center), and its reconstruction with discrete encoding of 64 tokens (bottom). In Alien, the player is the dark blue character, and the enemies are the large colored sprites. With 16 tokens per frame, the autoencoder often erases the player, switches colors, and misplaces rewards. When increasing the amount of tokens, it properly reconstructs the frame.
Table displays the final performance of IRIS trained with 64 tokens per frame in three games. Interestingly, even though the world model is more accurate, the performance in Alien only increases marginally (+36%). This observation suggests that Alien poses hard reinforcement learning problem, evidenced by the low performance of other baselines in that game. On the contrary, IRIS greatly benefits from having more tokens per frame for Asterix (+121%) %) and BankHeist (+432%)
Table 7: Returns on Alien, Asterix, and BankHeist with 64 tokens per frame instead of 16.
Game
Random
Human
SimPLe
CURL
DrQ
SPR
IRIS (16 tokens)
IRIS (64 tokens)
Alien
227.8
7127.7
616.9
711.0
865.2
841.9
420.0
570.0
Asterix
210.0
8503.3
1128.3
567.2
763.6
962.5
853.6
1890.4
BankHeist
14.2
753.1
34.2
65.3
232.9
345,4
53.1
282.5
19

---

Published as a conference paper at ICLR 2023
F BEYOND THE SAMPLE-EFFICIEN SETTING
IRIS can be scaled up by increasing the number of tokens used to encode frames, adding capacity the model, taking more optimization steps per environment steps, using more data. In this experiment, we investigate data scaling properties by increasing the number of environment steps from 100k 10M. However, maintain training time within our computational resources, we lower the ratio of optimization steps per environment steps from 1:1 to 1:50. As consequence, the results of this experiment at 100k frames would be worse than those reported the paper.
Table 8: Increasing the number of environment steps from 100k to 10M.
Game
Random
Human
(100k)
(10M)
Alien
227.8
7127.7
420.0
Amidar
5.8
1719.5
143.0
213.4
Assault
222.4
742.0
1524.4
9355.6
210.0
8503.3
853.6
6861.0
HankHeist
14.2
753.1
53.1
921.6
BattleZone
2360.0
37187.5
13074.0
34562.5
Boxing
0.1
Breakout
1.7
30.5
83.7
493,9
ChopperCommand
811.0
7387.8
1565.0
9814.0
Crazy CrunyClimher
10780.5
35829.4
59324.2
111068.8
DemonAttack
152.1
1971.0
2034.4
96218.6
Freeway
Frostbite
65.2
4334.7
259.1
290.3
Gopher
257.6
2412.5
2236.1
97370.6
Hero
1027.0
30826.4
7037.4
19212.0
Jameshond
29.0
302.8
462.7
5534.4
Kangamoo
52.0
3035.0
838.2
1793.8
Krull
1598.0
2665.5
6616.4
7344.0
Kunsteraster
258.5
22736.3
21759.8
39643.8
MsPacman
307.3
6951.6
999.1
1233.0
Pong
-20.7
14.6
14.6
21.0
PrivateEye
24.9
69571.3
100.0
100.0
Obert
163.9
13455,0
745.7
4012.1
RoadRunner
11.5
7845.0
9614.6
30609.4
Senguest
68.4
42054.7
661.3
1815.0
UpNDown
533.4
#Superhuman (t)
Mean (1)
0.000
1,000
1.046
7.488
Median (?)
0.000
1,000
0.289
1 1.207
0.000
1.000
0,501
2.239
Optimality Gap (1)
1.000
0.000
0,512
0.282
Table 8 illustrates that increasing the number of environment steps from 100k to 10M drastically improves performance for most games, providing evidence that IRIS could be scaled up beyond the sample-efficient regime. On some games, more data only yields marginal improvements, most likely due to hard exploration problems or visually challenging domains that would benefit from higher number of tokens to encode frames (Appendix E
20

---

Published as a conference paper at ICLR 2023
G COMPUTATIONAL RESOURCES
For each Atari environment, we repeatedly trained IRIS with 5 different random seeds. We ran our experiments with 8 Nvidia A00 40GB GPUs. With two Atari environments running on the same GPU, training takes around 7 days, resulting in an average of 3.5 days per environment.
SimPLe (Kaiser al. 2020), the only baseline that involves learning in imagination, trains for weeks with PIOO GPU on a single environment. As for SPR (Schwarzer et al. (021), the strongest baseline without lookahead search, it trains notably fast in 4.6 hours with a PIO0 GPU.
Regarding baselines with lookahead search, MuZero (Schrittwieser et al. 2020) originally used 40 TPUs for 12 hours to train single Atari environment, Ye et al. (2021) train both EfficientZero and their reimplementation of MuZero in 7 hours with RTX 3090 GPUs. EfficientZero's implementation relies distributed infrastructure with CPU and GPU threads running parallel, and C++/Cython implementation of MCTS. By contrast, IRIS and the baselines without lookahead search rely on straightforward single GPU single CPU implementations.
H EXPLORATION IN FREEWAY
The reward function in Freeway sparse since the agent is only rewarded when it completely crosses the road. In addition, bumping into cars will drag down, preventing it from smoothly ascending the highway. This poses an exploration problem for newly initialized agents because a random policy will almost surely never obtain non-zero reward with 100k frames budget.
Figure 10: A game of Freeway. Cars will bump the player down, making it very unlikely to cross the road and be rewarded for random policies.
The solution to this problem actually straightforward and simply requires stretches of time when the UP action is oversampled. Most Atari 100k baselines fix the issue with epsilon-greedy schedules and argmax action selection, where at some point the network configuration will be such that the UP action is heavily favored. In this work, we opted for the simpler strategy of having fixed epsilon-greedy parameter and sampling from the policy. However, we lowered the sampling temperature from to 0.01 for Freeway, in order to avoid random walks that would not be conducive to learning in the early stages training. As consequence, once received its first few rewards through exploration, IRIS was able to internalize the sparse reward function in its world model.
21


Reviews:
[{"review_id": "cgecSdMDrO", "reviewer": "Reviewer_fmDi", "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": "The submission benchmarks a transformer-based approach to background planning in the Atari 100k benchmark. The submission finds that its approach outperforms existing methods that have benchmarked on Atari 100k, excluding those that use decision-time planning.", "summary_of_the_review": "The submission shows a qualitatively novel approach to MBRL can, to some extent, be successful. My main gripe is that the submission shows these results in a domain that is, in my opinion, not well benchmarked; thus, IRIS's strength as a MBRL algorithm is a bit unclear to me. Nevertheless, I think the evidence presented in the submission is sufficiently strong to merit acceptance as is.", "strength_and_weaknesses": "### Strengths\n\n- Loosely speaking, existing successful MBRL algorithms are modifications of Dreamer or MuZero. Showing that another framework for MBRL can be successful is a valuable contribution.\n- The submission largely follows the guidelines for benchmarking laid out by the statistical precipice paper.\n\n### Weaknesses\n\n- At least in my opinion, Atari 100k is not particularly well benchmarked. I would say that the only \"good\" algorithm that has been benchmarked by experimenters incentivized to tune the algorithm to maximize performance is EfficientZero. Thus, the extent to which IRIS is good is a bit unclear to me. It would be interesting to see how IRIS performs on the standard Atari benchmark (or, inversely, how something like Dreamer compares to IRIS on Atari 100k).\n- The statistical precipice paper suggests that results on Atari 100k are reliable (under appropriate metrics) with as few as 10 runs; the submission only uses 5 runs. My gut feeling is that the margin of improvement under various metrics seems substantial enough that it would probably continue to hold under a more reliable number of runs, but it would be good to actually show this.\n\n---\n\nAddendum:\nI also concur with reviewer t9Kq: 1) the submission would benefit from additional attention to related work (such as [1],[2],[3]) and 2) additional ablations.\n\n### Comments on decision-time planning:\n\nThe submission argues: \"Moreover, IRIS could be combined with MCTS, both in imagination and in the real environment. Therefore, methods involving lookahead search should not be seen as direct competitors but rather as potential extensions to learning-only methods.\" In principle, this is of course true. However, as a matter of practice, I am not sure it is as clear. I am not aware of any examples of a non-MuZero-like architecture successfully utilizing decision-time planning. It is plausible to me that algorithms like Dreamer and IRIS, which perform well in a background planning regime, may not enjoy much benefit from decision-time MCTS.\n\n### Comments on superhuman performance:\n\n> Most notably, human experts were surpassed by deep RL algorithms in a multitude of arcade (Mnih et al., 2015; Schrittwieser et al., 2020; Hafner et al., 2021), real-time strategy (Vinyals et al., 2019; Berner et al., 2019), board (Silver et al., 2016; 2018; Schrittwieser et al., 2020) and imperfect information (Schmid et al., 2021; Brown et al., 2020a) games.\n\nI find this sentence is misleading. Deep RL algorithms have not surpassed human experts in most Atari games, as is clearly evidenced by the human world record metric. They have also not surpassed human experts in StarCraft (AlphaStar is only grandmaster level) or DOTA (OpenAI5 played a restricted version of the game and was found to be reliably exploitable by humans).", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "VYn0-yL2MqN", "reviewer": "Reviewer_CZra", "summary": "", "questions": "", "limitations": "", "rating": 8, "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 work proposes IRIS, which uses a world model to train agents. The world model consists of a discrete autoencoder and an autoregressive Transformer.\n\nThe discrete autoencoder $(E, D)$ consists of an encoder $E$, which converts an input image to tokens, and a decoder, which turns tokens back to an image. Given previous tokens and actions, the autoregressive Transformer then models (1) the transitions to next token; (2) reward; (3) termination of the episode.\n\nThe world model is then used to train an actor-critic method using standard reinforcement learning (RL) objectives.\n\nThe authors conducted experiments on Atari 100k benchmark. After 100k actions (about 2 hours of human gameplay experience), the proposed IRIS method can outperform several other baselines including SimPLe, CURL, and DrQ, and it outperforms human players on 10 out of 26 games.\n\nThe authors also show results of the good prediction ability of Transformer world models in terms of accurately predicting next status of Atari game images, and rewards. Some drawbacks of the world models, such as inability of predicting some games with multiple layers are also discussed.", "summary_of_the_review": "Overall, this work studies an interesting idea of using Transformer as world models in model-based RL, and the promising experimental results support the proposed methods well.\n\n**Update after rebuttal**:\n\nI would like to thank the authors for the feedback, which addressed my main concerns. I increased my score accordingly.", "strength_and_weaknesses": "**Strength**:\n\n1. Using Transformers as world models in model-based RL seems novel.\n2. The results are promising.\n\n**Weaknesses**:\n\n1. Although it is emphasized in the paper that sample efficiency should be a main consideration for model-based RL methods, I think training for a longer time and show that the proposed methods can achieve good results on most games would also be worth considering to verify the proposed models.\n2. The objectives of training the autoencoder and Transformer is missing in the paper and in the appendix (only mentioned by a few sentences). It would be better to also include those settings in the appendix.", "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": {}}, {"review_id": "M-BcYohum4", "reviewer": "Reviewer_t9Kq", "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 proposes a miniGPT-like transformer architecture for learning the world model of RL agents in POMDP environments. This World Model's training data is derived from the present policy model's interplay with the real world. The input and output images are then represented in the VQGAN-style and utilized to train the global model. Finally, the policy function of the RL Agent is trained using the fictitious world model data. The authors demonstrate the efficacy of this method on the Atari 100k benchmark, which achieves a mean human\nnormalized score of 1.046, and outperforms humans on 10 out of 26 games with two-hours training data.", "summary_of_the_review": "Overall, I think this article is worth reading and being published, as it explores the possibilities of using Transformer to learn World Model. Although I am not a Big Fan of the World model, which breaks down reinforcement learning into learning a simulater for the environment and then learning it, I still endorse the paper's conclusion that Transformer-based architectures (Vit, Swin, etc.) can encode image inputs into a series that can be easily understood and learned tokens.\n\nHowever, I still hope the author can give more in-depth understanding, such as how the encoding of the world-model itself accounts for this task, and whether the pure necessity of the world-model itself is sufficient (e.g., the same backbone structure of the A2C agent can be used, etc.)\n------------------------------------------------------------------------------------------------------\nI have raise my score to 8, while there are still many details I'd like to know about this paper, overall, I'm inclined to think it deserves to be published.", "strength_and_weaknesses": "\n[Pros]\n* The paper is well-writen and clear to read. the effieciecency and effacacy are good on the atari100k benchmark.\n* The code is clear and easy to follow.\n\n[Cons]\n* Noveltiy is insufficient (Some disucussion to [1][2][3] would be appreciated)\n* Increasing the number of ablations and analyses to improve comprehension of the contributions of various designed parts.(e.g. different tokenizer, backbone)\n\n[1] TRANSDREAMER: REINFORCEMENT LEARNING WITH TRANSFORMER WORLD MODELS https://arxiv.org/pdf/2202.09481.pdf \n[2] Reinforcement Learning with Action-Free Pre-Training from Videos, https://arxiv.org/pdf/2203.13880.pdf\n[3] Online Decision Transformer, https://arxiv.org/pdf/2202.05607.pdf\n\n\n*Some typos: sec2.2 blue arrow-> purple? arrow", "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": {}}, {"review_id": "6Tcp3D1UJM", "reviewer": "Reviewer_Knzj", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "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": "Authors train very impressive agents for the 100K Atari benchmark. The method consists in representing Atari frames with discrete tokens and training a transformer-based world model and a game agent. The transformer model is trained in a supervised way. The agent is trained entirely using frames from the world model. ", "summary_of_the_review": "Authors train very impressive agents for the 100K Atari benchmark.  Benchmark results are excellent but algorithmic, and modeling novelty is limited.", "strength_and_weaknesses": "Strengths:\n\n1. Simple and easy-to-understand learning algorithm.\n2. Well-thought-through reuse of standard modeling and algorithmic components.\n3. Outstanding results on Atari 100K benchmark, in particular, bearing in mind the limited compute and expressivity of models.\n\nWeaknesses:\n\n1. Models and agents are trained per game. I am curious what would happen if one tries to train a multi-game world model and agent in the spirit of Gato https://arxiv.org/abs/2205.06175 or Multi-game decision transformers https://arxiv.org/abs/2205.15241. \n2. Nit: the paper may benefit from using more modest language. I mean, in particular, the “drastically different architecture” claim or not very elegant comparisons to look-ahead algorithms.\n\nA bug that may turn into a feature:\n\nThe method needs “more tokens” for games that require the detailed representation of images (as authors put it: “Another kind of games difficult to simulate are mazes with moving enemies, such as MsPacman, BankHeist, and Alien”). This characteristic seems to be a limitation but may lead to an interesting long-context benchmark. Some environments may require more frames and details, leading to a new dataset that can be used to benchmark long-context transformer models.", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%