because 1 ensures that each processor state is in a subtree that matches the
subtree of the memory element it wishes to read.

As an important detail, our recursion proceeds in parallel, delegating the
shallower processor tree to a child PSAM process. This ensures that the child
process terminates before the parent process completes its own recursive call.
This, in turn, ensures that the dereference of
$m_{\newmcodes@\mathsf{shallow}},p_{\newmcodes@\mathsf{shallow}}$ (line 12)
does not illegally dereference the return pointer of an incomplete child
process (see Section 5).

Once both recursive calls are complete, we have two memory trees and two
processor trees, which we pairwise merge with calls to $\uplus$.

If the processor tree is a leaf node, then we are ready to dispatch a single
PRAM instruction. To do so, we first case analyze the
$\newmcodes@\mathsf{memory}$ tree. Note that it is impossible that the memory
tree is a $\newmcodes@\mathsf{Branch}$ node, because by 1 the processor tree
is always deeper than the memory tree. Thus, the memory tree can either hold
the single RAM element required by the current processor $\rho$, or it can be
empty. The latter case corresponds to a case where the processor $\rho$ wishes
to access a memory cell that has not yet been accessed by any processor. In
this latter case, the processor will read the all zeros word.

From here, we call $\newmcodes@\mathsf{handle-pram-instruction}$ with the
processor state and the accessed data element. The processor returns an
element to be written back to memory and up to two subsequent processor
states. We encode these values as trees, then we return the fresh
memory/processor trees.

#### Cleaning up.

Once we compute the fixed point of $\newmcodes@\mathsf{step}$, the PRAM
simulation is finished. However, recall that to achieve a _clean_ PSAM
procedure (Definition 5), our program must read all unread data.

Cleaning up is straightforward. Our processor tree is, by the fact that we
have finished the simulation, trivially empty, and to clean the memory tree,
we simply traverse it from root to leaves in a recursive procedure, without
writing back any data. This cleanup imposes at most constant factor overhead,
as the PRAM memory cannot have accessed more than $W(n)$ leaves.

###  7.3 Step Complexity

In the following, we assume that the processor tree has $p$ leaves; i.e.,
there are $p$ active processors.

#### Work.

We claim that $\newmcodes@\mathsf{step}$ uses at most $p\cdot O(\log n)$ total
PSAM work. To see this, first note that both the memory tree and the processor
tree have $O(\log n)$ depth; this is ensured by the fact that the number of
processors and the highest memory address are each at most polynomial in $n$.
If we for now ignore calls to $\uplus$, we see that each time we call
$\newmcodes@\mathsf{step}$ on an internal node, we expend constant work (and
time) before recursing down both subtrees. Our recursion only follows paths
through the tree to locations where processors reside, so the total number of
recursive calls is at most $p\cdot O(\log n)$.

The calls to $\uplus$ are more difficult to analyze, but they similarly can be
seen to consume at most $p\cdot O(\log n)$ total work. One way to see this is
as follows: For the processor tree, we merge together $O(p)$ trees resulting
from calls to $\newmcodes@\mathsf{encode-path}$, where each merged tree is a
“path tree” with one leaf and $O(\log n)$ internal nodes. Our
$\newmcodes@\mathsf{step}$ procedure merges these trees together in some order
resulting from $\newmcodes@\mathsf{steps}$’s recursion tree. However, instead
suppose we were to merge these trees together one at a time, maintaining an
accumulator tree and merging each path tree into this accumulator. Each such
merge would cost at most $O(\log n)$ work, simply due to the small size of the
path tree. Thus, this strategy uses at most $p\cdot O(\log n)$ work.

Now, $\newmcodes@\mathsf{step}$ does not merge path trees into an accumulator
in this way; it merges trees at each recursive call site. However, we claim
that $\newmcodes@\mathsf{step}$ uses _strictly less work_ than the above
accumulator approach. Indeed, the merge of two trees is strictly smaller than
those two trees alone, since the merged tree will join together some internal
nodes. Hence, merging the processor tree consumes at most $p\cdot O(\log n)$
work. This same argument holds for the memory tree.

Hence, $\newmcodes@\mathsf{step}$ consumes at most $p\cdot O(\log n)$ work.

#### Time.

We claim that $\newmcodes@\mathsf{step}$ computes root pointers of its output
trees within $O(\log n)$ PSAM time. This follows almost immediately from the
fact that both the memory and processor tree have logarithmic depth.

There is one non-trivial component to this argument: we must argue that calls
to $\uplus$ can be pipelined. Indeed, recall that $\uplus$ takes $O(\log n)$
to completely compute its output, but it computes its root pointer within
_constant_ time.

Formally, we can establish an inductive hypothesis that
$\newmcodes@\mathsf{step}$ outputs its root pointers within $O(\log n)$ time.
In the base case this holds, since $\newmcodes@\mathsf{handle-pram-
instruction}$ is assumed to take at most $O(\log n)$ time and since it takes
$O(\log n)$ steps to encode a tree path (indeed, the root of the encoded path
is available within $O(1)$ time). In the inductive case, we recursively call
$\newmcodes@\mathsf{step}$ twice in parallel, taking $O(\log n)$ time by the
inductive hypothesis. We conclude with calls to $\uplus$, which return their
root in $O(1)$ time. Hence, the general case incurs only constant additive
time on top of its recursive calls. As the depth of the tree is at most
logarithmic, the induction holds. Hence, $\newmcodes@\mathsf{step}$ computes
its output within $O(\log n)$ time.

##  8 Concluding Remarks

By combining our results from Sections 6 and 7, we attain a simulation of CRCW
PRAM, and our simulation is achieved using cyclic circuits, a simple and
easily implemented model. Surprisingly, the simulation incurs only
polylogarithmic overhead in terms of both work and time.

This demonstrates feasibility of powerful parallel machines, at least in
theory. Of course, substantial effort is needed to determine if this
theoretical feasibility can translate to practical outcomes. At the least, we
believe our result shows that cyclic circuits are far more interesting than
previously thought. While prior works investigated cyclic circuits, none
showed a connection between this simple model and PRAM.

#### An open question.

Can our simulation of PRAM by cyclic circuits be improved? Our simulation
achieves $O(\log^{4}n)$ work overhead and $O(\log^{3}n)$ runtime overhead, and
it is not obvious how to do significantly better. We incur $O(\log n)$
overhead from word size, $O(\log^{2}n)$ overhead from the dynamic permutation
network, and $O(\log n)$ overhead from the simulation of PRAM by PSAM. Thus,
improving our result in a modular way would require either (1) changing the
circuit model, (2) improving dynamic permutation networks, or (3) improving
the simulation of PRAM by PSAM; none of these improvements are obvious. Of
course, it might also be possible to mix the concerns of the dynamic
permutation with PRAM-by-PSAM simulation, or by applying some completely
different approach; these ideas are also not clear.

In terms of negative results, our only current insight is that there is a
trivial $\Omega(\log n)$ lower bound on work overhead. This bound comes simply
from “unfairness” of the comparison between the considered PRAM and cyclic
circuits: circuit wires hold individual bits while PRAM manipulates
$\Theta(\log n)$-bit words. To show a $\Omega(\log n)$ lower bound, we can
consider simulating a PRAM program that, e.g., forces the cyclic circuit to
take as input $O(n)$ elements and write them out in some input-specified
order. To achieve generality, the circuit must fully read the $O(n\cdot\log
n)$-bit input.

One might expect that this permutation problem would allow us to strengthen
the lower bound to $\Omega(\log^{2}n)$ work overhead, since it is well known
that comparison-based sorting requires $\Theta(n\cdot\log n)$ comparisons.
However, trying to apply this insight immediately runs into well-known open
questions regarding the complexity of sorting, see e.g. [FHLS19, ALS22].
Namely, it might be possible to sort with $o(n\cdot\log n)$ Boolean gates, and
this prevents us from establishing this stronger lower bound.

## References

  * [ABB00] Umut A Acar, Guy E Blelloch, and Robert D Blumofe.  The data locality of work stealing.  In Proceedings of the twelfth annual ACM Symposium on Parallel Algorithms and Architectures, 2000.
  * [AKS83] Miklós Ajtai, János Komlós, and Endre Szemerédi.  An $O(n\log n)$ sorting network.  In Proceedings of the fifteenth annual ACM Symposium on Theory of Computing, 1983.
  * [ALS22] Gilad Asharov, Wei-Kai Lin, and Elaine Shi.  Sorting short keys in circuits of size $o(n\log n)$.  SIAM Journal on Computing, 2022.
  * [Bat68] K. E. Batcher.  Sorting networks and their applications.  In Proceedings of the April 30–May 2, 1968, Spring Joint Computer Conference. Association for Computing Machinery, 1968.
  * [BCP16] Elette Boyle, Kai-Min Chung, and Rafael Pass.  Oblivious parallel ram and applications.  In Theory of Cryptography, 2016.
  * [Ben64] Václad E Beneš.  Optimal rearrangeable multistage connecting networks.  Bell System Technical Journal, 1964.
  * [BFGS20] Guy E Blelloch, Jeremy T Fineman, Yan Gu, and Yihan Sun.  Optimal parallel algorithms in the binary-forking model.  In Proceedings of the 32nd ACM Symposium on Parallelism in Algorithms and Architectures, 2020.
  * [BKS17] Paul Beame, Paraschos Koutris, and Dan Suciu.  Communication steps for parallel query processing.  Journal of the ACM, 2017.
  * [CC96] Wang-Jiunn Cheng and Wen-Tsuen Chen.  A new self-routing permutation network.  IEEE Transactions on Computers, 1996.
  * [CKP+93] David Culler, Richard Karp, David Patterson, Abhijit Sahay, Klaus Erik Schauser, Eunice Santos, Ramesh Subramonian, and Thorsten von Eicken.  Logp: Towards a realistic model of parallel computation.  SIGPLAN Not., 1993.
  * [CLR08] Henri Casanova, Arnaud Legrand, and Yves Robert.  Parallel algorithms.  CRC Press, 2008.
  * [CLRS22] Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein.  Introduction to algorithms.  MIT press, 2022.
  * [DC89] Patrick W Dymond and Stephen A Cook.  Complexity theory of parallel time and hardware.  Information and Computation, 1989.
  * [FHLS19] Alireza Farhadi, MohammadTaghi Hajiaghayi, Kasper Green Larsen, and Elaine Shi.  Lower bounds for external memory integer sorting via network coding.  In Proceedings of the 51st Annual ACM SIGACT Symposium on Theory of Computing, STOC 2019. Association for Computing Machinery, 2019.
  * [FW78] Steven Fortune and James Wyllie.  Parallelism in random access machines.  In Proceedings of the Tenth Annual ACM Symposium on Theory of Computing, STOC ’78. Association for Computing Machinery, 1978.
  * [GR89] Alan Gibbons and Wojciech Rytter.  Efficient parallel algorithms.  Cambridge University Press, 1989.
  * [HKO23] David Heath, Vladimir Kolesnikov, and Rafail Ostrovsky.  Tri-state circuits: A circuit model that captures ram.  In Advances in Cryptology – CRYPTO, 2023.
  * [KR90] Richard M. Karp and Vijaya Ramachandran.  Chapter 17 - parallel algorithms for shared-memory machines.  In Algorithms and Complexity, Handbook of Theoretical Computer Science. Elsevier, 1990.
  * [LO13] Steve Lu and Rafail Ostrovsky.  How to garble ram programs?  In Advances in Cryptology – EUROCRYPT, 2013.
  * [Mal93] S. Malik.  Analysis of cyclic combinational circuits.  In Proceedings of 1993 International Conference on Computer Aided Design (ICCAD), 1993.
  * [MSB12] Michael Mendler, Thomas R Shiple, and Gérard Berry.  Constructive boolean circuits and the exactness of timed ternary simulation.  Formal Methods in System Design, 2012.
  * [PF79] Nicholas Pippenger and Michael J. Fischer.  Relations among complexity measures.  Journal of the ACM, 1979.
  * [RB12] Marc D. Riedel and Jehoshua Bruck.  Cyclic boolean circuits.  Discrete Applied Mathematics, 2012.
  * [Rie04] Marc D Riedel.  Cyclic combinational circuits.  California Institute of Technology, 2004.
  * [Riv77] Rivest.  The necessity of feedback in minimal monotone combinational circuits.  IEEE Transactions on Computers, 1977.
  * [SBT96] Thomas R Shiple, Gérard Berry, and Hervé Touati.  Constructive analysis of cyclic circuits.  In Proceedings ED&TC European Design and Test Conference. IEEE, 1996.
  * [Sho60] Robert Allen Short.  A theory of relations between sequential and combinational realizations of switching functions.  Stanford University, 1960.
  * [Sto92] Leon Stok.  False loops through resource sharing (logic cad).  In 1992 IEEE/ACM International Conference on Computer-Aided Design. IEEE, 1992.
  * [SV84] Larry Stockmeyer and Uzi Vishkin.  Simulation of parallel random access machines by circuits.  SIAM Journal on Computing, 1984.
  * [Tar55] Alfred Tarski.  A lattice-theoretical fixpoint theorem and its applications.  Pacific Journal of Mathematics, 1955.
  * [Val90a] Leslie G. Valiant.  A bridging model for parallel computation.  Communications of the ACM, 1990.
  * [Val90b] Leslie G Valiant.  General purpose parallel architectures.  In Algorithms and Complexity. Elsevier, 1990.
  * [Vis84] Uzi Vishkin.  A parallel-design distributed-implementation (pddi) general-purpose computer.  Theoretical Computer Science, 1984.
  * [Wak68] Abraham Waksman.  A permutation network.  Journal of the ACM, 1968.
  * [Yao86] Andrew Chi-Chih Yao.  How to generate and exchange secrets.  In Symposium on Foundations of Computer Science. IEEE, 1986.
  * [ZCSM07] Yunquan Zhang, Guoliang Chen, Guangzhong Sun, and Qiankun Miao.  Models of parallel computation: a survey and classification.  Frontiers of Computer Science in China, 2007.


# Efficient World Models with Context-Aware Tokenization

Vincent Micheli  Eloi Alonso  François Fleuret

###### Abstract

Scaling up deep Reinforcement Learning (RL) methods presents a significant
challenge. Following developments in generative modelling, model-based RL
positions itself as a strong contender. Recent advances in sequence modelling
have led to effective transformer-based world models, albeit at the price of
heavy computations due to the long sequences of tokens required to accurately
simulate environments. In this work, we propose $\Delta$-iris, a new agent
with a world model architecture composed of a discrete autoencoder that
encodes stochastic deltas between time steps and an autoregressive transformer
that predicts future deltas by summarizing the current state of the world with
continuous tokens. In the Crafter benchmark, $\Delta$-iris sets a new state of
the art at multiple frame budgets, while being an order of magnitude faster to
train than previous attention-based approaches. We release our code and models
at https://github.com/vmicheli/delta-iris.

Artificial Intelligence, Machine Learning, Generative Modelling, World Models,
Transformers, Autoencoders, Reinforcement Learning, Decision making, ICML

##  1 Introduction

Deep Reinforcement Learning (RL) methods have recently delivered impressive
results (Ye et al., 2021; Hafner et al., 2023; Schwarzer et al., 2023) in
traditional benchmarks (Bellemare et al., 2013; Tassa et al., 2018). In light
of the evermore complex domains tackled by the latest generations of
generative models (Rombach et al., 2022; Achiam et al., 2023), the prospect of
training agents in more ambitious environments (Kanervisto et al., 2022) may
hold significant appeal. However, that leap forward poses a serious challenge:
deep RL architectures have been comparatively smaller and less sample-
efficient than their (self-)supervised counterparts. In contrast, more
intricate environments necessitate models with greater representational power
and have higher data requirements.

Model-based RL (MBRL) (Sutton & Barto, 2018) is hypothesized to be the key for
scaling up deep RL agents (LeCun, 2022). Indeed, world models (Ha &
Schmidhuber, 2018) offer a diverse range of capabilities: lookahead search
(Schrittwieser et al., 2020; Ye et al., 2021), learning in imagination
(Sutton, 1991; Hafner et al., 2023), representation learning (Schwarzer et
al., 2021; D’Oro et al., 2023), and uncertainty estimation (Pathak et al.,
2017; Sekar et al., 2020). In essence, MBRL shifts the focus from the RL
problem to a generative modelling problem, where the development of an
accurate world model significantly simplifies policy training. In particular,
policies learnt in the imagination of world models are freed from sample
efficiency constraints, a common limitation of RL agents that is magnified in
complex environments with slow rollouts.

Recently, the iris agent (Micheli et al., 2023) achieved strong results in the
Atari 100k benchmark (Bellemare et al., 2013; Kaiser et al., 2020). iris
introduced a world model composed of a discrete autoencoder and an
autoregressive transformer, casting dynamics learning as a sequence modelling
problem where the transformer composes over time a vocabulary of image tokens
built by the autoencoder. This approach opened up avenues for future model-
based methods to capitalize on advances in generative modelling (Villegas et
al., 2022; Achiam et al., 2023), and has already been adopted beyond its
original domain (comma.ai, 2023; Hu et al., 2023). However, in its current
form, scaling iris to more complex environments is computationally
prohibitive. Indeed, such an endeavor requires a large number of tokens to
encode visually challenging frames. Besides, sophisticated dynamics may
require to store numerous time steps in memory to reason about the past,
ultimately making the imagination procedure excessively slow. Hence, under
these constraints, maintaining a favorable imagined-to-collected data ratio is
practically infeasible.

$z^{1}_{1}$$\dots$$z^{K_{I}}_{1}$$z^{1}_{2}$$\dots$$z^{K_{I}}_{2}$$x_{1}$$\hat{x}_{1}$$x_{2}$$\hat{x}_{2}$$x_{0}$$a_{0}$$z^{1}_{1}$$\dots$$z^{K}_{1}$$x_{1}$$a_{1}$$z^{1}_{2}$$\dots$$z^{K}_{2}$$x_{1}$$x_{2}$$x_{0}$$a_{0}$$a_{1}$$a_{2}$$\hat{x}_{1}$$\hat{x}_{2}$EncoderDecoder
Figure 1: Discrete autoencoder of iris (Micheli et al., 2023) (left) and
$\Delta$-iris (right). iris encodes and decodes frames independently, meaning
that $z_{t}$ has to carry all the information necessary to reconstruct
$x_{t}$. On the other hand, $\Delta$-iris’ encoder and decoder are conditioned
on past frames and actions, thus $z_{t}$ only has to capture what has changed
and that cannot be inferred from actions, i.e. the stochastic delta. This
conditioning scheme enables us to drastically reduce the number of tokens
required to encode a frame with minimal loss ($K\ll K_{I}$), which is critical
to speed up the autoregressive transformer that predicts future tokens.
$\tilde{x}_{0}$$a_{0}$$z^{1}_{1}$$z^{2}_{1}$$\dots$$z^{K-1}_{1}$$z^{K}_{1}$$\tilde{x}_{1}$$a_{1}$$z^{1}_{2}$$z^{2}_{2}$$\dots$$z^{K-1}_{2}$$z^{K}_{2}$$\tilde{x}_{2}$$a_{2}$$z^{1}_{3}$$z^{2}_{3}$$\dots$$z^{K-1}_{3}$$z^{K}_{3}$$\hat{z}^{1}_{1}$$\hat{z}^{2}_{1}$$\dots$$\hat{z}^{K-1}_{1}$$\hat{z}^{K}_{1}$$\hat{z}^{1}_{2}$$\hat{z}^{2}_{2}$$\dots$$\hat{z}^{K-1}_{2}$$\hat{z}^{K}_{2}$$\hat{z}^{1}_{3}$$\hat{z}^{2}_{3}$$\dots$$\hat{z}^{K-1}_{3}$$\hat{z}^{K}_{3}$$\hat{r}_{0}$$\hat{d}_{0}$$\hat{r}_{1}$$\hat{d}_{1}$$\hat{r}_{2}$$\hat{d}_{2}$$x_{0}$$x_{1}$$x_{2}$
Figure 2: Unrolling dynamics over time. At each time step (separated by dashed
lines), the GPT-like autoregressive transformer $G$ predicts the
$\Delta$-tokens for the next frame, as well as the reward and a potential
episode termination. Its input sequence consists of action tokens,
$\Delta$-tokens, and

I

-tokens, namely continuous image embeddings that alleviate the need to attend to past $\Delta$-tokens for world modelling. More specifically, an initial frame $x_{0}$ is embedded into

I

-token $\tilde{x_{0}}$. From $\tilde{x_{0}}$ and $a_{0}$, $G$ predicts the reward $\hat{r}_{0}$, episode termination $\hat{d}_{0}\in\\{0,1\\}$, and in an autoregressive manner $\hat{z}_{1}=(\hat{z}_{1}^{1},\dots,\hat{z}_{1}^{K})$, the $\Delta$-tokens for the next frame. Note that, during the imagination procedure, the next frame (stripped box) is computed by the decoder $D$ based on previous frames, actions, and the $\Delta$-tokens generated by $G$, i.e. $x_{1}=D(x_{0},a_{0},\hat{z}_{1})$.

In the present work, we introduce $\Delta$-iris, a new agent capable of
scaling to visually complex environments with lengthier time horizons.
$\Delta$-iris encodes frames by attending to the ongoing trajectory of
observations and actions, effectively describing stochastic deltas between
time steps. This enriched conditioning scheme drastically reduces the number
of tokens to encode frames, offloads the deterministic aspects of world
modelling to the autoencoder, and lets the autoregressive transformer focus on
stochastic dynamics. Nonetheless, substituting the sequence of absolute image
tokens with a sequence of $\Delta$-tokens makes the task of the autoregressive
model more arduous. In order to predict the next transition, it may only
reason over previous $\Delta$-tokens, and thus faces the challenge of
integrating over multiple time steps as a way to form a representation of the
current state of the world. To resolve this issue, we modify the sequence of
the autoregressive model by interleaving continuous  I -tokens, that summarize
successive world states with frame embeddings, and discrete $\Delta$-tokens.

In the Crafter benchmark (Hafner, 2022), $\Delta$-iris exhibits favorable
scaling properties: the agent solves 17 out of 22 tasks after 10M frames of
data collection, supersedes DreamerV3 (Hafner et al., 2023) at multiple frame
budgets, and trains 10 times faster than iris. In addition, we include results
in the sample-efficient setting with Atari games. Through experiments, we
provide evidence that $\Delta$-iris learns to disentangle the deterministic
and stochastic aspects of world modelling. Moreover, we conduct ablations to
validate the new conditioning schemes for the autoencoder and transformer
models.

##  2 Method

We consider a Partially Observable Markov Decision Process (pomdp) (Sutton &
Barto, 2018). The transition, reward, and episode termination dynamics are
captured by the conditional distributions $p(x_{t+1}\mid x_{\leq t},a_{\leq
t})$ and $p(r_{t},d_{t}\mid x_{\leq t},a_{\leq t})$, where
$x_{t}\in\mathcal{X}=\mathbb{R}^{3\times h\times w}$ is an image observation,
$a_{t}\in\mathcal{A}=\\{1,\dots,A\\}$ a discrete action, $r_{t}\in\mathbb{R}$
a scalar reward, and $d_{t}\in\\{0,1\\}$ indicates episode termination. The
reinforcement learning objective is to find a policy $p_{\pi}(a_{t}\mid
x_{\leq t},a_{<t})$ that maximizes the expected sum of rewards
$\mathbb{E}_{\pi}[\sum_{t\geq 0}\gamma^{t}r_{t}]$, with discount factor
$\gamma\in(0,1)$.

Learning in imagination (Sutton, 1991; Sutton & Barto, 2018) consists of 3
stages that are repeated alternatively: experience collection, world model
learning, and policy improvement. Strikingly, the agent learns behaviours
purely within its world model, and real experience is only leveraged to learn
the environment dynamics.

In the vein of iris (Micheli et al., 2023), our world model is composed of a
discrete autoencoder (Van Den Oord et al., 2017) and an autoregressive
transformer (Vaswani et al., 2017; Radford et al., 2019), albeit with new
conditioning schemes and architectures. We first expose iris’ world model in
Section 2.1, then present $\Delta$-iris’ autoencoder and autoregressive model
in Sections 2.2 and 2.3, respectively. Finally, we describe the policy
improvement phase in Section 2.4. Appendix A gives a detailed breakdown of
model architectures and hyperparameters.

###  2.1 Background: iris

High-dimensional images are converted into tokens with a discrete autoencoder
$(E_{I},D_{I})$ (Van Den Oord et al., 2017). The encoder
$E_{I}:\mathbb{R}^{h\times w\times 3}\rightarrow\\{1,\dots,N_{I}\\}^{K_{I}}$
maps an input image $x_{t}$ into $K_{I}$ tokens from a vocabulary of size
$N_{I}$. The discretization is done by picking the index of the vector in the
vocabulary embedding table that is closest to the encoder output
$y_{t}\in\mathbb{R}^{K_{I}\times d}$. The $K_{I}$ tokens are then decoded back
into an image with
$D_{I}:\\{1,\dots,N_{I}\\}^{K_{I}}\rightarrow\mathbb{R}^{h\times w\times 3}$.
This discrete autoencoder is trained with $L_{1}$ reconstruction, perceptual
(Esser et al., 2021) and commitment losses (Van Den Oord et al., 2017)
computed on collected frames.

The transformer $G_{I}$ models the environment dynamics by operating over an
input sequence of image and action tokens
$(z_{0}^{1},\dots,z_{0}^{K_{I}},a_{0},z_{1}^{1},\dots,z_{1}^{K_{I}},a_{1},\dots,z_{t}^{1},\dots,z_{t}^{K_{I}},a_{t})$.
Image and action tokens are embedded with learnt lookup tables. At each time
step, $G_{I}$ predicts the transition, reward, and termination distributions:
$p_{G_{I}}(\hat{z}_{t+1}|z_{\leq t},a_{\leq t})\text{ with
}\hat{z}_{t+1}^{k}\,{\sim}\,p_{G_{I}}(\hat{z}_{t+1}^{k}|z_{\leq t},a_{\leq
t},z_{t+1}^{<k})$, $p_{G_{I}}(\hat{r}_{t}|z_{\leq t},a_{\leq t})$, and
$p_{G_{I}}(\hat{d}_{t}|z_{\leq t},a_{\leq t})$. The model is trained with a
cross-entropy loss on segments sampled from past experience.

At a high level, the autoencoder builds a vocabulary of image tokens to encode
each frame, and the transformer captures the environment dynamics by
autoregressively composing the vocabulary over time. As a result, this world
model is capable of attending to previous time steps to make its predictions,
and models the joint law of future latent states.

###  2.2 Disentangling deterministic and stochastic dynamics

iris (Micheli et al., 2023) encodes frames independently, making no assumption
about temporal redundancy within trajectories. One major drawback of this
general formulation is that, in environments with visually challenging frames,
a large number of tokens is required to encode frames losslessly.
Consequently, computations with the dynamics model become increasingly
prohibitive, as the attention mechanism scales quadratically with sequence
length. Therefore, limiting computation under such a trade-off may result in
degraded performance (Micheli et al. (2023) app. E)

One possible solution to achieve fast world modelling with minimal loss is to
condition the autoencoder on previous frames and actions. Intuitively,
encoding a frame given previous frames consists in describing what has
changed, the delta, between successive time steps. In many environments, the
delta between frames is often much simpler to describe than the frames
themselves. As a matter of fact, when the transition function is
deterministic, adding previous actions to the conditioning of the decoder
results in a world model, without the need to encode any information between
time steps. However, most environments of interest feature stochastic
dynamics, and apart from aleatoric uncertainty, architectural limitations such
as the agent’s memory may induce additional epistemic uncertainty. Hence, the
delta between two time steps usually consists of deterministic and stochastic
components.

For instance, an agent moving from one square to another in a grid-like
environment when pressing movement keys can be seen as a deterministic
component of the transition. On the other hand, the sudden apparition of an
enemy in a nearby square is a random event. Interestingly, only the stochastic
features of a transition should be encoded, and the autoencoder could directly
learn to model the deterministic dynamics, which do not require the
expressivity and ability to handle multimodality of an autoregressive model.
Therefore, when autoenconding frames by conditioning on previous frames and
actions, a frame encoding may only consist of a handful of $\Delta$-tokens,
instead of a large number of image tokens describing frames independently.

Section 3.4 provides empirical evidence that $\Delta$-iris’ autoencoder learns
to encode frames in such fashion, and Figure 1 illustrates the new
conditioning scheme of the autoencoder.

$\Delta$-tokens sampled randomly

$\Delta$-tokens sampled by the autoregressive transformer

$t=0$ $t=4$ $t=5$ $t=9$ $t=10$ $t=12$

Figure 3: Evidence of dynamics disentanglement. Two trajectories are imagined
with different ways of generating $\Delta$-tokens. In the top trajectory,
$\Delta$-tokens are sampled randomly. In the bottom trajectory, the
autoregressive transformer predicts future $\Delta$-tokens. The same starting
frame ($t=0$) and sequence of actions are used. With random $\Delta$-tokens,
the deterministic aspects of the dynamics (layout, movement, items, crafting)
are still properly modelled, but the stochastic dynamics (mobs, health
indicators) become problematic. For instance, the agent successfully cuts down
a tree between $t=4$ and $t=5$, and uses wood planks to build a crafting table
between $t=10$ and $t=12$. We observe that these dynamics are modelled in the
same way whether $\Delta$-tokens are sampled randomly or not. However, in the
top trajectory, large quantities of cows appear and disappear from the screen
incoherently, whereas the bottom trajectory does not display such erratic
patterns. This experiment shows that $\Delta$-iris encodes stochastic deltas
between time steps with $\Delta$-tokens, and its decoder handles the
deterministic aspects of world modelling. Appendix 9 contains additional
examples.

More formally, for any set $\mathcal{Y}$, we denote
$\mathcal{S}_{n}(\mathcal{Y})=\bigcup_{i=1}^{n}\mathcal{Y}^{i}$ the set of
tuples of elements from $\mathcal{Y}$ of maximum length $n$, and
$\mathcal{S}(\mathcal{Y})=\mathcal{S}_{\infty}(\mathcal{Y})$. Let
$\mathcal{Z}=\\{1,\dots,N\\}$ a vocabulary of discrete tokens. Given past
images and actions $(x_{0},a_{0},\dots,x_{t-1},a_{t-1})$, the encoder
$E:\mathcal{S}(\mathcal{X}\times\mathcal{A})\times\mathcal{X}\rightarrow\mathcal{Z}^{K}$
converts an image $x_{t}$ into $z_{t}=(z_{t}^{1},\dots,z_{t}^{K})$, a sequence
of $K$ discrete $\Delta$-tokens. The encoder is parameterized by a
Convolutional Neural Network (cnn) (LeCun et al., 1989). Actions are embedded
with a learnt lookup table and concatenated channel-wise with frames. We use
vector quantization (Van Den Oord et al., 2017; Esser et al., 2021) with
factorized and normalized codes (Yu et al., 2021) to discretize the encoder’s
continuous outputs. The cnn decoder
$D:\mathcal{S}(\mathcal{X}\times\mathcal{A})\times\mathcal{Z}^{K}\rightarrow\mathcal{X}$
reconstructs an image $\hat{x}_{t}$ from past frames, actions and
$\Delta$-tokens $(x_{0},a_{0},\dots,x_{t-1},a_{t-1},z_{t})$. Action and
$\Delta$-tokens are embedded with learnt lookup tables, and concatenated
channel-wise with feature maps obtained by forwarding frames through an
auxiliary cnn.

The discrete autoencoder is trained on previously collected trajectories with
a weighted combination of $L_{1}$, $L_{2}$ and max-pixel (Anand et al., 2022)
reconstruction losses, as well as a commitment loss (Van Den Oord et al.,
2017). The codebook is updated with an exponential moving average (Razavi et
al., 2019) and we use a straight-through estimator (Bengio et al., 2013) to
enable backpropagation.

###  2.3 Modelling stochastic dynamics

While it should be possible to predict future $\Delta$-tokens, given a
starting image, past actions and $\Delta$-tokens, we found this task much more
difficult than simply predicting future image tokens, given past image tokens
and actions, as in iris.

To better understand why this is the case, let us consider another example: in
a grid environment, $\Delta$-tokens may describe the unpredictable movement of
an enemy, randomly jumping from one square to another at every time step.
Based on the initial enemy location and after only a few time steps, it
becomes increasingly difficult to predict if the enemy and the agent are
located on the same square, which could trigger a battle and make the enemy
disappear. Indeed, situating the two entities involves reasoning about the
initial observation, and integrating over all of the previous action and
$\Delta$-tokens, which may have a complex dependence structure.

To address this problem, we alter the sequence of the dynamics model by
interleaving continuous  I -tokens, in reference to MPEG’s  I -frames
(Richardson, 2004), and discrete $\Delta$-tokens.  I -tokens alleviate the
need of integrating over past $\Delta$-tokens to form a representation of the
current state of the world, i.e. they deploy a “soft” Markov blanket for the
prediction of the next $\Delta$-tokens.

Table 1: Returns, number of parameters, and frames collected per second (FPS) for the methods considered. We compute FPS as the total number of environment frames collected divided by the training duration. $\Delta$-iris outperforms DreamerV3 for larger frame budgets, and is 10x faster than iris (64 tokens). Method | Return @1M | Return @5M | Return @10M | #Parameters | FPS
---|---|---|---|---|---
$\Delta$-iris | 7.7 (0.5) |  15.4 (0.4) |  16.1 (0.1) | 25M | 20
DreamerV3 XL |  9.2 (0.3) | 14.2 (0.2) | 15.1 (0.3) | 200M | 30
iris (64 tokens) | 5.5 (0.7) | - | - | 48M | 2
$\Delta$-iris w/o  I -tokens | 6.6 (0.2) | 10.4 (0.5) | 12.6 (0.8) | 24M | 22
DreamerV3 M | 6.2 (0.5) | 12.6 (0.7) | 13.7 (0.8) | 37M | 40
iris (16 tokens) | 4.4 (0.1) | - | - | 50M | 6
Figure 4: Returns at multiple frame budgets in the Crafter benchmark.
$\Delta$-iris achieves higher returns than DreamerV3 beyond 3M frames, and
surpasses iris for all frame budgets considered. Removing

I

-tokens from the input sequence of the autoregressive transformer significantly hurts performance.

$\Delta$-iris 4 tokens iris 16 tokens

Figure 5: Bottom $1\%$ test frames autoencoded by $\Delta$-iris (4 tokens) and
iris (Micheli et al., 2023) (16 tokens). Each token takes a value in
$\\{1,2,\dots,1023,1024\\}$, i.e. $\Delta$-iris encodes frames with
$4\times\log_{2}(1024)=40$ bits while iris uses 160 bits. Original frames,
reconstructions, and errors are respectively displayed in the top, middle, and
bottom rows. Even in the worst instances, $\Delta$-iris makes only minor
errors, whereas iris fails to accurately reconstruct frames. These errors
severely hamper the agent’s performance, as it purely learns behaviours from
frames generated by its autoencoder.

Autoregressive transformer with  I -tokens

Autoregressive transformer without  I -tokens

Figure 6: Trajectories imagined with (top) and without (bottom)

I

-tokens. In the top trajectory, we observe more than 30 seconds of gameplay generated by $\Delta$-iris’ world model. A wide variety of mechanics have been internalized: scrolling, chopping down trees, building a crafting table, mining iron, crafting pickaxes, etc. However, removing

I

-tokens from the sequence of the autoregressive transformer makes the task of predicting future $\Delta$-tokens drastically harder as evidenced by the agent glitching through walls and water in the bottom trajectory. These mistakes ultimately hinder the policy improvement phase, since the agent will reinforce behaviours in a world that does not properly reflect its environment.

We obtain  I -tokens by forwarding frames through an auxiliary cnn at each
time step. They are not produced by a discrete autoencoder. Since  I -tokens
are not predicted by the model but rather enrich its conditioning, there are
no incentives to include a lossy discretization operator or to optimize a
reconstruction loss. Instead, they are optimized end-to-end with the learning
objectives of the dynamics model. With this improved conditioning, the
dynamics model perceives the ongoing trajectory with a mixture of continuous
and discrete representations, while making its predictions autoregressively in
a discrete space.

Figure 2 displays the input sequence of the dynamics model and the quantities
it predicts. Given a sequence of past  I -tokens, action tokens, and
$\Delta$-tokens
$(\tilde{x}_{0},a_{0},z_{1}^{1},\dots,z_{1}^{K},\dots,\tilde{x}_{t-1},a_{t-1},z_{t}^{1},\dots,z_{t}^{k})$,
the dynamics model $G$ outputs a categorical distribution on $\mathcal{Z}$ for
the next $\Delta$-token
$\hat{z}_{t}^{k+1}\,{\sim}\,p_{G}(\hat{z}_{t}^{k+1}|\tilde{x}_{<t},z_{<t},a_{<t},z_{t}^{\leq
k})$. It also predicts distributions for rewards
$p_{G}(\hat{r}_{t}|\tilde{x}_{\leq t}z_{\leq t},a_{\leq t})$ and episode
terminations $p_{G}(\hat{d}_{t}|\tilde{x}_{\leq t},z_{\leq t},a_{\leq t})$.

$G$ is parameterized by a stack of transformer encoder layers with causal
self-attention (Vaswani et al., 2017; Radford et al., 2019). It is trained
with a cross-entropy loss for transition and termination predictions, and we
follow DreamerV3 (Hafner et al., 2023) in using discrete regression with two-
hot targets and symlog scaling for reward prediction (Imani & White, 2018).

###  2.4 Policy improvement

During the policy improvement phase, the policy $\pi$ learns in the
imagination pomdp of its world model, composed of the autoencoder $(E,D)$ and
the dynamics model $G$.

At time step $t$, the policy observes a reconstructed image observation
$\hat{x}_{t}$ and samples action $a_{t}\sim\pi(a_{t}|\hat{x}_{\leq t})$. The
world model then predicts the reward $\hat{r}_{t}$, the episode end
$\hat{d}_{t}$, and the next observation $\hat{x}_{t+1}=D(\hat{x}_{\leq
t},\hat{a}_{\leq t},\hat{z}_{\leq t},\hat{z}_{t+1})$, with $\hat{z}_{t+1}\sim
p_{G}(\hat{z}_{t+1}|\hat{x}_{\leq t},\hat{a}_{\leq t},\hat{z}_{\leq t})$. The
imagination procedure is initialized with a real observation $x_{0}$ sampled
from past experience, and is rolled out for $H$ steps. The procedure stops if
an episode termination is predicted before reaching the imagination horizon.

We employ to a large extent the actor-critic training method used for iris
(Micheli et al., 2023). A value baseline is trained to predict
$\lambda$-returns (Sutton & Barto, 2018) with the same discrete regression
objective as for reward prediction. The policy optimizes the reinforce with
value baseline (Sutton & Barto, 2018) learning objective over imagined
trajectories. Exploration is encouraged by adding an entropy maximization term
to the policy’s objective.

##  3 Experiments

In our experiments, we consider the Crafter benchmark (Hafner, 2022) to
illustrate $\Delta$-iris’ ability to scale to a visually rich environment with
large frame budgets. Besides, we also include Atari 100k games (Bellemare et
al., 2013; Kaiser et al., 2020) in Appendix C to showcase the performance and
speed of our agent in the sample-efficient setting.

We introduce the Crafter benchmark and baselines in Section 3.1. Then, we
present our results in Section 3.2. Finally, in Sections 3.3 and 3.4, we
propose qualitative experiments to validate $\Delta$-iris’ world model
architecture, and better our understanding of how the model represents
information.

###  3.1 Benchmark and baselines

Crafter (Hafner, 2022) is a procedurally generated environment, inspired by
the video game Minecraft, with visual inputs, a discrete action space and non-
deterministic dynamics. By incorporating mechanics from survival games and a
technology tree, this benchmark evaluates a broad range of agent capabilities
such as generalization, exploration, and credit assignment. During each
episode, the agent’s goal is to solve as many tasks as possible, e.g. slaying
mobs, crafting items, and managing health indicators.

Regarding baselines, we consider two model-based RL agents learning in
imagination: iris (Micheli et al., 2023) and DreamerV3 (Hafner et al., 2023).
We run several variants: iris (16 tokens), encoding frames with $K_{I}=16$
tokens, iris (64 tokens), encoding frames with $K_{I}=64$ tokens, and
configurations of DreamerV3 of different sizes, namely DreamerV3 XL and
DreamerV3 M. To demonstrate the importance of  I -tokens, we also run
$\Delta$-iris without  I -tokens in the sequence of the transformer, i.e. $G$
only operates over the first frame as well as actions and $\Delta$-tokens.

We keep a fixed imagined-to-collected data ratio of 64 to balance speed and
performance. Our experiments run on a Nvidia A100 40GB GPU, with 5 seeds for
all methods and ablations. We evaluate each run by computing the average
return over 256 test episodes every 1M frames. Note that we stop the iris
experiments before 10M frames because they are prohibitively slow.

###  3.2 Results

Table 1 exhibits key metrics and Figure 4 displays learning curves. After 10M
frames of data collection, $\Delta$-iris solves on average 17 out of 22 tasks,
setting a new state of the art for the Crafter benchmark. Beyond the 3M frames
mark, $\Delta$-iris consistently achieves higher returns than DreamerV3,
although DreamerV3 is better suited for the smallest frame budgets. A key
difference between the two methods is that $\Delta$-iris does not leverage the
representations of its world model for policy learning, which may be
especially useful in the scarce data regime. As our main objective is to
develop world model architectures that scale to complex environments and
larger frame budgets, we leave this exploration to future work. $\Delta$-iris
outperforms iris for all frame budgets considered, while training an order of
magnitude faster. Finally, removing  I -tokens from the sequence of the
dynamics model drastically hurts performance.

We believe that achieving higher returns at the 10M frames cap poses a hard
exploration problem. Indeed, three of the missing four tasks require crafting
new tools in the presence of a nearby crafting table and furnace. Discovering
these tools with a naive exploration strategy is highly unlikely, and we have
observed only a few occurrences of those events throughout training runs.

With too few training samples, the world model is unable to internalize these
new mechanics and reflect them during the imagination procedure. We
hypothesize that a biased data sampling procedure (Kauvar et al., 2023) could
be the key to unlock the missing achievements.

###  3.3 World model analysis

In Section 3.2, we validated our design choices for $\Delta$-iris with RL
experiments. However, downstream RL performance is an imperfect proxy for the
quality of a world model due to many possible confounding factors, e.g. the
choice of the RL algorithm, entangled world model and policy architectures, or
the continual learning loop. In this section, we directly focus on the
abilities of the world model.

Figure 5 illustrates the bottom $1\%$ autoencoded test frames with and without
conditioning the autoencoder on the ongoing trajectory (i.e. reconstructions
with $\Delta$-iris vs iris). With as few as 4 tokens per frame, $\Delta$-iris’
autoencoder is able to encode frames with minimal loss. On the other hand,
without access to previous frames and actions, and even with 16 tokens, iris’
autoencoder produces poor reconstructions.

Figure 6 displays trajectories imagined to illustrate whether crucial
mechanics have been internalized by the world model, when including  I -tokens
in the sequence of the autoregressive transformer or not. We observe that,
with  I -tokens, a multitude of game mechanics are well understood, but in the
absence of  I -tokens the world model is unable to simulate key concepts.
Appendix B includes additional quantitative results.

###  3.4 Evidence of dynamics disentanglement

In Section 2.2, we argued that, by design, $\Delta$-iris’ encoder describes
stochastic deltas between timesteps with $\Delta$-tokens. In the present
section, we propose to exhibit this phenomenon.

We pick a starting frame and a sequence of actions, and predict two different
trajectories with the world model. In one case, we sample future
$\Delta$-tokens randomly. In the other case, $\Delta$-tokens are produced by
the autoregressive transformer. We consider a scenario where the agent
collects wood then builds a crafting table in Figure 3. Appendix F displays
two other scenarios where the agent explores its surroundings, and where it
moves down then stands still.

We observe that, even when sampling $\Delta$-tokens randomly, the
deterministic aspects of the dynamics are properly modelled: grid layout,
agent movement, wood level increasing, crafting table appearing, etc. On the
other hand, stochastic dynamics become problematic: skeletons and cows
appearing and disappearing, food and water indicators decreasing too early,
unlikely quantities of enemies and objects, etc. These observations confirm
that $\Delta$-iris encodes stochastic deltas between time steps with
$\Delta$-tokens, and its decoder handles the deterministic aspects of world
modelling.

##  4 Related Work

### World Models and imagination

With Dyna, Sutton (1991) introduced the idea of learning behaviours in the
imagination of a world model. Ha & Schmidhuber (2018) went beyond the tabular
setting and proposed a new world model architecture, composed of a variational
autoencoder (Kingma & Welling, 2013) and a recurrent network (Hochreiter &
Schmidhuber, 1997; Gers et al., 2000), capable of simulating simple visual
environments. Following this breakthrough, multiple generations of Dreamer
agents (Hafner et al., 2020, 2021, 2023) were developed, with DreamerV2 being
the first imagination-based agent to outperform humans in Atari games, and
DreamerV3 being the first world model architecture applicable to a wide range
of domains without any specific tuning. DreamerV2 learns in the imagination of
a world model combining a convolutional autoencoder with a recurrent state-
space model (RSSM) (Hafner et al., 2019). The key modifications that enabled
DreamerV2 to improve over the original Dreamer agent were categorical latents
and KL balancing between prior and posterior estimates. DreamerV3 builds upon
DreamerV2 with more universal design choices such as symlog scaling of rewards
and values, combining free bits (Kingma et al., 2016) with KL balancing,
return scaling for static entropy regularization, and architectural novelties
for model scaling. Variants of Dreamer such as TransDreamer (Chen et al.,
2022) and STORM (Zhang et al., 2023) have also been explored, where
transformers replace the recurrent network in the RSSM for dynamics
prediction.

A potential limitation of RSSM-like architectures is that they do not model
the joint distribution of future latent states, and instead predict product
laws. One way to mitigate this discrepancy between the predicted distributions
and the distributions of interest is to encourage factorized distributions
(Hafner et al., 2023). On the other hand, autoregressive architectures
(Micheli et al., 2023) do model the joint distribution and do not require to
enforce independence, which may result in a more expressive model.

### Trajectory and video autoencoders

The idea of encoding frames with respect to past frames predates modern deep
learning, and is at the origin of efficient video compression algorithms, such
as MPEG (Richardson, 2004). In recent years, multiple works have implemented
variants of this approach. Ozair et al. (2021) propose an offline version of
MuZero (Schrittwieser et al., 2020) equipped with an autoregressive
transformer that performs search over trajectory-level discrete latent
variables and actions. Phenaki (Villegas et al., 2022) is a text-to-video
model composed of a spatio-temporal discrete autoencoder and a masked
bidirectional transformer. TECO (Yan et al., 2022) is an action-conditional
video prediction model composed of a discrete frame autoencoder conditioned on
the previous frame, a temporal autoregressive transformer, and a spatial
MaskGit (Chang et al., 2022). While these methods also encode frames by
conditioning on past frames, their dynamics models purely operate over
discrete tokens, and do not leverage continuous tokens to alleviate the need
to integrate over multiple time steps in order to make the next prediction.

Hafner et al. (2019) acknowledge that modelling stochastic dynamics may be
difficult, as it would involve remembering information from previous time
steps. The authors propose to solve this problem by carrying a “deterministic”
state over time via a recurrent network, at the core of their RSSM. We make a
similar observation, and further show that this task is still difficult even
when past information does not have to be carried by a recurrent state, as a
transformer can attend to all previous $\Delta$-tokens. Hence, it is not only
a memory problem, but also a modelling one. Here, we address this issue in a
manner that is compatible with autoregressive transformers, namely by
injecting continuous  I -tokens in the sequence of the dynamics model.

##  5 Conclusion

We introduced $\Delta$-iris, a new model-based agent relying on an efficient
world model architecture to simulate its environment and learn new behaviours.
$\Delta$-iris features a discrete autoencoder that encodes the stochastic
aspects of world modelling with discrete $\Delta$-tokens, and an
autoregressive transformer leveraging continuous  I -tokens to model
stochastic dynamics.

Through experiments, we showed the ability of our agent to scale to the
challenging Crafter benchmark, as well as its sample efficiency in Atari100k.
Finally, we illustrated how its world model internalized environment dynamics,
and conducted ablations to validate our proposed design choices.

In its current form, $\Delta$-iris uses the same number of tokens to encode
stochastic dynamics at each time step. However, the reality of most
environments is such that periods of low uncertainty are quickly followed by
moments of high randomness. Therefore, an improved version of the world model
could possibly predict dynamically various numbers of tokens based on the
current context. Besides, leveraging the internal representations of the world
model could potentially result in a lightweight and more robust policy.

## Impact Statement

The deployment of autonomous agents in real-world applications raises safety
concerns. Agents learning new behaviours may harm individuals and damage
property. With world models, we lower the amount of time spent interacting
with the real world and thus mitigate risks. In this work, we propose a world
model architecture that is amenable to scaling up to complex environments,
where accurate simulations are even more critical given the usually higher
stakes.

## Acknowledgements

We would like to thank Adam Jelley, Bálint Máté, Daniele Paliotta, Maxim
Peter, Youssef Saied, Atul Sinha, and Alessandro Sordoni for insightful
discussions and comments. Vincent Micheli was supported by the Swiss National
Science Foundation under grant number FNS-187494.

## References

  * Achiam et al. (2023) Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Agarwal et al. (2021) Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A. C., and Bellemare, M.  Deep reinforcement learning at the edge of the statistical precipice.  _Advances in neural information processing systems_ , 34, 2021.
  * Anand et al. (2022) Anand, A., Walker, J. C., Li, Y., Vértes, E., Schrittwieser, J., Ozair, S., Weber, T., and Hamrick, J. B.  Procedural generalization by planning with self-supervised world models.  In _International Conference on Learning Representations_ , 2022.
  * Bellemare et al. (2013) Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 47:253–279, 2013.
  * Bengio et al. (2013) Bengio, Y., Léonard, N., and Courville, A.  Estimating or propagating gradients through stochastic neurons for conditional computation.  _arXiv preprint arXiv:1308.3432_ , 2013.
  * Chang et al. (2022) Chang, H., Zhang, H., Jiang, L., Liu, C., and Freeman, W. T.  Maskgit: Masked generative image transformer.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 11315–11325, 2022.
  * Chen et al. (2022) Chen, C., Wu, Y.-F., Yoon, J., and Ahn, S.  Transdreamer: Reinforcement learning with transformer world models.  _arXiv preprint arXiv:2202.09481_ , 2022.
  * comma.ai (2023) comma.ai.  commavq, 2023.  URL https://github.com/commaai/commavq.
  * D’Oro et al. (2023) D’Oro, P., Schwarzer, M., Nikishin, E., Bacon, P.-L., Bellemare, M. G., and Courville, A.  Sample-efficient reinforcement learning by breaking the replay ratio barrier.  In _The Eleventh International Conference on Learning Representations_ , 2023.
  * Esser et al. (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.
  * Gers et al. (2000) Gers, F. A., Schmidhuber, J., and Cummins, F.  Learning to forget: Continual prediction with LSTM.  _Neural Computation_ , 12(10):2451–2471, 2000.
  * Ha & Schmidhuber (2018) Ha, D. and Schmidhuber, J.  Recurrent world models facilitate policy evolution.  _Advances in neural information processing systems_ , 31, 2018.
  * Hafner (2022) Hafner, D.  Benchmarking the spectrum of agent capabilities.  In _International Conference on Learning Representations_ , 2022.
  * Hafner et al. (2019) Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J.  Learning latent dynamics for planning from pixels.  In _International conference on machine learning_ , pp. 2555–2565. PMLR, 2019.
  * Hafner et al. (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.
  * Hafner et al. (2021) Hafner, D., Lillicrap, T. P., Norouzi, M., and Ba, J.  Mastering atari with discrete world models.  In _International Conference on Learning Representations_ , 2021.
  * Hafner et al. (2023) Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T.  Mastering diverse domains through world models.  _arXiv preprint arXiv:2301.04104v1_ , 2023.
  * Hochreiter & Schmidhuber (1997) Hochreiter, S. and Schmidhuber, J.  Long short-term memory.  _Neural Computation_ , 9(8):1735–1780, 1997.
  * Hu et al. (2023) Hu, A., Russell, L., Yeo, H., Murez, Z., Fedoseev, G., Kendall, A., Shotton, J., and Corrado, G.  Gaia-1: A generative world model for autonomous driving.  _arXiv preprint arXiv:2309.17080_ , 2023.
  * Imani & White (2018) Imani, E. and White, M.  Improving regression performance with distributional losses.  In _International conference on machine learning_ , pp. 2157–2166. PMLR, 2018.
  * Kaiser et al. (2020) Kaiser, Ł., Babaeizadeh, M., Miłos, P., Osiński, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., et al.  Model based reinforcement learning for atari.  In _International Conference on Learning Representations_ , 2020.
  * Kanervisto et al. (2022) Kanervisto, A., Milani, S., Ramanauskas, K., Topin, N., Lin, Z., Li, J., Shi, J., Ye, D., Fu, Q., Yang, W., Hong, W., Huang, Z., Chen, H., Zeng, G., Lin, Y., Micheli, V., Alonso, E., Fleuret, F., Nikulin, A., Belousov, Y., Svidchenko, O., and Shpilman, A.  Minerl diamond 2021 competition: Overview, results, and lessons learned.  In _Proceedings of the NeurIPS 2021 Competitions and Demonstrations Track_ , Proceedings of Machine Learning Research, 2022.
  * Kapturowski et al. (2019) Kapturowski, S., Ostrovski, G., Quan, J., Munos, R., and Dabney, W.  Recurrent experience replay in distributed reinforcement learning.  In _International conference on learning representations_ , 2019.
  * Kauvar et al. (2023) Kauvar, I., Doyle, C., Zhou, L., and Haber, N.  Curious replay for model-based adaptation.  In _International Conference on Machine Learning_ , 2023.
  * Kingma & Welling (2013) Kingma, D. P. and Welling, M.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Kingma et al. (2016) Kingma, D. P., Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I., and Welling, M.  Improved variational inference with inverse autoregressive flow.  _Advances in neural information processing systems_ , 29, 2016.
  * LeCun (2022) LeCun, Y.  A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27.  _Open Review_ , 62, 2022.
  * LeCun et al. (1989) LeCun, Y., Boser, B., Denker, J. S., Henderson, D., Howard, R. E., Hubbard, W., and Jackel, L. D.  Backpropagation applied to handwritten zip code recognition.  _Neural computation_ , 1(4):541–551, 1989.
  * Micheli et al. (2023) Micheli, V., Alonso, E., and Fleuret, F.  Transformers are sample-efficient world models.  In _International Conference on Learning Representations_ , 2023.
  * Mnih et al. (2015) Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al.  Human-level control through deep reinforcement learning.  _Nature_ , 518(7540):529–533, 2015.
  * Mnih et al. (2016) Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T., Harley, T., Silver, D., and Kavukcuoglu, K.  Asynchronous methods for deep reinforcement learning.  In _International conference on machine learning_ , pp. 1928–1937. PMLR, 2016.
  * Ozair et al. (2021) Ozair, S., Li, Y., Razavi, A., Antonoglou, I., Van Den Oord, A., and Vinyals, O.  Vector quantized models for planning.  In _International Conference on Machine Learning_ , pp. 8302–8313. PMLR, 2021.
  * Pathak et al. (2017) Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T.  Curiosity-driven exploration by self-supervised prediction.  In _International conference on machine learning_ , pp. 2778–2787. PMLR, 2017.
  * Radford et al. (2019) Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I.  Language models are unsupervised multitask learners, 2019.
  * Razavi et al. (2019) Razavi, A., van den Oord, A., and Vinyals, O.  Generating diverse high-fidelity images with vq-vae-2.  _Advances in neural information processing systems_ , 32, 2019.
  * Richardson (2004) Richardson, I. E.  _H. 264 and MPEG-4 video compression: video coding for next-generation multimedia_.  John Wiley & Sons, 2004.
  * Rombach et al. (2022) Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 10684–10695, 2022.
  * Schrittwieser et al. (2020) Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., Lillicrap, T. P., and Silver, D.  Mastering atari, go, chess and shogi by planning with a learned model.  _Nature_ , 588(7839):604–609, 2020.
  * Schwarzer et al. (2021) Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., and Bachman, P.  Data-efficient reinforcement learning with self-predictive representations.  In _International Conference on Learning Representations_ , 2021.
  * Schwarzer et al. (2023) Schwarzer, M., Ceron, J. S. O., Courville, A., Bellemare, M. G., Agarwal, R., and Castro, P. S.  Bigger, better, faster: Human-level atari with human-level efficiency.  In _International Conference on Machine Learning_ , 2023.
  * Sekar et al. (2020) Sekar, R., Rybkin, O., Daniilidis, K., Abbeel, P., Hafner, D., and Pathak, D.  Planning to explore via self-supervised world models.  In _International Conference on Machine Learning_ , pp. 8583–8592. PMLR, 2020.
  * Sutton (1991) Sutton, R. S.  Dyna, an integrated architecture for learning, planning, and reacting.  _ACM Sigart Bulletin_ , 1991.
  * Sutton & Barto (2018) Sutton, R. S. and Barto, A. G.  _Reinforcement Learning: An Introduction_.  A Bradford Book, Cambridge, MA, USA, 2018.
  * Tassa et al. (2018) Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., et al.  Deepmind control suite.  _arXiv preprint arXiv:1801.00690_ , 2018.
  * Van Den Oord et al. (2017) Van Den Oord, A., Vinyals, O., et al.  Neural discrete representation learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Villegas et al. (2022) Villegas, R., Babaeizadeh, M., Kindermans, P.-J., Moraldo, H., Zhang, H., Saffar, M. T., Castro, S., Kunze, J., and Erhan, D.  Phenaki: Variable length video generation from open domain textual description.  _arXiv preprint arXiv:2210.02399_ , 2022.
  * Yan et al. (2022) Yan, W., Hafner, D., James, S., and Abbeel, P.  Temporally consistent video transformer for long-term video prediction.  _arXiv preprint arXiv:2210.02396_ , 2022.
  * Ye et al. (2021) Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Y.  Mastering atari games with limited data.  _Advances in neural information processing systems_ , 34, 2021.
  * Yu et al. (2021) Yu, J., Li, X., Koh, J. Y., Zhang, H., Pang, R., Qin, J., Ku, A., Xu, Y., Baldridge, J., and Wu, Y.  Vector-quantized image modeling with improved vqgan.  _arXiv preprint arXiv:2110.04627_ , 2021.
  * Zhang et al. (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.

##  Appendix A Architectures and hyperparameters

###  A.1 Discrete autoencoder

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\times 64$
Layers | 5
Residual blocks per layer | 2
Channels in convolutions per layer | $[64,64,128,128,256]$
Downsampling after layer $n$ | $[1,0,1,1,0]$
Past actions embedding channels | 4
Decoder past frames embedder arch. | Same as encoder
Decoder past frames embedder output feature map size | $8\times 8\times 8$
Conditioning time steps | $1$
$L_{1}$ loss weight | 0.1
$L_{2}$ loss weight | 1.0
Max-pixel loss weight | 0.01
Commitment loss weight | 0.02
Table 3: Embedding table and latent state hyperparameters. Hyperparameter | Value
---|---
Vocabulary size (N) | 1024
Tokens per frame (K) | 4
Latent feature map size | $64\times 8\times 8$
Pre-discretization token size | $64\times 4\times 4$
Token embedding dimension | 64
Codebook moving average coefficient | 0.99

In early experiments, we used a transformer instead of a cnn for the
architecture of the autoencoder. It had a much longer context size of twenty
time steps. Although the transformer-based autoencoder performed better than
its cnn counterpart on static datasets, we observed that the cnn would learn
faster than the transformer in the continual learning setup. Besides, for the
sake of simplicity, we decreased the initial conditioning of the cnn
autoencoder from four time steps to one time step, as the slight increase in
reconstruction losses did not significantly hinder agent performance. These
observations are largely environment-dependent, thus the context size or the
architecture of the autoencoder should most likely be adapted accordingly.

###  A.2 Autoregressive transformer

Table 4: Transformer hyperparameters. Hyperparameter | Value
---|---
Timesteps | 21
Embedding dimension | 512
Layers | 3
Attention heads | 8
Weight decay | 0.01
I -token frame embedder arch. | Same as encoder with halved channels per layer

###  A.3 Actor-Critic

We tie the weights of the actor and critic, except for the last layer. The
actor-critic takes as input a frame, and forwards it through a convolutional
neural network (LeCun et al., 1989) followed by an LSTM cell (Hochreiter &
Schmidhuber, 1997; Gers et al., 2000; Mnih et al., 2016). For the cnn, we use
the same architecture as the encoder, except that we halve the number of
channels per layer. The dimension of the LSTM hidden state is 512.

Before starting the imagination procedure ($H=15$) from a given frame, we
burn-in (Kapturowski et al., 2019) the 5 previous frames to initialize the
hidden state. The discount factor $\gamma$ is 0.997, the parameter for
$\lambda$-returns is set to 0.95, and the coefficient for the entropy
maximization term is 0.001. Targets for value estimates are produced by a
moving average of the critic network, with update parameter 0.995 (Mnih et
al., 2015)

###  A.4 Training loop and shared hyperparameters

Table 5: Training loop and shared hyperparameters.

Hyperparameter | Value
---|---
Epochs | 1000
Environment steps first epoch | 100000
Environment steps per epoch | 10000
# Collection epochs | 990
Collection epsilon-greedy | 0.01
Training steps per epoch | 500

Hyperparameter | Value
---|---
Autoencoder batch size | 32
Transformer batch size | 32
Actor-critic batch size | 86
Learning rate | 1e-4
Optimizer | Adam
Max gradient norm | 10.0

As mentioned in Section 2, the world model and policy are trained with
temporal segments sampled from past experience. We use a count-based sampling
procedure over the entire history of episodes, i.e. the likelihood that a
given episode is chosen to produce the next sample is inversely proportional
to the number of times it was previously used. We raise inverse counts to the
power of 5 to further limit the bias towards older episodes.

##  Appendix B Impact of design choices on key world modelling metrics

Metrics are computed on a held-out test set after training various world
models on a dataset consisting of 10M frames collected by a $\Delta$-iris
agent throughout its training.

Table 6: Left: Impact of removing past frames and actions from the
conditioning of the autoencoder ($\Delta$-iris $\rightarrow$ iris). Right:
Impact of removing

I

-tokens from the conditioning of the autoregressive transformer.

Method |  $L_{2}$ loss
---|---
$\Delta$-iris (4 tokens) | 0.000185
iris (64 tokens) | 0.001715
iris (16 tokens) | 0.007496

Method | Next token loss (CE) | Reward loss (CE)
---|---|---
$\Delta$-iris | 1.57 | 0.108
$\Delta$-iris w/o  I -tokens | 1.73 | 0.135

Table 7: Impact of discarding the auxiliary max-pixel loss. Method |  $L_{2}$ loss | Max-pixel loss
---|---|---
$\Delta$-iris | 0.000185 | 0.018
$\Delta$-iris w/o max-pixel loss | 0.000178 | 0.031

##  Appendix C Atari 100k

The Atari 100k benchmark (Kaiser et al., 2020) features Atari games (Bellemare
et al., 2013) with diverse mechanics. The specificity of this benchmark is the
hard constraint on the number of interactions, namely one hundred thousand per
environment. Compared to the standard Atari benchmark, this constraint results
in a dramatic drop in real-time experience, from 900 hours to 2 hours.

Regarding baselines, we consider four model-based RL agents learning in
imagination: SimPLe (Kaiser et al., 2020), DreamerV3 (Hafner et al., 2023),
storm (Zhang et al., 2023), and iris (Micheli et al., 2023). We note that the
current best performing methods for Atari 100k resort to other approaches,
such as lookahead search for EfficientZero (Ye et al., 2021), or self-
supervised representation learning with periodic resets for BBF (Schwarzer et
al., 2023).

The usual metric of interest is the HNS, the human-normalized score, based on
the performance of human players with similar experience. A negative HNS
indicates worse than random performance whereas an HNS above 1 signifies
superhuman performance. We evaluate $\Delta$-iris by computing an average over
100 episodes collected at the end of training for each game (5 seeds). For the
baselines, we report the published results.

Table 8 displays returns across games and aggregate metrics (Agarwal et al.,
2021). $\Delta$-iris achieves higher aggregate metrics than iris, while
training in 26 hours, a 5-fold speedup.

Table 8: Returns on the 26 games of Atari 100k after 2 hours of real-time experience, and human-normalized aggregate metrics. Game | Random | Human | SimPLe | DreamerV3 | storm | iris |  $\Delta$-iris (ours)
---|---|---|---|---|---|---|---
Alien | 228 | 7128 | 617 | 959 | 984 | 420 | 391
Amidar | 6 | 1720 | 74 | 139 | 205 | 143 | 64
Assault | 222 | 742 | 527 | 706 | 801 | 1524 | 1123
Asterix | 210 | 8503 | 1128 | 932 | 1028 | 854 | 2492
BankHeist | 14 | 753 | 34 | 649 | 641 | 53 | 1148
BattleZone | 2360 | 37188 | 4031 | 12250 | 13540 | 13074 | 11825
Boxing | 0 | 12 | 8 | 78 | 80 | 70 | 70
Breakout | 2 | 31 | 16 | 31 | 16 | 84 | 302
ChopperCommand | 811 | 7388 | 979 | 420 | 1888 | 1565 | 1183
CrazyClimber | 10781 | 35829 | 62584 | 97190 | 66776 | 59324 | 57864
DemonAttack | 152 | 1971 | 208 | 303 | 165 | 2034 | 533
Freeway | 0 | 30 | 17 | 0 | 34 | 31 | 31
Frostbite | 65 | 4335 | 237 | 909 | 1316 | 259 | 279
Gopher | 258 | 2413 | 597 | 3730 | 8240 | 2236 | 6445
Hero | 1027 | 30826 | 2657 | 11161 | 11044 | 7037 | 7049
Jamesbond | 29 | 303 | 101 | 445 | 509 | 463 | 309
Kangaroo | 52 | 3035 | 51 | 4098 | 4208 | 838 | 2269
Krull | 1598 | 2666 | 2205 | 7782 | 8413 | 6616 | 5978
KungFuMaster | 259 | 22736 | 14863 | 21420 | 26182 | 21760 | 21534
MsPacman | 307 | 6952 | 1480 | 1327 | 2674 | 999 | 1067
Pong | -21 | 15 | 13 | 18 | 11 | 15 | 20
PrivateEye | 25 | 69571 | 35 | 882 | 7781 | 100 | 103
Qbert | 164 | 13455 | 1289 | 3405 | 4523 | 746 | 1444
RoadRunner | 12 | 7845 | 5641 | 15565 | 17564 | 9615 | 10414
Seaquest | 68 | 42055 | 683 | 618 | 525 | 661 | 827
UpNDown | 533 | 11693 | 3350 | 9234 | 7985 | 3546 | 4072
#Superhuman | 0 | N/A | 1 | 9 | 10 | 10 | 11
Mean | 0.00 | 1.00 | 0.33 | 1.10 | 1.27 | 1.05 | 1.39
Interquartile Mean | 0.00 | 1.00 | 0.13 | 0.50 | 0.64 | 0.50 | 0.65

##  Appendix D Crafter scores and individual success rates

Table 9: Crafter scores, i.e. geometric mean of success rates. Method | Crafter score @1M | Crafter score @5M | Crafter score @10M
---|---|---|---
$\Delta$-iris | 9.30 | 39.67 | 42.47
$\Delta$-iris w/o  I -tokens | 5.85 | 14.39 | 25.92
iris (64 tokens) | 6.66 | - | -
Figure 7: Individual success rates after collecting 1M frames. Figure 8:
Individual success rates after collecting 10M frames.

##  Appendix E Baselines

DreamerV3 results were obtained with commit 8fa35f8. We used the standard
configuration for Crafter, and set the run.train_ratio variable controlling
the imagined-to-collected data ratio to 64. Note that a new version of
DreamerV3 was recently released in April 2024. This update includes additional
and broadly applicable novelties for world model and policy learning.

iris results were obtained with commit ac6be40. For the training loop and
shared hyperparameters, we picked the same values as in Table 5. We increased
the dimension and attention heads of the transformer from 256 and 4 to 512 and
8, respectively. Finally, we used a replay buffer with a capacity of 1M
frames.

##  Appendix F Evidence of dynamics disentanglement

$\Delta$-tokens sampled randomly

$\Delta$-tokens sampled by the autoregressive transformer

$t=0$ $t=4$ $t=5$ $t=9$ $t=10$ $t=12$

$\Delta$-tokens sampled randomly

$\Delta$-tokens sampled by the autoregressive transformer

$t=0$ $t=4$ $t=5$ $t=9$ $t=10$ $t=12$

Figure 9: Two additional examples of dynamics disentanglement, as discussed in
Section 3.4 and Figure 3.


$|\log\mathbb{E}^{V_{N}}_{N,\beta}\exp(t\mathrm{Fluct}(\xi_{R,\alpha}))|\leq\frac{C|t|R^{d-2}}{\alpha^{3}}+Ct^{2}\left(\frac{R^{d-4}}{\alpha^{4}}+\frac{R^{d-2}}{\alpha}\right)+\frac{Ct^{4}R^{d-8}}{\alpha^{8}}+\frac{C|t|R^{d}}{N^{2/d}}.$
(5.32)

for all $|t|\leq C^{-1}\alpha^{2}R^{2}$. In particular for $t=1$ and
$\alpha^{3}R^{2}\leq N^{2/d}$, we have

$|\log\mathbb{E}^{V_{N}}_{N,\beta}\exp(\mathrm{Fluct}(\xi_{R,\alpha}))|\leq
C\frac{|t|R^{d-2}}{\alpha^{3}}.$ (5.33)

###### Proof.

A more general estimate on a non-blown up scale, i.e. the scale with
interstitial distance of order $N^{-1/d}$, and with the thermal equilibrium
measure $\mu_{\theta}$ in place of the equilibrium measure
$\mu_{\mathrm{eq}}$, is [Ser20, Theorem 1]. We slightly transform the estimate
by using $\beta\geq C^{-1}$, changing the interstitial length scale to $O(1)$,
and plugging in the specifics of $\xi_{R,\alpha}$. We also use (see [AS19,
Theorem 1])

$\|\mu_{\theta}-\mu_{\mathrm{eq}}\|_{L^{\infty}(\mathbb{R}^{d})}\leq\frac{C}{N^{2/d}}$

to replace the thermal equilibrium measure by $\mu_{\mathrm{eq}}$, generating
the $C|t|R^{d}N^{-2/d}$ term in (5.32). ∎

The following proposition shows that one can find a screening region of excess
positive charge whenever the absolute discrepancy is large in a ball. It is a
natural consequence of rigidity for fluctuations of smooth linear statistics.

######  Proposition 5.7.

Fix $\delta\in(0,d)$ and $R\geq C$ for a large enough $C>0$ and
$\alpha\in(0,1]$. For any $T\geq 1$, if $\mathrm{Fluct}(\xi_{R-\alpha
R,\alpha})\leq\frac{T}{10}R^{\delta}$, we can find $k\in\mathbb{Z}^{\geq 0}$
and a constant $C_{6}>0$ such that

$\mathrm{Disc}(\mathrm{Ann}_{[R-\alpha_{k}R,R]}(z))\geq\mathrm{Disc}(B_{R}(z))-\frac{T}{2}R^{\delta}$

for $\alpha_{k}=\alpha-C_{6}^{-1}R^{-d+\delta}k$ with
$C_{6}^{-1}R^{-d+\delta}k\leq\alpha_{k}\leq\alpha$. If instead
$\mathrm{Fluct}(\xi_{R,\alpha})\geq-\frac{T}{10}R^{\delta}$, we can also find
$\alpha_{k}$ as above with

$\mathrm{Disc}(\mathrm{Ann}_{[R,R+\alpha_{k}R]}(z))\geq-\mathrm{Disc}(B_{R}(z))-\frac{T}{2}R^{\delta}.$

###### Proof.

First, note that

$\mathrm{Disc}(B_{s+\varepsilon}(z))-\mathrm{Disc}(B_{s}(z)))=\int_{\mathrm{Ann}_{[s+\varepsilon,s)}(z)}\mathrm{fluct}(dx)$

whenever $\mathrm{fluct}$ has no atoms on $\partial B_{s}(z)$ or $\partial
B_{s+\varepsilon}(z)$. Thus, we have by integration in spherical coordinates
that

$\displaystyle\mathrm{Fluct}(\xi_{R-\alpha R,\alpha})$
$\displaystyle=\int_{0}^{R}\xi_{R-\alpha
R,\alpha}(s)d(\mathrm{Disc}(B_{s}(z)))=-\int_{0}^{R}\frac{d}{ds}\xi_{R-\alpha
R,\alpha}(s)\mathrm{Disc}(B_{s}(z))ds$
$\displaystyle=\mathrm{Disc}(B_{R}(z))-\int_{R-\alpha
R}^{R}\frac{d}{ds}\xi_{R-\alpha
R,\alpha}(s)\left(\mathrm{Disc}(B_{s}(z))-\mathrm{Disc}(B_{R}(z))\right)ds.$

Assuming now $\mathrm{Fluct}(\xi_{R-\alpha
R,\alpha})\leq\frac{T}{10}R^{\delta}$, by the mean value theorem, there must
exist $s\in(R-\alpha R,R)$ such that
$\mathrm{Disc}(B_{s}(z))-\mathrm{Disc}(B_{R}(z))\leq\frac{T}{10}R^{\delta}-\mathrm{Disc}(B_{R}(z))$.
This means that

$\mathrm{Disc}(\mathrm{Ann}_{[s,R)}(z))\geq\mathrm{Disc}(B_{R}(z))-\frac{T}{10}R^{\delta}.$

For any $k$ such that $R-\alpha R\leq R-\alpha_{k}R<s$, we have

$\displaystyle\mathrm{Disc}(\mathrm{Ann}_{(R-\alpha_{k}R,R)}(z))$
$\displaystyle\geq\mathrm{Disc}(\mathrm{Ann}_{[s,R)}(z))-\mu_{\mathrm{eq}}^{N}(\mathrm{Ann}_{(R-\alpha_{k},s]}(z))$
$\displaystyle\geq\mathrm{Disc}(\mathrm{Ann}_{[s,R)}(z))-CR^{d-1}|s-\alpha_{k}R|.$

We choose $k$ such that $|s-(R-\alpha_{k}R)|\leq C^{-1}R^{-d+1+\delta}$ for a
large enough constant $C$ to conclude.

If we instead assume
$\mathrm{Fluct}(\xi_{R,\alpha})\geq-\frac{T}{10}R^{\delta}$, then since

$\mathrm{Fluct}(\xi_{R,\alpha})=\mathrm{Disc}(B_{R}(z))-\int_{R}^{R+\alpha
R}\frac{d}{ds}\xi_{R,\alpha}(s)\left(\mathrm{Disc}(B_{s}(z))-\mathrm{Disc}(B_{R}(z))\right)ds,$

we can find $s\in(R,R+\alpha R)$ with

$\mathrm{Disc}(\mathrm{Ann}_{[R,s)}(z))\geq-\frac{T}{10}R^{\delta}-\mathrm{Disc}(B_{R}(z)).$

Choosing $\alpha_{k}$ such that $s\leq R+\alpha_{k}R\leq R+\alpha R$ and
$|s-(R+\alpha_{k}R)|\leq C^{-1}R^{-d+1+\delta}$ is sufficient to conclude. ∎

We are now ready to prove Theorem 6.

###### Proof of Theorem 6.

Let $\mathbb{P}=\mathbb{P}^{V_{N}}_{N,\beta}$. We choose $\alpha=R^{-\lambda}$
for $\lambda=2/5$ and consider discrepancies in $B_{R}(z)$ of size
$TR^{\delta}(1+\mathbf{1}_{d=2}\log R)$ for $\delta=d-4/5$ and $T$
sufficiently large. We will however write a mostly generic argument in terms
of $\lambda$ and $\delta$ and insert the specific values later. We will
consider the case of positive discrepancy first.

If $\mathrm{Fluct}(\xi_{R-\alpha R,R})\leq\frac{T}{10}R^{\delta}$, we can use
Proposition 5.7 to find $\alpha_{k}\in(0,\alpha]$ such that
$\mathrm{Disc}(\mathrm{Ann}_{[R-\alpha_{k}R,R]}(z))\geq\frac{T}{2}R^{\delta}(1+\mathbf{1}_{d=2}\log
R)$. By a union bound, we have

$\displaystyle\mathbb{P}(\\{\mathrm{Disc}(B_{R}(z))\geq TR^{\delta}\log R\\})$
(5.34) $\displaystyle\leq\mathbb{P}(\\{\mathrm{Fluct}(\xi_{R-\alpha
R,\alpha})>\frac{T}{10}R^{\delta}\\})+C\alpha
R^{d-\delta}\sup_{\alpha_{k}}\mathbb{P}(\\{\mathrm{Disc}(\mathrm{Ann}_{[R-\alpha_{k}R,R]}(z))\geq\frac{T}{2}R^{\delta}\log
R\\}),$

where the supremum is over all $\alpha_{k}=\alpha-C^{-1}R^{-d+\delta}k$,
$k\in\mathbb{Z}$ and $\alpha_{k}\in[0,\alpha)$. Note that we have
$\alpha_{k}\geq C^{-1}R^{-d+\delta}\gg R^{-1}$ always. Next, we bound the
supremum in (5.34).

Let $\alpha^{\prime}\in[C^{-1}R^{-d+\delta},\alpha]$ and let $\Omega$ be the
$\alpha^{\prime}$-thin annulus $\mathrm{Ann}_{[R-\alpha^{\prime}R,R]}(z)$. We
bound

$\mathbb{P}(\\{\mathrm{Disc}(\Omega)\geq\frac{T}{2}R^{\delta}(1+\mathbf{1}_{d=2}\log
R)\\})\leq\mathbb{P}(X(\Omega)\geq\mu(\Omega)+\rho|\Omega|)$

for $\rho=C^{-1}T(\alpha^{\prime})^{-1}R^{-d+\delta}(1+\mathbf{1}_{d=2}\log
R)$. Applying Proposition 5.6 shows

$\mathbb{P}(\\{\mathrm{Disc}(\Omega)\geq\frac{T}{2}R^{\delta}(1+\mathbf{1}_{d=2}\log
R)\\})\leq
e^{-c(\alpha^{\prime})^{2/3}R^{2/3+\delta}T}+e^{-c(\alpha^{\prime}R)^{d/3}R^{2+2\delta-2d}T^{2}}+e^{-(\alpha^{\prime}R)^{d+2}}$
(5.35)

whenever $T(\alpha^{\prime})^{-1/3}R^{\delta+2/3-d}\geq
TR^{\delta+2/3-d+\lambda/3}$ is large and we can estimate

$\|\mu_{\mathrm{eq}}^{N}-q\|_{L^{\infty}(B_{2R}(z))}\leq
C_{2}^{-1}TR^{\delta-d+\lambda}$

for a constant $q$. The latter condition happens when
$R^{1+d-\delta-\lambda}\ll N^{1/d}$.

Considering the other term in (5.34), we apply Theorem 8 with $|t|=C^{-1}$ and
Chebyshev’s inequality to bound

$\mathbb{P}(\\{\mathrm{Fluct}(\xi_{R-\alpha
R,\alpha})>\frac{T}{10}R^{\delta}\\})\leq
e^{C^{-1}(-\frac{T}{10}R^{\delta}+CR^{d-2+3\lambda})}.$ (5.36)

Note that $\alpha^{3}R^{2}=R^{4/5}\leq N^{2/d}$ so that (5.33) applies. We
thus apply our argument to parameters $\lambda$ and $\delta$ such that

$\delta+\frac{2}{3}-d+\frac{\lambda}{3}\geq 0,\quad\delta\geq d-2+3\lambda.$
(5.37)

One can check that the smallest choice of $\delta$ is $\delta=d-4/5$ with
$\lambda=2/5$. With these choices, choosing $R\leq N^{\frac{5}{7d}}$
guarantees that we can approximate $\mu_{\mathrm{eq}}^{N}$ by a constant
sufficiently well.

Finally, we estimate the RHS of (5.35) and (5.36) and plug them into (5.34).
Note that $\alpha^{\prime}\geq C^{-1}R^{-4/5}$, so the RHS of (5.35) is
bounded by

$e^{-cR^{d-10/15}T}+e^{-cR^{2/5+d/15}T^{2}}+e^{-R^{(d+2)/5}}.$

The factor $\alpha R^{d-\delta}$ within (5.34) can be absorbed into the above
at the cost of a constant factor within the exponent.

This finishes the one half of the proof of Theorem 6. The proof of the lower
bound on $\mathrm{Disc}(B_{R}(z))$ is nearly identical, except we use
fluctuation bounds on $\xi_{R,\alpha}$ to find a screening region outside of
$\partial B_{R}(z)$ with positive discrepancy. ∎

## References

  * [AB13] Gérard Ben Arous and Paul Bourgade.  Extreme gaps between eigenvalues of random matrices.  The Annals of Probability, 41(4):2648 – 2681, 2013.
  * [AKM19] Scott Armstrong, Tuomo Kuusi, and Jean-Christophe Mourrat.  Quantitative stochastic homogenization and large-scale regularity, volume 352.  Springer, 2019.
  * [Ame18] Yacin Ameur.  Repulsion in Low Temperature $\beta$-Ensembles.  Communications in Mathematical Physics, 359(3):1079–1089, May 2018\.
  * [Ame21] Yacin Ameur.  A localization theorem for the planar Coulomb gas in an external field.  Electronic Journal of Probability, 26(none):1 – 21, 2021.
  * [AR22] Yacin Ameur and José Luis Romero.  The planar low temperature Coulomb gas: separation and equidistribution, 2022.
  * [AS19] Scott Armstrong and Sylvia Serfaty.  Thermal approximation of the equilibrium measure and obstacle problem, 2019.
  * [AS21] Scott Armstrong and Sylvia Serfaty.  Local laws and rigidity for Coulomb gases at any temperature.  The Annals of Probability, 49(1):46 – 121, 2021.
  * [BBNY19] Roland Bauerschmidt, Paul Bourgade, Miika Nikula, and Horng-Tzer Yau.  The two-dimensional coulomb plasma: quasi-free approximation and central limit theorem.  Advances in Theoretical and Mathematical Physics, 23(4):841–1002, 2019.
  * [BMP21] Paul Bourgade, Krishnan Mody, and Michel Pain.  Optimal local law and central limit theorem for $\beta$-ensembles.  arXiv:1806.01831, 2021.
  * [Bou22] Paul Bourgade.  Extreme gaps between eigenvalues of Wigner matrices.  Journal of the European Mathematical Society, 24(8):2823 – 2873, 2022.
  * [CHM18] Djalil Chafaï, Adrien Hardy, and Mylène Maïda.  Concentration for Coulomb gases and Coulomb transport inequalities.  Journal of Functional Analysis, 275(6):1447–1483, 2018.
  * [Cif06] Orion Ciftja.  Monte Carlo study of Bose Laughlin wave function for filling factors 1/2, 1/4 and 1/6.  Europhysics Letters, 74(3):486, 2006.
  * [CSA20] Gabriel Cardoso, Jean-Marie Stéphan, and Alexander G Abanov.  The boundary density profile of a Coulomb droplet. freezing at the edge.  Journal of Physics A: Mathematical and Theoretical, 54(1):015002, Dec 2020.
  * [CW03] Orion Ciftja and Carlos Wexler.  Monte Carlo simulation method for Laughlin-like states in a disk geometry.  Phys. Rev. B, 67:075304, Feb 2003.
  * [DVJ08] D. J. Daley and D. Vere-Jones.  An Introduction to the Theory of Point Processes.  Probability and its Applications. Springer New York, 2nd edition, 2008\.
  * [EKYY12] László Erdős, Antti Knowles, Horng-Tzer Yau, and Jun Yin.  Spectral statistics of erdős-rényi graphs ii: Eigenvalue spacing and the extreme eigenvalues.  Communications in Mathematical Physics, 314(3):587–640, Sep 2012\.
  * [ESY10] László Erdős, Benjamin Schlein, and Horng-Tzer Yau.  Wegner estimate and level repulsion for Wigner random matrices.  International Mathematics Research Notices, 2010(3):436–479, 2010\.
  * [FTW19] Renjie Feng, Gang Tian, and Dongyi Wei.  Small gaps of GOE.  Geometric and Functional Analysis, 29(6):1794–1827, Dec 2019.
  * [FW21] Renjie Feng and Dongyi Wei.  Small gaps of circular $\beta$-ensemble.  The Annals of Probability, 49(2):997 – 1032, 2021.
  * [GLM80] Ch. Gruber, Ch. Lugrin, and Ph. A. Martin.  Equilibrium properties of classical systems with long-range forces. BBGKY equation, neutrality, screening, and sum rules.  Journal of Statistical Physics, 22(2):193–236, Feb 1980.
  * [JLM93] B Jancovici, Joel L Lebowitz, and G Manificat.  Large charge fluctuations in classical Coulomb systems.  Journal of statistical physics, 72(3):773–787, 1993.
  * [KK16] S. A. Khrapak and A. G. Khrapak.  Internal energy of the classical two- and three-dimensional one-component-plasma.  Contributions to Plasma Physics, 56(3-4):270–280, 2016.
  * [Leb21] Thomas Leblé.  The two-dimensional one-component plasma is hyperuniform, 2021.
  * [LRY19] Elliott H. Lieb, Nicolas Rougerie, and Jakob Yngvason.  Local incompressibility estimates for the Laughlin phase.  Communications in Mathematical Physics, 365(2):431–470, Jan 2019\.
  * [LS18] Thomas Leblé and Sylvia Serfaty.  Fluctuations of two dimensional coulomb gases.  Geometric and Functional Analysis, 28:443–508, 2018.
  * [Mar88] Ph. A. Martin.  Sum rules in charged fluids.  Rev. Mod. Phys., 60:1075–1127, Oct 1988.
  * [MW66] N. D. Mermin and H. Wagner.  Absence of ferromagnetism or antiferromagnetism in one- or two-dimensional isotropic Heisenberg models.  Phys. Rev. Lett., 17:1133–1136, Nov 1966.
  * [MY80] Ph. A. Martin and T. Yalcin.  The charge fluctuations in classical Coulomb systems.  Journal of Statistical Physics, 22(4):435–463, Apr 1980.
  * [NS15] Simona Rota Nodari and Sylvia Serfaty.  Renormalized energy equidistribution and local charge balance in 2D Coulomb systems.  International Mathematics Research Notices, 2015(11):3035–3093, 2015.
  * [NTV17] Hoi Nguyen, Terence Tao, and Van Vu.  Random matrices: tail bounds for gaps between eigenvalues.  Probability Theory and Related Fields, 167(3):777–816, Apr 2017\.
  * [OR20] Alessandro Olgiati and Nicolas Rougerie.  Stability of the Laughlin phase against long-range interactions.  Archive for Rational Mechanics and Analysis, 237(3):1475–1515, 2020\.
  * [PS17] Mircea Petrache and Sylvia Serfaty.  Next order asymptotics and renormalized energy for Riesz interactions.  Journal of the Institute of Mathematics of Jussieu, 16(3):501–569, 2017.
  * [Rou22a] Nicolas Rougerie.  The classical jellium and the Laughlin phase.  arXiv preprint arXiv:2203.06952, 2022.
  * [Rou22b] Nicolas Rougerie.  On the stability of Laughlin’s fractional quantum Hall phase.  arXiv preprint arXiv:2208.02619, 2022.
  * [RS16] Nicolas Rougerie and Sylvia Serfaty.  Higher-dimensional Coulomb gases and renormalized energy functionals.  Communications on Pure and Applied Mathematics, 69(3):519–605, 2016\.
  * [RSY14] N. Rougerie, S. Serfaty, and J. Yngvason.  Quantum Hall phases and plasma analogy in rotating trapped Bose gases.  Journal of Statistical Physics, 154(1):2–50, Jan 2014.
  * [RY15a] Nicolas Rougerie and Jakob Yngvason.  Incompressibility estimates for the Laughlin phase.  Communications in Mathematical Physics, 336(3):1109–1140, Jun 2015\.
  * [RY15b] Nicolas Rougerie and Jakob Yngvason.  Incompressibility estimates for the Laughlin phase, part ii.  Communications in Mathematical Physics, 339(1):263–277, Oct 2015\.
  * [RY18] Nicolas Rougerie and Jakob Yngvason.  The Laughlin liquid in an external potential.  Letters in Mathematical Physics, 108:1007–1029, 2018.
  * [Ser20] Sylvia Serfaty.  Gaussian fluctuations and free energy expansion for Coulomb gases at any temperature, 2020.
  * [SJ12] Dai Shi and Yunjiang Jiang.  Smallest gaps between eigenvalues of random matrices with complex Ginibre, Wishart and universal unitary ensembles, 2012.
  * [Tao13] Terence Tao.  The asymptotic distribution of a single eigenvalue gap of a Wigner matrix.  Probability Theory and Related Fields, 157(1):81–106, Oct 2013\.
  * [TV11] Terence Tao and Van Vu.  Random matrices: Universality of local eigenvalue statistics.  Acta Mathematica, 206(1):127 – 204, 2011.

Eric Thoma
Courant Institute, New York University.
Email<EMAIL_ADDRESS>