Masked Generative Priors Improve World Models Sequence Modelling Capabilities

## Abstract

Abstract Deep Reinforcement Learning (RL) has become the leading approach for creating artificial agents in complex environments. Model-based approaches, which are RL methods with world models that predict environment dynamics, are among the most promising directions for improving data efficiency, forming a critical step toward bridging the gap between research and real-world deployment. In particular, world models enhance sample efficiency by learning in imagination, which involves training a generative sequence model of the environment in a self-supervised manner.
Recently, Masked Generative Modelling has emerged as a more efficient and superior inductive bias for modelling and generating token sequences. Building on the Efficient Stochastic Transformer-based World Models (STORM) architecture, we replace the traditional MLP prior with a Masked Generative Prior (e.g., MaskGIT Prior) and introduce GIT-STORM.
We evaluate our model on two downstream tasks: reinforcement learning and video prediction. GIT-STORM demonstrates substantial performance gains in RL tasks on the Atari 100k benchmark.
Moreover, we apply Transformer-based World Models to continuous action environments for the first time, addressing a significant gap in prior research. To achieve this, we employ a state mixer function that integrates latent state representations with actions, enabling our model to handle continuous control tasks. We validate this approach through qualitative and quantitative analyses on the DeepMind Control Suite, showcasing the effectiveness of Transformer-based World Models in this new domain.
Our results highlight the versatility and efficacy of the MaskGIT dynamics prior, paving the way for more accurate world models and effective RL policies.

## 1 Introduction

Deep Reinforcement Learning (RL) has emerged as the premier method for developing agents capable of navigating complex environments. Deep RL algorithms have demonstrated remarkable performance across a diverse range of games, including arcade games (Mnih et al., ; Schrittwieser et al., ; Hafner et al., ; ), real-time strategy games (Vinyals et al., ; OpenAI, ), board games (Silver et al., ; ; Schrittwieser et al., ), and games with imperfect information (Schmid et al., ). Despite these successes, data efficiency remains a significant challenge, impeding the transition of deep RL agents from research to practical applications. Accelerating agent-environment interactions can mitigate this issue to some extent, but it is often impractical for real-world scenarios. Therefore, enhancing sample efficiency is essential to bridge this gap and enable the deployment of RL agents in real-world applications (Micheli et al., ). (^†^†^1 Delft University of Technology, NL. ^2 Bangladesh University of Engineering and Technology, BD. ^3 The University of Tokyo, JP. ^4 Mila, University of Montreal, CA. ^5 National University of Singapore, SG. Corresponding author: c.meo@tudelft.nl)

Model-based approaches (Sutton & Barto, ) represent one of the most promising avenues for enhancing data efficiency in reinforcement learning.
Specifically, models which learn a “world model” (Ha & Schmidhuber, ) have been shown to be effective in improving sample efficiency.
This involves training a generative model of the environment in a self-supervised manner. These models can generate new trajectories by continuously predicting the next state and reward, enabling the RL algorithm to be trained indefinitely without the need for additional real-world interactions.

Figure: Figure 1: Overview of our proposed GIT-STORM method. (Left) The MaskGIT prior introduced to model the dynamics of the environment. The bidirectional transformer (Devlin et al., ) combines the hidden state given by the autoregressive transformer and the masked posterior $z_{t}\circ m_{t}$ to produce the prior corresponding to the next timestep. (Right) MLP prior originally used in STORM.
Refer to caption: x1.png

However, the effectiveness of RL policies trained in imagination hinges entirely on the accuracy of the learned world model. Therefore, developing architectures capable of handling visually complex and partially observable environments with minimal samples is crucial.
Following Ha & Schmidhuber (), previous methods have employed recurrent neural networks (RNN) to model the dynamics of the environment Hafner et al. (; ; ). However, as RNNs impede parallelized computing due to their recurrent nature, some studies (Micheli et al., ; Robine et al., ; Zhang et al., ) have incorporated autoregressive transformer architectures (Vaswani et al., ) which have been shown to be effective across various domains, such as language (Devlin et al., ; Radford et al., ; Brown et al., ; Raffel et al., ), images (Dosovitskiy et al., ; He et al., ; Liu et al., ), and offline RL (Janner et al., ; Chen et al., ). For instance, IRIS (Micheli et al., ) utilize discrete autoencoders (Oord et al., ) to map raw pixels into a smaller set of image tokens to be used as the input to the world model, achieving superhuman performance in ten different environments of the Atari 100k benchmark (Kaiser et al., ).

**Table 1: Comparison between an MLP prior and a spatial MaskGIT prior for video dynamics using Fréchet Video Distance (FVD).**
|  | FVD $(\downarrow)$ |  |
| --- | --- | --- |
| Method | DMLab | SSv2 |
| TECO w/ MLP prior | 153 | 228 |
| TECO w/ MaskGIT prior | 48 | 199 |

However, autoregressive transformers often suffer from hallucinations (Ji et al., ), where predicted states of the environment are unfeasible, deteriorating the agent’s learning process. Additionally, their unidirectional generation process limits the ability to fully capture global contexts (Lee et al., ). To address these issues, TECO (Yan et al., ) introduces MaskGIT (Chang et al., ) prior $p_{\phi}({z}_{t+1}\mid h_{t})$, using a draft-and-revise algorithm to predict the next discrete representations in the sequence in video generation task. Interestingly, STORM shows that the latent representations $z_{t}$ have the biggest impact on the sequence modelling capabilities of the world model. Moreover, to the best of our knowledge, transformer-based world models have not yet been applied to continuous action environments (e.g., DeepMind Control Suite (DMC) (Kaiser et al., )). The primary challenge lies in the reliance on categorical latent states, which are often ill-suited for representing continuous actions. Addressing this gap is critical for extending the applicability of transformer-based world models to a broader range of tasks.

In this paper, we introduce GIT-STORM, a novel world model inspired by STORM (Zhang et al., ), which leverages the MaskGIT prior to enhance world model sequence modelling capabilities. Building on insights from Yan et al. (), we demonstrate the superior performance of the MaskGIT prior over an MLP prior in predicting video dynamics, as evidenced by results in the DMLab (Beattie et al., ) and SSv2 (Goyal et al., ) datasets (Table 1). Here we summarize the main contributions of this work:

C1: We propose GIT-STORM, a novel world model that enhances STORM (Zhang et al., ) with a MaskGIT prior network for improved sequence modelling. Our model achieves state-of-the-art results on the Atari 100k benchmark, outperforming methods like DreamerV3 (Hafner et al., ) and IRIS, with comprehensive ablation studies showing the impact of discrete representation quality on downstream RL tasks.

C2: We bridge the gap between transformer-based world models and continuous control tasks by using a State Mixer function that effectively combines categorical latent representations with continuous actions, enabling effective learning in continuous action spaces. Through rigorous evaluation on the DMC benchmark, we provide an in-depth analysis of the strengths and limitations of the proposed GIT-STORM model.

This paper marks a key step forward in extending transformer-based world models to more complex and diverse environments.

## 2 Related Works

### 2.1 Model-based RL: World Models

Model-based RL has been a popular paradigm of reinforcement learning. With the advent of neural networks, it has become possible to model high-dimensional state spaces and thus, use model-based RL for environments with high-dimensional observations such as RGB images. In the last few years, based on PlaNet (Hafner et al., ), Hafner et al. proposed the Dreamer series (Hafner et al., ; ; ), a class of algorithms that learn the latent dynamics of the environment using a recurrent state space model (RSSM), while learning behavioral policy in the latent space. Currently, DreamerV3 (Hafner et al., ) has been shown to work across multiple tasks with a single configuration, setting the state-of-the-art across different benchmarks. The actor and critic in DreamerV3 learn from abstract trajectories of representations predicted by the world model.

With the advent of transformers (Vaswani et al., ) in sequence modelling and the promise of scaling performance across multiple tasks with more data, replacing the traditional RSSM backbones with transformer-based backbones has become a very active research direction.
IRIS (Micheli et al., ) was one of the first transformer-based world model approaches, using VQVAE (Oord et al., ) to tokenize the observations and an autoregressive transformer that performs sequence modelling in the token space. Although IRIS obtains human-level performance on 10 out of 26 Atari games, its actor-critic operates in the RGB pixel space, making it almost 14x slower than DreamerV3. In contrast, methods such as TWM (Robine et al., ) and STORM (Zhang et al., ), use latent actor-critic input space. The proposed GIT-STORM employs it as well, as we believe it is the most promising direction to overcome sample efficiency constraints. TWM introduces a balanced sampling scheme and employs TransformerXL (Dai et al., ). More recently, STORM updated DreamerV3 by utilizing the transformer backbone. All aforementioned transformer-based world models use an MLP head to model a dynamics prior which is used to predict the discrete representation of the following timestep. In contrast, introduced by TECO (Yan et al., ), we employ a MaskGIT (Chang et al., ) prior head, which enhances the sequence modelling capabilities of the world model. Table 2 compares various design aspects of different world models. Furthermore, besides STORM, all the mentioned transformer-based world models concatenate the discrete action to the extracted categorical latent representations. As a result, none of these methods is able to handle continuous actions. In contrast, combining latent representations and actions with a state mixer, we successfully train STORM and GIT-STORM on a challenging continuous action environment (i.e., DMC).

**Table 2: Comparison between the proposed GIT-STORM and relevant world models. AC stands for Actor Critic, OneHot for OneHotCategorical.**
| Module | DreamerV3 | IRIS | TWM | STORM | GIT-STORM |
| --- | --- | --- | --- | --- | --- |
| (Hafner et al., ) | (Micheli et al., ) | (Robine et al., ) | (Zhang et al., ) | (ours) |  |
| Latent space | [OneHot, Hidden] | VQ Codes | OneHot | OneHot | OneHot |
| Dynamics Model | RSSM | Transformer | TransformerXL | Transformer | Transformer |
| Dynamics Prior | MLP | MLP | MLP | MLP | MaskGIT |
| AC Input Space | [Latent, Hidden] | RGB | Latent | [Latent, Hidden] | [Latent, Hidden] |
| Experience Sampling | Uniform | Uniform | Balanced | Uniform | Uniform |

### 2.2 Masked Modelling for Visual Representations and Generation

Inspired by the Cloze task (Taylor, ), BERT (Devlin et al., ) proposed a masked language model (MLM) pre-training objective that led to several state-of-the-art results on a wide class of natural language tasks. Following the success of BERT, Masked Autoencoders (MAEs) (He et al., ) learn to reconstruct images with masked patches during the pre-training stage. The learned representations are then used for downstream tasks. Zhang et al. () similarly, improves upon a BERT-like masking objective for its non-autoregressive generation algorithm.

The most relevant to our work is MaskGIT (Chang et al., ), a non-autoregressive decoding approach that consists of a bidirectional transformer model, trained by learning to predict randomly masked visual tokens.
By leveraging a bidirectional transformer (Devlin et al., ), it can better capture the global context across tokens during the sampling process. Furthermore, training on masked token prediction enables efficient, high-quality sampling at a significantly lower cost than autoregressive models.
MaskGIT achieves state-of-the art performance on ImageNet dataset and achieves a 64$\times$ speed-up on autoregressive decoding. The MaskGIT architecture has been applied to various tasks, such as video generation (Yan et al., ; Yu et al., ; ) and multimodal generation (Mizrahi et al., ). For example, Yan et al. () proposes TECO, a latent dynamics video prediction model that uses MaskGIT to model the prior for predicting the next timestep discrete representations, enhancing the sequence modelling of a backbone autoregressive transformer. Inspired by TECO, we adopt the use of MaskGIT prior for the world model, enhancing the sequence modelling capabilities, crucial for enabling and improving the agent policy learning behavior.

Further discussion of related works can be found in Appendix B.

## 3 Method

Following DreamerV3 (Hafner et al., ) and STORM (Zhang et al., ), we define our framework as a partially observable Markov decision process (POMDP) with discrete timesteps, $t\in\mathbb{N}$, scalar rewards, $r_{t}\in\mathbb{R}$, image observations, $o_{t}\in\mathbb{R}^{h\times w\times c}$, and discrete actions. $a_{t}\in\{1,\ldots,m_{a}\}$. These actions are governed by a policy, $a_{t}\sim\pi(a_{t}\mid o_{1:t},a_{1:t-1})$, where $o_{1:t}$ and $a_{1:t-1}$ represent the previous observations and actions up to timesteps $t$ and $t-1$, respectively. The termination of each episode is represented by a Boolean variable, $c_{t}\in\{0,1\}$. The goal is to learn an optimal policy, $\pi$, that maximizes the expected total discounted rewards, $\mathbb{E}_{\pi}\left[\sum_{t=1}^{\infty}\gamma^{t-1}r_{t}\right]$, where $\gamma\in[0,1]$ serves as the discount factor. The learning process involves two parallel iterative phases: learning the observation and dynamics modules (World Model) and optimizing the policy (Agent).

In this section, we first provide an overview of the dynamics module of GIT-STORM. Then, we describe our dynamics prior head of the dynamics module, inspired by MaskGIT (Chang et al., ) (Figure 1). Finally, we explain the imagination phase using GIT-STORM, focusing on the differences between STORM and GIT-STORM.
We follow STORM for the observation module and DreamerV3 for the policy definition, which are described in Appendix A.1 and A.2, respectively.

### 3.1 Overview: Dynamics Module

The dynamics module receives representations from the observation module and learns to predict future representations, rewards, and terminations to enable planning without the usage of the observation module (imagination). We implement the dynamics module as a Transformer State-Space Model (TSSM).
Given latent representations from the observation module, $z_{t}$, and actions, $a_{t}$, the dynamics module predicts hidden states, $h_{t}$, rewards, $\hat{r}_{t}$, and episode termination flags, $\hat{c}_{t}$ $\in\{0,1\}$ as follows,

$$ $\displaystyle\zeta_{t}$ $\displaystyle=g_{\theta}(z_{t},a_{t})$ (State Mixer) $\displaystyle h_{t}$ $\displaystyle=f_{\theta}(\zeta_{1:t})$ (Autoregressive Transformer) $\displaystyle{z}_{t+1}$ $\displaystyle\sim p_{\phi}({z}_{t+1}\mid h_{t})$ (Dynamics Prior Head) (1) $\displaystyle\hat{r}_{t}$ $\displaystyle\sim p_{\phi}(\hat{r}_{t}\mid h_{t})$ (Reward Head) $\displaystyle\hat{c}_{t}$ $\displaystyle\sim p_{\phi}(\hat{c}_{t}\mid h_{t})$ (Termination Head) $$

The world model is optimized to minimize the objective,

$$ $\mathcal{L}(\phi)=\frac{1}{BT}\sum_{n=1}^{B}\sum_{t=1}^{T}\left[\mathcal{L}_{ \text{rew}}(\phi)+\mathcal{L}_{\text{term}}(\phi)+\beta_{1}\mathcal{L}_{\text{ dyn}}(\phi)+\beta_{2}\mathcal{L}_{\text{rep}}(\phi)\right]$ (2) $$

where $\beta_{1},\beta_{2}$ are loss coefficients and $\mathcal{L}_{\text{rew}}(\phi),\mathcal{L}_{\text{term}}(\phi),\mathcal{L}_{
\text{rep}}(\phi),\mathcal{L}_{\text{dyn}}(\phi)$ are reward, termination, representation, and dynamics losses, respectively.
We use the symlog two-hot loss described in  Hafner et al. () as the reward loss. The termination loss is calculated as cross-entropy loss, $c_{t}\log\hat{c}_{t}+(1-c_{t})\log(1-\hat{c}_{t})$. In the following section, we define the dynamics prior in Eq. 1, as well as representation loss, $\mathcal{L}_{\text{rep}}$, and dynamics loss, $\mathcal{L}_{\text{dyn}}$.

### 3.2 Dynamics Prior Head: MaskGIT Prior

Given the expressive power of MaskGIT (Chang et al., ), we propose enhancing the dynamics module in the world model by replacing the current MLP prior with a MaskGIT prior, as shown in Figure 1.
Given the posterior, $z_{t}$, and a randomly generated mask, $m\in\{0,1\}^{N}$ with $M=\lceil\gamma N\rceil$ masked values where $\gamma=\cos\left(\frac{\pi}{2}t\right)$, the MaskGIT prior $p_{\phi}(z_{t+1}\;|\;h_{t})$ is defined as follows.

First, the hidden states, $h_{t}$, are concatenated with the masked latent representations, $z_{t}\circ m_{t}$, where $\circ$ indicates element-wise multiplication. Despite $h_{t}$ being indexed by $t$, it represents the output of the $f_{\theta}$ and thus encapsulates information about the subsequent timestep. Consequently, the concatenation of $z_{t}$ and $h_{t}$ integrates information from both the current and the next timestep, respectively.
A bidirectional transformer is then used to learn the relationships between these two consecutive representations, producing a summary representation, $\xi_{t}$.
Finally, logits are computed as the dot product (denoted as $\odot$ in Figure 1) between the MaskGIT embeddings, which represent the masked tokens, and $\xi_{t}$.
This dot product can be interpreted as a similarity distance between the embeddings and $\xi_{t}$. Indeed, from a geometric perspective, both cosine similarity and the dot product serve as similarity metrics, with cosine similarity focusing on the angle between two vectors, while the dot product accounts for both the angle and the magnitude of the vectors. Therefore, by optimizing the MaskGIT prior, this dot product aligns the embeddings with $\xi_{t}$, thereby facilitating and improving the computation of logits. In contrast, when using the MLP prior, the logits are generated as the output of an MLP that only takes $h_{t}$ as input. This approach requires the model to learn the logits space and their underlying meaning without any inductive bias, making the learning process more challenging.

During training, we follow the KL divergence loss of DreamerV3 (Hafner et al., ), which consists of two KL divergence losses which differ in the stop-gradient operator, $\operatorname{sg}(\cdot)$, and loss scale. We account for the mask tokens in the posterior and define $\mathcal{L_{\mathrm{dyn}}}$ and $\mathcal{L_{\mathrm{rep}}}$ as,

$$ $\displaystyle\mathcal{L_{\mathrm{dyn}}}(\phi)$ $\displaystyle\doteq\max\bigl{(}1,\operatorname{KL}\!\big{[}\operatorname{sg}(q _{\phi}(z_{t}\;|\;x_{t}))\circm_{t}\;\big{\|}\;\hskip 13.77771ptp_{\phi}(z_{t} \;|\;h_{t-1})\hphantom{)}\big{]}\bigr{)}$ (3) $\displaystyle\mathcal{L_{\mathrm{rep}}}(\phi)$ $\displaystyle\doteq\max\bigl{(}1,\operatorname{KL}\!\big{[}\hskip 13.77771ptq_ {\phi}(z_{t}\;|\;x_{t})\hphantom{)}\circm_{t}\;\big{\|}\;\operatorname{sg}(p_{ \phi}(z_{t}\;|\;h_{t-1}))\big{]}\bigr{)}$ (4) $$

where $m_{t}$ is multiplied element-wise with the posterior, eliminating the masked tokens from the loss.

#### Sampling.

During inference, since MaskGIT has been trained to model both unconditional and conditional probabilities, we can sample any subset of tokens per sampling iteration. Following Yan et al. (), we adopt the Draft-and-Revise decoding scheme introduced by Lee et al. () to predict the next latent state (Algorithm 1 and 2).
During the draft phase, we initialize a partition $\bm{\Pi}$ which contains $T_{\text{draft}}$ disjointed mask vectors $\mathbf{m}$ of size $(\text{latent dim}\div T_{\text{draft}})$, which together mask the whole latent representation. Iterating through all mask vectors in $\bm{\Pi}$, the resulting masked representations are concatenated with the hidden states $h_{t}$ from Eq. 1 and fed to the MaskGIT prior head that computes the logits of the tokens correspondent to $h_{t}$ and $\mathbf{m}^{i}$.
Such logits are then used to sample the new tokens that replace the positions masked by $\mathbf{m}^{i}$.
During the revise phase, the whole procedure is repeated $\Gamma$ times.
As a result, when sampling the new tokens, the whole representation is taken into account, resulting in a more consistent and meaningful sampled state.

### 3.3 State Mixer for Continuous Action Environments

When using a TSSM as the dynamics module, the conventional approach has been to concatenate discrete actions with categorical latent representations and feed this sequence into the autoregressive transformer. However, this method is ineffective for continuous actions, as one-hot categorical representations or VQ-codes (Oord et al., ) are poorly suited for representing continuous values. To overcome this limitation, we repurpose the state mixer function $g_{\theta}(\cdot)$ introduced in STORM, which combines the latent representation and the action into a unified mixed representation $\zeta_{t}$. This approach allows for the integration of both continuous and discrete actions with latent representations, enabling the application of TSSMs to environments that require continuous action spaces.

Figure: Algorithm 1 Draft-and-Revise decoding scheme

Figure: Algorithm 2 UPDATE of $\bm{z}$

### 3.4 Imagination Phase

Instead of training the policy by interacting with the environment, model-based approaches use the learned representation of the environment and plan in imagination (Hafner et al., ). This approach allows sample-efficient training of the policy by propagating value gradients through the latent dynamics. The interdependence between the dynamics generated by the world model and agent’s policy makes the quality of the imagination phase crucial for learning a meaningful policy. The imagination phase is composed of two phases, conditioning phase and the imagination one. During the conditioning phase, the discrete representations $z_{t}$ are encoded and fed to the autoregressive transformer. The conditioning phase gives context for the imagination one, using the cached keys and values (Yan et al., ) computed during the conditioning steps.

Differently from STORM, which uses a MLP prior to compute the next timestep representations, we employ MaskGIT to accurately model the dynamics of the environment. By improving the quality of the predicted trajectories, the agent is able to learn a superior policy. Moreover, while STORM applies KV caching only on the conditioning phases, we extend its use to the whole imagination phase as in IRIS (Micheli et al., ), significantly improving the related sample efficiency.

Figure: Figure 2: (Left) Human normalized mean, across the Atari 100k benchmark. GIT-STORM outperforms all other baselines. (Middle) Human normalized median. TWM achieves the highest median value of $51\%$. (Right) IQM. GIT-STORM outperforms all other baselines.
Refer to caption: x2.png

## 4 Experiments

In this section, we analyse the performance of GIT-STORM and its potential limitations by exploring the following questions: (a) How does the MaskGIT Prior affect TSSMs learning behavior and performances on related downstream tasks (e.g., Model-based RL and Video Prediction tasks)? (b) Can Transformer-based world models learn to solve tasks on continuous action environments when using state mixer functions?

### 4.1 Experimental Setup

To evaluate and analyse the proposed method, we consider both discrete and continuous actions environments, namely Atari 100k benchmark (Kaiser et al., ) and DeepMind Control Suite (Tassa et al., ) respectively. On both environments, we conduct both RL and video prediction tasks.

#### Benchmark and baselines.

Atari 100k benchmark consists of 26 different video games with discrete action space. The constraint of 100k interactions corresponds to a total of 400k frames used for training, as frame skipping is set to 4. For RL task on Atari 100k benchmark, we compare GIT-STORM against one model-free method, SimPLe (Kaiser et al., ), one RSSM, DreamerV3 (Hafner et al., ), and three TSSM models (i.e., IRIS (Micheli et al., ), TWM (Robine et al., ), and STORM (Zhang et al., )).
DMC benchmark consists of 18 control tasks with continuous action space. We restrict the models to be trained with only 500k interactions (1M frames) by setting frame skipping to 2. For RL task on DMC benchmark, we compare our model against SAC (Haarnoja et al., ), CURL (Laskin et al., ), DrQ-v2 (Yarats et al., ), PPO (Schulman et al., ), DreamerV3 (Hafner et al., ), and STORM (Zhang et al., ). We trained GIT-STORM on 5 different seeds.
For video prediction tasks, we compare GIT-STORM with STORM only to understand how the MaskGIT Prior affects the visual quality of predicted frames and its influence on the policy training.

Extended details of the baselines for both benchmarks can be found in Appendix I.

#### Evaluation metrics.

Proper evaluation of RL algorithms is known to be difficult due to both the stochasticity and computational requirements of the environments  (Agarwal et al., ). To provide an accurate evaluation of the models, we consider a series of metrics to assess the performances of the considered baselines on across the selected experiments. We report human normalized mean and median as evaluation metrics, aligning with prior literature. We also report interquartile Mean (IQM), Optimality Gap, Performance Profiles (scores distributions), and Probability of Improvement (PI), which provide a statistically grounded perspective on the model evaluation (Agarwal et al., ).

For video prediction task, we report two metrics: Fréchet Video Distance (FVD) (Unterthiner et al., ) to evaluate visual quality of the predicted frames, and perplexity (Jelinek et al., ) measure of the predicted tokens to evaluate the token utilization by the dynamics prior head. We use the trained agent to collect ground truth episodes and use the world model to predict the frames. We report the FVD over 256 videos which are conditioned on the first 8 frames to predict 48 frames.

A full description of these metrics can be found in Appendix J.

### 4.2 Results on Discrete Action Environments: Atari 100k

Figure: Figure 3: Probability of Improvement of the mentioned baselines and GIT-STORM in the Atari 100k benchmark (Left) and DMC benchmark (Right). The results represent how likely it is for GIT-STORM to outperform other baselines.
Refer to caption: x3.png

#### RL task.

Figure 2 summarizes the human normalized mean and median, and IQM score. The full results on individual environments can be found in the Appendix due to space limitations (Table 5). We can see that while TWM and DreamerV3 present a higher human median than GIT-STORM (TWM: $51\%$, DreamerV3: $49\%\rightarrow$ GIT-STORM: $42.6\%$), GIT-STORM dominates in terms of human mean (TWM: $96\%$, DreamerV3: $104\%\rightarrow$ GIT-STORM: $112.6\%$). In terms of IQM, a more robust and statistically meaningful metric, GIT-STORM significantly outperforms the related baselines (DreamerV3: $0.501$, IRIS: $0.502\rightarrow$ GIT-STORM: $0.522$).

Figure 3 (Left) compares PI against the baselines. Noticeably, GIT-STORM presents $PI>0.5$ for all baselines, which indicates that, from a probabilistic perspective GIT-STORM would outperform each baseline on a random task Y from Atari 100k with a probability greater than $0.5$.
Figure 8 illustrates the Optimality Gap, while  Figure 9 presents the fraction of runs with score $>\tau$ for different human normalized scores; both confirm the trends observed so far. Moreover, a closer look to  Table 5 reveals that GIT-STORM presents an optimality gap of $0.500$, marginally beating DreamerV3, which reports $0.503$ and significantly outperforming all other baselines.

#### Video Prediction task.

Table 4 shows video prediction results on selected Atari 100k environments. The table shows that GIT-STORM presents, on average, lower FVD and higher perplexity than STORM (e.g., in Freeway, STORM: 105.45, 33.15 $\rightarrow$ GIT-STORM: 80.33, 67.92, respectively). Figure 5 shows several video prediction results on each environment. For example, on Boxing, we can see that GIT-STORM is able to predict more accurately into the future. The differences in the other two games are smaller, as the player in each game has a much smaller dimension. We think GIT-STORM achieves higher perplexity because the learned agent can collect more diverse episodes.

**Table 3: FVD and perplexity comparisons of STORM and GIT-STORM on selected Atari 100k environments.**
| Game | FVD ($\downarrow$) | Perplexity ($\uparrow$) |  |  |
| --- | --- | --- | --- | --- |
| STORM | GIT-STORM | STORM | GIT-STORM |  |
| Boxing | 1458.32 | 1580.32 | 49.24 | 54.95 |
| Hero | 381.16 | 354.16 | 10.55 | 30.25 |
| Freeway | 105.45 | 80.33 | 33.15 | 67.92 |

### 4.3 Results on Continuous Action Environments: DeepMind Control Suite

#### RL task.

Figure 4 summarizes the human normalized mean and median, and IQM score. The full results on individual environments can be found in the Appendix (Table 6). Although DreamerV3 outperforms all other models on average, Table 6 shows that GIT-STORM presents state-of-the-art scores on two environments, Walker Stand and Quadruped Run. Compared to STORM, GIT-STORM consistently and significantly outperforms across the whole benchmark in terms of human median and mean (STORM: 31.50, 214.50 $\rightarrow$ GIT-STORM: 475.12, 442.10, respectively). For PI, GIT-STORM achieves $PI>0.5$ than STORM, PPO, and SAC (e.g., GIT-STORM: 0.75, 0.60 and 0.63, over STORM, PPO and SAC, respectively) ( Figure 3 (Right)).

#### Video Prediction task.

Table 4 shows video prediction results on selected DMC environments. The table shows that our model achieves lower FVD and higher perplexity than STORM for all environments. The video prediction results in  Figure 6 show that although both models fail to capture the dynamics accurately, GIT-STORM generates marginally better predictions, leading to higher perplexity as well.

Figure: Figure 4: Comparison of human normalized mean (left) and median (right) on DMC benchmark.
Refer to caption: x5.png

Figure: Figure 5: Generated trajectories of STORM and GIT-STORM on selected Atari 100k environments. The model uses the first 8 frames as context and then generates the following 48 frames.
Refer to caption: extracted/5959547/Figures/atari_vis.png

Figure: Figure 6: Generated trajectories of STORM and GIT-STORM on selected DMC environments. The model uses the first 8 frames as context and then generates the following 48 frames.
Refer to caption: extracted/5959547/Figures/dmc_vis_1.png

## 5 Discussion

The proposed GIT-STORM uses a Masked Generative Prior (MaskGIT) to enhance the world model sequence modelling capabilities. Indeed, as discussed in the introduction, high quality and accurate representations are essential to guarantee and enhance agent policy learning in imagination. Remarkably, the proposed GIT-STORM is the only world model, among the ones that use uniform sampling and latent actor critic input space, that is able to achieve non-zero reward on the Freeway environment (e.g., DreamerV3: 0, STORM: 0 $\rightarrow$ GIT-STORM: 13). Indeed, both STORM and IRIS resorted in ad-hoc solutions to get positive rewards, such as changing the sampling temperature (Micheli et al., ) and using demonstration trajectories (Zhang et al., ). Such result, together with the quantitative results on the Atari 100k and DMC benchmarks, clearly answer question (a) -
the presented MaskGIT prior improves the policy learning behavior and performance on downstream tasks (e.g., Model-based RL and Video Prediction) of TSSMs. Moreover, the FVD and perplexity comparisons in Table 4 and Table 4 suggest that GIT-STORM has better predictive capabilities, learns a better dynamics module, and presents more accurate imagined trajectories (Figure 5, Figure 6). Similarly to image synthesis (Chang et al., ) and video prediction (Yu et al., ) tasks, we show how using masked generative modelling is a better inductive bias to model the prior dynamics of discrete representations and improve the downstream usefulness of world models on RL tasks. Furthermore, the MaskGIT Prior can be used in any sequence modelling backbone that uses categorical latent representations (e.g., VideoGPT (Yan et al., ), IRIS (Micheli et al., )), positioning itself as a very versatile approach. In this work we do not apply a MaskGIT prior on top of IRIS only because of computational and time constraints - IRIS requires 168h of training on a V100 GPU for a single run (Zhang et al., ) .

Noticeably, the quantitative results on DMC benchmarks answer question (b) - It is possible to train TSSMs when using a mixer function to combine categorical representations and continuous actions. Indeed, both STORM and GIT-STORM are able to learn meaningful policies within the DMC benchmark. Remarkably, GIT-STORM outperforms STORM with an substantial margin, while using exactly the same policy learning algorithm. Interestingly,  Figure 14 presents an ablation of the used state mixer function, revealing that the overall learning behavior highly depends on the used inductive bias. Surprisingly, the simplest one (e.g., concatenation of $z_{t}$ and $a_{t}$) is the only one that works meaningfully. We leave the exploration of better inductive biases (e.g., imposing specific information bottlenecks (Meo et al., )) to improve the state mixer function as future work.

#### Limitations and Future Work.

The current implementation has been validated on environments that do not require extensive training steps (e.g., ProcGen (Cobbe et al., ), Minecraft (Kanitscheider et al., )) to be trained. We keep as a future work the validation of GIT-STORM on ProcGen and Minecraft environments. As suggested by  Yan et al. (), using a MaskGIT prior could benefit the world model learning behavior in a visually challenging environment like Minecraft. From a technical point of view, one of the main limitations of the proposed world model is that we use only one iteration for the Draft-and-Revise decoding scheme (Lee et al., ). Indeed, while using one iteration speeds up training and evaluation, we do not fully exploit the advantages of this decoding scheme. As a result, in environments like Pong or Breakout, which present small objects (e.g., white or red balls, respectively), using a masked generative approach can lead to filtering such objects out, degrading the downstream performances in these environments. The main reason is that the presented decoding scheme scales exponentially with the number of iterations. We leave as future work the definition of a decoding scheme that scales more efficiently with the number of iterations.

## 6 Conclusion

The motivation for this work stems from the need to improve the quality and accuracy of world models representations in order to enhance agent policy learning in challenging environments. Inspired by (Yan et al., ), we conducted experiments using the TECO framework on video prediction tasks with DMLab and SSv2 (Goyal et al., ) datasets. Replacing an MLP prior with a MaskGIT (Chang et al., ) prior significantly improved the sequence modelling capabilities and the related performance on the video prediction downstream task. Building upon these insights, we proposed GIT-STORM, which employs a MaskGIT Prior to enhance the sequence modelling capabilities of world models, crucial to improve the policy learning behavior (Micheli et al., ). Moreover, through the use of a state mixer function, we successfully combined categorical latent representations with continuous actions, and learned meaningful policies on the related environments. We validated the proposed approach on the Atari 100k and the DMC benchmarks. Our quantitative analysis showed that GIT-STORM on average outperforms all baselines in the Atari 100k benchmark while outperforming STORM with a significant margin on the DMC benchmark. Although our approach does not beat the state-of-the-art in the DMC benchmark, the presented quantitative and qualitative evaluations led to the conclusion that masked generative priors (e.g., MaskGIT Prior) improve world models sequence modelling capabilities and the related downstream usefulness.

## Appendix A GIT-STORM Framework

Figure: Figure 7: GIT-STORM End-to-End pipeline. Similar to STORM (Zhang et al., ), GIT-STORM performs sequence modelling using an autoregressive transformer, which predicts future stochastic latents, $z_{t}$, reward, $r_{t}$ and termination, $c_{t}$. In contrast with STORM, GIT-STORM uses a Masked Generative Prior to model the dynamics of the environment.
Refer to caption: x6.png

Similar to previous TSSM-based world models (Micheli et al., ; Chen et al., ; Zhang et al., ), GIT-STORM consists of a world model with two modules, VAE-based observation module and autoregressive dynamics module, and a policy trained in the latent space. Figure 7 describes the world model architecture of GIT-STORM. In the following sections, we provide details of the observation module and policy.

### A.1 Observation Module

Following STORM, the observation module is a variational autoencoder (VAE) (Kingma & Welling, ), which encodes observations, $o_{t}$, into stochastic latent representations, $z_{t}$, and decodes back the latents to the image space, $\hat{o}_{t}$:

$$ $\displaystyle\text{Observation encoder: }z_{t}\sim q_{\phi}(z_{t}\mid o_{t})$ (5) $\displaystyle\text{Observation decoder: }\hat{o}_{t}=p_{\phi}(z_{t})$ (6) $$

The observations are encoded using a convolutional neural network (CNN) encoder (LeCun et al., ) which outputs the logits used to sample from a categorical distribution. The distribution head applies an unimix function over the computed logits to prevent the probability of selecting any category from being zero (Sullivan et al., ). Since the sampled latents lack gradients, we use the straight-through gradients trick (Bengio et al., ) to preserve them. The decoder, modeled using a CNN, reconstructs the observation from the latents, $z_{t}$. While the encoder is updated using gradients coming from both observation and dynamics modules, the decoder is optimized using only the Mean Squared Error (MSE) between input and reconstructed frames:

$$ $\mathcal{L}_{\text{Observation Model}}=\text{MSE}(o_{t},\hat{o}_{t})$ (7) $$

### A.2 Policy Learning

Following the model-based RL research landscape (DreamerV3; Hafner et al., ) we cast the agent policy learning framework using the actor-critic approach (Mnih et al., ). The agent actor-critic is trained purely from agent state trajectories $s_{t}=[z_{t},h_{t}]$ generated by the world model.
The actor aims to learn a policy that maximizes the predicted sum of rewards and the critic aims to predict the distribution of discounted sum of rewards by the current actor:

$$ $\text{Actor: }a_{t}\sim\pi_{\theta}(a_{t}|s_{t})\ ,\ \text{Critic: }V_{\psi}(s _{t})\approx\mathbb{E}_{\pi_{\theta},p_{\phi}}\left[\sum_{\tau=0}^{\infty} \gamma^{\tau}r_{t+\tau}\right],$ (8) $$

where $\gamma$ is a discount factor.

We follow the setup of STORM (Zhang et al., ) and DreamerV3 (Hafner et al., ) to train the agent. First, a random trajectory is sampled from the replay buffer to compute the initial state of the agent. Then, using the sampled trajectory as context, the world model and actor generate a trajectory of imagined model states, $s_{1:L}$, actions, $a_{1:L}$, rewards, $\hat{r}_{1:L}$, and termination flags, $\hat{c}_{1:L}$, where $L$ is the imagination horizon. To estimate returns that consider rewards beyond the prediction horizon, we compute bootstrapped $\lambda$-returns (Sutton & Barto, ; Hafner et al., ) defined recursively as follows:

$$ $G_{l}^{\lambda}=\hat{r}_{l}+\gamma\hat{c}_{l}\left[(1-\lambda)V_{\psi}(s_{l+1} )+\lambda V_{l+1}^{\lambda}\right]\ ,\ G_{L}^{\lambda}=V_{\psi}(s_{L})$ (9) $$

To stabilize training and prevent the model from overfitting, we regularize the critic towards predicting the exponential moving average (EMA) of its own parameters. The EMA of the critic is updated as,

$$ $\psi_{l+1}^{\text{EMA}}=\sigma\psi_{l}^{\text{EMA}}+(1-\sigma)\psi_{l},$ (10) $$

where $\sigma$ is the decay rate. As a result, the critic learns to predict the distribution of the return estimates using the following maximum likelihood loss:

$$ $\mathcal{L}_{\psi}=\frac{1}{BL}\sum_{n=1}^{B}\sum_{l=1}^{L}\left[(V_{\psi}(s_{ l})-\text{sg}(G_{l}^{\lambda}))^{2}+(V_{\psi}(s_{l})-\text{sg}(V_{\psi^{\text{ EMA}}}(s_{l})))^{2}\right],$ (11) $$

The actor learns to choose actions that maximize return while enhancing exploration using an entropy regularizer (Williams & Peng, ; Hafner et al., ). Reinforce estimator  (Williams, ) is used for actions, resulting in the surrogate loss function:

$$ $\mathcal{L}_{\theta}=\frac{1}{BL}\sum_{n=1}^{B}\sum_{l=1}^{L}\left[-\text{sg} \left(\frac{G_{l}^{\lambda}-V_{\psi}(s_{l})}{\max(1,S)}\right)\ln\pi_{\theta}( a_{l}|s_{l})-\eta H(\pi_{\theta}(a_{l}|s_{l}))\right],\\$ (12) $$

where $\text{sg}(\cdot),H(\cdot)$ are stop gradient operator and entropy, respectively, and $\eta$ is a hyperparameter coefficient of the entropy loss. When training the actor, the rewards are computed between the range from the $5^{\text{th}}$ to the $95^{\text{th}}$ percentile and smoothed out by using an EMA to be robust to outliers. Therefore, the normalization ratio $S$ is,

$$ $S=\text{EMA}(\text{percentile}(G_{l}^{\lambda},95)-\text{percentile}(G_{l}^{ \lambda},5)).$ (13) $$

## Appendix B Extended Related Works: Video Prediction Modelling

Video prediction, a fundamental task in computer vision, aims to generate or predict sequences of future frames based on conditioning past frames. The downstream tasks of video prediction modelling span a wide range of domains, showcasing its significance in different fields, such as autonomous driving (Hu et al., ), robot navigation (DeSouza & Kak, ) controllable animation (Mahapatra & Kulkarni, ), weather forecasting (Bi et al., ; Meo et al., ), and model based reinforcement learning (Hafner et al., ; ; ; ; Zhang et al., ; Micheli et al., ).
Video prediction modelling is known for its sample inefficiency, which poses significant challenges in learning accurate and reliable models in a feasible time (Ming et al., ). To address this, recent advancements have introduced spatio-temporal state space models, which typically consist of a feature extraction component coupled with a dynamics prediction module. These models aim to understand and predict the evolution of video frames by capturing both spatial and temporal relationships. Notable examples include NUWÄ (Wu et al., ) and VideoGPT (Yan et al., ) which respectively use 2D and 3D convolutional layers to extract the latent representations and an autoregressive transformer to perform sequence modelling in the latent space. Moreover, TECO (Yan et al., ) introduces the use of MaskGIT (Chang et al., ) prior to improve the accuracy of the predicted discrete latents and uses a 1D convolution to enhance temporal consistency. Furthermore, VideoPoet (Kondratyuk et al., ), which is able to handle multiple modalities and perform a variety of tasks besides video prediction.

## Appendix C Full Results on RL Task

In this section we report and present the full evaluation and comparison on the two RL benchmark environments, Atari 100k (Kaiser et al., ) and DMC (Tassa et al., ).
Table 5 and Table 6 are the results on Atari 100k and DMC, respectively.

**Table 5: Evaluation on the $26$ games in the Atari $100$k benchmark. We report mean scores as well as aggregated human normalized mean and median, Interquantile Mean (IQM), and Optimality Gap. Following the conventions of Hafner et al. (), scores that are the highest or within $5\%$ of the highest score are highlighted in bold.**
| Game | Rand | Hum | SimPLe | TWM | IRIS | DreamerV3 | STORM | GIT-STORM |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  | reported | reported | reported | reproduced | reproduced | ours |  |
| Alien | 228 | 7128 | 617 | 675 | 420 | 804 | 1364 | 1145 |
| Amidar | 6 | 1720 | 74 | 122 | 143 | 122 | 239 | 181 |
| Assault | 222 | 742 | 527 | 683 | 1524 | 642 | 707 | 967 |
| Asterix | 210 | 8503 | 1128 | 1116 | 854 | 1190 | 865 | 811 |
| Bank Heist | 14 | 753 | 34 | 467 | 53 | 752 | 375 | 503 |
| Battle Zone | 2360 | 37188 | 4031 | 5068 | 13074 | 11600 | 10780 | 9470 |
| Boxing | 0 | 12 | 8 | 78 | 70 | 71 | 80 | 81 |
| Breakout | 2 | 30 | 16 | 20 | 84 | 24 | 12 | 12 |
| Chopper Command | 811 | 7388 | 979 | 1697 | 1565 | 680 | 2293 | 2048 |
| Crazy Climber | 10780 | 35829 | 62584 | 71820 | 59234 | 86000 | 54707 | 55237 |
| Demon Attack | 152 | 1971 | 208 | 350 | 2034 | 203 | 229 | 223 |
| Freeway | 0 | 30 | 17 | 24 | 31 | 0 | 0 | 13 |
| Frostbite | 65 | 4335 | 237 | 1476 | 259 | 1124 | 646 | 582 |
| Gopher | 258 | 2413 | 597 | 1675 | 2236 | 4358 | 2631 | 8562 |
| Hero | 1027 | 30826 | 2657 | 7254 | 7037 | 12070 | 11044 | 13351 |
| Jamesbond | 29 | 303 | 101 | 362 | 463 | 290 | 552 | 471 |
| Kangaroo | 52 | 3035 | 51 | 1240 | 838 | 4080 | 1716 | 1601 |
| Krull | 1598 | 2666 | 2204 | 6349 | 6616 | 7326 | 6869 | 7011 |
| Kung Fu Master | 256 | 22736 | 14862 | 24555 | 21760 | 19100 | 20144 | 24689 |
| Ms Pacman | 307 | 6952 | 1480 | 1588 | 999 | 1370 | 2673 | 1877 |
| Pong | -21 | 15 | 13 | 19 | 15 | 19 | 8 | 6 |
| Private Eye | 25 | 69571 | 35 | 87 | 100 | 140 | 2734 | 2225 |
| Qbert | 164 | 13455 | 1289 | 3331 | 746 | 1875 | 2986 | 3924 |
| Road Runner | 12 | 7845 | 5641 | 9109 | 9615 | 14613 | 12477 | 17449 |
| Seaquest | 68 | 42055 | 683 | 774 | 661 | 571 | 525 | 459 |
| Up N Down | 533 | 11693 | 3350 | 15982 | 3546 | 7274 | 7985 | 10098 |
| Human Mean ($\uparrow$) | 0% | 100% | 33% | 96% | 105% | 104% | 94.7% | 112.6% |
| Human Median ($\uparrow$) | 0% | 100% | 13% | 51% | 29% | 49% | 35.7% | 42.6% |
| IQM ($\uparrow$) | 0.00 | 1.00 | 0.130 | 0.459 | 0.501 | 0.502 | 0.426 | 0.522 |
| Optimality Gap ($\downarrow$) | 1.00 | 0.00 | 0.729 | 0.513 | 0.512 | 0.503 | 0.528 | 0.500 |

**Table 6: Evaluation on the DeepMind Control Suite benchmark. We report scores under visual inputs at 1M frames as well as aggregated human normalized mean and median. Following the conventions of Hafner et al. (), scores that are the highest or within $5\%$ of the highest score are highlighted in bold.**
| Task | SAC | CURL | PPO | DrQ-v2 | DreamerV3 | STORM | GIT-STORM |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Acrobot Swingup | 5.1 | 5.1 | 2.3 | 128.4 | 210.0 | 12.2 | 2.1 |
| Cartpole Balance | 963.1 | 979.0 | 507.3 | 991.5 | 996.4 | 208.9 | 567.0 |
| Cartpole Balance Sparse | 950.8 | 981.0 | 890.4 | 996.2 | 1000.0 | 15.2 | 790.9 |
| Cartpole Swingup | 692.1 | 762.7 | 259.9 | 858.9 | 819.1 | 124.8 | 452.2 |
| Cartpole Swingup Sparse | 154.6 | 236.2 | 0.0 | 706.9 | 792.9 | 0.6 | 97.3 |
| Cheetah Run | 27.2 | 474.3 | 95.5 | 691.0 | 728.7 | 137.7 | 552.5 |
| Cup Catch | 163.9 | 965.5 | 821.4 | 931.8 | 957.1 | 735.5 | 841.5 |
| Finger Spin | 312.2 | 877.1 | 121.4 | 846.7 | 818.5 | 753.8 | 787.0 |
| Finger Turn Easy | 176.7 | 338.0 | 311.0 | 448.4 | 787.7 | 307.3 | 334.1 |
| Finger Turn Hard | 70.5 | 215.6 | 0.0 | 220.0 | 810.8 | 1.4 | 148.6 |
| Hopper Hop | 3.1 | 152.5 | 0.3 | 189.9 | 369.6 | 0.0 | 193.6 |
| Hopper Stand | 5.2 | 786.8 | 6.6 | 893.0 | 900.6 | 0.0 | 664.6 |
| Pendulum Swingup | 560.1 | 376.4 | 5.0 | 839.7 | 806.3 | 0.0 | 0.0 |
| Quadruped Run | 50.5 | 141.5 | 299.7 | 407.0 | 352.3 | 46.2 | 396.6 |
| Quadruped Walk | 49.7 | 123.7 | 107.1 | 660.3 | 352.6 | 55.4 | 445.4 |
| Reacher Easy | 86.5 | 609.3 | 705.8 | 910.2 | 898.9 | 72.7 | 222.4 |
| Reacher Hard | 9.1 | 400.2 | 12.6 | 572.9 | 499.2 | 24.3 | 12.3 |
| Walker Run | 26.9 | 376.2 | 32.7 | 517.1 | 757.8 | 387.2 | 427.6 |
| Walker Stand | 159.3 | 463.5 | 163.8 | 974.1 | 976.7 | 934.8 | 954.8 |
| Walker Walk | 38.9 | 828.8 | 96.0 | 762.9 | 955.8 | 758.0 | 854.7 |
| Median | 78.5 | 431.8 | 101.5 | 734.9 | 808.5 | 31.5 | 475.12 |
| Mean | 225.3 | 504.7 | 211.9 | 677.4 | 739.6 | 214.5 | 442.1 |

Figure: Figure 8: Optimality Gap of the mentioned baselines and GIT-STORM. It shows the amount by each algorithm fails to achieve human performance, where lower values are better.
Refer to caption: x7.png

Figure: Figure 9: Performance profiles of the considered baseline and our proposed model GIT-STORM. The graph presents the fraction of games above a certain human normalized score.
Refer to caption: x8.png

Figure: Figure 10: More generated trajectories of STORM and GIT-STORM on selected DMC environments. The model uses the first 8 frames as context and then generates the following 48 frames.
Refer to caption: extracted/5959547/Figures/dmc_vis_2.png

## Appendix D Training Curves

In this section, we provide the training curves of GIT-STORM for both Atari 100k and DMC benchmark.

Figure: Figure 11: Performance evaluation across all the checkpoints for the Atari 100k benchmark. The solid line represents the average over $5$ seeds while the fill area is defined in terms of maximum and minimum values corresponding to each checkpoint.
Refer to caption: x9.png

Figure: Figure 12: Performance evaluation across all the checkpoints for the DMC benchmark for GIT-STORM and STORM. The solid line represents the average over $5$ seeds while the fill area is defined in terms of standard deviation values corresponding to each checkpoint.
Refer to caption: extracted/5959547/Figures/dmc_benchmark_line.png

## Appendix E Ablation study

### E.1 Dimensions of Dynamic Prior Head

In order to find the best configuration for the MaskGIT prior, we conduct experiments on three different environments with different embedding and vocabulary dimensions corresponding to the bidirectional transformer. While the performance of different configurations varies between environments, we find that a bigger embedding size achieves higher scores on average as seen in Figure 13.

As shown in DreamerV3  (Hafner et al., ), the model achieves better performance as it increases in the number of trainable parameters. Thus, to provide a fair comparison with STORM, we restrict the transformer corresponding to the MaskGIT prior to a similar number of parameters as the MLP prior defined in STORM.

Figure: Figure 13: Different MaskGIT configurations for the Bidirectional Transformer embedding size. Bigger embedding sizes achieve better results. Three different seeds were used for this experiment.
Refer to caption: x10.png

### E.2 VQ-VAE vs One Hot Categorical

The world model state in model-based RL is represented in terms of a latent representation based on raw observations from the environment. However, there is no clear consensus on the representation of the latent space, with SimPLE  (Kaiser et al., ) using a Binary-VAE, IRIS  (Micheli et al., ) using a VQ-VAE while DreamerV3  (Hafner et al., ), STORM  (Zhang et al., ) and TWM  (Robine et al., ) employ a Categorical-VAE.

While recent methods show empirically the advantages of a Categorical-VAE in Atari environments, there is no comprehensive study on different latent space representations. Thus, Table 7 provides a comparison between a VQ-VAE and Categorical-VAE latent representation in the context of GIT-STORM, motivating our choice of latent space. The comparison is performed on three environments with different levels of complexity in terms of visual representations.

**Table 7: Comparison between a VQ-VAE and Categorical-VAE latent representation for the world model state on three Atari 100k environments.**
| Game | VQ-VAE | One Hot Categorical |
| --- | --- | --- |
| Boxing | 0 | 81 |
| Hero | 0 | 13351 |
| MsPacman | 255 | 1877 |

In order to keep the comparison between the two representations accurate, we scale down the VQ-VAE to only $32$ codebook entries, each consisting of $32$ dimensions, matching the size of the one-hot categorical representation of $32$ categories with $32$ classes each. While the VQ-VAE in IRIS  (Micheli et al., ) uses a considerably bigger vocabulary and embedding size, we believe the additional number of parameters introduced provide a biased estimation of the representation capabilities of the latent space. Moreover, we notice that the VQ-VAE approach introduces a significant overhead in terms of training and sampling time. Table 7 shows that the VQ-VAE latent representations collapse and fail to learn a meaningful policy. In contrast the categorical representation achieves impressive results with the same compute budget.

### E.3 State Mixer Ablation

As described in Sec. 3.1, latent representations $z_{t}$ and actions $a_{t}$ are mixed using a state mixer function $g(\cdot)$. To understand the affect of different mixing strategies for the underlying task, we compare three different mixing functions in the DMC benchmark: (1) concatenation, (2) concatenation followed by attention and (3) cross attention between state and actions. Figure 14 illustrates the results. Surprisingly, we find that the simple approach works the best for the tasks – concatenation of state and action significantly outperforms the attention-based approaches in the chosen tasks.

Figure: Figure 14: Comparison between different state and action mixing strategies tested in the DMC environments. We find that simple concatination works the best for the chosen tasks.
Refer to caption: extracted/5959547/Figures/dmc_mixer_benchmark_line.png

## Appendix F Video Prediction Downstream Task: TECO

In order to assess the capabilities of the MaskGIT prior in modelling latent dynamics across different tasks, we consider video generation tasks as a representative study. More specifically, we consider Temporally Consistent Transformer for Video Generation (TECO)  (Yan et al., ) on DeepMind Lab (DMLab) (Beattie et al., ) and Something-Something v.2 (SSv2) (Goyal et al., ) datasets. TECO uses a spatial MaskGIT Prior to generate the state corresponding to the next timestep. Table 1 highlights the importance of the prior network and supports our earlier results on the Atari 100k benchmark. Indeed, when replacing the MaskGit prior network with an MLP one with the same number of parameters, the FVD  (Unterthiner et al., ) on both DMLab and SSv2 datasets significantly increases, going from $48$ to $153$ and from $199$ to $228$ in the DMLab and SSv2 datasets respectively.

## Appendix G Hyperparameters

**Table 8: Hyperparameters regarding the dynamics module, training settings and environment. We use the same hyperparameters as STORM (Zhang et al., ) to focus our experiments on the MaskGIT prior.**
| Hyperparameter | Symbol | Value |
| --- | --- | --- |
| Transformer layers | $K$ | 2 |
| Transformer feature dimension | $D$ | 512 |
| Transformer heads | - | 8 |
| Dropout probability | $p$ | 0.1 |
| World model training batch size | $B_{1}$ | 16 |
| World model training batch length | $T$ | 64 |
| Imagination batch size | $B_{2}$ | 1024 |
| Imagination context length | $C$ | 8 |
| Imagination horizon | $L$ | 16 |
| Update world model every env step | - | 1 |
| Update agent every env step | - | 1 |
| Environment context length | - | 16 |
| Gamma | $\gamma$ | 0.985 |
| Lambda | $\lambda$ | 0.95 |
| Entropy coefficiency | $\eta$ | $3\times 10^{-4}$ |
| Critic EMA decay | $\sigma$ | 0.98 |
| Optimizer | - | Adam (Kingma & Ba, ) |
| World model learning rate | - | $1.0\times 10^{-4}$ |
| World model gradient clipping | - | 1000 |
| Actor-critic learning rate | - | $3.0\times 10^{-5}$ |
| Actor-critic gradient clipping | - | 100 |
| Gray scale input | - | False |
| Frame stacking | - | False |
| Frame skipping | - | 4 (max over last 2 frames) |
| Use of life information | - | True |
| MaskGIT Transformer layers | - | 4 |
| MaskGIT Transformer feature dimension | - | 128 |
| MaskGIT Transformer heads | - | 8 |
| MaskGIT Dropout probability | - | 0.0 |
| Mask Schedule | - | cosine |
| Draft Rounds | $T_{draft}$ | 1 |
| Revise Rounds | $T_{revise}$ | 1 |
| Repetitions | $M$ | 1 |

**Table 9: Specific structure of the image encoder used in GIT-STORM (ours) and STORM (Zhang et al., ). The size of the modules is omitted and can be derived from the shape of the tensors. ReLU refers to the rectified linear units used for activation, while Linear represents a fully-connected layer. Flatten and Reshape operations are employed to alter the indexing method of the tensor while preserving the data and their original order. Conv denotes a CNN layer (LeCun et al., ), characterized by $\text{kernel}=4$, $\text{stride}=2$, and $\text{padding}=1$. BN denotes the batch normalization layer (Ioffe & Szegedy, ).**
| Submodule | Output tensor shape |
| --- | --- |
| Input image ($o_{t}$) | $3\times 64\times 64$ |
| Conv1 + BN1 + ReLU | $32\times 32\times 32$ |
| Conv2 + BN2 + ReLU | $64\times 16\times 16$ |
| Conv3 + BN3 + ReLU | $128\times 8\times 8$ |
| Conv4 + BN4 + ReLU | $256\times 4\times 4$ |
| Flatten | $4096$ |
| Linear | $1024$ |
| Reshape (produce $\mathcal{Z}_{t}$) | $32\times 32$ |

**Table 10: Structure of the image decoder. DeConv denotes a transpose CNN layer (Zeiler et al., ), characterized by $\text{kernel}=4$, $\text{stride}=2$, and $\text{padding}=1$.**
| Submodule | Output tensor shape |
| --- | --- |
| Random sample ($z_{t}$) | $32\times 32$ |
| Flatten | $1024$ |
| Linear + BN0 + ReLU | $4096$ |
| Reshape | $256\times 4\times 4$ |
| DeConv1 + BN1 + ReLU | $128\times 8\times 8$ |
| DeConv2 + BN2 + ReLU | $64\times 16\times 16$ |
| DeConv3 + BN3 + ReLU | $32\times 32\times 32$ |
| DeConv4 (produce $\hat{o}_{t}$) | $3\times 64\times 64$ |

**Table 11: Action mixer $\zeta_{t}=g_{\theta}(z_{t},a_{t})$. Concatenate denotes combining the last dimension of two tensors and merging them into one new tensor. The variable $A$ represents the action dimension, which ranges from $3$ to $18$ across different games. $D$ denotes the feature dimension of the Transformer. LN is an abbreviation for layer normalization (Ba et al., ).**
| Submodule | Output tensor shape |
| --- | --- |
| Random sample ($z_{t}$), Action ($a_{t}$) | $32\times 32$, $A$ |
| Reshape and concatenate | $1024+A$ |
| Linear1 + LN1 + ReLU | $D$ |
| Linear2 + LN2 (output $e_{t}$) | $D$ |

**Table 12: Positional encoding module. $w_{1:T}$ is a learnable parameter matrix with shape $T\times D$, and $T$ refers to the sequence length.**
| Submodule | Output tensor shape |
| --- | --- |
| Input ($e_{1:T}$) | $T\times D$ |
| Add ($e_{1:T}+w_{1:T}$) |  |
| LN |  |

**Table 13: Transformer block. Dropout mechanism (Srivastava et al., ) can prevent overfitting.**
| Submodule | Module alias | Output tensor shape |
| --- | --- | --- |
| Input features (label as $x_{1}$) |  | $T\times D$ |
| Multi-head self attention | MHSA | $T\times D$ |
| Linear1 + Dropout($p$) |  |  |
| Residual (add $x_{1}$) |  |  |
| LN1 (label as $x_{2}$) |  |  |
| Linear2 + ReLU | FFN | $T\times 2D$ |
| Linear3 + Dropout($p$) | $T\times D$ |  |
| Residual (add $x_{2}$) | $T\times D$ |  |
| LN2 | $T\times D$ |  |

**Table 14: MLP settings. A 1-layer MLP corresponds to a fully-connected layer. 255 is the size of the bucket of symlog two-hot loss (Hafner et al., ).**
| Module name | Symbol | MLP layers | Input/ MLP hidden/ Output dimension |
| --- | --- | --- | --- |
| Reward head | $p_{\phi}$ | 3 | D/ D/ 255 |
| Termination head | $p_{\phi}$ | 3 | D/ D/ 1 |
| Policy network | $\pi_{\theta}(a_{t}|s_{t})$ | 3 | D/ D/ A |
| Critic network | $V_{\psi}(s_{t})$ | 3 | D/ D/ 255 |

## Appendix H Computational Resources

Throughout our experiments, we make use of NVIDIA A100 and H100 GPUs for both training and evaluation on an internal cluster, a summary of which can be found in Table 15. For the Atari 100k benchmark, we find that each individual experiment requires around 20 hours to train. For the video prediction tasks, DMLab requires 3 days of training on 4 NVIDIA A100 GPUs. For DMC Vision tasks, we used H100 GPUs to sample from 16 environments concurrently, which reduced our training time to only 8 hours for 1M steps. Compared to this, using A100 for one environment takes 7 days. We acknowledge that the research project required more computing resources than the reported ones, due to preliminary experiments and model development.

**Table 15: Summary of resources used in experiments.**
| Experiment type | GPU Type | # of Days to train |
| --- | --- | --- |
| Atari100k | 1x A100 | 20 hours |
| DMLab | 4x A100 | 3 days |
| DMC Vision | 1x H100 | 8 hours |

## Appendix I Baselines

To assess our approach downstream capabilities on Atari 100k we select the following baselines: SimPLe (Kaiser et al., ) trains a policy using PPO (Schulman et al., ) leveraging a world model represented as an action-conditioned video generation model;
TWM (Robine et al., ) uses a transformer-based world model that leverages a Transformer-XL architecture and a replay buffer which uses a balanced sampling scheme (Dai et al., ). IRIS (Micheli et al., ), that uses a VideoGPT (Yan et al., ) based world model; DreamerV3 (Hafner et al., ), a general algorithm which achieves SOTA results on a multitude of RL benchmarks. Lastly, we consider STORM (Zhang et al., ), an efficient algorithm based on DreamerV3 that uses the transformer architecture for the world model.
Since Hafner et al. () shows that the replay buffer size is a scaling factor, to present a fair comparison we reproduce DreamerV3, which uses a replay buffer of 1M samples by default and full precision variables for the Atari 100k benchmark, using a replay buffer of $100K$ samples and half precision variables, consistent with our approach. Moreover, since STORM does not follow the evaluation protocol proposed in Agarwal et al. (), after setting reproducible seeds, we reproduce STORM on the Atari 100k benchmark using the code released by the authors, and report the results as a result of running the released code.

For DMC Suite, we consider several state-of-the-art algorithms. Soft Actor-Critic (SAC) (Haarnoja et al., ) is a popular algorithm for continuous control tasks, known for its data efficiency due to the use of experience replay. However, SAC often requires careful tuning, particularly for the entropy coefficient, and its performance can degrade when handling high-dimensional input spaces (Hafner, ). Another baseline is Proximal Policy Optimization (PPO) (Schulman et al., ), a widely-used RL algorithm recognized for its robustness and stability across a range of tasks. Additionally, we include DrQ-v2 (Yarats et al., ) and CURL (Laskin et al., ), both of which are tailored for visual environments. These methods leverage data augmentation to improve the robustness of learned policies, making them highly effective in scenarios where pixel-based observations dominate. Finally, we consider DreamerV3, which is the current state-of-the-art in this environment.

## Appendix J Metrics

In order to meaninfgully evaluate the considered baselines we follow the protocol suggested in Agarwal et al. (), which proposes the following metrics for a statistically grounded comparison:

- •
Human Normalized Score: To account for the discrepancies between raw score ranges in Atari games, and at the same time comparing the algorithm’s capabilities with the human benchmark, the human normalized score is used to assess the performance of an algorithm on a specific environment. The Human Normalized Score is defined as $\frac{agent_{score}-random_{score}}{human_{score}-random_{score}}$.
- •
Human Mean: The Human Mean is an aggregate metric used to assess the performance across the whole Atari benchmark. The mean is computed using the Human Normalized Score for each environment, as previously defined.
- •
Human Median: Similar to the Human Mean, the Human Median is an aggregate metric across the Atari benchmark that is insensitive to high-score environments, which instead harm the statistical significance of the Human mean. According to  Agarwal et al. (), both the Human Mean and Human Median are necessary to assess the performance of an algorithm in Atari.
- •
Interquantile Mean (IQM): Interquantile Mean is a popular statistical tool that only considers $50\%$ of the results, effectively ignoring the lowest and highest performing environments. IQM aims to address the shortcomings of the Human Mean by ignoring outliers, while being more statistically significant than the Human Median, which only considers a single value.
- •
Performance profiles (score distributions): Considering the variety of score ranges across different Atari environments, some of which may be heavy-tailed or contain outliers, point or interval estimates provide an incomplete picture with respect to an algorithm’s performance. Performance profiles aim to alleviate this issues by revealing performance variability across tasks more significantly than interval and point estimates, like the Human Mean and Human Median.
- •
Optimality Gap: The Optimality Gap represents another alternative to the Human Mean, and accounts for how much the algorithm fails to meet a minimum Human Normalized Score of $\gamma=1$. The metric considers $\gamma$ as the desirable target and does not account for values greater than it. In the context of the Atari benchmark, $\gamma=1$ represents the human performance. Using the Optimality Gap, the algorithms are compared without taking in consideration super-human performance, which is considered irrelevant.
- •
Probability of Improvement: Instead of treating algorithm’s comparison as a binary decision (better or worse), the Probability of Improvement, indicates a probability corresponding to how likely it is for algorithm $X$ to outperform algorithm $Y$ on a specific task.

For sequence modelling and video prediction task, we use the following metrics:

- •
Perplexity: Perplexity is mathematically defined as the exponentiated average negative log-likelihood of a sequence. Given a sequence of categorical representations $z_{0},z_{1},\dots,z_{t}$, the perplexity of $z$ is computed as:
$\text{PPL}(z)=\exp\left\{-\frac{1}{t}\sum_{i=1}^{t}\log p_{\phi}(z_{i}\mid z_{
<i})\right\}.$
Here, $\log p_{\theta}(z_{i}\mid z_{<i})$ is the log-likelihood of the $i$-th token, conditioned on the preceding tokens $z_{<i}$, according to the model. In this context, perplexity serves as a measure of the model’s ability to predict the tokenized representations of images in a sequence.
- •
Fréchet Video Distance (FVD): Introduced in (Unterthiner et al., ), FVD is a metric designed to evaluate the quality of video generation models. It builds on the idea of the widely-used Fréchet Inception Distance (FID), which is applied to assess the quality of generated images, but extends it to video by incorporating temporal dynamics. FVD is particularly effective for comparing the realism of generated videos with real video data, making it a crucial metric in video prediction and generation tasks.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - Could the authors elaborate on why GIT-STORM occasionally does not surpass STORM and the conditions where improvements are only minor? Understanding this would clarify the contextual efficacy of the MaskGIT prior.
- Regarding the reported state-of-the-art claim, Table 6 suggests that DrQ-v2 outperforms GIT-STORM in some highlighted environments. Could the authors comment why they claim GIT-STORM provides SOTA results on these? It is not the case, right?
- What is the rationale for improving STORM over directly utilizing DreamerV3, which appears to perform better in many scenarios? Or put differently: why would one care to improve STORM with the proposed modifications when there is DreamerV3 and I could just use it or improve over DreamerV3? 

I am open to increase my score once there is clarity on these questions.
soundness: 2
strengths: The empirical evaluation spans discrete and continuous action benchmarks, providing a robust assessment of GIT-STORM’s performance. The reported results demonstrate that GIT-STORM not only improves sample efficiency in RL tasks but also enhances video prediction quality, particularly in the Atari 100k benchmark, aligning well with the study's objectives. Moreover, the paper is well-written with a clear structure, providing a good experience as a reader. Extending the transformer-based world models to continuous action tasks also poses a sufficient novelty and broadens the utility of these models in RL and video prediction applications.
summary: The paper, Masked Generative Priors Improve World Models Sequence Modelling Capabilities, introduces GIT-STORM, an extension of the STORM architecture, incorporating MaskGIT as a dynamics prior to enhance sequence modeling in world models. The authors address two main gaps in previous research: the limitation of transformer-based world models in continuous action environments and the inadequacies of prior methods, like STORM, in capturing effective state representations. Through experiments on Atari 100k (discrete) and DeepMind Control Suite (continuous), GIT-STORM demonstrates improvements in RL and video prediction, suggesting that Masked Generative Priors could be a powerful inductive bias for world models, supporting broader applicability across diverse RL tasks and environments.
weaknesses: It remains unclear why GIT-STORM does not consistently outperform STORM across all benchmarks or why it fails to close the performance gap with DreamerV3 in environments beyond Atari 100k. The paper does not fully explain the conditions under which GIT-STORM’s improvements are more marginal, suggesting a need for clearer insights into the impact of individual architectural components.

The paper claims state-of-the-art results for GIT-STORM on select environments, yet Table 6 seems to indicate that DrQ-v2 outperforms GIT-STORM on two environments (where the authors claim they are better?). Clarifying the conditions under which GIT-STORM achieves these results or adjusting the claim would help ensure consistency and accuracy in presenting the model's achievements.

The proposed approach for handling continuous action spaces is promising, yet lacks a comprehensive empirical analysis. Additional studies on more diverse continuous control tasks could provide stronger validation of the state mixer function's effectiveness and the broader applicability of the model in continuous settings. Most importantly, the modifications from STORM to GIT-STORM are extensive, involving MaskGIT, state mixer, policy adjustments from DreamerV3, and an observation module from STORM. The compounded modifications make it difficult to discern the exact contribution of each component to the reported performance improvements. A more focused ablation study could is required to isolate the impact of each modification.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - Results on the Freeway task have very high variance according to Figure 10. How many out of the five runs does GIT-STORM actually achieve non-zero performance? 
- The most challenging aspect of learning the Freeway task is obtaining the first successful trajectory, which I believe is more related to the exploration strategy than state predictions, given the sparse rewards. How does GIT-STORM assist the agent in exploring more efficiently? Is this strategy stable, or are the successful trajectories obtained by random seeds?
- Why would the pendulum swingup task fail for both STORM and GIT-STORM? DreamerV2, DreamerV3 and TransDreamer can learn this task fairly easily.
- The experiment results in Table 5 and Figure 10 appear inconsistent. For instance, the Gopher score reported in Table 5 is 8562, but the last point in Figure 10 shows a performance of around 2500. Do these two results use different metrics?
- Could you add the learning curves of STORM or DreamerV3 to Figure 10 for a better comparison, considering that you have reproduced these results?
soundness: 2
strengths: - The motivation of incorporating a MaskGIT prior into the STORM architecture is clear.
- The proposed method is straightforward and easy to reproduce.
- MaskGIT can effectively improve the video prediction quality of STORM,  indicating applicability of GIT-STORM to more complicated tasks.
summary: The paper introduces GIT-STORM, which incorporates three modifications to the base algorithm STORM: a MaskGIT prior that replaces the MLP dynamics head, a draft-and-revise decoding scheme for enhanced consistency, and a state mixer for continuous action environments. Experimental results demonstrate that GIT-STORM surpasses STORM on both Atari 100K and DMC benchmarks. Video prediction results indicate that this improvement is attributed to more accurate representations.
weaknesses: - The paper contains a misstatement in its contributions. The authors claim that they "apply transformer-based world models to continuous action environments for the first time". This claim is inaccurate, as TransDreamer[1] can also be applied to continuous action environments. The authors are evidently aware of this paper, given that they have cited it in this work.
- The state-mixer design is not properly addressed. If the authors claim this part of their contribution, they should either elaborate on the design, or provide empirical results to show the superiority of this method. Based on the overlapping tasks, TransDreamer appears to have better performance than GIT-STORM+state-mixer on the continuous control benchmark DMC.
- The experimental results in Atari 100K only demonstrate marginal improvement. The gain over STORM seems to primarily originate from the gopher task alone, which contains inconsistent results, as detailed in the questions section.

[1] Chen et al. TransDreamer: Reinforcement Learning with Transformer World Models.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. On lines 307-309, I think STORM uses KV caching in both the conditioning phase and the imagination phase, see [here](https://github.com/weipu-zhang/STORM/blob/e0b3fd44320d7e213ec905c673ad3f35b61b89f4/sub_models/world_models.py#L363). The `predict_next()` uses `forward_with_kv_cache()` for decoding.

2. Missing comma on line 214?

3. What's new in the proposed state mixer compared to the STORM's action mixer?

4. `Freeway` is a hard exploration environment, as the agent has to repeat the `up` operation many times to get a first reward, which is a rare event for a random policy. Without the first reward, the value space is all zero and the policy would be further optimized as a uniform random policy. STORM, IRIS, and DIAMOND have different tricks that can mitigate such an issue. But what is the underlying reason for GIT-STORM to reach a non-zero result? I think this is not related to the improved decoding or world modelling quality since DreamerV3 and STORM (w/o traj) could also produce a nearly perfect reconstruction and prediction on `Freeway`.

5. For the `Quadruped Run` in Figure 6, I wonder if it's too small (compared to Figure 4 in [DreamerV3](https://arxiv.org/pdf/2301.04104)).

6. Lines 529-530, "Replacing...", the order is reversed.
soundness: 3
strengths: 1. This paper clearly distinguishes itself from previous work, with good comparison and illustration.

2. One-hot categorical latent is widely used in recent model-based RL, yet the research on it is insufficient. This paper provides a novel view of it.

3. This paper bridges the gap of the lack of evaluation of transformer-based world models on continuous control tasks.
summary: This paper proposed to replace the MLP head with MaskGIT prior in STORM to achieve a higher quality of latent generation, and therefore achieve better performance on the Atari100k benchmark.
This paper also bridges the gap of the lack of evaluation of transformer-based world models on continuous control tasks.
weaknesses: 1. The motivation and effect of using MaskGIT head in world models are unclear.
Is there any evidence that the world models would have hallucinations, and how could a MaskGIT head mitigate such issues?
How to distinguish if the improved performance (both RL and FVD) comes from more parameters or the MaskGIT prior?

    There should be some further investigation into the mechanism of the MaskGIT head. Such as:

    (a) What's the difference between the latent variables (or distributions) generated with the MLP head and MaskGIT head?

    (b) This MaskGIT head looks like a skip/residual connection from $z_{t}$ to $\hat{z}_{t+1}$, would this reduce the KL divergence in the training or imagination?

    (c) These are sample questions. An investigation like this would improve the soundness and contribution of this paper.

2. Section 2.1 could be more concise, as these are not quite related to the key contributions and are frequently repeated in each of the model-based RL papers.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: There are also some minor questions:

1. In Line 309, why KV cache can improve sample efficiency? Do you mean computational efficiency?
2. To my knowledge, perplexity is a metric whose lower values mean better. However, in Table 3, higher perplexity is marked as better.
3. In Figure 6, the quadruped agents are too small in the images. This work seems to have used an unusual camera setting for these tasks.

If the authors well address my concerns, I am willing to improve my rating.
soundness: 2
strengths: To my knowledge, MaskGIT models, with their strong expressiveness, are not yet utilized for world models in the MBRL community.
summary: This paper proposes GIT-STORM, which utilizes a MaskGIT model instead of an MLP for the prior head in world models (based on STORM). It also makes minor modifications (a state mixer) to support continuous actions. Experiments are done on Atari100k and DMC benchmarks, considering both policy learning and video prediction performance. GIT-STORM outperforms its base method STORM.
weaknesses: 1. The illustration and descriptions of the model are confusing. Can authors provide more insights for their specific designs?
   - In Figure 1 (left), it seems that GIT-STORM uses masked $z_t$ as inputs for reconstructing $z_{t+1}$. This is strange since, in the original MaskGIT, we mask and reconstruct masked target tokens. Similarly, I think it is more reasonable to mask $z_{t+1}$ as inputs. 
   - In Figure 1 (left), there is no $\xi_t$ but only $\eta_t$. 
   - Also, the dot product seems to be a commonly used trick that ties weights for embedding and linear layer before Softmax. If so, relevant literature should be cited.
   - The Draft-and-Revise decoding scheme, if not proposed by this work, should be moved into a preliminary section. 
2. The contribution to supporting continuous actions is overclaimed (as 'for the first time'). In fact, concatenating or summating continuous inputs with hidden states is a too straightforward approach in current VLA models (e.g., OpenVLA for inputting continuous visual representations) and action-conditioned video prediction models (e.g., iVideoGPT for inputting continuous actions).
3. The performance of GIT-STORM on DMC is outperformed by its base method, DreamerV3.

Meta-review: While GIT-STORM shows some promising results in improving world model sequence modeling through masked generative priors, several critical concerns remain insufficiently addressed. The authors' claim of being the first to apply transformer-based world models to continuous actions was overstated, given the existence of TransDreamer. Additionally, there is concerning ambiguity in the MaskGIT head design and sampling process, particularly regarding the gap between training and inference time behavior. The performance gains over STORM appear modest and heavily depend on specific environments, while being still outperformed by DreamerV3 on DMC tasks. Though the authors provided detailed responses, their explanations remain unconvincing, suggesting the work needs further development before publication. I recommend rejecting this paper.

Decision: Reject