Accurate and Efficient World Modeling with Masked Latent Transformers

Thank you for your positive feedback and valuable returns. Please find below our response to the concerns that you raised in the review.

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.

We 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.

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.

While the architecture of EMERALD may appear similar, previous approaches for image and video generation differs in several major points:

• 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.

• 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). In contrast, EMERALD uses latent encoder/decoder networks to project spatial features to temporal feature vectors.

• 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.

• 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.

We 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.

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.

Weakness 1

In 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.

Weakness 2

We 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.

Additional 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.

Weakness 3

EMERALD 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.

Concerning 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:

• #FLOPs World Model Forward (number of FLOPs to process the observations and predict next state, rewards and episode continuations)
• #FLOPs World Model Imagination (number of FLOPs to imaginate H=15 steps in the future)
• #FLOPs Env step (number of FLOPs to process the observations and sample actions for N=16 parallel environments during exploration)

As 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.

The 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.

Figure 3 is not clear as it differs from equations (1) in the world model overview in terms of inputs maskGIT equation in (1) should have z_{t+1} as the prediction? doesnt make sense to predict z_t from z_t?

No, 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}$.

in the world model overview it takes as input h and z but figure 3 shows as taking in z and o?

The 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.

Is Latent Dec the same as the Decoder?

No, 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).

Thank you for reviewing our paper. Please find below our response to the concerns and questions that you raised in the review.

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?

The 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.

For instance, in Breakout, the IRIS series, which achieves higher reconstruction accuracy, significantly outperforms other methods, whereas EMERALD does not exhibit a similar advantage.

Yes 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.

Providing more detailed experiments on latent trajectory prediction would further strengthen the contribution of this work.

As 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}$):

Rebuttal Table 1:

The 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:

Rebuttal Table 2:

We 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.

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.

Weakness 1:

The 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.

The following Table summarizes the results after 10M environment steps over 3 seeds:

As 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.

Weakness 2:

We 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.

Given 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!

the proposed spatial latent space structure appears relatively straightforward.

We 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.

Q1:

Yes, 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.

Q2:

We 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.

Q3:

The 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.

in Table 3, DreamerV3 with RSSM (Line 1) outperforms TSSM (Line 3).

The 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).

Q5:

Thanks for pointing this out! The reference to DreamerV2 will be cited accordingly.