Accurate and Efficient World Modeling with Masked Latent Transformers
Authors: Maxime Burchi
Radu Timofte
Venue: ICML 2025 poster (2025)
Keywords: model-based reinforcement learning
maskGIT
transformer network
crafter
Primary area: reinforcement_learning
TL;DR: We introduce EMERALD, a Transformer model-based reinforcement learning algorithm using a spatial latent state with MaskGIT predictions to generate accurate trajectories in latent space and improve the agent performance.

Abstract: The Dreamer algorithm has recently obtained remarkable performance across diverse environment domains by training powerful agents with simulated trajectories. However, the compressed nature of its world model's latent space can result in the loss of crucial information, negatively affecting the agent's performance. Recent approaches, such as $\Delta$-IRIS and DIAMOND, address this limitation by training more accurate world models. However, these methods require training agents directly from pixels, which reduces training efficiency and prevents the agent from benefiting from the inner representations learned by the world model. In this work, we propose an alternative approach to world modeling that is both accurate and efficient. We introduce EMERALD (Efficient MaskEd latent tRAnsformer worLD model), a world model using a spatial latent state with MaskGIT predictions to generate accurate trajectories in latent space and improve the agent performance. On the Crafter benchmark, EMERALD achieves new state-of-the-art performance, becoming the first method to surpass human experts performance within 10M environment steps. Our method also succeeds to unlock all 22 Crafter achievements at least once during evaluation.

Accurate and Efficient World Modeling with Masked Latent Transformers
Maxime Burchi'
Radu timofte'
Abstract
Crafter Achievements Score (10M environment steps)
Leam agent from reconstructed images Spatial and temparal world model state
The Dreamer algorithm has recently obtained remarkable performance across diverse environ- ment domains by training powerful agents with simulated trajectories. However, the compressed nature of its world model's latent space can re- sult in the loss of crucial information, negatively affecting the agent's performance. Recent ap- proaches, such as A-IRIS and DIAMOND, ad- dress this limitation by training more accurate world models. However, these methods require training agents directly from pixels, which re- duces training efficiency and prevents the agent from benefiting from the inner representations learned by the world model. In this work, we propose an alternative approach to world model- ing that is both accurate and efficient. We intro- duce EMERALD (Efficient MaskEd latent iRAns- former worLD model), world model using a spatial latent state with MaskGIT predictions generate accurate trajectories in latent space and improve the agent performance. On the Crafter benchmark, EMERALD achieves new state-of- the-art performance, becoming the first method to surpass human experts performance within 10M environment steps. Our method also succeeds to unlock all 22 Crafter achievements least once during evaluation.
60%
Temporal only
58.1
40%-
34.9
39.6
42.5
20%
0%
Dreamenva
Dreamerl'3 Xt
4:R (Kai 2024)
Expents
EMESALO (Ours)
Collected Frames Per Second (FP5 on RTX 3090)
38
30-
27
23
15
DreamerV3
DrnamerV3 xt
AJRS (KOAL 2024)
EMERALD (Ours)
Figure 1: Achievements score and collected Frames Per Second (FPS) of recently published model-based methods on the Crafter benchmark. EMERALD is first method to exceed the performance of human experts on the benehmark. The world model uses spatial latent state with MaskGIT predictions to improve performance and training efficiency, A-IRIS proposed an accurate world model but suffers from lower training efficiency due to autoregressive decoding in latent space and agent learning from reconstructed images.
1. Introduction
Model-basee reinforcement learning has attracted inereasing research attention in recent years (Hafner et al., 2019; 2020; 2021; 2023; Schrittwieser et al., 2020). The growing com- putational capabilities hardware systems have allowed re- searchers make significant progress, training world mod- els from nigh-dimensional observations like videos (Yan et al., 2023) using deep neural networks (LeCun al., 2015)
function approximations, World models (Sutton, 1991; Ha & Schmidhuber, 2018) summarize the experience of an agent into predictive function that can be used instead of the real environment to learn complex behaviors. Having access to a model of the environment enables the agent to simulate multiple plausible trajectories in parallel, improv- ing generalization, sample efficiency, and via planning.
'Computer Vision Lab, CAIDAS IFI. University of Würzburg, Germany. Correspondence to: Maxime Burchi <maxime.burchi@e www.wuuwwurur..............
Among these approaches, the Dreamer algorithm (Hafner et al., 2020; 2021; 2023) has achieved remarkable perfor- mance various environments by training a world model to generate imaginary trajectories in latent space. However, while the world model has no difficulty in simulating simple
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Accurate and Efficient World Modeling with Masked Latent Transformers
4.
9999 9
4444
99
4
4/4
1
21
21
1-1
112
(a) EMERALD
(b) Dreamer (M)
Figure 2: Comparison of EMERALD image reconstruction with DreamerV3. We show the 5 frames with highest recon- struction emrrr among batch of 1024 test observations. The top row indicates original images, the middle row shows reconstructed images and the bottom row shows reconstruction error. We observe that EMERALD achieves near-perfect reconstruction, with errors resulting mostly from player orientation and textures. Dreamer V3 fails to perceive crucial details like diamonds and skeleton arrows, which diminishes the agent's 's perception capacity and negatively impacts its performance.
visual environments such as DeepMind Control (Tassa et al., 2018) tasks or Atari games (Bellemare et al., 2013), some inaccuracies appear when modeling more complex environ- ments like Crafter (Hafner, 2022) and MineRL (Guss et al., 2019). As shown in Figure 2b, the world model has diffi- culty perceiving important details on Crafter. This failure to learn accurate representations for the agent leads to lower performance.
approach was later adapted by TECO (Yan et al., 2023), demonstrating that MaskGIT-based latent space predictions could enable powerful video generation models. Inspired by these advances, we apply MaskGIT to model-based re- inforcement learning, allowing the agent to better perceive crucial environment details while improving decoding effi- ciency. Furthermore, the agent benefits from internal repre- sentations learned by the world model, such as long-term memory, enabling it to track important objects and effec- tively map the environment. As shown in Figure 1, EMER- ALD sets record on the Crafter benchmark, being the first method to exceed human experts performance within 10M environment steps with score of 58.1 %. Our method also succeeds to unlock all 22 Crafter achievements at least once during evaluation. Additionally, we report results on the commonly used Atari 100k benchmark (Kaiser et al., 2020) to demonstrate the general efficacy of EMERALD on Atari games that do not necessarily require the use of spatial latents to achieve near perfect reconstruction.
To remedy this problem, recent approaches have proposed more accurate world models. A-IRIS (Micheli et al., 2024) proposed to encode stochastic deltas for changes in the ob- servation. Concurrently, DIAMOND (Alonso et al., 2024) proposed a diffusion-based world model using U-Net ar- chitecture. Both of these approaches proposed new world model alternatives to improve the reconstruction quality of generated trajectories and increase the agent performance. However, they also require learning the agent from recon- struction images, which impact learning efficiency and does not allow the agent to benefit from inner representations learned by the world model such as long-term memory. This significantly limits the potential of such methods for perception and training efficiency.
2. Related Works
2.1. Model-based Reinforcement Learning
In this work, we propose an alternative approach to world modeling that is both accurate and efficient. We intro- duce EMERALD, Transformer model-based reinforce- ment learning algorithm using a spatial latent state with MaskGIT predictions to generate accurate trajectories in latent space and improve the agent performance. Origi- nally designed for vector-quantized image generation (Van Den Oord al., 2017), MaskGIT (Chang et al., 2022) im- proves decoding speed and generation quality compared to autoregressive sequential decoding (Esser et al., 2021). This
Model-based reinforcement learning approaches use a model of the environment simulate agent trajectories, improving generalization, sample efficiency, and decision- making via planning. Following the success of deep neural networks for learning function approximations, researchers proposed to learn world models using gradient backpropa- gation, allowing the development of powerful agents with limited data. One of the earliest model-based algorithms applied to image data SimPLe (Kaiser et al., 2020), which proposed a world model for Atari games in pixel space using
2

---

Accurate and Efficient World Modeling with Masked Latent Transformers
Table 1: Comparison between EMERALD and other recent model-based approaches. Latent (2a) refers to image latent representations carrying spatial information while hidden (h₂) refers to world model hidden state carrying temporal information. EMERALD uses spatial latent state to improve world model accuracy. It also trains the agent latent space to increase efficiency and benefit from world model inner representations. Efficient decoding is performed using MaskGIT.
Attributes
IRIS
D3eamerV3
A-IRIS
DIAMOND
EMERALD (Ours)
World Model
Transformer
GRU Latent
Transformer
U-Net
Transformer
Tokens Latent Representation
Latent (4 x VQ-VAE
Latent (2 2)
N/A
Latent (4 x 4)
Categorical-VAn
VQ-VAE
N/A
Categorical-Ve -erg
Decoding
Sequential
N/A
Sequential
EDM
MaskGIT
Agent State (s²)
Image
Latent, hidden
Image
Image
Latent (4 x 4), hidden
a convolutional autoencoder architecture. The world model predicts future frames and environment rewards based on past observations and selected actions, enabling the training of Proximal Policy Optimization (PPO) agent (Schulman et al., 2017) with simulated trajectories.
ages into discrete tokens and an autoregressive Transformer to predict future tokens. TransDreamer (Chen al., 2021) proposed to replace Dreamer's 's RSSM with Transformer State-Space Model (TSSM) using masked self-attention to imagine future trajectories. TWM (Robine et al., 2023) (Transformer-based World Model) proposed a similar ap- proach, encoding states, actions and rewards as distinct successive input tokens for the autoregressive Transformer. More recently, STORM (Zhang et al., 2024) achieved results comparable DreamerV3 with better training efficiency on the Atari 100k benchmark.
Concurrently, PlaNet (Hafner et al., 2019) introduced Re- current State-Space Model (RSSM) incorporating Gated Recurrent Units (GRUs) (Cho et al., 2014) to learn a world model in latent space, planning using model predictive con- trol. PlaNet learns convolutional variational autoencoder (VAE) (Kingma& & Welling, 2014) with a pixel reconstruc- tion loss to encode observations into stochastic state repre- sentations. The RSSM then predicts future stochastic states and environment rewards based on previous stochastic and deterministic recurrent states. Following the success of PlaNet on DeepMind Visual Control tasks, Dreamer (Hafner et al., 2020) improved the algorithm by learning an actor and a value network from the world model hidden repre- sentations. DreamerV2 (Hafner et al., 2021) extended the algorithm to Atari games, replacing Gaussian latents by cat- egorical latent states using straight-through gradients (Ben- gio et al., 2013) to improve performance. Lastly, Dream- erV (Hafner et al., 2023) mastered diverse domains us- ing the same hyper-parameters with a set of architectural changes to stabilize learning across tasks. The agent uses symlog predictions for the reward and value function to address the scale variance across domains. The networks also employ layer normalization (Ba et al., 2016) to im- prove robustness and performance while scaling to larger model sizes. It stabilizes policy learning by normalizing the returns and value function using an Exponential Moving Average (EMA) of the returns percentiles. With these modi- fications, 33eamerV3 outperformed specialized model-free and
2.2. Accurate World Modeling in Pixel Space
Another line of work focuses on designing accurate world models to train agents from reconstructed trajectories in pixel space. Analogously to SimPLe, the agent's policy and value functions are trained from image reconstruction instead of world model hidden state representations. This requires learning auxiliary encoder networks for the policy and value functions. Micheli et al. (2024) proposed A- IRIS, encoding stochastic deltas between time steps using previous action and image as conditions for the encoder and decoder. This increased VQ-VAE compression ratio and image reconstruction capabilities, achieving state-of-the-art performance on the Crafter (Hafner, 2022) benchmark.
Meanwhile, DIAMOND (Alonso et al., 2024) introduced a diffusion- based (Sohl-Dickstedd et al., 2015) world model with EDM decoding (Karras et al., 2022) generate high- quality trajectories and improve the agent performance. The paper demonstrated that diffusion can successfully be ap- plied to world modeling. requiring as few as 3 denoising steps for imagination. The method was evaluated on the Atari 100k benchmark showing better reconstruction quality and improved performance compared to DreamerV3. How- ever, was not applied to a memory-demanding benchmark like Crafter that requires both good memory and visual per- ception to achieve strong results. The authors disclosed that initial experiments on the Crafter benchmark did not achieve good performance due to the absence of an effective
algorithms in a wide range of benchmarks.
Concurrent approaches have proposed to use Transformer- based world model to improve the hidden representations and memory capabilities compared to RNNs. IRIS (Micheli al., 2023) first proposed world model composed of VQ-VAE (Van Den Oord et al., 2017) to convert input im-
3

---

Accurate and Efficient World Modeling with Masked Latent Transformers
Fes1
E
Ge-
Bet
Cete:
Part)
Rree
Ren
Gashed Kays Villun
Temporal Masked Transforee Network
Gushol Koys
and
Terporaal Masked Transformer Network
Lattett
Lanent
Laten
Laten'
Latent
hy
Ap-11
Ap-a
Rs
Eroocer
Decoder
Encookr
Depoder
Encader
Decoder
Frender
Pg-12
6:12
World Model Learning
(b) Actor-Critic Learning
Figure 3: Efficient masked latent Transformer-bant based world model. EMERALD spatial latent state 2ς and temporal hidden state h2 to model the environment accurately and effectively. The world model predictions are made using spatial MaskGIT predictor network to increase decoding speed while maintaining accuracy. Actor-critic learning is performed by imagining trajectories latent space which allows the agents to benefit from world model inner representations.
memory module for the world model. We find that DIA- MOND has difficulties predicting future frames in Crafter, generating hallucinations. In this work, we propose to im- prove performance and training efficiency by learning accurate world model in latent space. Table compares the architectural details of recent model-based approaches with our proposed method.
Inspired by the success of MaskGIT and similar approaches such as draft-and-revie (Lee al., 2022) for image gener- ation, Yan et al. (2023) proposed TECO (Temporally Con- sistent Video Transformer). TECO is video generation model using MaskGIT with draft-and-revise decoding to predict future frames in latent space. They showed that us- ing MaskGit prior allows for not just faster but also higher quality sampling compared to an autoregressive sequential prior. Using pre-trained VQ-GAN (Esser et al., 2021) vec- quantizer, the algorithm demonstrated strong memory capabilities and generation quality. In this work, we propose to apply MaskGIT to model-based reinforcement learning, We train world model that is both accurate and efficient without requiring pre-trained discrete representations.
2.3. Scheduled MaskGIT Predictions in Latent Space
Parallel tokens prediction with refinements was intro- duced by MaskGIT (Masked Generative Image Trans- former) (Chang et al., 2022) as an alternative to sequen- tial decoding for vector quantized image generation (Van Den Oord et al., 2017; Ramesh et al., 2021; Esser et al., 2021). MaskGIT proposed to replace autoregressive de- coding with scheduled parallel decoding of masked tokens to improve generation quality and significantly decrease decoding time. During training, the method samples de- proceeds to mask N HW] tokens where > cos(" r) follows cosine schedule and H and W correspond to the height and width of the latent space in tokens. At decoding time, the model sampled to progressively predict all the tokens. This is done by selecting the most probable tokens and masking the remaining tokens for the next decoding step.
Meanwhile, the use of MaskGIT as prior for model-based reinforcement learning was concurrently explored by GIT- STORM (Meo et al., 2024). GIT-STORM proposed to ap- ply MaskGIT decoding using draft-and-revise strategy to the recently proposed STORM model. However, we note that the motivation behind GIT-STORM is different from our work: Similarly to image and video generation works (Chang et al., 2022; Yan et al., 2023), EMERALD uses MaskGIT as an alternative to sequential token decod- ing in order improve decoding efficiency for spatial la- tent spaces. In contrast, the GIT-STORM paper applied MaskGIT to the vector latent space of STORM in order to
coding times т E [0, 1) uniformly.

---

Accurate and Efficient World Modeling with Masked Latent Transformers
Autoregressive
Sequential Decoding
192
Schaduled Paraliel
Decoding
9.0
1
2.2
2.2
2
2/
1=0
1=1
2=2
1=3
1=4
1=5
1=6
t=7 1
Figure 4: Comparison of scheduled parallel decoding EMERALD vs. autoregressive sequential decoding used by IRIS and A-IRIS. We illustrate the H x W G latent space of our method with only 4 of the 32 groups for better clarity. Sequential decoding predicts one token at a time, significantly impacting efficiency. In contrast, EMERALD uses parallel predictions with scheduled refinements, reducing decoding time while preserving the coherence of predicted tokens.
improve the quality of sampling. Key differences also lie in the architecture of the MaskGIT network. GIT-STORM per- forms attention on the 32 feature tokens the vector latent space while EMERALD first concatenates the tokens along the feature dimension and performs attention on spatial posi- tions similarly to the original MaskGIT paper (Chang et al., 2022).
3.1. World Model Overview
Analogously to the Dreamer algorithm (Hafner et al., 2023) and other recent approaches learning a world model latent space (Robine et al., 2023; Zhang et al., 2024), we encode imput image observations o into hidden representations us- ing convolutional VAE with categorical latents. To balance prediction accuracy and efficiency, our world model adopts carefully designed architecture with spatial latent states 2 and temporal hidden states hy The spatial latent states are first projected into temporal feature vectors for efficient processing by the world model. temporal Transformer network with masked self-attention then generates the tem- poral hidden states he by modeling long-range dependencies. These temporal states subsequently projected back into spatial features for MaskGIT predictions. The trainable components of the world model are the following:
3. Method
We introduce EMERALD, a Transformer model-based re- inforcement learning algorithm using a spatial latent state with MaskGIT predictions generate accurate trajectories in latent space and improve the agent performance. EMER- ALD can be separated into three main components: The Transformer world model that learns to simulate the envi- ronment in latent space. A critic network that Iearns to estimate the sum of future rewards. And an actor network that learns select actions that maximize future rewards estimated by the critic network. Fach component is trained concurrently by sampling trajectories from replay buffer of past experiences. This section describes the architecture and optimization process of our proposed masked latent Transformer world model. Additionally, we provide a brief overview of the imagination process for actor-critic learning. Figure provides a visual overview of world model learn- ing and imagination process during the actor-critic learning phase.
Encoder Network:
~ ( |ot)
Transformer Network:
h2 = fo(211-1,
MaskGIT Predictor:
2. pa pal
l1, mask))
(1)
Decoder Network:
6, ~ Pà h2,21)
Reward Predictor:
A. ps (fr h 2t)
Continue Predictor:
è, ps (c he, 2x)
The spatial latent states and temporal hidden states are com- bined to form the world model state s, = [x,, hy], which integrates both spatial and temporal information. Latent state predictions are made using spatial MaskGIT network,
5

---

Accurate and Efficient World Modeling with Masked Latent Transformers
while separate networks predict the environment reward ét and episode continuation éz
start
1.0
= cos r)
1 step
Mask schedule y(r)
0.8
Encoder Network While previous approaches project encoder features to a compressed vector latent space to simplify world model predictions, our method leverages a spatial latent space to improve world model accuracy and the agent performance. The spatial feature represen- tations are first projected categorical distribution logits E H DDD which are then reorganized into multiple fea- ture groups G for sampling. Instead of representing entire latent vector with a single embedding, the token represen- tations combine embeddings from multiple feature groups. This grouping mechanism was originally used by Dream- e2 (Hafner et al., 2021) to improve training stability using flexible latent representation. We sample discrete stochas- tic states 2, IR xxxxxxxXXD
0.6
2 step
0.4
0.2
3 step
0.0
0.0
0.2
0.6 Decoding time r
0.4
0.8
1.0
Figure 5: Cosine mask schedule. During training, we uni- formly sample decoding times T between and 1. During imagination, the world model samples all masked tokens and refines N HWHWG] tokens with lower probability, We illustrate the schedule used for S = 3 decoding steps.
from the encoder, which serves spatial latent states for our world model.
The use of spatial latent space allows to effectively im- prove the accuracy of the world model without significantly increasing the amount of trainable parameters. The world model stochastic state capacity can also be increased by using a larger number of groups and/or token dimensions. However, this results in significantly larger number of parameters due to linear projections. In contrast, the spatial latent space benefits from weight sharing, which provides useful position bias for the world model.
minimizing the KL divergence with the unmasked latent state as follows:
Lmask = [L[ sg(g5(>; o.)) ( h2
2t mask)
]
(2)
We use the stop gradient operator sg(.) prevent the gradi- ents of targets from being backpropagated.
In order to stabilize learning, the predictor also uses a dy- namics loss with regularization (Hafner et al., 2023). A linear layer is optimized to predict the target latent states from upsampled hidden states. This allows guide the predictions of the MaskGIT predictor in early training. The regularization terms stabilize learning by training the en- coder representations to become more predictable. Both loss terms are scaled with loss weights Bäyn = 0.5 and Breg = 0.1, respectively:
Temporal Masked Transformer Network EMERALD uses Transformer State-Space Model (TSSM) (Chen et al., 2021) to learn long-range memory dependencies and pre- diet future observations more accurately. The Transformer network uses masked self-attention with relative positional encodings (Dai et al., 2019), which simplifies the use of the world model during imagination and evaluation. We also use truncated backpropagation through time (TBTT) (Pašukonis et al., 2023) to preserve information over time. This requires the sampling of training trajectories sequentially in order to access past memory, passing cached attention keys and values from one batch to the next. The Transformer net- work outputs hidden states h, E Txx,, carrying temporal information.
Ldgn Bdgn (L(so(( oc)) || p hb) + Breg KL gp(zz lo)) (g(ps(2 hk))
(3)
Decoder Network The decoder network receives both spatial latent states 2; and temporal hidden states h4 as inputs. This helps the encoder to learn more compressed representations by conditioning on past context. The two representations are projected and concatenated along the channel dimension for decoding. The reconstruction loss Lyee learns hidden representations for the world model by reconstructing input visual observations as follows:
Spatial MaskGIT Predictor The MaskGIT predictor net- work uses a Transformer architecture with spatial attention to model relationships between spatial tokens. After spatial upsampling by the latent decoder network, the temporal hid- den states serve conditioning inputs accurately predict the next spatial tokens. To train the MaskGIT predictor, we uniformly sample decoding times E [0, 1) and use a cosine schedule to mask N = HyWWG tokens from the latent space: zmaas = 24 @ m, with mask m E {0, H
Lioc 11 0z
(4)
H W xG
Reward and Continuation Predictors Lyew and Leon train the world model to predict environment rewards and
The MaskGIT predictor learns to predict masked tokens by
6

---

Accurate and Efficient World Modeling with Masked Latent Transformers
EMERALD
Human Experts
Delta-IRIS
DreamerV3-XL
100
100
Success Rate (%)
Success Rate (%)
75
1
50
0.1
25
0.01
Collect Coal
Collect Diamond
Collect Drink
Collect fron
Collect Sapting
Collect Stone Collect Wood
Defeat Skeleton
Defeat Zombie
Eat Cow
Make Iron Pickaxe
Eat Plant
Make Iron Sword
Make Stone Pickaxe
Make Stone Sword
Make Wood Pickaxe
Make Wood Sword
Place Furnace
Place Plant
Place Stone
Place Table
Wake up
Score
Figure 6: Achievements success rates over 256 evaluation episodes after training 10M environment steps. EMERALD succeeds to solve all achievements at least once. Our method also achieves new state-of-the-an performance, becoming the first method to surpass human experts performance in terms of achievements score.
episode continuation flags, which are used to compute the returns of imagined trajectories during the imagination phase. We adopt the symlog cross-entropy loss from Dream- erV3 (Hafner et al., 2023), which scales and transforms rewards into twohot encoded targets to ensure robust learn- ing across games with different reward magnitudes:
4. Experiments
In this section, we describe our experiments on the Crafter benchmark. We show the results obtained by EMERALD in Table 2. We also perform an ablation study on the world model architecture section 4.3. Finally, we analyze the impact of the number of decoding steps world model predictions in section 4.4. Additional comparison the Atari 100k benchmark with other model-based methods can be found in the appendix A.7.
Trrew SymlogCrossFant(hh,.. rt)
(5a)
Leon BinaryCrossEnt(q cx)
(5b)
Complete Training Objective Given an input batch con- taining B sequences of T image observations 01.,, actions a1.7 rewards T1:7 and episode continuation flags C1.7, the world model parameters (Ø) are optimized to minimize the following loss function:
4.1. Crafter Benchmark
The Crafter benchmark was proposed in Hafner (2022) to evaluate a wide range of general abilities within a single en- vironment. Crafter inspired by the video game Minecraft, It features visual inputs, discrete action space of 17 dif- ferent actions and non-deterministic environment dynamics. The goal of the agent is to achieve as many achievements possible among list of 22 achievements! The bench- mark evaluates the capacity of the agent to use long-term memory and accurate visual perception to collect resources and craft items while surviving in a hostile environment. The agents receive a single reward of +1 for each unlocked achievement. It also perceives a reward of -0.1 for every health point lost and reward of +0.1 for every health point that is regenerated. The benchmark compares the achieve- ments score, which is computed as the geometric mean of success rates for all 22 achievements. This metric prioritizes general agents that unlock wide range of achievements over those that unlock a small number of achievements very frequently. The mean episode return can also be used to
T(O))
Σ t=1
Lmaask + Tedgn + Terewo Trcon + Treee (6)
T
3.2. World Model Imagination Process
The agent critic and actor networks are trained with imagi- nary trajectories generated from the world model. Learning takes place entirely in latent space, which allows the agent to process large batch sizes and increase generalization. We flatten the model states of the sampled sequences along the sam-
batch and time dimensions to generate Bimg = ple trajectories using the world model. The self-attention keys and values features computed during the world model training phase are cached be reused during the agent be- havior learning phase and preserve past context. As shown in Figure 3b, the world model imagines H = 15 steps into the future using the Transformer network and the dynamics network head, selecting actions by sampling from the actor network categorical distribution. We detail the behavior learning process more detail in the appendix A.6.
Score (%) = exp EN In(1 s.)) - where 8; [0; 100) are the N = 22 achievement success rates.

---

Accurate and Efficient World Modeling with Masked Latent Transformers
t=0
3=3
1=4
5
t=6
t=0
t=3
t=4
t=5
f=
S=1
S=3
S=8
2-2
Figure 7: Impact of the number of decoding steps on world model prediction accuracy and efficiency. The randomness of the Crafter environment requires the world model make refinements the predicted tokens avoid the sampling of tokens with contradictory representations. White rectangle boundaries show hallucinations due to incoherent predictions made during imagination. We use S = 3 decoding steps during imagination to achieve good prediction accuracy and efficiency.
directly compare the metric optimized during training.
at least once when evaluating on 256 episodes. The agent progressively learns to master levels of the technology tree, leading to the collection of diamonds.
Table 2: Crafter benchmark comparison (10M env frames). We show achievement scores and environment mean returns aggregated over 5 different seeds. Following Micheli et al. (2024), we also compute FPS as the total number of envi- ronment frames collected divided by the training duration.
4.3. Ablation Study
We study the impact using spatial latent states and a Transformer-baas world model on performance and train- ing efficiency. Table 3 shows the performance obtained by each ablation, applying one modification at time. We first experiment replacing the DreamerV3 vector latent space by a spatial latent space with MaskGIT predictions. We find that the use of a spatial latent state significantly helps to reduce reconstruction error and increase performance. This improves world model accuracy by encoding important environment details for the agent in the latent space. We then compare the use of Transformer-base world model with masked self-attention with the recurrent-based world model of D3eamerV3. We find that attention helps to further improve performance by conditioning representations on a longer context. The capacity of self-attention to model temporal relationships without recurrence makes the Trans- former architecture highly effective at capturing historical context. The use of Transformer-basee architecture also improves training efficiency by not requiring recurrent pro- cessing of hidden states during the world model forward.
Method
Score (%)
Return
#Params FPS
Human Experts
50.5
14.3 + 2.3
DreamerV3 (M)
34.9
14.0 + 0.4
37M
38
D3eamerV3 (XI.)
39.6
15.4 + 0.1
200M
23
A-IRIS
42.5
0.1
25M
12
EMERALD (Ours)
58.1
16.8 + 0.6
30M
27
4.2. Results
Table 2 compares the performance of EMERALD with A- IRIS and DreamerV3 on the Crafter benchmark after train- ing 10M environment steps. We show achievements scores, mean episode returns and total number of parameters. Anal- ogously to Micheli et al. (2024), we also compare the num- ber of collected Frames Per Second (FPS) using single RTX 3090 GPU for training. EMERALD achieves a score of 58,1% being the first algorithm to exceed human experts performance on the benchmark. EMERALD's s enhanced perception enables the agent detect objects and enemies more effectively, leading to increased survival time and im- proved performance. Our method also benefits from reduced training time compared to AIIIS and DreamerV3 XI. The MaskGIT prior reduces decoding time during imagination while preserving the coherence of predicted tokens. Figure 6 compares the success rates of individual achievements. We find that EMERALD successfully solves all achievements
4.4. Choice of the number of decoding steps
While parallel predictions allow the world model to signifi- cantly reduce decoding time, this can also generate unclear predictions with unmatching sets of tokens. The random- ness of the Crafter environment requires the world model to make refinements on the predicted tokens to avoid the sampling of tokens with contradictory representations. To
8

---

Accurate and Efficient World Modeling with Masked Latent Transformers
Table 3: Ablation study on world model architecture and latent space size (10M environment frames), We compare model
performance, training efficiency and reconstruction quality. The performance metrics are aggregated over seeds.
World Model Architecture Latent Space Size
MaskGIT
Score (%)
Return
#Params
FPS
Pixel L2 loss
Dreamer V3 (M) RSSM
32
No
34.9
14.0 + 0.4
37M
38
0.000522
D3eamerV3 (S) RSSM
4 x x 32
Yes
40.4
15.8 + 0.1
28M
20
0.000231
32
No
31.8
13.1 + 0.5
30M
44
0.000890
EMERALD TSSM
4 X 4 X 32
Yes
58.1
16.8 + 0.6
30M
27
0.000241
duration, which impacts training efficiency. We performed corresponding ablation study the impact of the number of imagination decoding steps on final performance over seeds. Table 4 shows the impact on final performance when using different numbers of decoding steps for imagination. We find that the decrease of prediction accuracy has a notice- able impact on final performance. The decrease of average accuracy in world model predictions leads to the generation of less accurate trajectories for the actor and critic networks. Our experiments using and 8 decoding steps achieves higher returns and achievement scores compared using single decoding step simple linear head for prediction.
0.7
80
Rollout duration (second)
Prediction accuracy (%)
0.6
75
0.5
70
0.4
S-16 steps
65
S-8 8 steps
0.3
5=3 steps
0.2
5=1 step
60
No MaskGIT
0.1
2
12 Num pred steps in future
10
14
Table 4: Impact of the number of world model decoding
Figure 8: World model prediction accuracy. We show the to- ken prediction accuracy for world model rollout aggregated over 5 seeds for several numbers of decoding steps.
steps S on final performance and training efficiency.
#Decoding Steps
Score (%)
Return
FPS
No MaskGIT
51.6
16.1 + 0.7
33
S 1 step
53.8
16.1 = 0.5
33
avoid incoherent predictions, we follow MaskGIT by using multiple decoding steps during the imagination phase. Fig- ure 7 shows the impact of the number of decoding steps on prediction quality. We find that using single decoding step, without revising the prediction, can generate blurry predic- tions. These unmatching representations in latent space cause the world model to make incoherent predictions for successive time steps. Using a small number of decoding steps helps to revise the prediction and limit this effect while maintaining good prediction efficiency. Therefore, set
S = steps
58.1
16,8 = 0.6
27
S = 8 steps
55.1
16.5 + 0.6
23
5. Conclusion
We propose EMERALD, a Transformer model-based rein- forcement learning algorithm using spatial latent state with MaskGIT predictions generate accurate trajectories in la- tent space and improve the agent performance. EMERALD achieves new state-of- oftte--f--hhearr performance on the Crafter benchmark with budget of 10M environment steps, becom- ing the first method to exceed human experts performance. We study the impact of latent space size and find that the use of a spatial latent state helps the world model perceive en- vironment details that are crucial to the agent. Additionally, demonstrate that Transformer-baat based architecture out- performs a recurrent-based approach by leveraging attention to model long-range dependencies more effectively. Based on these findings, we hope this work will inspire researchers to further explore the impact of masked latent world models in more complex environments, such as Minecraft.
S = 3 decoding steps in experiments.
Figure 8 shows the average accuracy of predictions for dif- ferent numbers of decoding steps at imaginations time (No MaskGIT designates predictions made by the linear head learned by the dynamics loss Ldya). The accuracy is aver- aged of the 5 EMERALD seeds and computed by comparing the target future states with predicted states during rollout, conditioned the correct sequence of future actions. We find that using less than 3 decoding steps during imagination results in a drop of accuracy. We also compare the rollout time in seconds required to imaginate 15 time steps in the future. Using a larger number of decoding steps can lead to a small increase accuracy but also results in longer rollout

---

Accurate and Efficient World Modeling with Masked Latent Transformers
Acknowledgments
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-12888 2021.
This work was partly supported by The Alexander von Hum- boldt Foundation (AvH).
Impact Statement
Guss, W. H.,, Codel, C., Hofmann, K., Houghton, B., Kuno, N., Milani, S., Mohanty, S., Liebana, D. P., Salakhutdi- nov, R., Topin, N., et al. The minerl 2019 competition on sample efficient reinforcement learning using human priors. arXiv preprint PrXiv:190410007 2019.
The development of autonomous agents for real-world ap- plications introduces various safety and environmental con- cerns. In real-world scenarios, an agent might cause harm to individuals and damage to its surroundings during both training and deployment. Although using world models dur- ing training can mitigate these risks by allowing policies to be learned through simulation, some risks may still persist. This statement aims to inform users of these potential risks and emphasize the importance of Al safety in the application of autonomous agents to real-world scenarios.
Ha, D. and Schmidhuber, J. Recurrent world models facil- itate policy evolution. Advances in Neural Information Processing Systems, 31, 2018.
Hafner, D. Benchmarking the spectrum of agent capabilities. In International Conference on Learning Representations, 2022.
Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D.,
References
Lee, H., and Davidson, J. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2565, PMLR, 2019.
Alonso, E., Jelley, A., Micheli, V., Kanervisto, A., Storkey, A. J., Pearce, T., and Fleuret, F. Diffusion for world mod- eling: Visual details matter in atari. Advances in Neural Information Processing Systems, 37:58757-58791, 2024.
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
Ba, L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint
2016.
Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M.
Hafner, D., Lillicrap. T., Norouzi, M., and Ba, J. Master- ing atari with discrete world models. In International Conference on Learning Representations, 2021.
The arcade learning environment: An evaluation plat- form for general agents. Journal of artificial intelligence research, 47:253-211 279, 2013.
Hafner, D., Pasukonis. J., Ba, J., and Lillicrap. T. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104, 2023.
Bengio, Y., Léonard, N., and Courville, A. Estimating or propagating gradients through stochastic neurons for con- ditional computation. arXiv preprint 1rXiv:130..3433 2013.
Kaiser, L., Babaeizadeh, M., Milos, P., Osinski, B., Camp- bell, R. H., Czechowski, K., Erhan, D., Finn, C., Koza- kowski, P., Levine, S., et al. Model-based reinforcement learning for atari. In International Conference on Learn- ing Representations, 2020.
Chang, H., Zhang, H., Jiang, L., Liu, C., and Freeman, W. T. Maskgit: Masked generative image transformer. In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315-11325, 2022.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion- generative models. Advances in Neural Information Processing Systems, 35: 26565-26577 2022.
Chen, C., Yoon, J., Wu, Y.-F., and Ahn, S. Transdreamer: Reinforcement learning with transformer world models. In Deep RL Workshop NeurIPS 2021, 2021.
Cho, K., van Merrienboer, B., Gulcehre, C., Bougares, F., Schwenk, H., and Bengio, Y. Learning phrase represen- tations using rnn encoder-decoder for statistical machine translation. In Conference on Empirical Methods in Nat- ural Language P Processing (EMNLP 2014), 2014.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International Conference on Learning Repre- sentations, 2014.
LeCun, Y., Bengio, Y., and Hinton, G. Deep learning. nature, 521(7553):436 444, 2015.
Dai, Z., Yang, Z.,, Yang, Y., Carbonell, J. G., Le, Q., and Salakhutdinov, R. Transformer-xl: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computa- tional Linguistics, pp. 2978 2988, 2019.
Lee, D., Kim, C., Kim, S., Cho, M., and HAN, W. S. Draft- and-revise: Effective image generation with contextual rq- transformer. Advances in Neural Information Processing Systems, 35:30127-30138, 2022.
10

---

Accurate and Efficient World Modeling with Masked Latent Transformers
Meo, C., Lica, M., Ikram, Z., Nakano, A., Shah, V., Di- dolkar, A. R., Liu, D., Goyal, A., and Dauwels, J. Masked generative priors improve world models sequence mod- elling capabilities. arXiv preprint arXiv:2410.07836, 2024.
Van Den Oord, A., Vinyals, O., al. Neural discrete rep- resentation learning. Advances in Neural Information Processing Systems, 30, 2017.
Wang. S., Liu, S., Ye, W., You, J., and Gao, Y. Efficientzero v2: Mastering discrete and continuous control with lim- ited data. In International Conference Machine Learn- ing, pp. 51041-51062. PMLR, 2024.
Micheli, V., Alonso, E., and Fleuret, F. Transformers are sample-efficient world models. In International Confer- ence on Learning Representations, 2023.
Williams, R. J. Simple statistical gradient- algo- rithms for connectionist reinforcement learning. Machine learning, 8:229-256, 1992.
Micheli, V., Alonso, E., and Fleuret, F. Efficient world mod- els with context-aware tokenization. In International Con- ference Machine Learning, pp. 2024.
PMLR,
Yan, W., Hafner, D., James, S., and Abbeel, P. Temporally consistent transformers for video generation. In Inter- national Conference on Machine Learning, pp. 39062- 39098. PMLR, 2023.
Pašukonis, J., Lillicrap, T. P., and Hafner, Evaluating long-term memory in 3d mazes. In International Confer- ence on Learning Representations, 2023.
Zhang, W., Wang, G., Sun, J., Yuan, Y., and Huang, G. Storm: Efficient stochastic transformer based world mod- els for reinforcement learning Advances in Neural Infor- mation Processing Systems, 36, 2024.
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Rad- ford, A., Chen, M., and Sutskever, I. Zero-shot text- to-image generation. In International Conference on Machine Learning, pp. 8821-8831. PMLR, 2021.
Robine, J., Höftmann, M., Uelwer, T., and Harmeling, S. Transformer-based world models are happy with 100k interactions. In International Conference on Learning Representations, 2023.
Schritwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D.. Graepel, T., et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839): 604 609, 2020.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprini arXiv: /001 0777770037777 2017.
Schwarzer, M., Ceron, J. S. O., Courville, A., Bellemare, M. G., Agarwal, R., and Castro, P. Bigger, better, faster: Human-level atari with human-level efficiency. In International Conference on Machine Learning, pp. 30365-30380. PMLR, 2023.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International Conference on Machine Learning, pp. 2256-2265. PMLR, 2015.
Sutton, R. S. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 264) 163, 1991.
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.

---

Accurate and Efficient World Modeling with Masked Latent Transformers
A. Appendix
A.I. Latent Space Comparison
(a) DreamerV3's Vector Latent Space
Encodor
Decoder
119
1144
(b) EMERALD's Spatial Latent Space
Figure 9: Comparison between the standard DreamerV3 vector latent space and EMERALD spatial latent space. We illustrate the size of each grouped latent space in tokens, where each token is of dimension D/G. The compressed nature of the DreamerV3 vector latent space E IR (D/(G) can result in the loss of crucial information, negatively impacting the agent's performance. In contrast, the use of spatial latent state 2, E (xx//C) improves world model accuracy and provides further information to the agent.
21
2g
MLP
Spatial MaskGIT
Masking
Latent Dec
Latent Dec
omask
h
h
MLP Head Prediction
(b) MaskGIT Decoding
Figure 10: Comparison between standard MLP head and MaskGIT for decoding. The MLP head is very efficient and samples predictions without refinements, which can lead to ambiguous predictions with tokens having unmatching representations. MaskGIT limits this effect by progressively sampling and refining the prediction to avoid incoherent representations.
12

---

Accurate and Efficient World Modeling with Masked Latent Transformers
A.2. Success Rates Comparison
Table 5: Success rates of each method on the Crafter benchmark (1OM environment steps).
Achievement
Human Experts
DreamerV3 (M)
DreamerV3 (XL)
A-IRIS
EMERALD
Collect Coal
86.0%
77.3%
86.8%
90.6%
91.6%
Collect Diamond
12.0%
0.0%
0.0%
0.0%
0.5%
Collect Drink
92.0%
93.0%
96.2%
98.4%
95.8%
Collect Iron
53.0%
12.1%
32.6%
76.6%
71.2%
Collect Sapling
67.0%
86.3%
97.0%
96.9%
99.9%
Collect Stone
100.0%
95.4%
98.2%
98.4%
97.9%
Collect Wood
100.0%
99.9%
99.9%
100.0%
99.8%
Defeat Skeleton
31.0%
50.5%
74.3%
81.3%
82.3%
Defeat Zombie
84.0%
87.6%
92.4%
96.9%
92.3%
Eat Cow
89.0%
84. 1%
94.9%
98.4%
90.5%
Eat Plant
8.0%
0.4%
0.8%
0.0%
0.1%
Make Iron Pickaxe
26.0%
0.0%
0.0%
0,0%
41.5%
Make Iron Sword
22.0%
0.0%
0.0%
1.6%
41.6%
Make Stone Pickaxe
78.0%
72.0%
84.2%
95.3%
93.7%
Make Stone Sword
78.0%
80.3%
91.0%
95.3%
83.5%
Make Wood Pickaxe
100.0%
96.5%
98.6%
98.4%
99.2%
Make Wood Sword
45.0%
90.0%
96.9%
98.4%
98.8%
Place Furnace
32.0%
90.0%
96.0%
95.3%
96.2%
Place Plant
24.0%
86.3%
97.0%
96.9%
99.9%
Place Stone
90.0%
94.8%
98.2%
98.4%
97.4%
Place Table
100.0%
99.3%
99.9%
98.4%
99.4%
Wake Up
73.0%
92.2%
98.0%
89.1%
92.2%
Score
50.5%
34.9%
39.6%
42.5%
58.1% %
13

---

Accurate and Efficient World Modeling with Masked Latent Transformers
A.3. Crafter Achievement Curves
Reward
Length
Collect Coal
Collect Diamond
1500
1000
10
1
500
0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
1e7
1e7
1e7
1e7
Collect Drink
Collect Iron
Collect Sampling
Collect Stone
15
300
100
200
10-
100
0.0
0.5
0.5
0.0
0.5
1e7
1e7
1e7
1e7
Collect Wood
Defeat Skeleton
Defeat Zombie
Eat Cow
10
10
5
0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
1e7
1e7
1e7
1e7
Eat Plant
Make Iron Pickaxe
Make Iron Sword
Make Stone Pickaxe
10
0.5
0.5
1e7
1e7
1e7
1e7
Make Stone Sword
Make Wood Pickaxe
Make Wood Sword
Place furaae
10-
20
5
5-
10:
0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
1e7
1e7
1e7
1e7
Place Plant
Place Stone
Place Table
Wake Up
10
10
15
10
5-
0
0.5
0.5
0.5
1e7
1e7
1e7
1e7
Figure 11: Achievement counts of EMERALD with shaded min and max.
14

---

Accurate and Efficient World Modeling with Masked Latent Transformers
A.4. Model Architecture
Table 6: Architecture of the encoder and decoder networks. The size of submodules is omitted and can be derived from output shapes. We use layer normalization (LN) and SiLU activation layer inside residual blocks.
Decoder submodules
Output shape
Input latent state (2ç)
x 4 32 x 32
Encoder submodules
Output shape
Reshape (label as x:))
x 4 x 1024
Input image (ot)
64 x 64 x 3
Input hidden state (hz)
512
Conv strided
32 x 32 32
Linear Reshape
4 X 4 x 128
Residual block
32 x 32 32
Conv (label as tbh)
x x 256
Conv strided
16 x 16 x 64
Residual block
16 x 16 64
Concat v2 and Th.
4 4 1280
Conv strided
x 128
Conv
4 X X 256
Residual block
8 x 8 x 128
Residual block
4 X X 256
Conv strided
4 4 x 256
Conv transposed
8 x 8 X 128
Residual block
x 4 x 256
Residual block
8 x x 128
Conv
4 X 4 x 1024
Conv transposed
16 16 x 64
Reshape + Softmax
4 x 4 x 32 x 32
Residual block
16 16 64
Sample One Hot (outputs 24)
4 x x 32 x 32
Conv transposed
32 32
32
Residual block
32
32
32
Conv transposed (outputs o,)
64 x 64
3
Table Transformer network. The latent states 20:T-1 and one-hot encoded actions d0:T-1 E T'A are combined using an action mixer network (Zhang al., 2024). The features are processed by the Transformer network compute temporal hidden states hi.r.
Submodule
Module alias
Output shape
Inputs latent states (20:T-1)
T x 7 32 32
Reshape Conv
Latent Encoder
T x 4 4 128
Flatten Linear
T x 512
Coneat actions (Go:T-1)
T (512 + A)
Linear LN + SiLU
Action Mixer
512
Linear LN
T x 512
Transformer block x Outputs hidden states (h1:T)
Transformer Network
x 512
15

---

Accurate and Efficient World Modeling with Masked Latent Transformers
Table 8: MaskGIT predictor. The latent states are masked and concatenated with spatial hidden states. The features are processed by the Transformer network to predict masked tokens &g
Submodule
Module alias
Output shape
Inputs hidden states (h₂) Linear + Reshape
512
Latent Decoder
4 x 128
Conv
X 256
Concat masked latent states (z+)
X 4 x 1536
Flatten + Linear
Latent Embedding
16 x 256
Add pos embeddings
16 x 256
Transformer block x Kmosk
16 256
Linear
MaskGIT Predictor
16 x 1024
Reshape Softmax
X 4 X 32 x 32
Sample + One Hot (outputs 2,)
x 4 x 32 x 32
Table 9: World model predictors and actor-critic networks architecture. Each network first projects spatial latent states 2, into feature vectors for concatenation with hidden states ht. Each Linear layer is followed by layer normalization and SiLU activation except for the last layer which outputs distribution logits.
Predictor submodules
Output shape
Inputs latent states (2+)
4 x x 32 x 32
Network
Output dim
Output Distribution
Reshape Conv
x 128
Reward predictor
255
Symlog Discrete
Flatten Linear
512
Continue predictor
1
Bernoulli
Concat hidden states (h+)
1024
Critic network
255
Symlog Discrete
Linear LN SiLU
512
Actor network
A
One hot Categorical
Linear + LN + SiLU
512
Linear prediction layer
Output dim
16

---

Accurate and Efficient World Modeling with Masked Latent Transformers
A.5. Model Hyperparameters
Table 10: EMERALD hyper-parameters.
Hyperparameter
Value
General
Image resolution
64 x 64
Batch size (B)
16
Sequence length (T)
64
Optimizer
Adam
Environment parallel instances
16
Collected frames per training Step
16
Replay buffer capacity
1M
World model
Latent space size (H x G)
x
Number categories per group
32
Temporal Transformer blocks
4
Temporal Transformer width
512
Spatial MaskGIT blocks
2
Spatial MaskGIT width
256
Number of attention heads
8
Dropout probability
0.1
Attention context length
64
Learning rate
10-4
Gradient clipping
1000
Actor-critic
Imagination horizon (H)
15
Number of decoding steps (S)
3
Return discount
0.997
Return lambda
0.95
Critic EMA Decay
0,98
Return normalization Momentum
0.99
Actor entropy Scale
3 10-4
Learning rate
3 10-5
Gradient clipping
100
17

---

Accurate and Efficient World Modeling with Masked Latent Transformers
A.6. Actor Critic Learning
Analogously world model predictor networks, the actor and critic networks are designed simple MLPs with projection layer for the spatial latent states. The two network have parameter vectors (0) and (t)), respectively.
Actor Network:
at (a V/ (v,l(t)
Critic Network:
Critic Learning Following DreamerV3, the critic network learns to minimize the symlog cross-entropy loss with discretized X-returns obtained from imagined trajectories with rewards and episode continuation flags predicted by the world model:
R A+t + rd++1 (a1 (8 ARA
+1
Va($H+1)
The critic does not use target network but relies on its own predictions for estimating rewards beyond the prediction horizon. This requires stabilizing the critic by adding regularizing term toward the outputs of its own EMA network Var. Equation defines the critic network loss:
H
Lerstic (1²) = SymlogCrossEin (oenttee,,, (vr, R.) + Crosss (v,, V (s:))
1
Σ
(9)
H
discrete returns regression
critic EMA regulatizer
Actor Learning The actor network learns to select actions that maximize the predicted returns using Reinforce (Williams, 1992) while maximizing the policy entropy to ensure sufficient exploration during both data collection and imagination. The actor network loss is defined as follows:
Lactor (0)
1 H
Σ
sg(A)) log me(ac st) ( st))
(10)
1=1
reinforce
entropy regularizer
Where A (R² Vo(s)))/ max(1, 8) defines advantages computed using normalized returns. The returns scaled using exponentially moving average statistics of their 5th and 95th batch percentiles to ensure stable learning across all Atari games:
= EMA(Per(FF) 95) - Per RR, 5), momentum = =0.99)
(11)
18

---

Accurate and Efficient World Modeling with Masked Latent Transformers
A.7. Atari 100k Benchmark
The commonly used Atari 100k benchmark was proposed in Kaiser et al. (2020) evaluate reinforcement learning agents on Atari games in low data regime. The benchmark includes 26 Atari games with budget of 400k environment frames, amounting to 100k interactions between the agent and the environment using the default action repeat setting. This amount of environment steps corresponds to about two hours (1.85 hours) of real-time play, representing similar amount of time that a human player would need to achieve reasonably good performance.
We evaluate our method on the benchmark to assess EMERALD's performance on environments that do not necessarily require the use of spatial latents to achieve near perfect reconstruction. We also demonstrate improved training efficiency compared Á-IRIS and DIAMOND. Following preceding works,
human-normalizee metries and compare the mean and median returns across all 26 games. The numan-normalized scores are computed for each game using the scores achieved by human player and the scores obtained by random policy: normed score agent score-random score
human score-random
The current state-of sifie-of-theaatee tht-the-rtt art is held by EfficientZero V2 (Wang et al.. 2024), which uses Monte-Cart Tree Search to select the best action at every time step. Another recent notable work is BBF (Schwarzer et al., 2023), a model-free agent using learning techniques that are orthogonal to our work such as periodic network resets and improve performance. In this work, ensure fair comparison, we compare our method with model-based upproaches that do not utilize look-ahead search techniques.
annealing to
We compare EMERALD with Diamond (Alonso et al., 2024), AIIRI (Micheli et al., 2024) and DreamerV3 (Hafner et al., 2023) as well as other mrodll-based approaches like STORM (Zhang et al., 2024), IRIS (Micheli et al., 2023) and SimPLe (Kaiser et al., 2020). We find that EMERALD achieves comparable aggregated performance as A-IRIS and Diamond while offering higher training efficiency. When using a RTX 3090 GPU for training, EMERALD requires only 17 hours to reach 400k environment steps on the Breakout game while A-IRIS and DIAMOND need 27 hours and 75 hours, respectively.
Table 11: Agent scores and Human-normalizze metrics on the 26 games the Atari 100k benchmark. We show average scores over seeds. Bold numbers indicate best performing method for each game.
SimPl.e
D3eamerV3
STORM
IRIS
DIAMOND
EMERALD
Alien
228
7128
617
675
420
959
984
391
744
651
Amidar
1720
74
122
143
139
205
64
226
129
Assault
222
742
527
683
1524
706
801
1123
1526
798
Asterix
8503
1128
1116
854
1028
2492
3698
1045
Bank Heist
14
753
34
467
53
649
641
1148
20
927
Bartle Zone
2360
37188
403
5068
13074
12250
13540
11825
15800
Boxing
12
8
78
70
78
80
70
71
Breakoat
20
31
16
302
133
62
Chopper Command
811
7388
979
1697
1565
420
1888
1183
1370
990
Crazy Climber
10780
35829
62584
71820
59324
97190
66776
57854
99168
75380
Demon Attack
208
165
533
288
Freeway
0
30
17
24
31
0
34
31
33
31
Frustbite
65
279
Gopher
258
2412
597
1675
2236
3730
8240
6445
5898
14702
Hero
1027
30826
2657
7254
7037
11161
11044
7049
5622
7655
James Rond
29
303
100
362
463
445
509
309
427
195
Kangaroo
52
3035
51
1240
838
4098
4208
2269
5382
8780
Krull
1598
6349
7600
Kung Fu Muster
258
22736
14862
24555
21760
21420
26182
21534
18714
22822
999
Pong
-21
13
19
15
18
11
20
20
20
Private Eye
25
69571
35
87
100
882
7781
103
114
100
Qhert
164
13455
1289
3331
746
3405
4522
1444
4499
1245
Road Runner
12
7845
5641
9107
9615
15565
17564
10414
20673
6620
Seaquest
68
42055
661
468
Up N Down
533
11693
3350
15982
3546
7600
7985
4072
3856
$227
Superhuman
N/A
10
10
11
11
Normed Mean (%)
100
33
96
105
112
127
139
146
134
Normed Median (%)
19


Reviews:
[{"review_id": "y4nSLypgir", "reviewer": "Reviewer_mGS9", "summary": "This paper introduces EMERALD, a world modeling approach that balances accuracy and efficiency. EMERALD leverages spatial latent states and MaskGIT-based prediction to generate precise trajectories in the latent space. By improving the perception of critical environmental details, EMERALD enhances the quality of imagined rollouts, ultimately boosting agent performance. Empirical evaluations on the Crafter benchmark demonstrate that EMERALD outperforms existing methods by generating high-fidelity latent trajectories.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper. Please find below our response to the concerns and questions that you raised in the review.\n\n> in the ablation study, while EMERALD's advantage over RSSM-based frameworks does not stem from higher reconstruction fidelity, the paper does not discuss whether it results from reduced dynamics or representation error. For example, does EMERALD exhibit fewer hallucinations over longer imagination rollouts?\n\nThe ablation study in Table 3 shows that the initial performance of world models can be improved by using a spatial latent space and using a Transformer-based world model. On the increase of performance due to architecture change (line 2 VS line 4): Our spatial and temporal TSSM uses self-attention which allows the model to easily perceive past information while the RSSM uses a recurrent state with a limited capacity. To further understand the reason behind the observed results, we performed further studies on the latent state predictions of both world model alternatives. We compared the token prediction accuracy over 5 seeds for both world models when predicting future states. We find that EMERALD achieves an average accuracy of 81.51% correctly predicted tokens for next state prediction against 66.27% when using a RSSM (line 2). We also analyzed attention maps of EMERALD's TSSM and found that the world model learns to attend to all positions seen during training, up to the context of 64 time steps. We conclude that our proposed spatial and temporal TSSM has a positive impact on representation capability and prediction accuracy, which leads to better final performance. Concerning hallucinations, we did not observe an increase of hallucinations  when using a RSSM. However, given the lower token prediction accuracy, we notice that the model can sometimes have difficulties to predict futures that are coherent with past context.\n\n> For instance, in Breakout, the IRIS series, which achieves higher reconstruction accuracy, significantly outperforms other methods, whereas EMERALD does not exhibit a similar advantage.\n\nYes we are aware that $\\Delta$-IRIS achieves strong results on Breakout, the method uses a max pixel reconstruction loss which focuses on reducing the error on the pixel with maximum error. This is very helpful for games like Breakout or Pong where the ball object is crucial for achieving strong results. We observed that performance in Breakout is very linked to the correct reconstruction of the ball. EMERALD facilitates reconstruction and reaches final mean scores superior to 200 for some of the seeds, but we did not use the max pixel loss. \n\n> Providing more detailed experiments on latent trajectory prediction would further strengthen the contribution of this work.\n\nAs explained earlier, we performed further studies on the prediction of the world model in latent spaces. We also computed the average accuracy (%) of predictions for different numbers of decoding steps at imaginations time (No MaskGIT designates predictions made by the MLP head learned by the dynamics loss $L_{dyn}$):\n\n\n**Rebuttal Table 1:**\n| Num pred steps in future: | 1 | 5 | 10 | 15 | Rollout duration (second) |\n|----------|:------------:|:------:|:------:|:------:|:------:|\n| No MaskGIT | 78.57% | 72.22% | 64.80% | 58.59% | 0.10 |\n| S=1 step | 78.52% | 72.19% | 65.75% | 59.92% | 0.14 |\n| S=3 steps | 81.51% | 74.32% | 68.09% | 62.61% | 0.22 |\n| S=8 steps | 82.69% | 75.34% | 68.47% | 62.82% | 0.42 |\n| S=16 steps | 82.80% | 75.56% | 68.51% | 62.64% | 0.73 |\n\nThe accuracy is averaged of the 5 EMERALD seeds and computed by comparing the target future states with predicted states during rollout, conditioned on the correct sequence of future actions. We find that using less than 3 decoding steps during imagination results in a drop of accuracy. We also compare the rollout time in seconds required to imaginate 15 time steps in the future. Using a larger number of decoding steps can lead to a small increase in accuracy but also results in longer rollout duration, which impacts training efficiency. We performed a corresponding ablation to study the impact of the number of imagination decoding steps on final performance over 5 seeds:\n\n**Rebuttal Table 2:**\n| \\#Decoding Steps | Score (\\%) | Return | FPS | \n|----------|:------------:|:------:|:------:|\n| No MaskGIT | 51.6 | 16.1 $\\pm$ 0.7 | 33 |\n| S = 1 step | 53.8 | 16.1 $\\pm$ 0.5 | 33 |\n| S = 3 steps (EMERALD) | 58.1 | 16.8 $\\pm$ 0.6 | 27 |\n| S = 8 steps | 55.1 | 16.5 $\\pm$ 0.6 | 23 | \n\nWe find that the decrease of prediction accuracy has a noticeable impact on final performance. The decrease of average accuracy in world model predictions leads to the generation of less accurate trajectories for the actor and critic networks. Our experiments using 3 and 8 decoding steps achieves higher returns and achievement scores compared to using a single decoding step or a simple MLP head for prediction.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "They may have potential for real-world applications that require high-precision reconstruction and prediction (e.g., robotics, automated exploration).", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "This paper constructs an accurate and efficient world model using a Masked Latent Transformer, advancing the state-of-the-art in Transformer-based world models. Improving prediction accuracy is crucial for generating more realistic imagined rollouts, which in turn facilitates more effective training in imagination. The approach of jointly predicting the next spatial token based on the current spatial latent state and temporal latent state is well-motivated, as it explicitly enhances the world model’s awareness of the current state, leading to more precise predictions. Unlike most Transformer-based world models, EMERALD incorporates temporal hidden states during decoding, rather than relying solely on spatial tokens for reconstruction. This design improves reconstruction accuracy compared to purely spatial-token-based methods. The proposed algorithm significantly outperforms DreamerV3 (RSSM-based) and the IRIS series (Transformer-based) on the Crafter benchmark. Furthermore, the method introduces parallel prediction with scheduled refinements, substantially reducing decoding time while preserving the coherence of predicted tokens. The experimental results on Crafter are compelling.\n\nThe core contribution of this paper lies in generating more precise and efficient latent trajectories, rather than merely improving reconstruction accuracy. However, in the ablation study, while EMERALD's advantage over RSSM-based frameworks does not stem from higher reconstruction fidelity, the paper does not discuss whether it results from reduced dynamics or representation error. For example, does EMERALD exhibit fewer hallucinations over longer imagination rollouts?\n\nAdditionally, in the Atari benchmark experiments (Appendix), we observe that EMERALD's performance is inconsistent in environments where high-precision prediction and reconstruction are critical. For instance, in Breakout, the IRIS series, which achieves higher reconstruction accuracy, significantly outperforms other methods, whereas EMERALD does not exhibit a similar advantage. Providing more detailed experiments on latent trajectory prediction would further strengthen the contribution of this work.", "other_comments_or_suggestions": "No", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JPsqMuHK9W", "reviewer": "Reviewer_cbAq", "summary": "This paper proposes a world model architecture in which spatial latent states are predicted using a MaskGIT predictor. Experiments are conducted on the Crafter benchmark, achieving superhuman performance.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback on our paper. Please find below our response to the concerns and questions that you raised in the review.\n\n> Weakness 1:\n\nThe Crafter benchmark evaluates a wide range of general abilities (survival, memory) and was used by $\\Delta$-IRIS to evaluate its method. It provides an adequate challenge for developing accurate world models with long term memory capacity. For further benchmarking, we performed experiments on the Craftax benchmark (https://openreview.net/forum?id=hg4wXlrQCV). Craftax was proposed at the ICML 2024 conference to provide a more challenging alternative to Crafter. We crop and pad the images to obtain 128x128 pixels inputs. Given the larger resolution, we also add a strided convolution layer to all world models considered. \n\nThe following Table summarizes the results after 10M environment steps over 3 seeds:\n\n| Method | Score (\\%) | Return | \\#Params | FPS |\n|----------|:------------:|:------:|:------:|:------:|\n| DreamerV3 (M) | 2.4 | 13.5 $\\pm$ 1.3 | 37M | 27 |\n| DreamerV3 (XL) | 2.6 | 15.7 $\\pm$ 0.3 | 200M | 18 |\n| EMERALD (Ours) | 3.0 | 16.6 $\\pm$ 0.3 | 30M | 20 |\n\nAs for Crafter, we find that EMERALD achieves faster convergence and better performance compared to DreamerV3. $\\Delta$-IRIS experiments are underway but may not be completed on time given the longer associated training time.\n\n> Weakness 2:\n\nWe observe that the use of MaskGIT as prior for model-based RL was indeed explored concurrently to our work in the GITSTORM paper. GITSTORM was recently peer reviewed at the ICLR 2025 conference but reviewers suggested that the work needed further development before publication. The paper proposed to apply MaskGIT decoding using a draft and revise strategy to the recently proposed STORM model. However, we note that the motivation behind GITSTORM is different from our work: Similarly to image and video generation works, EMERALD uses MaskGIT as an alternative to sequential token decoding in order to improve decoding efficiency. In contrast, the GITSTORM paper applies MaskGIT to the vector latent space of STORM to improve the quality of sampling. Key differences also lie in the architecture of the MaskGIT network. GITSTORM performs attention on G=32 token positions while EMERALD first concatenates the 32 group tokens along the feature dimension and performs attention on HxW=16 spatial positions. \n\nGiven the relation of GITSTORM to our work and despite its recent refusal at the ICLR 2025 conference, we are nevertheless ready to discuss it and highlight the key differences in motivation and architecture with EMERALD in the related works section!\n\n> the proposed spatial latent space structure appears relatively straightforward. \n\nWe propose a novel TSSM network that processes both spatial and temporal information in a carefully designed manner to increase accuracy while maintaining efficiency. EMERALD's TSSM embodies a classical TSSM to process temporal-only relationships and a spatial MaskGIT predictor to process spatial-only relationships. We see our proposed spatial and temporal MaskGIT TSSM as a serious contribution to world modeling that balances prediction accuracy and efficiency.\n\n> Q1:\n\nYes, we explain that training agents directly from pixels prevents the agent from benefiting from inner representations learned by the world model. This requires learning additional image encoders to learn separate representations that may not be as effective for the agent. The self-supervised objectives of the world model learn both compressed representations of observations using the reconstruction loss and memory representations of past observations by predicting future latent states. \n\n> Q2:\n\nWe choose DreamerV3 M to compare with a world model having a similar amount of training parameters. DreamerV3 XL achieves a lower reconstruction error with a L2 error of 0.000359. However, despite the increased number of parameters, DreamerV3 XL still does not achieve better reconstruction than EMERALD and effectuates similar mistakes such as  predicting different blocks or not perceiving mobs.\n\n> Q3:\n\nThe motivation for our paper is the development of an accurate and efficient alternative to recently proposed $\\Delta$-IRIS and DIAMOND, which requires generating trajectories in pixel space. As illustrated in Figure 4, we propose to use a spatial latent space and to replace sequential decoding by MaskGIT decoding to further improve training efficiency. We provide a response to your question and new studies in reviewer mGS9’s dedicated response.\n\n> in Table 3, DreamerV3 with RSSM (Line 1) outperforms TSSM (Line 3).\n\nThe parameter choice of EMERALD is aligned with DreamerV3 Small. For a more accurate comparison, we also performed a 5 seeds experiment for DreamerV3 (S). The agent achieves a return of 11.6 $\\pm$ 0.7 and an achievement score of 22.7, which is lower than (line 3).\n\n> Q5:\n\nThanks for pointing this out! The reference to DreamerV2 will be cited accordingly.", "claims_and_evidence": "Partially. See Q1&2.", "methods_and_evaluation_criteria": "Partially. See Weakness 1.", "theoretical_claims": "N/A. No theoretical results are provided.", "experimental_designs_or_analyses": "Yes. See Q2&3.", "relation_to_broader_scientific_literature": "N/A.", "essential_references_not_discussed": "Yes. A key contribution of this paper is using MaskGIT for latent prior prediction, but this technique was previously introduced in *GITSTORM* (https://arxiv.org/abs/2410.07836), which is not cited or discussed.", "other_strengths_and_weaknesses": "Strengths: \n\n1. Achieving superhuman performance on the Crafter benchmark.\n\nWeaknesses:\n\n1. Limited Benchmarking: The primary experiments focus on Crafter, with only supplementary Atari results in the appendix.\n2. Novelty: As the use of MaskGIT for latent prior prediction was already (or concurrently) explored in GITSTORM, it should be discussed to highlight the difference with EMERALD. Furthermore, the proposed spatial latent space structure appears relatively straightforward. TSSM has already been validated as effective for modeling long-range dependencies in prior work.\n3. Unclear motivation and missing ablation study of MaskGIT design: See Q3.", "other_comments_or_suggestions": "See below.", "questions_for_authors": "1. The paper claims that \"training agents directly from pixels\" prevents the agent from benefiting from inner representations. Could the authors provide clearer evidence for this? In contrast, DIAMOND (a diffusion-based world model for training agents from pixels) currently sets the state-of-the-art on the widely used Atari100k benchmark.\n2. In Figure 2, why is DreamerV3 M used for comparison instead of DreamerV3 XL? Given that DV3 M has a relatively small recurrent state dimension (1024) compared to XL (4096), this could directly impact reconstruction quality, especially in visually demanding environments like Crafter.\n3. EMERALD can better perceive crucial environment details due to its more expressive spatial latent states. However, what is the necessity of the MaskGIT predictor? No ablation study is provided in Section 4.3 to justify its role. Additionally, in Table 3, DreamerV3 with RSSM (Line 1) outperforms TSSM (Line 3). What could be the reason for this?\n4. Writing suggestions:\n   - A figure (or a paragraph) is needed to explicitly highlight the difference between proposed spatial latent spaces ($H\\times W \\times G \\times (D/G)$) and DreamerV2/V3 ($G \\times (D/G)$).\n   - Table 3 should have clearer notation indicating which experiments use MaskGIT prediction.\n5. Equation (3) originates from the KL balancing loss in DreamerV2 and should be properly cited.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9dJD3oDrjz", "reviewer": "Reviewer_d7bP", "summary": "The paper mainly focuses on an approach towards world modeling where the prediction of dynamics is done by a spatial maskGIT. This results in significant improvements on the Crafter benchmark when compared to other models, and performs well on Atari. This also results in improved efficiency over an existing approach aimed at learning strong dynamics by doing reconstruction in the raw pixel space.\n\n## Update after Rebuttal\nThe authors successfully responded to my questions and the questions of the other reviewers. I feel confident in my score of a 4 and in the broader merit of the paper. The issues with the FLOP reporting seem slightly concerning, especially since delta iris has much lower FLOPs per env step than dreamerv3/EMERALD. However, the FLOPs regarding world model rollout now make sense and I trust that the authors will fix this issue for the final manuscript.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper and providing valuable suggestions. Please find below our response to the concerns and questions that you raised in the review.\n\n> Weakness 1\n\nIn Table 3, the lines 2 and 4 compare the use of a RSSM or TSSM for world modeling when using a spatial latent space. Using EMERALD TSSM increases final performance from 15.8 to 16.8. The number of training FPS also increases from 20 to 27. This is because the RSSM is recurrent, which slows down training while the TSSM can process the temporal elements in parallel during the world model learning phase. The use of attention with the TSSM also increases final performance by providing stronger memory representations for the agent. When comparing the accuracy of world model predictions in percent of correctly predicted future tokens, EMERALD reaches an accuracy of 81.51% correctly predicted tokens against 66.27% when using a RSSM. This can be explained by the fact that some future predictions require having access to past information that the RSSM has difficulties to accumulate in the limited GRU recurrent state.\n\n> Weakness 2\n\nWe performed an additional ablation comparing the number of decoding steps used during imagination and the use of a MLP head for prediction. Please refer to Rebuttal Table  1 and Table 2 in reviewer mGS9’s response.\n\nAdditional experiments on the Craftax benchmark reveals that using a simple MLP head instead of MaskGIT leads to the severe accumulation of hallucinations after a certain number of imagination steps. \n\n> Weakness 3\n\nEMERALD achieves better performance on games where tiny details are crucial. This is helpful in games like Pong or Breakout where our method reaches final scores superior to 200 for some of the seeds. In contrast, STORM suffers from higher reconstruction error, impacting results for these games.\n\nConcerning FLOPs (#multiply-and-add operations generated by matrix multiplies), it can be a relevant metric for the efficiency of world models that is indeed hardware agnostic. However, FLOPs do not take into account the ability of the Transformer architecture to process temporal information in parallel while recurrent units require processing each element sequentially during the world model learning phase. We nevertheless compute the amount of FLOPs for each world model using the standard batch size of 16 and sequence length of 64:\n- \\#FLOPs World Model Forward (number of FLOPs to process the observations and predict next state, rewards and episode continuations)\n- \\#FLOPs World Model Imagination (number of FLOPs to imaginate H=15 steps in the future)\n- \\#FLOPs Env step (number of FLOPs to process the observations and sample actions for N=16 parallel environments during exploration)\n\n| Method | \\#FLOPs World Model Forward (Billion) | \\#FLOPs World Model Imagination (Billion) | \\#FLOPs Env Step (Billion) |\n|----------|:------------:|:------:|:------:|\n| DreamerV3 (M) | 189 |  360 | 1.2 |\n| DreamerV3 (XL) | 812 | 2033 | 5.4 |\n| $\\Delta$-IRIS | 699 | 2273 | 0.7 |\n| EMERALD | 275 | 392 | 3.1 |\n\nAs expected, we find that a correlation can be made between FPS recorded during training and number of FLOPs. $\\Delta$-IRIS uses image tokens, which requires encoding observed frames to predict next latent states sequentially. This requires decoding predicted latent states to observations for each time step. The policy also predicts the next actions given reconstructed observations. This greatly increases the amount of FLOPs required for imagination compared to DreamerV3 and EMERALD.\n\nThe number of  RTX 3090 hours on the Atari 100k is provided in appendix A.6 (line 899-900). On the Crafter benchmark, EMERALD takes 100 hours to reach 10M env steps while DreamerV3 (M) and DreamerV3 (XL) take 75 hours and 120 hours, respectively. On the other hand, training one seed of $\\Delta$-IRIS takes 230 hours on one 3090 GPU.\n\n> Figure 3 is not clear as it differs from equations (1) in the world model overview in terms of inputs\nmaskGIT equation in (1) should have z_{t+1} as the prediction? doesnt make sense to predict z_t from z_t?\n\nNo, the equation in (1) is actually correct. The MaskGIT predictor takes the masked $z_{t}$ as input for predicting the unmasked $z_{t}$ target. For better clarity, we can update the formula to specify that the input is masked as $z_{t}^{mask}$.\n\n> in the world model overview it takes as input h and z but figure 3 shows as taking in z and o?\n\nThe maskGIT predictor takes $h_{t}$ and $z_{t}^{mask}$ as input. The dotted arrow in figure 3 actually refers to the use of $h_{t}$ as input for decoder reconstruction, not the input of observations $o_{t}$ to the MaskGIT predictor. \n\n> Is Latent Dec the same as the Decoder?\n\nNo, the decoder maps $h_{t}$ and $z_{t}$ to pixel observations predictions. On the other hand, the latent decoder network projects $h_{t}$ to spatial features for the MaskGIT predictor. You can find the detailed architecture of the two networks in appendix A.3 (Table 5 and Table 7).", "claims_and_evidence": "The claims made in the submission are decently supported by clear and convincing evidence. Although the results on the Crafter benchmark were great, outperforming all other existing world models, they did not generalize very well to other environments, particularly within the Atari 100k benchmark (table 10). Additionally, compared to DreamerV3 medium model, the proposed EMERALD is not more efficient. This is further exacerbated by the reporting of frames rather than FLOPs, which does not make the efficiency results hardware agnostic. Overall, the results do support the idea that EMERALD helps in environments with very difficult to predict dynamics/high dimensional states through great performance in crafter, but it would have been better to see more results in similar environments and more results regarding more interpretable efficiency metrics (FLOPs).", "methods_and_evaluation_criteria": "The crafter and atari evaluation criteria does make sense for the problem application, although the results would be more believable if they were gotten on more high dimensional state space environments.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "I verified the validity of all experimental designs as well as the ablations, I did not find any issues.", "relation_to_broader_scientific_literature": "The key contributions are directly related to Model Based Reinforcement Learning/World Models and their scalability, efficiency, and architectures. The contributions are also related to self-supervised learning in vision, particularly regarding the debate of direct pixel prediction vs latent prediction.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "### Weaknesses\n- Proposed approach is more expensive than dreamerv3 during inference because of maskgit\n\t- Transformer is more expensive than RNN during inference as well, it is unclear how much of a performance boost this causes from the experiments (it seems the results in Table 3 don't isolate this factor)\n        - It's likely a significant amount of the progress is from just using Transformer instead of an RSSM\n- It's not clear to me that its not better than DreamerV3 solely because of a more complex head on top of the dreamerv3 latent predictor. Some ablation experiments trying different latent predictors (i.e. MLP, transformer, etc) would have confirmed that MaskGIT is a good choice.\n- Seems to do worse/similar to other MBRL approaches like STORM on Atari 26\n\t- They need to report the number of RTX hours, and more importantly FLOPs, that EMERALD takes to train compared to other models", "other_comments_or_suggestions": "- line 142 meanwhile should be lowercase\n- 204 transformer world -> transformer world model\n- Figure 3 is not clear as it differs from equations (1) in the world model overview in terms of inputs\n- maskGIT equation in (1) should have z_{t+1} as the prediction? doesnt make sense to predict z_t from z_t?", "questions_for_authors": "It's still not clear how the maskgit portion works--according to equation (1) in the world model overview it takes as input h and z but figure 3 shows as taking in z and o?\n- Is Latent Dec the same as the Decoder?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kLpt2mYlPR", "reviewer": "Reviewer_mkE9", "summary": "This paper proposes EMERALD, a world model that can produce highly accurate rollouts. The architecture is similar to prior works on using transformers as world models, with the exception that it uses MaskGIT to do prediction rather than a naive raster-scan next token prediction scheme. The authors argue that MaskGIT allows the model to learn more accurate predictions as well as improve computational efficiency for rollouts. EMERALD is evaluated on crafter, which is a well known benchmark. The results show that EMERALD can outperform the current state of the art as well as allowing for faster computation.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and valuable returns. Please find below our response to the concerns that you raised in the review.\n\n> The results seem to suggest that EMERALD is only on par with the baseline models. This somewhat muddies the authors' claim that their method is better at prediction. I can believe that because of the simplicity of the Atari environments, current state-of-the-art already 'saturate' the performance. But this result should be more explicitly discussed in the main text to paint a fuller picture of the efficacy of the method.\n\nWe would like to remind that the motivation behind our work is the development of a novel world model that is both accurate and efficient. While EMERALD achieves results comparable with $\\Delta$-IRIS and DIAMOND on Atari 100k, the training time required is significantly reduced. Moreover, contrary to the Crafter benchmark, the Atari 100k benchmark does not require long term memory to solve different games. Only a few history frames are sufficient to achieve strong performance. Many games do not require the use of spatial latents to achieve near perfect reconstruction. But we choose to still evaluate our method on the benchmark for comparison.\n\n\n> I have slight concerns about the novelty of the work. It seems that this paper has taken an architecture developed for image generation and applied it on standard transformer world models. While this does improve the results, I am not entirely sure whether the amount of innovation here would be of much interest to the broader community.\n\nWhile the architecture of EMERALD may appear similar, previous approaches for image and video generation differs in several major points:\n\n- First, Image/Video generation approaches use pre-trained discrete representation (VQ-GAN) to generate tokens in latent space. EMERALD does not require pre-trained representations and learns discrete representations during training.\n\n- Second, TECO learns a second inner VQ-VAE from the pretrained VQ-GAN representations. This is done by encoding pre-trained representations into a further compressed latent space for the world model. Details can be found on the official implementation (https://github.com/wilson1yan/teco/blob/1e92c088965586835005bb1891d4616c2b7bfd5c/teco/models/teco.py#L133).\nIn contrast, EMERALD uses latent encoder/decoder networks to project spatial features to temporal feature vectors.\n\n- Image/Video generation approaches use vector quantized representations while EMERALD uses categorical representations with softmax sampling of the tokens. TECO uses a draft-and-revise decoding approach while we select the tokens with higher confidence at each decoding step.\n\n- Finally, we demonstrate that first: model-based agents can successfully be learned in a spatial latent space using MaskGIT predictions and second: EMERALD improves performance on the visually challenging Crafter benmark and commonly used Atari 100k benchmark.\n\nWe think that our work provides notable contributions to model-based reinforcement learning and that it presents promising results for the use of spatial-temporal world models with maskGIT predictions.", "claims_and_evidence": "The main claim is that using MaskGIT decoding helps with generating more accurate rollouts in less time. These claims are supported by experimental evidence. On crafter, EMERALD supersedes state-of-the-art results with less training time compared to the baselines.", "methods_and_evaluation_criteria": "The proposed method is sensible. Broadly speaking, the method builds on the Dreamer framework with a transformer world model and MaskGIT decoding scheme. All of these components are well established in the literature and combining these is a reasonable idea.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The main experiment results is the success rates of the method across different crafter tasks. In all but one tasks, EMERALD outperforms the baselines and human experts. Overall I find this set of results convincing. The FPS of the world models are also reported and, as the authors claimed, EMERALD can do prediction faster.", "relation_to_broader_scientific_literature": "This work situates within the literature of improving world models and RL for more complex tasks, and pushes the state-of-the-art in this domain. The main innovation here is the application of a more recent image generation method (namely MaskGIT) on existing architectures.", "essential_references_not_discussed": "None that I am aware of.", "other_strengths_and_weaknesses": "Strength: \n- The idea is relatively simple and the result is promising both in terms of performance and computational efficiency.\n- The writing and presentation is in general clear.\n\nWeakness:\n- I have slight concerns about the novelty of the work. It seems that this paper has taken an architecture developed for image generation and applied it on standard transformer world models. While this does improve the results, I am not entirely sure whether the amount of innovation here would be of much interest to the broader community.", "other_comments_or_suggestions": "None.", "questions_for_authors": "None.\n\n## Post rebuttal\nThe rebuttal clarified my questions. I believe the paper can be improved by a longer discussion on the results on the Atari games as well as new experiments during the rebuttal period. However, my concern on the novelty of the method (that the method is a simple combination of MaskedGIT and standard world models) remains. Overall I believe the results presented support the authors' claims. My recommendation for the acceptance of the paper remains.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The proposed method establishes a new state-of-the-art on the challenging Crafter environment. Some reviewers initially criticized the lack of variation in evaluation benchmarks. The authors addressed this and other concerns sufficiently, also providing additional experiments on DMC, where performance of DreamerV3 is roughly matched. 
Since two reviewers gave and maintained an accept rating, while the other two only weakly lean towards rejection, I recommend acceptance.
Decision: Accept (poster)