Improving Token-Based World Models with Parallel Observation Prediction (2024)

1 Introduction

Sample efficiency remains a central challenge in reinforcement learning (RL) due to the substantial data demands of successful RL algorithms(Mnih etal., 2015; Silver etal., 2016; Schrittwieser etal., 2020; Berner etal., 2019; Vinyals etal., 2019).One prominent model-based approach for addressing this challenge is known as world models. In world models, the agent’s learning is solely based on simulated interaction data produced by a learned model of the environment through a process called imagination.World models are gaining increasing popularity due to their effectiveness, particularly in visual environments(Hafner etal., 2023).

In recent years, attention-based sequence models, most notably the Transformer architecture (Vaswani etal., 2017), demonstrated unmatched performance in language modeling tasks (Devlin etal., 2019; Brown etal., 2020; Bubeck etal., 2023; Touvron etal., 2023).The notable success of these models when applied to sequences of discrete tokens sparked motivation to employ these architectures to other data modalities by learning appropriate token-based representations.In computer vision, discrete representations are becoming a mainstream approach for various tasks (vanden Oord etal., 2017; Dosovitskiy etal., 2021; Esser etal., 2021; Li etal., 2023).In RL, token-based world models were recently explored in visual environments (Micheli etal., 2023).The visual perception module in these methods is called a tokenizer, as it maps image observations to sequences of discrete symbols.This way, agent interaction is translated into a language-like sequence of discrete tokens, which are processed individually by the world model.

During imagination, to generate the tokens of the next observation with the auto-regressive model, the prediction is carried sequentially token-by-token.Effectively, this highly-sequential computation results in a severe bottleneck that pronouncedly hinders token-based approaches.Consequently, this bottleneck practically caps the length of the observation token sequences which in turn degrades performance.This limitation renders current token-based methods impractical for complex problems.

In this paper, we present Parallel Observation Prediction (POP), a novel mechanism that resolves the imagination bottleneck of token based world models (TBWMs).With POP, the enire next observation token sequence is generated in parallel during world model imagination.At its core, POP augments a Retentive Network (RetNet) sequence model (Sun etal., 2023) with a novel forward mode devised for retaining world model training efficiency.Additionally, we present REM (Retentive Environment Model), a TBWM agent driven by a POP-augmented RetNet architecture.

2 Method

Notations.We consider the Partially Observable Markov Decision Process (POMDP) settingwith image observations $\mathbf{o}_{t}\in\Omega\subseteq\mathbb{R}^{h\times w\times 3}$,discrete actions $a_{t}\in\mathcal{A}$,scalar rewards $r_{t}\in\mathbb{R}$,episode termination signals $d_{t}\in\{0,1\}$,dynamics $\mathbf{o}_{t+1},r_{t},d_{t}\sim p(\mathbf{o}_{t+1},r_{t},d_{t}|\mathbf{o}_{$,and discount factor $\gamma$.The objective is to learn a policy $\pi$ such that for every situation the output $\pi(a_{t}|\mathbf{o}_{\leq t},a_{<t})$ is optimal w.r.t. the expected discounted sum of rewards from that situation $\operatorname*{\mathbb{E}}[\sum_{\tau=0}^{\infty}\gamma^{\tau}R_{t+\tau}]$ under the policy $\pi$.

2.1 Overview

REM builds on IRIS (Micheli etal., 2023), and similar to most prior works on world models for pixel input(Hafner etal., 2021; Kaiser etal., 2020; Hafner etal., 2023), REM follows a $\mathcal{V}$-$\mathcal{M}$-$\mathcal{C}$ structure(Ha & Schmidhuber, 2018):a $\mathcal{V}$isual perception module that compresses observations into compact latent representations, a predictive $\mathcal{M}$odel that captures the environment’s dynamics, and a $\mathcal{C}$ontroller that learns to act to maximize return.Additionally, a replay buffer is used to store environment interaction data.An overview of REM’s training cycle is presented in Figure 2.A pseudo-code algorithm of REM is presented in Appendix A.2.

$\mathcal{V}$ - Tokenizer

We instantiate the visual perception component as a tokenizer, mapping input observations into latent tokens.Following (Micheli etal., 2023), the tokenizer is a VQ-VAE discrete auto-encoder(vanden Oord etal., 2017; Esser etal., 2021) comprised of an encoder, a decoder, and an embedding table.The embedding table $\mathcal{E}=\{\mathbf{e}_{i}\}_{i=1}^{N}\in\mathbb{R}^{N\times d}$ consists of $N$ trainable vectors.The encoder first maps an input image $\mathbf{o}_{t}$ to a sequence of $d$-dimensional latent vectors $(\mathbf{h}^{1}_{t},\mathbf{h}^{2}_{t},\cdots,\mathbf{h}^{K}_{t})$.Then, each latent vector $\mathbf{h}^{k}_{t}\in\mathbb{R}^{d}$ is mapped to the index of the nearest embedding in $\mathcal{E}$, , $z_{t}^{k}=\operatorname*{arg\,min}_{i}\|\mathbf{h}_{t}^{k}-\mathbf{e}_{i}\|$.Such indices are called tokens.For an input image $\mathbf{o}_{t}$, its latent token sequence is denoted as $\mathbf{z}_{t}=(z_{t}^{1},z_{t}^{2},\cdots,z_{t}^{K})$.To map a token sequence back to the input space, we first retrieve the embedding for each token and obtain a sequence $(\hat{\mathbf{h}}_{t}^{1},\hat{\mathbf{h}}_{t}^{2},\cdots,\hat{\mathbf{h}}_{t}$ where $\hat{\mathbf{h}}_{t}^{k}=\mathbf{e}_{z_{t}^{k}}$.Then, inverse to the encoding process, the decoder is responsible for mapping this sequence to a reconstructed observation $\hat{\mathbf{o}}_{t}$.

The tokenizer is trained on frames sampled uniformly from the replay buffer.Its optimization objective, architecture, and other details are deferred to Appendix A.1.1.

$\mathcal{M}$ - World Model

At the core of a world model is the component that captures the dynamics of the environment and makes predictions based on historical observations.Here, $\mathcal{M}$ is learned entirely in the latent token space, modeling the following distributions at each step $t$:

	Transition:	$\displaystyle p(\hat{\mathbf{z}}_{t+1}|\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_$		(1)
	Reward:	$\displaystyle p(\hat{r}_{t}|\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_{t},a_{t}),$		(2)
	Termination:	$\displaystyle p(\hat{d}_{t}|\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_{t},a_{t}).$		(3)

To map observation tokens to embedding vectors, $\mathcal{M}$ uses the code vectors $\mathcal{E}$ learned by the tokenizer $\mathcal{V}$.Note that $\mathcal{E}$ is not updated by $\mathcal{M}$.In addition, $\mathcal{M}$ maintains dedicated embedding tables for mapping actions and special tokens (detailed in Section 2.3) to continuous vectors.

$\mathcal{C}$ - Controller

REM’s actor-critic controller $\mathcal{C}$ is trained to maximize return entirely in imagination (Kaiser etal., 2020; Hafner etal., 2021; Micheli etal., 2023).$\mathcal{C}$ comprises of a policy network $\pi$ and a value function estimator $V^{\pi}$, and operates on latent tokens and their embeddings. In each optimization step, $\mathcal{M}$ and $\mathcal{C}$ are initialized with a short trajectory segment sampled from the replay buffer.Subsequently, the agent interacts with the world model for $H$ steps.At each step $t$, the agent plays an action sampled from its policy $\pi(a_{t}|\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_{t-1},a_{t-1},\mathbf{z}_{t})$.The world model evolves accordingly, generating $\hat{r}_{t}$, $\hat{d}_{t}$, and $\hat{\mathbf{z}}_{t+1}$ by sampling from the appropriate distributions (Eqn. (1-3)). The resulting trajectories are then used to train the agent.Following (Micheli etal., 2023), we adopted the actor-critic objectives of DreamerV2 (Hafner etal., 2021).We leave the full details of its architecture and optimization to Appendix A.1.3.

2.2 Retention Preliminaries

Similar to Transformers (Vaswani etal., 2017), a RetNet model (Sun etal., 2023) consists of a stack of layers, where each layer contains a multi-head Attention-like mechanism, called Retention, followed by a fully-connected network.A unique characteristic of the Retention mechanism is that it has a dual form of recurrence and parallelism, called “chunkwise”, for improved efficiency when handling long sequences.This form allows to split such sequences into smaller “chunks”, where a parallel computation takes place within chunks and a sequential recurrent form is used between chunks, as shown in Figure3.The information from previous chunks is summarized by a recurrent state $\mathbf{S}\in\mathbb{R}^{d\times d}$ maintained by the Retention mechanism.

Formally, consider a sequence of tokens $(x_{1},x_{2},\cdots,x_{m})$.In our RL context, this sequence is a token trajectory composed of observation-action sub-sequences $(z_{t}^{1},\cdots,z_{t}^{K},a_{t})$ we call blocks.As such trajectories are typically long, we split them into chunks of $B$ tokens, where $B=c(K+1)$ is a multiple of $K+1$ so that each chunk only contains complete blocks.Here, the hyperparameter $c$ can be tuned according to the size of the models, the hardware, and other factors to maximize efficiency.Let $\mathbf{X}=(\mathbf{x}_{1},\mathbf{x}_{2},\cdots,\mathbf{x}_{m})\in\mathbb{R}^$ be the $d$-dimensional token embedding vectors.The Retention output $\mathbf{Y}_{[i]}=\mathrm{Retention}(\mathbf{X}_{[i]},\mathbf{S}_{[i-1]},i)$ of the $i$-th chunk is given by

$$\mathbf{Y}_{[i]}=\left(\mathbf{Q}_{[i]}\mathbf{K}^{\top}_{[i]}\odot\mathbf{D}$$

(4)

where the bracketed subscript $[i]$ is used to index the $i$-th chunk, $\mathbf{Q}=\left(\mathbf{X}\mathbf{W}_{Q}\right)\odot\Theta$, $\mathbf{K}=\left(\mathbf{X}\mathbf{W}_{K}\right)\odot\bar{\Theta}$, $\mathbf{V}=\mathbf{X}\mathbf{W}_{V}$, and $\boldsymbol{\xi}\in\mathbb{R}^{B\times d}$ is a matrix with $\boldsymbol{\xi}_{ij}=\eta^{i+1}$.Here,$\mathbf{W}_{Q},\mathbf{W}_{K},\mathbf{W}_{V}\in\mathbb{R}^{d\times d}$ are learnable weights, $\eta$ is an exponential decay factor, the matrix $\mathbf{D}\in\mathbb{R}^{B\times B}$ combines an auto-regressive mask with the temporal decay factor $\eta$, and the matrices $\Theta,\bar{\Theta}\in\mathbb{C}^{m\times d}$ are for relative position embedding (see Appendix A.3).Note the chunk index $i$ argument of the Retention operator, which controls positional embedding information through $\Theta$.The chunkwise update rule of the recurrent state is given by

$$\mathbf{S}_{[i]}=(\mathbf{K}_{[i]}\odot\boldsymbol{\zeta})^{\top}\mathbf{V}_{[$$

(5)

where $\mathbf{S}_{[0]}=\mathbf{S}_{0}=0$, and $\boldsymbol{\zeta}\in\mathbb{R}^{B\times d}$ is a matrix with $\boldsymbol{\zeta}_{ij}=\eta^{B-i-1}$.On the right hand side of Equations 4 and 5, the first term corresponds to the computation within the chunk while the second term incorporates the information from previous chunks, encapsulated by the recurrent state.Further details about the RetNet architecture are deferred to Appendix A.3.

2.3 World Model Imagination

As the agent’s training relies entirely on world model imagination, the efficiency of the trajectory generation is critical.During imagination, predicting $\hat{\mathbf{z}}_{t+1}$ constitutes the primary non-trivial component and consumes the majority of processing time.In IRIS, the prediction of $\hat{\mathbf{z}}_{t+1}$ unfolds sequentially, as the model is limited to predicting only one token ahead at each step.This limitation arises since the identity of the next token, which remains unknown at the current step, is necessary for the prediction of later tokens.Thus, generating $H$ observations costs $KH$ sequential world model calls.This leads to poor GPU utilization and long computation time.

To overcome this bottleneck, POP maintains a set of $K$ dedicated prediction tokens $\mathbf{u}=(u_{1},\ldots,u_{K})$ together with their corresponding embeddings $\mathcal{E}_{\mathbf{u}}\in\mathbb{R}^{K\times d}$.To generate $\hat{\mathbf{z}}_{t+1}$ in one pass, POP simply computes the RetNet outputs starting from $\mathbf{S}_{t}$ using $\mathbf{u}$ as its input sequence, as illustrated in Figure4.Note that at imagination, the chunk size is limited to a single block, i.e., to $K+1$.Here, the notation $\mathbf{S}_{t}$ refers to the state that summarizes the first $t$ observation-action blocks.To obtain $\mathbf{S}_{t}$, we use RetNet’s chunkwise forward to summarize an initial context segment of blocks sampled from the replay buffer.Essentially, for every $t$, POP models the following distribution for next observation prediction:

$$p(\hat{\mathbf{z}}_{t+1}|\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_{t},a_{t},$$

with

$$p(\hat{\mathbf{z}}_{t+1}^{k}|\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_{t},a_{t},$$

It is worth noting that the tokens $\mathbf{u}$ are only intended for observation token predictions and are never employed in the update of the recurrent state.

This approach effectively reduces the total number of world model calls during imagination from $KH$ to $2H$, eliminating the dependency on the number of observation tokens $K$.In fact, POP provides an additional generation mode that further reduces the number of sequential calls to $H$.We defer the details on this mode to Appendix A.1.2.Also, by using a recurrent state that summarizes long history sequences, POP improves efficiency further, as the per-token prediction cost reduces.Effectively, POP offers improved scalability at the expense of a higher overall computational cost ($(2K+1)H$ compared to $(K+1)H$).Our approach add to existing evidence suggesting that enhanced scalability is often favorable, even at the expense of additional computational costs, with Transformers (Vaswani etal., 2017) serving as a prominent example.

2.4 World Model Training

While applying POP during imagination is fairly straightforward, it requires modification of the training data.Consider an input trajectory segment $(\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_{T},a_{T})$ sampled from the replay buffer.To make meaningful observation predictions at imagination, the model should be trained to predict $\mathbf{z}_{t}$ given $(\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_{t-1},a_{t-1},\mathbf{u})$, for each time step $t$ of every input segment.Hence, for every $t$, the input sequence should contain $\mathbf{u}$ at block $t$. However, replacing $\mathbf{z}_{t}$ with $\mathbf{u}$ in the original sequence is inadequate, as the prediction of future observations, rewards, and termination signals depends on $\mathbf{z}_{t}$.Thus, the standard approach of computing all outputs from the same input sequence is not viable, as in this case these two requirements contradict each other (Figure5).The challenge then lies in devising an efficient method for computing the outputs for all time steps in parallel.

To tackle this challenge, we first note that each trajectory prefix can be summarized into a single recurrent state.For example, for the first input chunk $(\mathbf{z}_{1},a_{1},\ldots,\mathbf{z}_{c},a_{c})$,$(\mathbf{z}_{1},a_{1})$ can be summarized into $\mathbf{S}_{[1,1]}$ and $(\mathbf{z}_{1},a_{1},\mathbf{z}_{2},a_{2})$ can be summarized into $\mathbf{S}_{[1,2]}$.Here, we use the subscript $[i,j]$ to conveniently refer to the $j$-th block within the $i$-th chunk (this notation is demonstrated in Figure 3), with $\mathbf{S}_{[i,0]}=\mathbf{S}_{[i-1]}$ and $\mathbf{S}_{[i,c]}=\mathbf{S}_{[i]}$.Thus, our plan is to first compute all states $\mathbf{S}_{[i,1]},\ldots,\mathbf{S}_{[i,c]}$ in parallel, and then predict all next observations from all $(\mathbf{S}_{[i,j]},\mathbf{u})$ tuples.

To compute all recurrent states $\mathbf{S}_{[i,j]}$ in parallel, a two-step computation is carried.First, intermediate states $\tilde{\mathbf{S}}_{[i,j]}$ are computed in parallel for all $j$ with

$$\tilde{\mathbf{S}}_{[i,j]}=\left(\mathbf{K}_{[i,j]}\odot\boldsymbol{\zeta}$$

(6)

where $\boldsymbol{\zeta}\in\mathbb{R}^{(K+1)\times d}$ is a matrix with $\boldsymbol{\zeta}_{ij}=\eta^{K-i}$.Then, each recurrent state is computed sequentially by

$$\mathbf{S}_{[i,j]}=\tilde{\mathbf{S}}_{[i,j]}+\eta^{K+1}\mathbf{S}_{[i,j-1]}.$$

(7)

As the majority of the computational burden lies in the first step, the sequential computation in the second step has minimal impact on the overall speedup.

Once we have all states ready, the output of $(\mathbf{S}_{[i,j]},\mathbf{u})$ for all $1\leq j\leq c$ is computed in parallel.Here, we stress that the existing Retention mechanism can only perform batched input computation with recurrent states $\mathbf{S}_{t}$ of the same time step $t$.This is due to the shared positional embedding information applied to every input sequence in the batch.To overcome this, we devise a mechanism which extends RetNet to support the batched computation of the $(\mathbf{S}_{[i,j]},\mathbf{u})$ tuples, while applying the appropriate positional encoding information.A pseudo code of our novel POP extension of RetNet is given in Algorithms 1 and 2.The latter presents the core of the mechanism (described above), while the former describes the higher level layer-by-layer computation with a final aggregation for combining the produced outputs.Figure 6 illustrates a simplified example of the POP Forward mechanism (Algorithms 1 and 2) for a single-layer model.For brevity, our pseudo code and illustrations only considers Retention layers, omitting other modules of RetNet (Appendix A.3).

1:Input: chunk size $1\leq c\leq H$,token embeddings $\mathbf{X}_{[i]}$ of chunk $i$,per-layer recurrent states $\{\mathbf{S}_{[i-1]}^{l}\}_{l=1}^{L}$.

2:Initialize $\mathbf{A}^{0}_{[i]}\leftarrow\mathbf{X}$

3:Initialize $\mathbf{B}^{0}_{[i,1]},\ldots,\mathbf{B}^{0}_{[i,c]}\leftarrow\mathcal{E}_{$

4:for$l=1$ to $L$do

See Also

How to fine-tune Llama3 using DPO on Brev 🤙Self-Attention Explained with Code

5:$\mathbf{A}^{l}_{[i]},\mathbf{B}^{l}_{[i]},\mathbf{S}_{[i]}^{l}\leftarrow\text{$

6:endfor

7:for$j=1$ to $c$do

8:$\mathbf{Y}_{[i,j]}\leftarrow\text{Concat}(\mathbf{B}^{L}_{[i,j]},\mathbf{A}^{L$

9:endfor

10:Return $\mathbf{Y},\{\mathbf{S}_{[i]}^{l}\}_{l=1}^{L}$

1:Input: Chunk latents $\mathbf{A}_{[i]}$, observation prediction latents $\mathbf{B}_{[i]}$, recurrent state $\mathbf{S}_{[i-1]}$, chunk index $i$.

2:$\mathbf{A}_{[i]}\leftarrow\mathrm{Retention}(\mathbf{A}_{[i]},\mathbf{S}_{[i-1$ (Eqn. 4)

3:Compute $\tilde{\mathbf{S}}_{[i,1]},\ldots,\tilde{\mathbf{S}}_{[i,c]}$ in parallel (Eqn. 6)

4:for$j=1$ to $c$ do

5:$\mathbf{S}_{[i,j]}\leftarrow\tilde{\mathbf{S}}_{[i,j]}+\eta^{K+1}\mathbf{S}_{[$ (Eqn. 7)

6:endfor

7:$\mathbf{S}_{[i]}\leftarrow\mathbf{S}_{[i,c]}$

8:$\mathbf{B}_{[i,j]}\leftarrow\mathrm{Retention}(\mathbf{B}_{[i,j]},\mathbf{S}_{$ in parallel for $j=1,\ldots,c$ (Eqn. 4)

9:Return $\mathbf{A}_{[i]},\mathbf{B}_{[i]},\mathbf{S}_{[i]}$

To train the world model, trajectory segments of $H$ steps from past experience are uniformly sampled from the replay buffer and translated into token sequences.These sequences are processed in chunks of $c$ observation-action blocks to produce the modeled distributions, as depicted in Figure 6.Optimization is carried by minimizing the cross-entropy loss of the transitions and termination outputs, and the appropriate loss of the reward outputs, depending on the task.For continuous rewards, the mean-squared error loss is used while for discrete ones cross-entropy is used instead.

			Non-Token-Based				Token-Based
Game	Random	Human	SimPLe	DreamerV3	TWM	STORM	IRIS	REM (ours)
Alien	227.8	7127.7	616.9	959.4	674.6	983.6	420.0	607.2
Amidar	5.8	1719.5	74.3	139.1	121.8	204.8	143.0	95.3
Assault	222.4	742.0	527.2	705.6	682.6	801.0	1524.4	1764.2
Asterix	210.0	8503.3	1128.3	932.5	1116.6	1028.0	853.6	1637.5
BankHeist	14.2	753.1	34.2	648.7	466.7	641.2	53.1	19.2
BattleZone	2360.0	37187.5	4031.2	12250.0	5068.0	13540.0	13074.0	11826.0
Boxing	0.1	12.1	7.8	78.0	77.5	79.7	70.1	87.5
Breakout	1.7	30.5	16.4	31.1	20.0	15.9	83.7	90.7
ChopperCommand	811.0	7387.8	979.4	410.0	1697.4	1888.0	1565.0	2561.2
CrazyClimber	10780.5	35829.4	62583.6	97190.0	71820.4	66776.0	59324.2	76547.6
DemonAttack	152.1	1971.0	208.1	303.3	350.2	164.6	2034.4	5738.6
Freeway	0.0	29.6	16.7	0.0	24.3	0.0	31.1	32.3
Frostbite	65.2	4334.7	236.9	909.4	1475.6	1316.0	259.1	240.5
Gopher	257.6	2412.5	596.8	3730.0	1674.8	8239.6	2236.1	5452.4
Hero	1027.0	30826.4	2656.6	11160.5	7254.0	11044.3	7037.4	6484.8
Jamesbond	29.0	302.8	100.5	444.6	362.4	509.0	462.7	391.2
Kangaroo	52.0	3035.0	51.2	4098.3	1240.0	4208.0	838.2	467.6
Krull	1598.0	2665.5	2204.8	7781.5	6349.2	8412.6	6616.4	4017.7
KungFuMaster	258.5	22736.3	14862.5	21420.0	24554.6	26182.0	21759.8	25172.2
MsPacman	307.3	6951.6	1480.0	1326.9	1588.4	2673.5	999.1	962.5
Pong	-20.7	14.6	12.8	18.4	18.8	11.3	14.6	18.0
PrivateEye	24.9	69571.3	35.0	881.6	86.6	7781.0	100.0	99.6
Qbert	163.9	13455.0	1288.8	3405.1	3330.8	4522.5	745.7	743.0
RoadRunner	11.5	7845.0	5640.6	15565.0	9109.0	17564.0	9614.6	14060.2
Seaquest	68.4	42054.7	683.3	618.0	774.4	525.2	661.3	1036.7
UpNDown	533.4	11693.2	3350.3	7567.1	15981.7	7985.0	3546.2	3757.6
#Superhuman (↑)	0	N/A	1	9	8	9	10	12
Mean (↑)	0.000	1.000	0.332	1.124	0.956	1.222	1.046	1.222
Median (↑)	0.000	1.000	0.134	0.485	0.505	0.425	0.289	0.280
IQM (↑)	0.000	1.000	0.130	0.487	0.459	0.561	0.501	0.673
Optimality Gap (↓)	1.000	0.000	0.729	0.510	0.513	0.472	0.512	0.482

3 Experiments

We follow most prior works on world models and evaluate REM on the widely-recognized Atari 100K benchmark (Kaiser etal., 2020) for sample-efficient reinforcement learning.The Atari 100K benchmark considers a subset of 26 Atari games.For each game, the agent is limited to 100K interaction steps, corresponding to 400K game frames due to the standard frame-skip of 4.In total, this amounts to roughly 2 hours of gameplay.To put in perspective, the original Atari benchmark allows agent to collect 200M steps, that is, 500 times more experience.

3.1 Results

On Atari, it is standard to use human-normalized scores (HNS) (Mnih etal., 2015), calculated as $\frac{\text{agent\_score}\ -\ \text{random\_score}}{\text{human\_score}\ -\$, rather than raw game scores.Here, the final score of each training run is computed as an average over 100 episodes collected at the end of training.In the work of (Agarwal etal., 2021), the authors found discrepancies between conclusions drawn from point estimate statistics such as mean and median and a more thorough statistical analysis that also considers the uncertainty in the results.Adhering to their established protocol and utilizing their toolkit, we report the mean, median, and interquantile mean (IQM) human-normalized scores, and the optimality gap, with 95% stratified bootstrap confidence intervals in Figure 7.Performance profiles are presented in Figure 8.Average scores of individual games are reported in Table1.

REM attains an IQM human normalized score of 0.673, outperforming all baselines.Additionally, REM improves over IRIS on 3 out of the 4 metrics (i.e., mean, optimality gap, and IQM), while being comparable in terms of its median score.Remarkably, REM achieves superhuman performance on 12 games, more than any other baseline (Table 1).REM also exhibits state-of-the-art scores on several games, including Assault, Boxing, and Chopper Command.These findings support our empirical claim that REM performs similarly or better than previous token-based approaches while running significantly faster.

3.2 Ablation Studies

To analyze the impact of different components of our approach on REM’s performance, we conduct a series of ablation studies.For each component, we compare the final algorithm to a version where the component of interest is disabled.Due to computational resource constraints, the evaluation is performed on a subset of 8 games from the Atari 100K benchmark using 5 random seeds for each game.This subset includes games with large score differences between IRIS and REM, as we are particularly interested in studying the impact of each component in these games.Concretely, this subset includes the games “Assult”, “Asterix”, “Chopper Command”, “Crazy Climber”, “Demon Attack”, “Gopher”, “Krull”, and “Road Runner”.We performed ablation studies on the following aspects: the POP mechanism, the latent space architecture of $\mathcal{C}$ and its action inputs, the latent resolution of $\mathcal{V}$, and the observation token embeddings used by $\mathcal{M}$.

The probability of improvement (Agarwal etal., 2021) and IQM human-normalized scores are presented in Figure 9.Figure 10 offers a comparison of the training times, juxtaposing REM with its efficiency-related ablations.

Analyzing POP

To study the impact of POP on REM’s performance, we replaced the POP-augmented RetNet of $\mathcal{M}$ with a vanilla RetNet.In this version, denoted as ”No POP”, the prediction of next observation tokens is performed sequentially token-by-token, as done in IRIS.

Our results suggest that POP retains the agent’s performance (Figure 9) while significantly reducing the overall computation time (Figure 10).In Appendix A.4, we provide additional results indicating that the world model’s performance are also retained.POP achieves lower total computation time by expediting the actor-critic learning phase, despite the increased computational cost implied by the observation prediction tokens.

Actor-Critic Architecture and Action Inputs

For $\mathcal{C}$, we considered an incremental ablation.First, we replaced the architecture of REM’s controller $\mathcal{C}$ with that of IRIS (denoted “$\mathcal{C}_{\text{IRIS}}$”).In contrast to REM, this version processes fully reconstructed pixel frames and does not incorporate action inputs.Formally, $\mathcal{C}_{\text{IRIS}}$ models $\pi(a_{t}|\hat{\mathbf{o}}_{\leq t}),V^{\pi}(\hat{\mathbf{o}}_{\leq t})$.In the second ablation, REM was modified so that only the action inputs of $\mathcal{C}$ were disabled.This ablation corresponds to $\pi(a_{t}|\hat{\mathbf{z}}_{\leq t}),V^{\pi}(\hat{\mathbf{z}}_{\leq t})$.

Our findings indicate that both the latent codes based architecture and the added action inputs contribute to the final performance of REM (Figure 9).Additionally, the latent codes based architecture of $\mathcal{C}$ leads to reduced computational overhead and shorter actor-critic learning times (Figure 10).

Tokenizer Resolution

Here, we compare REM to a version with a reduced latent resolution of $4\times 4$, similar to that of IRIS.The results in Figure 9 provides clear evidence that the latent resolution of the tokenizer has a significant impact on the agent’s performance.Our results demonstrates that POP enables REM to utilize higher latent resolutions while incurring shorter computation times than prior token-based approaches.

World Model Embeddings

In REM, $\mathcal{M}$ translates observation tokens to embedding vectors using the embedding table $\mathcal{E}$ learned by $\mathcal{V}$.These embeddings encode the visual information as learned by $\mathcal{V}$.In contrast, IRIS maintains a separate embedding table learned by the world model for that purpose.Here, the results in Figure 9 provide empirical evidence indicating that leveraging this encoded visual information leads to improved performance.In Appendix A.4, we provide additional evidence suggesting that the world model’s next-observation predictions are also improved.

4 Related Work

Model-based reinforcement learning (RL), with its roots in the tabular setting (Sutton, 1991), has been a focus of extensive research in recent decades.The deep RL agent of (Ha & Schmidhuber, 2018) leveraged an LSTM (Hochreiter & Schmidhuber, 1997) sequence model with a VAE (Kingma & Welling, 2014) to model the dynamics in visual environments, demonstrating that successful policies can be learned entirely from simulated data.This approach, commonly known as world models, was later applied to Atari games (Kaiser etal., 2020) with the PPO (Schulman etal., 2017) RL algorithm.Later, a series of works (Hafner etal., 2020, 2021, 2023) proposed the Dreamer algorithms, which are based on a recurrent state space model (RSSM) (Hafner etal., 2019) to model dynamics.The latest DreamerV3 was evaluated on a variety of challenging environments, providing further evidence of the promising potential of world models.In contrast to token-based approaches, where each token serves as a standalone sequence element, Dreamer encodes each frame as a vector of categorical variables, which are processed at once by the RSSM.

Following the success of the Transformer architecture (Vaswani etal., 2017) in language modeling (Brown etal., 2020), and motivated by their favorable scaling properties compared to RNNs, Transformer were recently explored in RL (Parisotto etal., 2020; Chen etal., 2021; Reed etal., 2022; Shridhar etal., 2023).World model approaches also adopted the Transformer architecture.(Micheli etal., 2023) blazed the trail for token-based world models with IRIS, representing agent trajectories as language-like sequences.By treating each observation as a sequence, its Transformer-based world model gains an explicit sub-observation attention resolution.Despite IRIS’s high performance, its imagination bottleneck results in a substantial disadvantage.

In addition to IRIS, non-token-based world models driven by Transformers were proposed.TWM (Robine etal., 2023) utilizes the Transformer-XL architecture (Dai etal., 2020) and a non-uniform data sampling.STORM (Zhang etal., 2023) proposes an efficient Transformer based world model agent which sets state-of-the-art result for the Atari 100K benchmark.STORM has a significantly smaller 2-layer Transformer compared to the 10-layer models of TWM and IRIS, demonstrating drastically reduced training times and improved agent performance.

5 Conclusions

In this work, we presented a novel parallel observation prediction (POP) mechanism augmenting Retention networks with a dedicated forward mode to improve the efficiency of token-based world models (TBWMs).POP effectively solves the imagination bottleneck of TBWMs and enables them to deal with longer observation sequences.Additionally, we introduced REM, a TBWM agent equipped with POP.REM is the first world model agent driven by the RetNet architecture.Empirically, we demonstrated the superiority of REM over prior TBWMs on the Atari 100K benchmark, rendering REM competitive with the state-of-the-art, both in terms of agent performance and overall run time.

Our work opens up many promising avenues for future research by making TBWMs practical and cost-efficient.One such direction could be to explore a modification of REM where the recurrent state of the world model summarizes the entire history of the agent.Similarly, a history-preserving RetNet architecture should be considered for the controller as well.Another promising avenue would be to leverage the independent optimization of the tokenizer to enable REM to use pretrained visual perception models in environments where visual data is abundant, for example, the real world.Such perceptual models could be trained at scale, and allow REM to store only compressed observations in its replay buffer, further improving its efficiency.Lastly, token-based methods for video generation tasks can benefit from using the POP mechanism for generating entire frames in parallel conditioned on the past context.We believe that this is an exciting avenue to explore with a potentially high impact.

Acknowledgements

This project has received funding from the European Union’s Horizon Europe Programme under grant agreement No. 101070568.

Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

References

• Agarwal etal. (2021)Agarwal, R., Schwarzer, M., Castro, P.S., Courville, A.C., and Bellemare, M.Deep reinforcement learning at the edge of the statistical precipice.In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J.W. (eds.), Advances in Neural Information Processing Systems, volume34, pp. 29304–29320. Curran Associates, Inc., 2021.URL https://proceedings.neurips.cc/paper_files/paper/2021/file/f514cec81cb148559cf475e7426eed5e-Paper.pdf.
• Ba etal. (2016)Ba, J.L., Kiros, J.R., and Hinton, G.E.Layer normalization.arXiv preprint arXiv:1607.06450, 2016.
• Berner etal. (2019)Berner, C., Brockman, G., Chan, B., Cheung, V., Debiak, P., Dennison, C., Farhi, D., Fischer, Q., Hashme, S., Hesse, C., Józefowicz, R., Gray, S., Olsson, C., Pachocki, J., Petrov, M., deOliveiraPinto, H.P., Raiman, J., Salimans, T., Schlatter, J., Schneider, J., Sidor, S., Sutskever, I., Tang, J., Wolski, F., and Zhang, S.Dota 2 with large scale deep reinforcement learning.CoRR, abs/1912.06680, 2019.URL http://arxiv.org/abs/1912.06680.
• Brown etal. (2020)Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J.D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D.Language models are few-shot learners.In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume33, pp. 1877–1901. Curran Associates, Inc., 2020.URL https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf.
• Bubeck etal. (2023)Bubeck, S., Chandrasekaran, V., Eldan, R., Gehrke, J., Horvitz, E., Kamar, E., Lee, P., Lee, Y.T., Li, Y., Lundberg, S.M., Nori, H., Palangi, H., Ribeiro, M.T., and Zhang, Y.Sparks of artificial general intelligence: Early experiments with GPT-4.CoRR, abs/2303.12712, 2023.doi: 10.48550/ARXIV.2303.12712.URL https://doi.org/10.48550/arXiv.2303.12712.
• Chen etal. (2021)Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I.Decision Transformer: Reinforcement Learning via Sequence Modeling.In Advances in Neural Information Processing Systems, volume18, 2021.
• Dai etal. (2020)Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q.V., and Salakhutdinov, R.Transformer-XL: Attentive language models beyond a fixed-length context.In ACL 2019 - 57th Annual Meeting of the Association for Computational Linguistics, Proceedings of the Conference, 2020.doi: 10.18653/v1/p19-1285.
• Devlin etal. (2019)Devlin, J., Chang, M.W., Lee, K., and Toutanova, K.BERT: Pre-training of deep bidirectional transformers for language understanding.In NAACL HLT 2019 - 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies - Proceedings of the Conference, volume1, 2019.
• Dosovitskiy etal. (2021)Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N.An image is worth 16x16 words: Transformers for image recognition at scale.In International Conference on Learning Representations, 2021.URL https://openreview.net/forum?id=YicbFdNTTy.
• Esser etal. (2021)Esser, P., Rombach, R., and Ommer, B.Taming transformers for high-resolution image synthesis.In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12873–12883, 2021.
• Ha & Schmidhuber (2018)Ha, D. and Schmidhuber, J.Recurrent world models facilitate policy evolution.In Advances in Neural Information Processing Systems 31, pp. 2451–2463. Curran Associates, Inc., 2018.URL https://papers.nips.cc/paper/7512-recurrent-world-models-facilitate-policy-evolution.https://worldmodels.github.io.
• Hafner etal. (2019)Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J.Learning latent dynamics for planning from pixels.In 36th International Conference on Machine Learning, ICML 2019, volume 2019-June, 2019.
• Hafner etal. (2020)Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M.Dream to control: Learning behaviors by latent imagination.In International Conference on Learning Representations, 2020.URL https://openreview.net/forum?id=S1lOTC4tDS.
• Hafner etal. (2021)Hafner, D., Lillicrap, T.P., Norouzi, M., and Ba, J.Mastering atari with discrete world models.In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.URL https://openreview.net/forum?id=0oabwyZbOu.
• Hafner etal. (2023)Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T.Mastering diverse domains through world models.arXiv preprint arXiv:2301.04104, 2023.
• Hendrycks & Gimpel (2017)Hendrycks, D. and Gimpel, K.Bridging nonlinearities and stochastic regularizers with gaussian error linear units, 2017.URL https://openreview.net/forum?id=Bk0MRI5lg.
• Hochreiter & Schmidhuber (1997)Hochreiter, S. and Schmidhuber, J.Long short-term memory.Neural computation, 9(8):1735–1780, 1997.
• Kaiser etal. (2020)Kaiser, Ł., Babaeizadeh, M., Miłos, P., Osiński, B., Campbell, R.H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., Mohiuddin, A., Sepassi, R., Tucker, G., and Michalewski, H.Model based reinforcement learning for atari.In International Conference on Learning Representations, 2020.URL https://openreview.net/forum?id=S1xCPJHtDB.
• Kingma & Welling (2014)Kingma, D.P. and Welling, M.Auto-encoding variational bayes.In 2nd International Conference on Learning Representations, ICLR 2014 - Conference Track Proceedings, 2014.
• Li etal. (2023)Li, T., Chang, H., Mishra, S., Zhang, H., Katabi, D., and Krishnan, D.Mage: Masked generative encoder to unify representation learning and image synthesis.In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2142–2152, 2023.
• Micheli etal. (2023)Micheli, V., Alonso, E., and Fleuret, F.Transformers are sample-efficient world models.In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.URL https://openreview.net/pdf?id=vhFu1Acb0xb.
• Mnih etal. (2015)Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A.A., Veness, J., Bellemare, M.G., Graves, A., Riedmiller, M., Fidjeland, A.K., Ostrovski, G., etal.Human-level control through deep reinforcement learning.nature, 518(7540):529–533, 2015.
• Parisotto etal. (2020)Parisotto, E., Song, F., Rae, J., Pascanu, R., Gulcehre, C., Jayakumar, S., Jaderberg, M., Kaufman, R.L., Clark, A., Noury, S., Botvinick, M., Heess, N., and Hadsell, R.Stabilizing transformers for reinforcement learning.In III, H.D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 7487–7498. PMLR, 13–18 Jul 2020.URL https://proceedings.mlr.press/v119/parisotto20a.html.
• Ramachandran etal. (2018)Ramachandran, P., Zoph, B., and Le, Q.V.Searching for activation functions, 2018.URL https://openreview.net/forum?id=SkBYYyZRZ.
• Reed etal. (2022)Reed, S., Zolna, K., Parisotto, E., Colmenarejo, S.G., Novikov, A., Barth-maron, G., Giménez, M., Sulsky, Y., Kay, J., Springenberg, J.T., Eccles, T., Bruce, J., Razavi, A., Edwards, A., Heess, N., Chen, Y., Hadsell, R., Vinyals, O., Bordbar, M., and deFreitas, N.A generalist agent.Transactions on Machine Learning Research, 2022.ISSN 2835-8856.URL https://openreview.net/forum?id=1ikK0kHjvj.Featured Certification, Outstanding Certification.
• Robine etal. (2023)Robine, J., Höftmann, M., Uelwer, T., and Harmeling, S.Transformer-based world models are happy with 100k interactions.In The Eleventh International Conference on Learning Representations, 2023.URL https://openreview.net/forum?id=TdBaDGCpjly.
• Schrittwieser etal. (2020)Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., Lillicrap, T., and Silver, D.Mastering Atari, Go, chess and shogi by planning with a learned model.Nature, 588(7839):604–609, dec 2020.ISSN 14764687.doi: 10.1038/s41586-020-03051-4.URL https://www.nature.com/articles/s41586-020-03051-4.
• Schulman etal. (2017)Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O.Proximal policy optimization algorithms.ArXiv, abs/1707.06347, 2017.URL https://api.semanticscholar.org/CorpusID:28695052.
• Shridhar etal. (2023)Shridhar, M., Manuelli, L., and Fox, D.Perceiver-actor: A multi-task transformer for robotic manipulation.In Liu, K., Kulic, D., and Ichnowski, J. (eds.), Proceedings of The 6th Conference on Robot Learning, volume 205 of Proceedings of Machine Learning Research, pp. 785–799. PMLR, 14–18 Dec 2023.URL https://proceedings.mlr.press/v205/shridhar23a.html.
• Silver etal. (2016)Silver, D., Huang, A., Maddison, C.J., Guez, A., Sifre, L., vanden Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., Dieleman, S., Grewe, D., Nham, J., Kalchbrenner, N., Sutskever, I., Lillicrap, T., Leach, M., Kavukcuoglu, K., Graepel, T., and Hassabis, D.Mastering the game of go with deep neural networks and tree search.Nature, 529:484–503, 2016.URL http://www.nature.com/nature/journal/v529/n7587/full/nature16961.html.
• Sun etal. (2023)Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F.Retentive network: A successor to transformer for large language models.arXiv preprint arXiv:2307.08621, 2023.
• Sutton (1991)Sutton, R.S.Dyna, an integrated architecture for learning, planning, and reacting.ACM SIGART Bulletin, 2(4), 1991.ISSN 0163-5719.doi: 10.1145/122344.122377.
• Sutton & Barto (2018)Sutton, R.S. and Barto, A.G.Reinforcement Learning: An Introduction.A Bradford Book, Cambridge, MA, USA, 2018.ISBN 0262039249.
• Touvron etal. (2023)Touvron, H., Martin, L., Stone, K.R., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D.M., Blecher, L., Ferrer, C.C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A.S., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I.M., Korenev, A.V., Koura, P.S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I., Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E.M., Subramanian, R., Tan, X., Tang, B., Taylor, R., Williams, A., Kuan, J.X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T.Llama 2: Open foundation and fine-tuned chat models.ArXiv, abs/2307.09288, 2023.URL https://api.semanticscholar.org/CorpusID:259950998.
• vanden Oord etal. (2017)vanden Oord, A., Vinyals, O., and kavukcuoglu, k.Neural discrete representation learning.In Guyon, I., Luxburg, U.V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume30. Curran Associates, Inc., 2017.URL https://proceedings.neurips.cc/paper_files/paper/2017/file/7a98af17e63a0ac09ce2e96d03992fbc-Paper.pdf.
• Vaswani etal. (2017)Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L.u., and Polosukhin, I.Attention is all you need.In Guyon, I., Luxburg, U.V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume30. Curran Associates, Inc., 2017.URL https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
• Vinyals etal. (2019)Vinyals, O., Babuschkin, I., Czarnecki, W.M., Mathieu, M., Dudzik, A., Chung, J., Choi, D.H., Powell, R., Ewalds, T., Georgiev, P., Oh, J., Horgan, D., Kroiss, M., Danihelka, I., Huang, A., Sifre, L., Cai, T., Agapiou, J.P., Jaderberg, M., Vezhnevets, A.S., Leblond, R., Pohlen, T., Dalibard, V., Budden, D., Sulsky, Y., Molloy, J., Paine, T.L., Gulcehre, C., Wang, Z., Pfaff, T., Wu, Y., Ring, R., Yogatama, D., Wünsch, D., McKinney, K., Smith, O., Schaul, T., Lillicrap, T., Kavukcuoglu, K., Hassabis, D., Apps, C., and Silver, D.Grandmaster level in StarCraft II using multi-agent reinforcement learning.Nature, 575(7782), 2019.ISSN 14764687.doi: 10.1038/s41586-019-1724-z.
• Ye etal. (2021)Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Y.Mastering Atari Games with Limited Data.In Advances in Neural Information Processing Systems, volume30, 2021.
• Zhang etal. (2023)Zhang, W., Wang, G., Sun, J., Yuan, Y., and Huang, G.STORM: Efficient stochastic transformer based world models for reinforcement learning.In Thirty-seventh Conference on Neural Information Processing Systems, 2023.URL https://openreview.net/forum?id=WxnrX42rnS.

Appendix A Appendix