T RANSFORMERS ARE S AMPLE -E FFICIENT W ORLD M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 codebase at this https URL.
For the review process, we provide the code and visualizations in the supplementary materials.
I NTRODUCTION
Deep Reinforcement Learning (RL) has become the dominant paradigm for developing competent agents in challenging environments.
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.
However, a common drawback of these methods is their extremely low sample efficiency.
Indeed, experience requirements range from months of gameplay for DreamerV2 (Hafner et al., 2021) in
Atari 2600 games (Bellemare et al., 2013b) to thousands of years for OpenAI Five in Dota2 (Berner et 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 to achieve 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., 2021), and learning in imagination (Ha & Schmidhuber, 2018; Kaiser et al., 2020; Hafner et 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-based 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 in imagination is DreamerV2 (Hafner et al., 2021), although it 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.
z 10 ,...,z Kaxˆ xˆ z 11 ,..., ˆ z Kaˆ xˆ dˆ r...
ˆ dˆ rˆ z 1 s ,..., ˆ z
Ks a s ˆ x s ˆ d s −ˆ r s −
G
G G
E
D π D π
D π t =t =...
t = s
( E, D, G ) .
The Transformer architecture (Vaswani et al., 2017) is now ubiquitous in Natural Language Processing (Devlin et al., 2019;
Radford et al., 2019; Brown et al., 2020b; Raffel et al., 2020), and is also gaining traction in Computer Vision (Dosovitskiy et al., 2021;
He et al., 2022), as well as in Offline
Reinforcement Learning (Janner et 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-based models are trained with highdimensional signals and a 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 a 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 et al., 2017) as a mapping from raw pixels to a much smaller amount 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.
In the present work, we introduce IRIS (Imagination with auto-Regression 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 a sequence modeling problem, where an autoencoder builds a language of image tokens and a Transformer composes that language over time.
With minimal tuning and a drastically different architecture, IRIS outperforms a line of recent methods (Kaiser et al., 2020; Hessel et al., 2018;
Laskin et al., 2020; Yarats et al., 2021; Schwarzer et al., 2021) for sample-efficient RL in the Atari 100k benchmark (Kaiser et al., 2020).
After only two hours of real-time experience, it achieves a 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.
M ETHOD
We formulate the problem as a Partially Observable Markov Decision Process ( POMDP ) with image observations x t ∈ R h × w × 3 , discrete actions a t ∈ { 1 , . . .
, A } , scalar rewards r t ∈ R , episode termination d t ∈ { 0 , 1 } , discount factor γ ∈ (0 , 1) , initial observation distribution ρ 0 , and environment dynamics x t +1 , r t , d t ∼ p ( x t +1 , r t , d t | x ≤ t , a ≤ t ) .
The reinforcement learning objective is to train a policy π that yields actions maximizing the expected sum of rewards
E π
[ (cid:80) t ≥ 0 γ t r t ] .
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.
; Hafner et al.
(2020; 2021), our agent learns to act exclusively within its model of the world, and we only make use of real experience to learn the environment dynamics.
We repeatedly perform the three following steps:
• collect_experience : gather experience in the real environment with the current policy.
• update_world_model : improve rewards, episode ends and next observations predictions.
• update_behavior : in imagination, improve the policy and value functions.
The world model is composed of a discrete autoencoder (Van Den Oord et al., 2017), to convert an image to tokens and back, and a GPT -like autoregressive Transformer (Vaswani et al., 2017;
Radford et al., 2019; Brown et al., 2020b), whose task is to capture environment dynamics.
Figureillustrates 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 A provides a comprehensive description of model architectures and hyperparameters.
Algorithm 1 summarizes the training protocol.
2. F ROM IMAGE OBSERVATIONS TO TOKENS
The discrete autoencoder ( E, D ) learns a symbolic language of its own to represent high-dimensional images as a small number of tokens.
The back and forth between frames and tokens is illustrated with green arrows in Figure 1.
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 onlygames 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 E : R h × w × 3 → { 1 , . .
. , N } K converts an input image x t into K tokens from a vocabulary of size N .
Let E = { e i } Ni =1 ∈
R
N
×
d be the corresponding embedding table of d -dimensional vectors.
The input image x t is first passed through a Convolutional Neural
Network ( CNN ) (LeCun et al., 1989) producing output y t ∈ R K × d .
We then obtain the output tokens z t =
( z 1 t , . .
. , z Kt ) ∈ { 1 , . .
, N } K as z kt = argmin
i (cid:13)(cid:13)
y kt − e i (cid:13)(cid:13) 2 , the index of the closest embedding vector in E (Van Den Oord et al., 2017; Esser et al., 2021).
Conversely, the CNN decoder
D : { 1 , . . .
, N } K → R h × w × 3 turns K tokens back into an image.
This discrete autoencoder is trained on previously collected frames, with an equally weighted combination of a L 2 reconstruction loss, a commitment loss (Van Den Oord et al., 2017; Esser et al., 2021), and a perceptual loss (Esser et al., 2021; Johnson et al., 2016; Larsen et al., 2016).
We use a straight-through estimator (Bengio et al., 2013) to enable backpropagation training.
M ODELING DYNAMICS
At a high level, the Transformer G captures the environment dynamics by modeling the language of the discrete autoencoder over time.
Its central role of unfolding imagination is highlighted with the blue arrows in Figure 1.
Specifically, G operates over sequences of interleaved frame and action tokens.
An input sequence ( z 10 , . . .
, z K 0 , a 0 , z 11 , . .
z K 1 , a 1 , . . .
, z 1 t , . . .
, z Kt , a t ) is obtained from the raw sequence ( x 0 , a 0 , x 1 , a 1 , . . .
, x t , a t ) by encoding the frames with E , as described in Section 2.1.
At each time step t , the Transformer models the three following distributions:
[Equation]
Note that the conditioning for the k -th token also includes z <kt +1 :
= ( z 1 t +1 , . . .
, z k − 1 t +1 ) , the tokens that were already predicted, i.e. the autoregressive process happens at the token level.
We train G in a self-supervised manner on segments of L time steps, sampled from past experience.
We use a cross-entropy loss for the transition and termination predictors, and a mean-squared error loss or a cross-entropy loss for the reward predictor, depending on the reward function.
2. L EARNING IN IMAGINATION
Together, the discrete autoencoder ( E, D ) and the Transformer G form a world model, capable of imagination.
The policy π , depicted with purple arrows in Figure 1, exclusively learns in this imagination MDP .
At time step t , the policy observes a reconstructed image observation ˆ x t and samples action a t ∼ π ( a t | ˆ x ≤ t ) .
The world model then predicts the reward ˆ r t , the episode end ˆ d t , and the next observation ˆ x t +
= D (ˆ z t +1 ) , with ˆ z t +1 ∼ p G (ˆ z t +1 |
z 0 , a 0 , ˆ z 1 , a 1 , . .
, ˆ z t , a t ) .
This imagination procedure is initialized with a observation x 0 sampled from past experience, and is rolled out for H
, the imagination horizon hyperparameter.
We stop if an episode end is predicted before reaching the horizon.
Figure 1 illustrates the imagination procedure.
As we roll out imagination for a fixed number of steps, we cannot simply use a Monte Carlo estimate for the expected return.
Hence, to bootstrap the rewards that the agent would get beyond a given time step, we have a value network V that estimates V (ˆ x t ) ≃ E π (cid:2) (cid:80) τ ≥ t γ τ − t ˆ r τ (cid:3) .
Many actor-critic methods could be employed to train π and V in imagination (Sutton & Barto, 2018;
Kaiser et al., 2020; Hafner et al., 2020).
For the sake of simplicity, we opt for the learning objectives and hyperparameters of DreamerV2 (Hafner et al., 2021), that delivered strong performance in Atari games.
Appendix B gives a detailed breakdown of the reinforcement learning objectives.
E XPERIMENTS
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 et al., 2020).
We present the benchmark and its baselines in Section 3.1.
We describe the evaluation protocol and discuss the results in Section 3.2.
Qualitative examples of the world model’s capabilities are given in Section 3.3.
[Table]
3. B ENCHMARK AND B ASELINES
Atari 100k consists of 26 Atari games (Bellemare et al., 2013a) with various mechanics, evaluating a wide range of agent capabilities.
In 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 benchmarked on Atari 100k.
SimPLe (Kaiser et al., 2020) trains a policy with PPO (Schulman et al., 2017) in a 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., 2021) enforces consistent representations of input images across augmented views and neighbouring time steps.
The aforementioned baselines carry additional techniques to improve performance.
Instead, IRIS does not rely on prioritized experience replay (Schaul et al., 2016), epsilon-greedy scheduling, or data augmentations.
We make a distinction between methods with and without lookahead search.
Indeed, algorithms such as Monte Carlo Tree Search (MCTS) (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.
On the contrary, IRIS and the baselines above do not require distributed infrastructures and multiple GPUs.
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.
MuZero (Schrittwieser et al., 2020) and EfficientZero (Ye et al., 2021) are the current standard for search-based methods.
MuZero leverages MCTS as a policy improvement operator, by unrolling
[Figure] [Figure]
[Figure] [Figure]
multiple hypothetical trajectories in the latent space of a world model.
EfficientZero improves upon
MuZero by introducing a self-supervised consistency loss, predicting returns over short horizons in one shot, and correcting off-policy trajectories with its world model.
3. R ESULTS
The human normalized score is the established measure of performance in Atari 100k.
It is defined as
[Figure]
_ random , where score_random comes from a random policy, and score_human is obtained from human players (Wang et al., 2016).
Table 1 displays returns across games and human-normalized aggregate metrics.
For MuZero and
EfficientZero, we report the averaged results published by Ye et al.
(2021) (3 runs).
We use results from the Atari 100k case study conducted by Agarwal 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 at 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 5 the human normalized scores with stratified bootstrapped confidence intervals for mean, median, and interquartile mean (IQM).
For finer comparisons, we also provide performance profiles and probabilities of improvement in Figure 6.
With the equivalent of only two hours of gameplay, IRIS achieves a superhuman mean score of 1.
(+70%), an IQM of 0.501 (+49%), an 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 a new state of the art for methods without lookahead search in the Atari 100k benchmark.
We also note that IRIS outperforms MuZero, although the latter was not designed for the sample-efficient regime.
In addition, performance profiles (Figure 6a) reveal that IRIS is on par with the strongest baselines for its bottom 50% of games, at which point it stochastically dominates (Agarwal et al., 2021; Dror et al., 2019) the other methods.
Similarly, the probability of improvement is 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.
(2021) observe 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 is 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 on a double exploration problem.
IRIS has to first discover a new aspect of the game for its world model to internalize it.
Only then may the policy rediscover and exploit it.
Figure 7 details this phenomenon in Frostbite and Krull , two games with multiple levels.
In summary, as long as transitions between levels do not depend on low-probability events, the double exploration problem does not hinder performance.
Another kind of games difficult to simulate are mazes with moving enemies, such as MsPacman ,
BankHeist , and Alien .
As discussed in Appendix E, this is due to the small number of tokens used to encode frames.
When we increase that amount, the world model is able to situate enemies, albeit at the cost of increased computation.
3. W ORLD M ODEL A NALYSIS
As IRIS learns behaviors entirely in its imagination, the quality of the world model is the cornerstone of our approach.
For instance, it is 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 performance.
Hence, no matter the amount of imagined trajectories, the agent will learn suboptimal policies if the world model is flawed.
While Section 3.2 provides a 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 3 depicts pixel-perfect predictions in Pong .
Finally, we illustrate in Figure 4 predictions for rewards and episode terminations, which are crucial to the reinforcement learning objective.
R ELATED W ORK L EARNING IN THE IMAGINATION OF WORLD MODELS
The idea of training policies in a learnt model of the world was first investigated in tabular environments (Sutton & Barto, 2018).
Ha & Schmidhuber (2018) showed that simple visual environments could be simulated with autoencoders and recurrent networks.
SimPLe (Kaiser et al., 2020) demonstrated that a PPO policy (Schulman et al., 2017) trained in a video prediction model outperformed humans in some Atari games.
Improving upon Dreamer (Hafner et al., 2020), DreamerV2 (Hafner et al., 2021) was 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-space model (RSSM)
(Hafner et al., 2019) for latent dynamics learning.
More recently, Chen et al.
(2022) explored a variant of DreamerV2 where a Transformer replaces the recurrent network in the RSSM.
R EINFORCEMENT L EARNING WITH T RANSFORMERS
Following spectacular advances in natural language processing (Manning & Goldie, 2022), the reinforcement learning community has recently stepped into the realm of Transformers.
Parisotto et al.
(2020) make the observation that the standard Transformer architecture is difficult to optimize with RL objectives.
The authors propose to 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-supervised learning objective.
The Trajectory Transformer (Janner et al., 2021) and the
Decision Transformer (Chen et al., 2021) represent offline trajectories as a static dataset of sequences.
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, the Decision Transformer can handle image inputs but it cannot be easily extended as a world model.
Ozair et al.
(2021) 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.
V IDEO GENERATION WITH DISCRETE AUTOENCODERS AND T RANSFORMERS
VQGAN (Esser et al., 2021) and DALL - E (Ramesh et al., 2021) use discrete autoencoders to 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 et al., 2021) models sequences of frames instead of a single frame for text conditional video generation.
VideoGPT (Yan et al., 2021) introduces video-level discrete autoencoders, and Transformers with spatial and temporal attention patterns, for unconditional and action conditional video generation.
C ONCLUSION
We introduced IRIS , an agent that learns purely in the imagination of a world model composed of a 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-hardened agents, IRIS opens a 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 it could probably leverage the internal representations of the world model.
Another exciting avenue of research would be to combine learning in imagination with MCTS.
Indeed, both approaches deliver impressive results, and their contributions to agent performance are orthogonal.
R EPRODUCIBILITY S TATEMENT
The different components and their training objectives are 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 F. Algorithm 1 makes explicit the interplay between components in the training loop.
In Section 3.2, we provide the source of the reported results for the 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.
E THICS S TATEMENT
The development of autonomous agents for real-world environments raises many safety and environmental concerns.
During its training period, an agent may cause serious harm to 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 a 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.
A M ODELS AND H YPERPARAMETERS A. D ISCRETE AUTOENCODER
Our discrete autoencoder is based on the implementation of VQGAN (Esser et al., 2021).
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 et al., 2016).
[Table] [Table]
Note that during experience collection in the real environment, frames still go through the autoencoder to keep the input distribution of the policy unchanged.
See Algorithm 1 for details.
A. T RANSFORMER
Our autoregressive Transformer is based on the implementation of minGPT (Karpathy, 2020).
It takes as input a sequence of L ( K + 1) tokens and embeds it into a L ( K + 1) × D tensor using an
A × D embedding table for actions, and a N × D embedding table for frames tokens.
This tensor is forwarded through M Transformer blocks.
We use GPT2-like blocks (Radford et al., 2019), i
[Table]
e. each block consists of a self-attention module with layer normalization of the input, wrapped with a residual connection, followed by a per-position multi-layer perceptron with layer normalization of the input, wrapped with another residual connection.
[Table]
A. A CTOR -C RITIC
The weights of the actor and critic are shared except for the last layer.
The actor-critic takes as input a 64 × 64 × 3 frame, and forwards it through a convolutional block followed by an LSTM cell (Mnih et al., 2016; Hochreiter & Schmidhuber, 1997; Gers et al., 2000).
The convolutional block consists of the same layer repeated four times: a 3x3 convolution with stride 1 and padding 1, a ReLU activation, and 2x2 max-pooling with stride 2.
The dimension of the LSTM hidden state is 512.
Before starting the imagination procedure from a given frame, we burn-in (Kapturowski et al., 2019)
the 20 previous frames to initialize the hidden state
[Table] [Table]
B A CTOR - CRITIC LEARNING OBJECTIVES
We follow Dreamer (Hafner et al., 2020; 2021) in using the generic λ -return, that balances bias and variance, as the regression target for the value network.
Given an imagined trajectory (ˆ x 0 , a 0 , ˆ r 0 , ˆ d 0 , . . .
, ˆ x H − 1 , a H − 1 , ˆ r H − 1 , ˆ d H − 1 , ˆ x H ) , the λ -return can be defined recursively as follows:
[Equation]
The value network V is trained to minimize L V , the expected squared difference with λ -returns over imagined trajectories.
[Equation]
Here, sg( · ) denotes the gradient stopping operation, meaning that the target is a constant in the gradient-based optimization, as classically established in the literature (Mnih et al., 2015; Hessel et al., 2018; Hafner et al., 2020).
As large amounts of trajectories are generated in the imagination MDP , we can use a straightforward reinforcement learning objective for the policy, such as REINFORCE (Sutton & Barto, 2018).
To reduce the variance of REINFORCE gradients, we use the value V (ˆ x t ) as a baseline (Sutton & Barto, 2018).
We also add a weighted entropy maximization objective to maintain a sufficient exploration.
The actor is trained to minimize the following REINFORCE objective over imagined trajectories:
[Equation]
[Table]
C O PTIMALITY GAP
[Figure]
[Table] [Figure] [Table]
[Figure]
D
IRIS A LGORITHM
Algorithm 1:
Procedure training_loop() : for epochs do collect_experience( steps_collect ) for steps_world_model do update_world_model() for steps_behavior
do update_behavior()
Procedure collect_experience( n ) : x 0 ← env.reset() for t = 0 to n − 1 do ˆ x t ← D ( E ( x t )) //
forward frame through discrete autoencoder
Sample a t ∼ π ( a t | ˆ x t ) x t +1 , r t , d t ← env.step( a t ) if d t = 1 then x t +1 ← env.reset()
D ← D ∪ { x t , a t , r t , d t } n − 1 t =
Procedure update_world_model() :
Sample { x t , a t , r t , d t } τ + L − 1 t = τ ∼ D
Compute z t := E ( x t ) and ˆ x t := D ( z t ) for t
= τ, . .
, τ + L −
Update E and D
Compute p G (ˆ z t +1 , ˆ r t , ˆ d t | z τ , a τ , . . .
, z t , a t ) for t =
τ, . .
, τ + L −
Update G
Procedure update_behavior() :
Sample x 0 ∼ D z 0 ← E ( x 0 ) ˆ x 0 ← D ( z 0 ) for t = 0 to H − 1 do
Sample a t ∼ π ( a t | ˆ x t )
Sample ˆ z t +1 , ˆ r t , ˆ d t ∼ p G (ˆ z t +1 , ˆ r t , ˆ d t | z 0 , a 0 , . . .
, ˆ z t , a t ) ˆ x t +1 ← D (ˆ z t +1 )
Compute V (ˆ x t ) for t = 0 , . . .
, H
Update π and V
E A UTOENCONDING FRAMES WITH VARYING AMOUNTS OF TOKENS
The length of the input sequence of G is determined by the number of tokens K used to encode a single frame and the number of timesteps L 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 Atari games where enemies and players are moving in mazes with rewards to collect.
Due to the high number of possible configurations, the discrete autoencoder struggles to properly encode frames with only K = 16 tokens.
Indeed, sometimes the player, its enemies, or rewards are not correctly reconstructed, which severely hinders agent performance.
In Figure we show that when increasing the number of tokens per frame to 64, the discrete autoencoder is perfectly capable of dealing with detailed environments such as Alien .
However, this increases the sequence length of G from 340 to 1300.
Therefore, with more computational resources,
IRIS would most likely improve in these settings
F C OMPUTATIONAL RESOURCES
For each Atari environment, we repeatedly trained IRIS with 5 different random seeds.
We ran our experiments with 8 Nvidia A100 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 et al., 2020), the only baseline that involves learning in imagination, trains forweeks with a P100 GPU on a single environment.
As for SPR (Schwarzer et al., 2021), the strongest baseline without lookahead search, it trains notably fast in 4.6 hours with a P100 GPU.
Regarding baselines with lookahead search, MuZero (Schrittwieser et al., 2020) originally used
TPUs for 12 hours to train in a single Atari environment.
Ye et al.
(2021) train both EfficientZero and their reimplementation of MuZero in 7 hours with 4 RTX 3090 GPUs.
EfficientZero’s implementation relies on a distributed infrastructure with CPU and GPU threads running in parallel, and a C++/Cython implementation of MCTS.
By contrast, IRIS and the baselines without lookahead search rely on straightforward single GPU / single CPU implementations.
G E XPLORATION IN F REEWAY
The reward function in Freeway is sparse since the agent is only rewarded when it completely crosses the road.
In addition, bumping into cars will drag it 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 a non-zero reward with a 100k frames budget.
The solution to this problem is 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 a fixed epsilon-greedy parameter and sampling from the policy.
However, we lowered the sampling temperature from 1 to 0.01 for Freeway, in order to avoid random walks that would not be conducive to learning in the early stages of training.
As a consequence, once it received its first few rewards through exploration, IRIS was able to internalize the sparse reward function in its world model.