Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient

## Abstract

Abstract Model-based reinforcement learning (RL) offers a solution to the data inefficiency that plagues most model-free RL algorithms. However, learning a robust world model often demands complex and deep architectures, which are expensive to compute and train. Within the world model, dynamics models are particularly crucial for accurate predictions, and various dynamics-model architectures have been explored, each with its own set of challenges. Currently, recurrent neural network (RNN) based world models face issues such as vanishing gradients and difficulty in capturing long-term dependencies effectively. In contrast, use of transformers suffers from the well-known issues of self-attention mechanisms, where both memory and computational complexity scale as O ​ ( n 2 ) O(n^{2}) , with n n representing the sequence length. To address these challenges we propose a state space model (SSM) based world model, specifically based on Mamba, that achieves O ​ ( n ) O(n) memory and computational complexity while effectively capturing long-term dependencies and facilitating the use of longer training sequences efficiently. We also introduce a novel sampling method to mitigate the suboptimality caused by an incorrect world model in the early stages of training, combining it with the aforementioned technique to achieve a normalised score comparable to other state-of-the-art model-based RL algorithms using only a 7 million trainable parameter world model. This model is accessible and can be trained on an off-the-shelf laptop. Our code is available at https://github.com/realwenlongwang/drama.git .

## 1 Introduction

Deep Reinforcement Learning (RL) has achieved remarkable success in various challenging applications, such as Go (Silver et al., , ), Dota (Berner et al., ), Atari (Mnih et al., ; Schrittwieser et al., ), and MuJoCo (Schulman et al., ; Haarnoja et al., ). However, training policies capable of solving complex tasks often requires millions of interactions, which can be impractical and poses a barrier to real-world applications. Thus, improving sample efficiency has become one of the most critical goals in RL algorithm development.

World models have shown promise in improving sample efficiency through an auto-generative process that produces artificial samples on which to train RL agents, a method referred to as model-based RL (Micheli et al., ; Robine et al., ; Zhang et al., ; Hafner et al., ). In this approach, interaction data is used to learn the environment dynamics using a sequence model, allowing the agent to train on artificial experiences generated by the resulting dynamics model instead of relying on real-world interactions. This approach shifts the problem from improving the policy directly using real samples (which is sample inefficient) to improving the accuracy of the world model to match the real environment (which is more sample efficient). However, model-based RL faces a well-known challenge: when the model is inaccurate due to limited observed samples, especially early in training, the learned policy can become biased towards suboptimal behaviour, and detecting model errors is difficult, if not impossible.

In sequence modelling, linear complexity is highly desirable because it allows models to efficiently process longer sequences without a dramatic increase in computational and memory resources. This is particularly important when training world models, which require efficient sequence modelling to simulate complex environments over long time horizons. Recurrent Neural Networks (RNNs), particularly advanced variants like Long Short-Term Memory (LSTM) and Gated Recurrent Units (GRU), offer linear complexity, making them computationally attractive for this task. However, RNNs still struggle with vanishing gradient issues and are inefficient in capturing long-term dependencies (Hafner et al., , ). More recently, transformer architectures, which have dominated natural language processing (Vaswani et al., ), quickly gained widespread popularity in fields such as image processing and offline RL following groundbreaking work in these areas (Dosovitskiy et al., ; Chen et al., ). The transformer structure has demonstrated its effectiveness in model-based RL as well (Micheli et al., ; Robine et al., ; Zhang et al., ). However, transformers suffer from both memory and computation complexity that scale as $O(n^{2})$, where $n$ is the sequence length, creating challenges for world models that require long sequences to simulate complex environments.

Currently, State Space Models (SSMs) are attracting significant attention for their ability to efficiently handle long-sequence problems with linear complexity. Among SSMs, Mamba has emerged as a strong competitor to transformer-based architectures in various fields, including natural language processing (Gu and Dao, ; Dao and Gu, ), computer vision (Zhu et al., ), and offline RL (Ota, ). Applying Mamba’s architecture to model-based RL is particularly appealing due to its linear memory and computational scaling with sequence length, while also effectively capturing long-term dependencies. Moreover, efficiently capturing environmental dynamics can reduce the likelihood that the behaviour policy is learned within an inaccurate world model, which we also address by incorporating a novel dynamic frequency-based sampling method.

In this paper, we make three key contributions:

- •
We introduce DRAMA, the first model-based RL agent built on the Mamba SSM, with Mamba-2 as the core of its architecture. We evaluate DRAMA on the Atari100k benchmark, demonstrating that it achieves performance comparable to other state-of-the-art algorithms while using only a 7 millions trainable parameter world model.
- •
Additionally, we compare the performance of Mamba-1 and Mamba-2, demonstrating that Mamba-2 achieves superior results as a dynamics model in the Atari100k benchmarks, despite it slightly limiting expressive power in order to enhance training efficiency.
- •
Finally, we propose a novel but straightforward sampling method, i.e., dynamic frequency-based sampling (DFS) to mitigate the challenges posed by imperfect dynamics models.

Figure: Figure 1: Drama world model architecture. Starting from sequence index $i$, the raw gaming frames are encoded into $\displaystyle{\bm{z}}_{i}$ and combined with the action $\displaystyle a_{i}$ as input to the Mamba blocks. The input channel dimension is divided by the head dimension $\displaystyle p$ to generate the deterministic recurrent state $\displaystyle d_{i}$. This recurrent state $\displaystyle d_{i}$ is used to predict the next embedding $\displaystyle\hat{{\bm{z}}}_{i+1}$, reward $\hat{r_{i}}$, and termination flag $\hat{t_{i}}$, which represent the outcomes based on the current frame and action. The decoder reconstructs the original frame from the encoded embeddings $\displaystyle{\bm{z}}_{i}$ rather than from the predicted embeddings $\displaystyle\hat{{\bm{z}}_{i}}$. The Mamba-2 block employs employs a semiseparable matrix structure, which can be decomposed into $q\times q$ submatrices, enabling more efficient computation and processing.
Refer to caption: x1.png

## 2 Method

We describe the problem as a Partially Observable Markov Decision Process (POMDP), where at each discrete time step $\displaystyle t$, the agent observes a high-dimensional image $\displaystyle{\bm{\mathsfit{O}}}_{t}\in{\mathbb{O}}$ rather than the true state $\displaystyle{s}_{t}\in{\mathbb{S}}$ with the conditional observation probability given by $\displaystyle p({\bm{\mathsfit{O}}}_{t}|{s}_{t})$. The agent selects actions from a discrete action set $\displaystyle{a}_{t}\in{\mathbb{A}}=\{0,1,\dots,n\}$. After executing an action $\displaystyle{a}_{t}$, the agent receives a scalar reward $\displaystyle{r}_{t}\in\mathbb{R}$, a termination flag $\displaystyle{e}_{t}\in[0,1]$, and the next observation $\displaystyle{\bm{\mathsfit{O}}}_{t+1}$. The dynamics of the MDP is described by the transition probability $\displaystyle p({s}_{t+1},{r}_{t}|{s}_{t},{a}_{t})$. The behaviour of the agent is determined by a policy $\displaystyle f({\bm{\mathsfit{O}}}_{t};{\bm{\theta}})$, parameterised by $\displaystyle{\bm{\theta}}$, where $\displaystyle f:{\mathbb{O}}\rightarrow{\mathbb{A}}$ maps the observation space to the action space. The goal of this policy is to maximise the expected sum of discounted rewards $\displaystyle\mathbb{E}\sum_{t}\gamma^{t}{r}_{t}$, given that $\displaystyle\gamma$ is a predefined discount factor.

Unlike model-free RL, model-based RL does not rely directly on real experiences to improve the policy $\displaystyle f({\bm{\mathsfit{O}}}_{t};{\bm{\theta}})$ (Sutton and Barto, ). Instead, it learns a world model $\displaystyle f({\bm{\mathsfit{O}}}_{t},a_{t};\omega)$ from actual experiences to capture the dynamics of the POMDP. The actual experiences are stored in a replay buffer, allowing them to be repeatedly sampled for training the world model. The world model consists of a variational autoencoder (VAE) (Kingma and Welling, ; Hafner et al., ), a dynamics model, and linear heads to predict rewards and termination flags. The details of our world model are discussed in Section 2.2.

Each time the world model has been updated, a batch of experiences is sampled from the replay buffer to initiate a process called ‘imagination’. Starting from an actual initial observation and using an action generated by the current behaviour policy, the dynamics model generates the next latent state. This process is repeated until the agent collects enough imagined samples to improve the policy. We explain this process in detail in Section 2.3.

### 2.1 State Space Modelling with Mamba

State space models (SSMs) are mathematical constructs inspired by control theory to represent the complete state of a system at a given point in time. These models map an input sequence to an output sequence $\displaystyle{\bm{x}}\in\mathbb{R}^{l}\rightarrow{\bm{y}}\in\mathbb{R}^{l}$, where $l$ denotes the sequence length. In structured SSMs, a hidden state $\displaystyle{\bm{H}}\in\mathbb{R}^{(n,l)}$ is used to track the sequence dynamics, as described by the following equations:

$$ $\begin{split}{\bm{H}}_{t}&={\bm{A}}{\bm{H}}_{t-1}+{\bm{B}}{x}_{t}\\ {y}_{t}&={\bm{C}}^{\intercal}{\bm{H}}_{t}\\ \end{split}$ (1) $$

where $\displaystyle{\bm{\mathsfit{A}}}\in\mathbb{R}^{(n,n)},{\bm{B}}\in\mathbb{R}^{(n,1)},{\bm{C}}\in\mathbb{R}^{(n,1)}\text{ and }{\bm{H}}_{t}\in\mathbb{R}^{(n,1)}$, of which $\displaystyle n$ represents the predefined dimension of the hidden state that remains invariant to the sequence length. To efficiently compute the hidden states, it is common to structure $\displaystyle{\bm{A}}$ as a diagonal matrix, as discussed in (Gu et al., ; Gupta et al., ; Smith et al., ; Gu and Dao, ). Additionally, selective SSMs, such as Mamba-1, extend the matrices $\displaystyle({\bm{A}},{\bm{B}},{\bm{C}})$ to be time-varying, introducing an extra dimension corresponding to the sequence length. The shapes of these time-varying matrices are $\displaystyle{\bm{\mathsfit{A}}}\in\mathbb{R}^{(T,N,N)},{\bm{B}}\in\mathbb{R}^{(T,N)},\text{and }{\bm{C}}\in\mathbb{R}^{(T,N)}$ (^1^11In Mamba-1, the time variation of $\displaystyle{\bm{A}}$ is influenced by a discretisation parameter $\displaystyle\Delta$. For more details, please refer (Gu and Dao, )).

Dao and Gu () introduced the concept of structured state space duality (SSD), which further restricts the diagonal matrix $\displaystyle{\bm{A}}$ to be a scalar multiple of the identity matrix, forcing all diagonal elements to be identical. To compensate for the resulting reduced expressive power, Mamba-2 introduces a multi-head technique, akin to attention, by treating each input channel as $\displaystyle p$ independent sequences. Unlike Mamba-1, which computes SSMs as a recurrence, Mamba-2 approaches the sequence transformation problem through matrix multiplication, which is more GPU-efficient:

$$ $\begin{split}{y}_{t}&={\bm{C}}_{t}^{\intercal}{\bm{H}}_{t}\\ {y}_{t}&=\sum_{i=0}^{t}{\bm{C}}_{t}^{\intercal}{\bm{A}}_{t:i}{\bm{B}}_{i}{x}_{i}\\ \end{split}$ (2) $$

where $\displaystyle{\bm{A}}_{t:i}$ is $\displaystyle{\bm{A}}_{t}{\bm{A}}_{t-1}\dots\bm{A}_{i+1}$. This allows the SSM to be formulated as a matrix transformation:

$$ $\begin{split}{\bm{y}}&=SSM({\bm{x}};{\bm{\mathsfit{A}}},{\bm{B}},{\bm{C}})={\bm{M}}x\\ {M}_{j,i}&\coloneq\begin{cases}{\bm{C}}_{t}^{\intercal}{\bm{A}}_{t:i}{\bm{B}}_{i}&\text{if }j\geq i\\ 0&\text{if }j<i\end{cases}\end{split}$ (3) $$

Mamba-2 reformulates the state-space equations as a single matrix multiplication by utilising semi-separable matrices (Vandebril et al., ; Dao and Gu, ), which is well known in computational linear algebra as shown by Figure 1. The matrix $\displaystyle{\bm{M}}$ can also be written as:

$$ $\begin{split}{\bm{M}}&={\bm{L}}\circ{\bm{C}}{\bm{B}}^{\intercal}\in\mathbb{R}^{\mathtt{(T,T)}}\\ {\bm{L}}&=\begin{bmatrix}1&\\ a_{1}&1&\\ a_{2}a_{1}&a_{2}&1\\ \vdots&\vdots&\ddots&\ddots\\ a_{\mathtt{T}-1}\dots a_{1}&a_{\mathtt{T}-1}\dots a_{2}&\dots&a_{\mathtt{T}-1}&1\\ \end{bmatrix}\end{split}$ (4) $$

where $\displaystyle a_{t}\in[0,1]$ is a input-dependent scalar. The matrix $\displaystyle{\bm{L}}$ connects the SSM mechanism with the causal self-attention mechanism by removing the softmax function and applying a mask matrix $\displaystyle{\bm{L}}$ to the ‘attention-like’ matrix. It is, in fact, equivalent to causal linear attention when all $\displaystyle a_{t}=1$.

As a result, Mamba-2 achieves 2-8 times faster training speeds than Mamba-1, while maintaining linear scaling with sequence length.

### 2.2 World Model Learning

Our world model has two main components: an auto-encoder and a dynamics model. Additionally it includes two MLP heads for reward and termination predictions. The architecture of the world model is illustrated in Figure 1.

#### 2.2.1 Discrete Variational Auto-encoder

The autoencoder extends the standard variational autoencoder (VAE) architecture (Kingma and Welling, ) by incorporating a fully-connected layer to discretise the latent embeddings, consistent with previous approaches (Hafner et al., ; Robine et al., ; Zhang et al., ). The raw observation is a three-dimensional image, $\displaystyle{\bm{\mathsfit{O}}}_{t}\in[0,255]^{(h,w,c)}$, at time step $t$. The encoder compresses the observation into a vector of discrete numbers, denoted as $\displaystyle{\bm{z}}_{t}\sim p({\bm{z}}_{t}|{\bm{\mathsfit{O}}}_{t})$. The decoder reconstructs the raw image, $\displaystyle\hat{{\bm{\mathsfit{O}}}}_{t}$, from $\displaystyle{\bm{z}}_{t}$. Gradients are passed directly from the decoder to the encoder using the straight-through estimator, bypassing the sampling operation during backpropagation (Bengio et al., ).

#### 2.2.2 Dynamics Model

The dynamics model simulates the environment in the latent variable space, $\displaystyle{\bm{z}}_{t}$, using a deterministic state variable, $\displaystyle{\bm{d}}_{t}$. Since we are employing SSMs like Mamba-1 and Mamba-2, this should not be confused with the hidden states typically used by SSMs to track dynamics. At each time step $t$, the next token in the sequence is determined by both the current latent variable, $\displaystyle{\bm{z}}_{t}$ and the current action $\displaystyle a_{t}$. To integrate these, we first concatenate them and project the result using a fully-connected layer before passing it to the dynamics model. Given a sequence length $\displaystyle l$, the deterministic state is derived from all previous latent variables and actions. The dynamics model can be expressed as:

$$ $\displaystyle\text{Dynamics model:}\hskip 51.21495pt{\bm{d}}_{t}=f({\bm{z}}_{t-l:t},a_{t-l:t};\omega)$ (5) $\displaystyle\text{Latent variable predictor:}\hskip 17.07164pt\hat{{\bm{z}}}_{t+1}\sim p(\hat{{\bm{z}}}_{t+1}|{\bm{d}}_{t};\omega)$ $$

We implement the dynamics model with Mamba-2 (Dao and Gu, ). Specifically, each time a batch of samples, denoted as $\displaystyle{\bm{\mathsfit{O}}}\in[0,255]^{(b,l,h,w,c)}$, is drawn from the experience buffer $\mathcal{E}$ , where b represents batch size,l the sequence length, and i, w, c the image height, width, and channel dimension respectively. After encoding, the batch will be compressed to $\displaystyle{\bm{\mathsfit{Z}}}\in\mathbb{R}^{(b,l,d)}$ where$d$ is the dimension of the latent variable. The latent variable passes through a linear layer with the action to produce the input $\displaystyle{\bm{\mathsfit{X}}}\in\mathbb{R}^{(b,l,d)}$ of the Mamba blocks. To fully leverage the parallel computational capabilities of GPUs, the training process must not be strictly sequential. As a result, the targets of the dynamic model are independent of its outputs, which contrasts with the approach used in DreamerV3.

Mamba-1 first transforms the input tensor $\displaystyle{\bm{\mathsfit{X}}}_{b,:l,d}$ into a sequence of hidden states $\displaystyle{\bm{\mathsfit{H}}}\in\mathbb{R}^{(b,l-1,n)}$ , which are then mapped back to the deterministic state sequence $\displaystyle{\bm{\mathsfit{D}}}_{b,:l,d}$ using time-varying parameters. Since the hidden states operate in a fixed dimension
$\displaystyle n$ (unlike standard attention mechanisms, where the state scales with the sequence length), Mamba-1 achieves linear computational complexity with respect to sequence length.

Mamba-2 applies a similar transformation but leverages matrix multiplication. The input tensor $\displaystyle{\bm{\mathsfit{X}}}$’s dimension $d$ is first split into $\displaystyle d/p$ heads, which are processed independently. The transformation matrix is a specially designed semiseparable lower triangular matrix, which can be decomposed into $\displaystyle q\times q$ blocks. Different types of blocks are designed for specific purposes, such as handling causal attention over short ranges and transforming the hidden states.

### 2.3 Behaviour Policy Learning

The behaviour policy is trained within the ‘imagination’, an auto-generative process driven by the dynamics model. Specifically, a batch of shape $\displaystyle(b_{img},l_{img})$ is sampled from the replay buffer, where $\displaystyle b$ starting points are sampled and $\displaystyle l_{img}$ consecutive steps are selected starting from each. Since the Mamba dynamics model is efficient at handling long sequences, we can leverage actual experiences to estimate a more informative hidden state for the ‘imagination’ process. The rollout begins from the last transition in the sequence, $l_{img}$, and continues for $\displaystyle h$ steps. Notably, the rollout does not stop when an episode ends, unlike the prior SSM-based meta-RL model (Lu et al., ) where the hidden state must be manually reset, as the Mamba-based dynamics model automatically resets the state at episode boundaries (Gu and Dao, ).

A key difference between Mamba-based and transformer-based world models in the ‘imagination’ process is that Mamba updates inference parameters independent of sequence length. This independence is crucial for accelerating the ’imagination’ process, a significantly time-consuming component in model-based RL. The behaviour policy’s state concatenates the prior discrete variable $\displaystyle\hat{{\bm{z}}}_{t}$ with the deterministic variable $\displaystyle{\bm{h}}_{t}$ to exploit the temporal information. While the behaviour policy utilises a standard actor-critic architecture, other on-policy algorithms can also be applied. In this work, we adopt the recommendations from (Andrychowicz et al., ) and adjust the loss functions and value normalisation techniques as described in (Hafner et al., ).

### 2.4 Dynamic Frequency-Based Sampling (DFS)

In model-based RL, the behaviour model often underestimates rewards due to inaccuracies in the world model, impeding exploration and error correction (Sutton and Barto, ). These inaccuracies are particularly common early in training when the model relies on limited data. Thus, we propose a sample-efficient method to address this issue, i.e., Dynamic Frequency-based Sampling (DFS).

The primary objective is to sample transitions that the world model has sufficiently learned to initiate ‘imagination’. To accomplish this, we introduce two vectors during training, each matching the length of the transition buffer $|\mathcal{E}|$.
For the world model, $\displaystyle{\bm{v}}=({v}_{1},{v}_{2},\ldots,{v}_{|\mathcal{E}|}),\text{where }{v}_{i}\in\mathbb{Z}^{+}\text{ for }i\in\{1,2,\ldots,|\mathcal{E}|\}$, tracks the number of the transition has been used to improve the world model. The consequencing sampling probability is denoted as, $(p_{1},p_{2},\ldots,p_{|\mathcal{E}|})=\texttt{softmax}(-{\bm{v}})$. For ‘imagination’, ${\bm{b}}=({b}_{1},{b}_{2},\ldots,{b}_{|\mathcal{E}|}),\text{where }{b}_{i}\in\mathbb{Z}^{+}\text{ for }i\in\{1,2,\ldots,|\mathcal{E}|\}$, counts the times of transition has been used to improve the behaviour policy. The resulting sampling probability is denoted as, $(p_{1},p_{2},\ldots,p_{|\mathcal{E}|})=\texttt{softmax}(f({\bm{v}},{\bm{b}})),\text{where }f({\bm{v}},{\bm{b}})={\bm{v}}-{\bm{b}}-\max(0,{\bm{v}}-{\bm{b}})$. Note that for behaviour policy training, DFS employs balanced sampling similar to (Robine et al., ). During training, two cases arise: 1) $\exists i\in|\mathcal{E}|$, ${v}_{i}\geq{b}_{i},f({v}_{i},{b}_{i})=0$, In this case, the transition has been trained more frequently with the world model than with the behaviour policy, suggesting that the world model is likely capable of making accurate predictions from this transition. 2) $\exists i\in|\mathcal{E}|,{v}_{i}<{b}_{i},f({v}_{i},{b}_{i})={v}_{i}-{b}_{i}$, indicating that the transition is either under-trained as a starting point for the world model generation process or has been over-fitted to the behaviour policy. Consequently, the probability of selecting this transition for behaviour policy training decreases. These two mechanisms ensure that ’imagination’ sampling favors transitions learned by the world model, while avoiding excessive determinism.

## 3 Experiments

In this work, the proposed DRAMA framework is implemented on top of the STORM infrastructure (Zhang et al., ). We evaluate the model using the Atari100k benchmark (Kaiser et al., ), which is widely used for assessing the sample efficiency of RL algorithms. Atari100k limits interactions with the environment to 100,000 steps (equivalent to 400,000 frames with 4-frame skipping). We present the benchmark and analyse our results in Section 3.1 . Ablation experiments and their analysis are provided in Section 3.2.

### 3.1 Results

**Table 1: Comparison of game performance metrics for various algorithms across multiple Atari games. For Freeway IRIS enhances exploration using a distinct set of hyperparameters, while STORM leverages offline expert knowledge. TWM reports the results with a 21.6M model while IRIS does not report the exact number of parameters, they use the same transformer embedding dimension and layer number as TWM plus a behaviour policy with CNN layers. Dreamer notably uses a 200M parameter model and achieves good results in a series of diverse tasks. STORM dose not report the number of trainable parameters.**
|  | Random | Human | PPO | SimPLe | SPR | TWM | IRIS | STROM | DreamerV3 | Drama |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Alien | 228.0 | 7128.0 | 276.0 | 617.0 | 842.0 | 675.0 | 420.0 | 984.0 | 1118.0 | 820.0 |
| Amidar | 6.0 | 1720.0 | 26.0 | 74.0 | 180.0 | 122.0 | 143.0 | 205.0 | 97.0 | 131.0 |
| Assault | 222.0 | 742.0 | 327.0 | 527.0 | 566.0 | 683.0 | 1524.0 | 801.0 | 683.0 | 539.0 |
| Asterix | 210.0 | 8503.0 | 292.0 | 1128.0 | 962.0 | 1117.0 | 854.0 | 1028.0 | 1062.0 | 1632.0 |
| BankHeist | 14.0 | 753.0 | 14.0 | 34.0 | 345.0 | 467.0 | 53.0 | 641.0 | 398.0 | 137.0 |
| BattleZone | 2360.0 | 37188.0 | 2233.0 | 4031.0 | 14834.0 | 5068.0 | 13074.0 | 13540.0 | 20300.0 | 10860.0 |
| Boxing | 0.0 | 12.0 | 3.0 | 8.0 | 36.0 | 78.0 | 70.0 | 80.0 | 82.0 | 78.0 |
| Breakout | 2.0 | 30.0 | 3.0 | 16.0 | 20.0 | 20.0 | 84.0 | 16.0 | 10.0 | 7.0 |
| ChopperCommand | 811.0 | 7388.0 | 1005.0 | 979.0 | 946.0 | 1697.0 | 1565.0 | 1888.0 | 2222.0 | 1642.0 |
| CrazyClimber | 10780.0 | 35829.0 | 14675.0 | 62584.0 | 36700.0 | 71820.0 | 59324.0 | 66776.0 | 83931.0 | 52242.0 |
| DemonAttack | 152.0 | 1971.0 | 160.0 | 208.0 | 518.0 | 350.0 | 2034.0 | 165.0 | 577.0 | 201.0 |
| Freeway | 0.0 | 30.0 | 2.0 | 17.0 | 19.0 | 24.0 | 31.0 | 34.0 | 0.0 | 15.0 |
| Frostbite | 65.0 | 4335.0 | 127.0 | 237.0 | 1171.0 | 1476.0 | 259.0 | 1316.0 | 3377.0 | 785.0 |
| Gopher | 258.0 | 2412.0 | 368.0 | 597.0 | 661.0 | 1675.0 | 2236.0 | 8240.0 | 2160.0 | 2757.0 |
| Hero | 1027.0 | 30826.0 | 2596.0 | 2657.0 | 5859.0 | 7254.0 | 7037.0 | 11044.0 | 13354.0 | 7946.0 |
| Jamesbond | 29.0 | 303.0 | 41.0 | 100.0 | 366.0 | 362.0 | 463.0 | 509.0 | 540.0 | 372.0 |
| Kangaroo | 52.0 | 3035.0 | 55.0 | 51.0 | 3617.0 | 1240.0 | 838.0 | 4208.0 | 2643.0 | 1384.0 |
| Krull | 1598.0 | 2666.0 | 3222.0 | 2205.0 | 3682.0 | 6349.0 | 6616.0 | 8413.0 | 8171.0 | 9693.0 |
| KungFuMaster | 258.0 | 22736.0 | 2090.0 | 14862.0 | 14783.0 | 24555.0 | 21760.0 | 26183.0 | 23920.0 | 17236.0 |
| MsPacman | 307.0 | 6952.0 | 366.0 | 1480.0 | 1318.0 | 1588.0 | 999.0 | 2673.0 | 1521.0 | 2270.0 |
| Pong | -21.0 | 15.0 | -20.0 | 13.0 | -5.0 | 19.0 | 15.0 | 11.0 | -4.0 | 15.0 |
| PrivateEye | 25.0 | 69571.0 | 100.0 | 35.0 | 86.0 | 87.0 | 100.0 | 7781.0 | 3238.0 | 90.0 |
| Qbert | 164.0 | 13455.0 | 317.0 | 1289.0 | 866.0 | 3331.0 | 746.0 | 4522.0 | 2921.0 | 796.0 |
| RoadRunner | 12.0 | 7845.0 | 602.0 | 5641.0 | 12213.0 | 9109.0 | 9615.0 | 17564.0 | 19230.0 | 14020.0 |
| Seaquest | 68.0 | 42055.0 | 305.0 | 683.0 | 558.0 | 774.0 | 661.0 | 525.0 | 962.0 | 497.0 |
| UpNDown | 533.0 | 11693.0 | 1502.0 | 3350.0 | 10859.0 | 15982.0 | 3546.0 | 7985.0 | 46910.0 | 7387.0 |
| Normalised Mean (%) | 0.0 | 100.0 | 11.0 | 33.0 | 62.0 | 96.0 | 105.0 | 127.0 | 125.0 | 105.0 |
| Normalised Median (%) | 0.0 | 100.0 | 3.0 | 13.0 | 40.0 | 51.0 | 29.0 | 58.0 | 49.0 | 27.0 |

We compare our model against several benchmarks across 26 Atari games. In Table 1, the ‘Normalised Mean’ refers to the average normalised score, calculated as: $\displaystyle(evaluated_score-random_score)/(human_score-random_score)$. For each game, we train DRAMA with 5 different seeds and track training performance using a running average of 5 episodes, as recommended by Machado et al. (), a practice also followed in related work (Hafner et al., ).

Despite utilising an extra-small world model (7M parameters, referred to as the XS model), we achieve performance comparable to IRIS and TWM. Furthermore, by employing a stronger auto-encoder and a larger SSM hidden state dimension (10M parameters, referred to as the S model), we demonstrate improved results in ablation experiments on a reduced set of games. However, we emphasise that our goal is not to achieve the highest benchmark ranking, but to illustrate that Mamba can serve as a solid foundation for the dynamics model in model-based RL.

Table 1 demonstrates that DRAMA, with Mamba-2 as the dynamics model, is both sample- and parameter-efficient. For comparison, Simple (Kaiser et al., ) trains a video prediction model to optimise a PPO agent (Schulman et al., ), while SPR (Schwarzer et al., ) uses a dynamics model to predict in latent space, enhancing consistency through data augmentation. TWM (Robine et al., ) employs a Transformer-XL architecture to capture dependencies among states, actions, and rewards, training a policy-based agent. This method incorporates short-term temporal information into the embeddings to avoid using the dynamics model during actual interactions. Similarly, IRIS (Micheli et al., ) uses a Transformer as its dynamics model, but generates new samples in image space, allowing pixel-level feature extraction for behaviour policies. DreamerV3 (Hafner et al., ), which employs an RNN-based dynamics model along with robustness techniques, achieves superhuman performance on the Atari100k benchmark using a 200M parameter model—20 times larger than ours. STORM (Zhang et al., ), which adopts many of DreamerV3’s robustness techniques while replacing the dynamics model with a transformer, reaches similar performance on the Atari100k benchmark as DreamerV3.

Drama excels in games like Boxing and Pong, where the player competes against an autonomous agent in simple, static environments, requiring a less intense auto-encoder. This strong performance indicates that Mamba-2 effectively captures both ball dynamics and the opponent’s position. Similarly, Drama performs well in Asterix, which benefits from its ability to predict object movements. However, Drama struggles in Breakout, where performance can be improved with a more robust auto-encoder in Figure 4. Additionally, Drama excels in games like Krull and MsPacman, which require longer sequence memory, but faces challenges in sparse reward games like Jamesbond and PrivateEye.

### 3.2 Ablation experiments

We selected a representative subset of games for our ablation experiments. Krull is a multi-scene game with dense rewards, while Boxing is a single-scene game featuring an AI-controlled opponent. Freeway is a sparse reward game that requires exploration, and Kangaroo demands multitasking and object identification for different actions.

#### 3.2.1 Dynamic Frequency-Based Sampling

In this experiment, we compare DFS with the uniform sampling method in a Mamba-2 based Drama. As shown in Figure 2, DFS demonstrates a significant advantage over uniform sampling in the games Freeway and Kangaroo, with a smaller advantage observed in Krull and Boxing. The ablation results further highlight the effectiveness of DFS in mitigating the suboptimality of the behaviour policy when learning within a flawed world model. This is especially evident in Freeway, where agents often become trapped in a passive policy, waiting for the game to end without taking any meaningful action.

Figure: Figure 2: Uniform Sampling vs. Dynamic Frequency-Based Sampling. DFS demonstrates particular effectiveness in Freeway and Kangaroo.
Refer to caption: x2.png

#### 3.2.2 Mamba-1 vs. Mamba-2

As mentioned in Sec 2.1, Mamba-2 imposes restrictions on $\displaystyle{\bm{A}}$ for efficiency. However, it remains an open question whether these constraints negatively affect the performance of SSMs, as previous studies have not offered comprehensive theoretical or empirical evidence on the matter (Dao and Gu, ). In response to this gap, we compare Mamba-2 and Mamba-1 as the backbone of the world model in model-based RL. Ablation experiments were conducted using DFS, with both Mamba-1 and Mamba-2 configured with the same default hyperparameters.

Figure 3 illustrates that Mamba-2 outperforms Mamba-1 in games Krull, Boxing and Freeway. In Krull, the player navigates through different scenes and solves various tasks. In the later stages, rescuing the princess while avoiding hits results in a significant score boost, while failure leads to a plateau in score. As shown, Mamba-1 experiences a score plateau in Krull, whereas Mamba-2 successfully overcomes this challenge, leading to higher performance. Note that Freeway is a sparse reward game requiring high-quality exploration. A positive training effect is achieved only by combining DFS with Mamba-2 without any additional configuration.

Figure: Figure 3: Mamba-1 vs. Mamba-2. Mamba2 has shown a superior performance over Mamba-1 three out of four games.
Refer to caption: x3.png

### 3.3 More trainable parameters

As model-based RL agents consist of multiple trainable components, tuning the hyperparameters for each part can be resource-intensive and is not the primary focus of this research. Previous work has demonstrated that increasing the neural network’s size often leads to stronger performance on benchmarks Hafner et al. (). In Figure 4, we demonstrate that Drama achieves overall better performance when using a more robust auto-encoder and a larger SSM hidden state dimension $n$. Notably, the S model exhibits significantly improved results in games like Breakout and BankHeist, where pixel-level information plays a crucial role.

Figure: Figure 4: S model vs. XS model. We adjusted the game set to emphasise the importance of recognising small objects. The S model features a more robust auto-encoder than the XS model, with additional filters and 3M more trainable parameters. In terms of performance, the S model significantly outperforms the XS model in Breakout and BankHeist. However, it underperforms in Kangaroo and shows comparable performance in ChopperCommand.
Refer to caption: x4.png

## 4 Related work

### 4.1 Model-based RL

The origin of model-based RL can be traced back to the Dyna architecture introduced by Sutton and Barto (), although Dyna selects actions through planning rather than learning. Notably, Sutton and Barto () also highlighted the suboptimality that arises when the world model is flawed, especially as the environment improves. The concept of learning in ’imagination’ was first proposed by Ha and Schmidhuber (), where a world model predicts the dynamics of the environment. Later, SimPLe (Kaiser et al., ) applied MBRL to Atari games, demonstrating improved sample efficiency compared to state-of-the-art model-free algorithms. Beginning with Hafner et al. (), the Dreamer series introduced a GRU-powered world model to solve a diverse range of tasks, such as Mujoco, Atari, Minecraft, and others (Hafner et al., , , ). More recently, inspired by the success of transformers in NLP, many MBRL studies have adopted transformer architectures for their dynamics models. For instance, IRIS (Micheli et al., ) encodes game frames as sets of tokens using VQ-VAE (Oord et al., ) and learns sequence dependencies with a transformer. In IRIS, the behavior policy operates on raw images, requiring an image reconstruction during the ’imagination’ process and an additional CNN-LSTM structure to extract information. TWM (Robine et al., ), another transformer-based world model, uses a different structure. It stacks grayscale frames and does not activate the dynamics model during actual interaction phases. However, its behaviour policy only has access to short-term temporal information, raising questions about whether learning from tokens that already include this short-term information could be detrimental to the dynamics model. STORM (Zhang et al., ), closely following DreamerV3, replaces the GRU with a vanilla transformer. Additionally, it incorporates a demonstration technique, populating the buffer with expert knowledge, which has shown to be particularly beneficial in the game Freeway.

### 4.2 Structure State space model based RL

Structured SSMs were originally introduced to tackle long-range dependency challenges, complementing the transformer architecture (Gu et al., ; Gupta et al., ). However, Mamba and its successor, Mamba-2, have emerged as powerful alternatives, now competing directly with transformers (Gu and Dao, ; Dao and Gu, ). Deng et al. () implemented an SSM-based world model, comparing it against RNN-based and transformer-based models across various prediction tasks. Despite this, SSM-based world models have yet to be tested in the context of model-based RL, including Mamba-1 and Mamba-2. Mamba-1 has recently been applied to offline RL, either with a standard Mamba-1 block (Ota, ) or a Mamba-attention hybrid model (Huang et al., ). Lu et al. () proposed applying modified SSMs to meta-RL, where hidden states are manually reset at episode boundaries. Since both Mamba-1 and Mamba-2 are input-dependent, such resets are unnecessary.

## 5 Conclusion

In conclusion, DRAMA, our proposed Mamba-based world model, addresses key challenges faced by RNN and transformer-based world models in model-based RL. By achieving $O(n)$ memory and computational complexity, our approach enables the use of longer training sequences. Furthermore, our novel sampling method effectively mitigates suboptimality during the early stages of training, contributing to a model that is both lightweight, with only 7 million trainable parameter world model, and accessible, being trainable on standard hardware. Overall, our method achieves a normalised score competitive with other state-of-the-art RL algorithms, offering a practical and efficient alternative for model-based RL systems. Although Drama enables longer training and inference sequences, it does not demonstrate a decisive advantage that would allow it to dominate other world models on the Atari100k benchmark. An interesting direction for future work is to explore specific tasks where longer sequences drive superior performance in model-based RL. Despite advances in world models, model-based RL still faces several challenges, such as long-horizon behaviour planning and learning, informed exploration, and the dynamics of jointly training the world model and behaviour policy. Another promising future direction is to investigate to what extent Mamba can help address these challenges.

## Appendix A Appendix

### A.1 Atari100k Learning Curves

Figure: Figure 5: Atari100k Learning Curve.
Refer to caption: x5.png

### A.2 Loss and Hyperparameters

#### A.2.1 Variational Auto-encoder

The hyperparameters shown in Table 2 correspond to the default model, also referred to as XS in Figure 4. For the S model, we simply double the number of filters per layer to obtain a stronger auto-encoder.

**Table 2: Hyperparameters for the auto-encoder.**
| Hyperparameter | Value |
| --- | --- |
| Frame shape (h, w, c) | (64, 64, 3) |
| Layers | 5 |
| Filters per layer (Encoder) | (16, 32, 48, 64, 64) |
| Filters per layer (Decoder) | (64, 64, 48, 32, 16) |
| Kernel | 5 |
| Act | SiLU |
| Batch Norm | Yes |

#### A.2.2 Mamba-1 and Mamba-2

Similar to the previous section, the values shown in Table 3 correspond to the default model. For the S model, we double the latent state dimension, allowing more relevant information to be stored in the recurrent state. In the Mamba-2 model, the enhanced architecture supports a larger latent state dimension without significantly increasing the training time.

**Table 3: Hyperparameters for Mamba-1 and Mamba-2. Except the head dimension is only for Mamba-2, the other hyperparameters are shared. The head number is 512/128 = 4.**
| Hyperparameter | Value |
| --- | --- |
| Hidden state dimension (d) | 512 |
| Layers | 2 |
| Latent state dimension (n) | 16 |
| RMS Norm | True |
| Act | SiLU |
| Mamba-2: Head dimension (p) | 128 |

#### A.2.3 Reward and termination prediction heads

Both the reward and termination flag predictors take the deterministic state output from the dynamic model to make their predictions. Due to the quality of the hidden state extracted by the dynamic model, a single fully connected layer is sufficient for accurate predictions.

**Table 4: Hyperparameters for reward and termination prediction heads.**
| Hyperparameter | Value |
| --- | --- |
| Hidden units | 256 |
| Layers | 1 |

The world model is optimized in an end-to-end and self-supervised manner on batches of shape $(b,l)$ drawn from the experience replay.

$$ $\mathcal{L}(\omega)=\mathbb{E}\left[\begin{aligned} \sum_{i=1}^{l}\underbrace{(O_{i}-\hat{O}_{i})^{2}}_{\text{reconstruction loss}}+\mathcal{L}_{dyn}(\omega)&+0.1*\mathcal{L}_{rep}(\omega)\\ &-\underbrace{\ln p(\hat{r}_{i}|d_{i};\omega)}_{\text{reward prediction loss}}-\underbrace{\ln p(\hat{t}_{i}|d_{i};\omega)}_{\text{termination prediction loss}}\end{aligned}\right]$ (6) $$

where

$$ $\begin{split}&\mathcal{L}_{\text{dyn}}(\omega)=\max\left(1,\text{KL}\left[\text{sg}(p({\bm{z}}_{i+1}|{\bm{\mathsfit{O}}}_{i+1};\omega))\ \middle\|\ q(\hat{{\bm{z}}}_{i+1}|d_{i};\omega)\right]\right)\\ &\mathcal{L}_{rep}(\omega)=\max\left(1,\text{KL}\left[p({\bm{z}}_{i+1}|{\bm{\mathsfit{O}}}_{i+1};\omega)\ \middle\|\ \text{sg}(q(\hat{{\bm{z}}}_{i+1}|d_{i};\omega))\right]\right)\\ \end{split}$ (7) $$

and sg$(\cdot)$ represents the stop gradient operation.

#### A.2.4 Actor Critic Hyperparameters

We adopt the behavior policy learning setup from DreamerV3 [Hafner et al., ] for simplicity and strong performance, as the behaviour policy model is not central to our main contribution.

**Table 5: Hyperparameters for the behaviour policy.**
| Hyperparameter | Value |
| --- | --- |
| Layers | 2 |
| Gamma | 0.985 |
| Lambda | 0.95 |
| Entropy coefficient | 3e-4 |
| Max gradient norm | 100 |
| Actor hidden units | 256 |
| Critic hidden units | 512 |
| RMS Norm | True |
| Act | SiLU |
| Batch size ($b_{img}$) | 1024 |
| Imagine context length ($l_{img}$) | 8 |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: - **Evidence of Long-Term Dependency Capture:** You mention the ability of DRAMA to capture long-term dependencies. Could you provide more direct evidence or experiments that demonstrate this capability? For example, have you considered tasks that specifically require long-term memory or conducted analyses that quantify the effective memory length of the model?
- **Scaling DRAMA for Direct Comparison:** Have you considered scaling up DRAMA to match the model size of DreamerV3 for a direct comparison? If not, could you explain the limitations (e.g., hardware constraints) that prevented this? Testing DRAMA with larger models on a few games could provide insights into its scalability and performance limits.  
- **Effectiveness of DFS:** Could you elaborate on how the dynamic frequency-based sampling (DFS) method compares to other sampling strategies in terms of its impact on learning efficiency and final performance? Including quantitative comparisons or ablation studies would be helpful.
 - **Hyperparameter Sensitivity and Trade-offs:** Can you provide insights into the trade-offs between model size and performance in DRAMA? Specifically, how does increasing the size of the Mamba-2 model or the autoencoder affect results across different games?
- **Code and Reproducibility:** Is there a plan to release the code and pretrained models for DRAMA to facilitate reproducibility and further research in this area?
soundness: 3
strengths: - **Originality:** The paper introduces the novel application of Mamba-2 SSMs within MB-RL, specifically as the dynamics model in the world model. This is a new approach that addresses the limitations of existing architectures like typical RNNs and transformers and it makes a lot of sense in my opinion.
- **Quality:** The authors provide thorough experimental evaluations on the Atari100k benchmark, demonstrating that DRAMA achieves competitive performance with significantly fewer parameters (at least in the tasks and model sizes tested). The inclusion of ablation studies comparing Mamba-1 and Mamba-2, as well as the impact of DFS, strengthens the empirical results.
- **Clarity:** The paper is well-written and structured, providing clear explanations of the methodology, including detailed descriptions of Mamba-2 and how it is integrated into the world model. The figures and tables effectively support the textual content although I would make some further stylistic improvements in Figure 1 to maximise clarity.
- **Practical impact:** By achieving comparable performance to state-of-the-art methods with a smaller and more computationally efficient model, the DRAMA method contributes to making MB-RL more accessible and practical, particularly in resource-constrained environments.
summary: This paper introduces DRAMA, a model-based reinforcement learning (MB-RL) agent that leverages the Mamba-2 architecture, a state space model (SSM), as its core dynamics architecture. Traditional MB-RL approaches often rely on recurrent neural networks (RNNs) or transformers for world modelling, which suffer from issues like vanishing gradients, difficulty in capturing long-term dependencies, and quadratic scaling of computational complexity with sequence length. DRAMA addresses these challenges by utilizing Mamba-2, which achieves linear computational and memory complexity while effectively capturing long-term dependencies.

Additionally, the authors propose a novel dynamic frequency-based sampling (DFS) method to mitigate the suboptimality arising from imperfect world models during early training stages. They evaluate DRAMA on the Atari100k benchmark, demonstrating that it achieves performance comparable to state-of-the-art algorithms using a significantly smaller world model (7 million parameters) that can be trained on standard hardware. The paper also includes ablation studies comparing Mamba-1 and Mamba-2, highlighting the superior performance of Mamba-2 despite its constrained expressive power.
weaknesses: - **Unsupported claims about capturing long-term dependencies:** While the authors claim repeatedly that Mamba-2 effectively handling long-term dependencies, the paper provides limited direct evidence or analysis to demonstrate this capability. Including experiments or analyses that specifically test and showcase the ability to capture long-term dependencies would strengthen the paper. For instance, a task designed to require long-term memory or metrics that quantify the model's ability to capture dependencies over long sequences could be included. In the current form of the paper, it is not clear to me what "long" really means.
- **Limited Comparison with Scaled Models:** The comparison with DreamerV3 is somewhat limited. While the authors emphasize parameter efficiency, it would be valuable to see how DRAMA performs when scaled up to match the model size of DreamerV3, even on a subset of games. This would help understand the limits and potential of their approach and whether the advantages of Mamba-2 persist at larger scales. If hardware limitations prevented this, a discussion of these constraints would be very helpful.
- **Marginal Performance Gains:** While DRAMA achieves comparable performance to existing methods, the improvements are not substantial across all games. Demonstrating scenarios where DRAMA significantly outperforms other approaches would strengthen the claims about its effectiveness.
- **DFS Method Clarity:** The explanation of the dynamic frequency-based sampling method could be more detailed. Providing more comprehensive comparisons with other sampling strategies would help in understanding its effectiveness. Additionally, including ablation studies that isolate the impact of DFS would clarify its contribution.
- **Hyperparameter Sensitivity:** The paper mentions that increasing the model size leads to better performance but does not deeply explore this aspect. Though I understand that extensive hyperparameter search might be too difficult given computational constraints, some analysis of how sensitive DRAMA is to hyperparameter choices, including model size, sequence length, and learning rates, would be very valuable for understanding its practical applicability, otherwise the paper looks relatively incomplete.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. At line 222, it is not clear what targets mean.
2. Section 2.3 is not described well. At line 240, when the ‘b’ starting points are sampled of length $l_{img}$, is it just picking random samples along $b_{img}$ sequences? Since the batch sampled from replay buffer is of length $l_{img}$, I am unsure what this additional sampling is doing? Why not just sample $b$ trajectories from the buffer?
3. What is $h_t$ in behavior policy learning? The deterministic variable is defined as $d_t$ in Eq 5.
4. While training Dreamer, the method uses the whole sequence of ($b_{img}, l_{img}$) to compute a good hidden state and uses all sampled to generate trajectories in the future. I am curious to know why only the last hidden state $l_{img}$ is used for learning the policy and not the whole sequence like Dreamer? (describe around line 243).
5. The terminal flag is defined as $e_t$ at line 93, whereas it is $t_i$ in description of Figure 1. Also, the description of Fig 1 uses $i$ for indexing the time and the Section 2 starts with $t$ as time index. The description of Fig 1 should match the notation in the main text.
6. Was any experiment conducted to see if the behavior model underestimates rewards especially with limited data? 
7. For the ablation presented in Sec 3.2.2, were both variants Mamba-1 and Mamba-2 based WM trained with DFS?
8. How does the proposed method compare with R2I [1]. Since they propose using SSMs, should it be included as another baseline?
9. [Typo] Employs is written twice at line 132.

#### References
[1] Samsami et al., Mastering Memory Tasks with World Models, ICLR’24.
soundness: 3
strengths: - The proposed idea of using SSM for WMs is interesting as they provide crucial benefits over training with Transformers and RNNs. 
- The proposed method achieves good performance with significantly fewer parameters (7M) when compared with baselines.
summary: This paper proposes using State-Space Models (SSMs) for learning World Models. Specifically, the architecture comprises an encoder to get discrete latents, a SSM module (Mamba-2) to estimate the dynamics which is used to predict the latent embedding of observation, reward value and termination flag. This world model is used to train a policy by imagination. The paper also uses a method to sample transitions from the replay buffer based on the number of times the transition is used to update WM and policy. Experiments conducted on Atari100K show that the proposed method Drama attains similar performance to IRIS and TWM with a much smaller model in terms of parameters.
weaknesses: - It is hard to articulate where the performance gains are coming from. Section 3.2.1 discusses that DFS provides an advantage over uniform sampling. Since DFS is agnostic to most baselines, it is important to see a comparison of either Drama with Uniform Sampling or baselines with DFS sampling to understand if the architecture is helping or the sampling.
- The paper is not well written and it is hard to understand the details and motivation behind the design choices. Questions 1-5 below expand on this. The paper can use a pseudocode to describe the behavior learning part.

---

confidence: 3
contribution: 3
initial_rating: 3
presentation: 3
questions: - [Q1] Could the decoder operate based on the output of Mamba-2, such that $d$ rather than $z$ serves as the input?
soundness: 3
strengths: - [S1] The paper combines established methods in a novel way, effectively addressing an existing gap in world model research.
- [S2] The proposed model is computationally efficient, requiring only 7 million trainable parameters, making it accessible.
summary: This paper presents a new world model for deep reinforcement learning that is based on the Mamba-2 architecture. In addition, the authors introduce dynamic frequency-based sampling, which leads to more accurate imaginations of the world model. The approach achieves competitive performance on the Atari 100k benchmark while maintaining a lightweight architecture with only 7 million trainable parameters.
weaknesses: - [W1] The extent to which the Mamba architecture contributes to the model's performance remains unclear. Specifically, it is unclear how DFS impacts scores across all games. Extending ablation study 3.2.1 to cover more games, or conducting a new study that replaces Mamba with an RNN or transformer, would clarify these contributions.
- [W2] While the paper emphasizes Mamba's computational efficiency, there is a lack of exact wall-clock training and inference times. The abstract claims the model can be trained on a standard laptop, so providing specific runtime metrics would substantiate this claim.
- [W3] Several presentation aspects should be improved:
  - The paragraph in lines 99–105 shifts from general statements about model-based RL to specific details about the paper's world model. This transition could lead readers to infer that every world model relies on a variational autoencoder or linear heads, which isn't necessarily the case. Additionally, other model-based RL methods exist that don't utilize world models, such as those using lookahead search.
  - Many figures contain small text that is difficult to read without zooming.
  - Minor suggestions:
    - Highlight the highest scores in Table 1 for easy reference.
    - Correct notations for all \hat{} terms (e.g., \hat{r}_t instead of \hat{r_t}).
    - Address minor notation errors: e.g., incorrect $A$ on lines 148, 154, 169 and incorrect $T$ on line 177
    - Variables on line 218 should be in math mode.
    - Consider changing "auto-generative" to "autoregressive" on line 238?
    - Revise line 264 to read "tracks the number of *times* the transition has been used."
- [W4] To strengthen the soundness of findings, additional evaluations on alternative benchmarks, such as the DeepMind Control Suite, would be valuable. That said, I understand this may be challenging to realize.

I would consider raising my scores if these issues were addressed.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: - What is the value of h?
- "A key difference between Mamba-based and transformer-based world models in the ‘imagination’ process is that Mamba updates inference parameters independent of sequence length."--can you explain it more?
soundness: 3
strengths: - The combination of MAMBA and Sequence-based world models is novel
- The authors demonstrate comparable results with a smaller model size
summary: The authors propose a MAMBA-based world model architecture, as opposed to previous transformer and RSSM based ones. They also compare between MAMBA-1 and MAMBA-2 for world models. Finally, they evaluate and ablate their technique on the Atari100k benchmark.
weaknesses: - Some results are unclear. For example, could you go into details on why the model is not doing well on the Breakout game? What is it producing? Some figures from the game can be also useful--not just this one, but the ones in which model does well too.
- It is not clear why imagination context length is kept to 8. I would suggest providing experiments, both involving time-complexity and performance, for different imagination context length. Also, explaining why this is done would be useful.

Meta-review: The paper presents a significant contribution by effectively incorporating Mamba architecture into model-based RL, achieving competitive performance with only 7M parameters. While reviewers raised concerns about long-term dependency claims and computational efficiency, the authors provided comprehensive responses with new experimental results demonstrating Mamba's effectiveness at long sequences, faster imagination time than transformers, and improved DFS performance in 11 games. The authors also conducted thorough ablations showing superior performance over small-scale DreamerV3. With all reviewers responding positively to these detailed responses, and given the practical impact of achieving state-of-the-art performance with significantly reduced parameters, this work warrants acceptance.

Decision: Accept (Poster)