Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient
Authors: Wenlong Wang
Ivana Dusparic
Yucheng Shi
Ke Zhang
Vinny Cahill
Venue: ICLR 2025 Poster (2025)
Keywords: Mamba
Model based reinforcement learning
Atari100k
Mamba-2
Primary area: reinforcement learning

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 requires complex and deep architectures, which are computationally expensive and challenging to train. Within the world model, sequence models play a critical role in accurate predictions, and various architectures have been explored, each with its own challenges. Currently, recurrent neural network (RNN)-based world models struggle with vanishing gradients and capturing long-term dependencies. Transformers, on the other hand, suffer from the quadratic memory and computational complexity of self-attention mechanisms, scaling as $O(n^2)$, where $n$ is the sequence length.

To address these challenges, we propose a state space model (SSM)-based world model, Drama, specifically leveraging Mamba, that achieves $O(n)$ memory and computational complexity while effectively capturing long-term dependencies and enabling efficient training with longer sequences. We also introduce a novel sampling method to mitigate the suboptimality caused by an incorrect world model in the early training stages. Combining these techniques, Drama achieves a normalised score on the Atari100k benchmark that is competitive with other state-of-the-art (SOTA) model-based RL algorithms, using only a 7 million-parameter world model. Drama is accessible and trainable on off-the-shelf hardware, such as a standard laptop. Our code is available at https://github.com/realwenlongwang/Drama.git.

Published as a conference paper at ICLR 2025
DRAMA: MAMBA-ENABLED MODEL-BASED REIN- FORCEMENT LEARNING IS SAMPLE AND PARAMETER EFFICIENT
Wenlong Wang, Ivana Dusparic, Yucheng Shi, Zhang & Vinny Cahill School of Computer Science and Statistics Trinity College Dublin, the University of Dublin
College Green, Dublin 2, Ireland [wangwl, ivana. duspariccccc shiy2, 2hangk2, @ted.
ABSTRACT
Model-based reinforcement learning (RL) offers solution to the data inefficiency that plagues most model-free RL algorithms. However, learning robust world model often requires complex and deep architectures, which are computationally expensive and challenging to train. Within the world model, sequence models play a critical role in accurate predictions, and various architectures have been ex- plored, each with its challenges. Currently, recurrent neural network (RNN)- based world models struggle with vanishing gradients and capturing long-term dependencies. Transformers, the other hand, suffer from the quadratic memory and computational complexity of self-attention mechanisms, scaling as O(??), where the sequence length.
To address these challenges, we propose a state space model (SSM)-based world model, Drama, specifically leveraging Mamba, that achieves O(n) memory and computational complexity while effectively capturing long-term dependen- cies and enabling efficient training with longer sequences. We also introduce a novel sampling method to mitigate the suboptimality caused by an incorrect world model in the early training stages. Combining these techniques, Drama achieves a normalised score on the Atari100k benchmark that is competitive with other state-of-the-art (SOTA) model-based RL algorithms, using only 7 million-parameter world model. Drama accessible and trainable on off-the- shelf hardware, such standard laptop. Our code is available at https: com t orr.tuubbooom/rraleeeeeoogggenng
INTRODUCTION
Deep Reinforcement Learning (RL) has achieved remarkable success various challenging appli- cations, such Go Siiverr al. 2016 2017), Dota (Berner al. 2019). Atari (Mnih et al. 2013 Schrittwieser et al. 2020). and MuJoCo (Schulman et al., 2017 Haarnoja et af. 2018). However, training policies capable of solving complex tasks often requires millions of environment interac- tions, which can be impractical and pose barrier to real-world applications. Thus, improving sample efficiency has become critical goal in RL algorithm development.
World models have shown promise in improving sample efficiency by generating artificial training samples through an autoregressive process, method referred as model-based RL (Micheli al.) 2023 Robine et al. 2023 Zhang et al.] (2233] Hafner et al. 2023). In this approach, interaction data is used to learn the environment dynamics using sequence model, allowing the agent to train on artificial experiences generated by the resulting sequence model instead of relying real-world interactions. This approach shifts the problem from improving the policy directly using real sam- ples (which is sample inefficient) to improving the accuracy of the world model match the real environment (which is more sample efficient). However, model-based RL faces a well-known chal- lenge: when the model is inaccurate due to limited observed samples, especially early in training, the policy eventually learned can converge to suboptimal behaviour, and detecting model errors is difficult, if not impossible.

---

Published as a conference paper at ICLR 2025
In sequence modelling, linear complexity (in sequence length) 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 effi- cient sequence modelling to simulate complex environments over long time horizons. RNNs, partic- ularly advanced variants like Long Short-Term Memory (LSTTM) and Gated Recurrent Units (GRU), offer linear complexity, making them computationally attractive for this task. However, RNNs still struggle with vanishing gradient issues and exhibit limitations in capturing long-term dependencies (Hafner et al. 2019 2023). More recently, transformer architectures, which have dominated natural language processing (NLP) (Vaswani et al. 2017), have gained traction in fields like image process- ing and offline RL following groundbreaking work in these areas (Dosovisskiy et al.) 202 1 Chen jet al., 2021). The transformer structure has demonstrated its effectiveness model-based RL as well (Micheli et al., [2023] Robine et al.] 2023] Zhang et al.] 2023). However, transformers suffer from both memory and computation complexity that scale as Onn"). where n is the sequence length, posing challenges for world models that require long sequence to simulate complex environments.
Currently, SSMs are attracting significant attention for their ability to efficiently model long- sequence problems with linear complexity. Among SSMs, Mamba has emerged
competitive alternative to transformer-base based architectures in various fields, including NLP (Gu & Dao 2024]
Dao & Gu 2024), computer vision (Zhu et al. 2024), and offline RL. (Lv al. 2024). Applying Mamba's s architecture to model-based RI. is particularly appealing due to its linear memory and computational scaling with sequence length, coupled with its ability capture long-term dependen- cies effectively. Moreover, efficiently capturing environmental dynamics can reduce the likelihood that the behaviour policy being learned within an inaccurate world model, a challenge we further ad- dress by incorporating a novel dynamic frequency-basse -bsed 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 achieves performance comparable to other SOTA algorithms while using only a 7 million trainable parameter world model.
We compare the performance of Mamba and Mamba-2 demonstrating that Mamba-2 achieves superior results as a sequence model in the Atari100k benchmarks, despite slightly limiting expressive power enhance training efficiency.
Finally, we propose novel but straightforward sampling method, dynamic frequency- based sampling (DFS), to mitigate the challenges posed by imperfect sequence models.
2 METHOD
We describe the problem as a Partially Observable Markov Decision Process (POMDP), where at each discrete time step i, the agent observes high-dimensional image O O rather than the true state 8 E S with the conditional observation probability given by p(O,(se) The agent selects actions from discrete action set a₂ A {0, 1,...,,,, n). After executing an action ae, the agent receives scalar reward rz E IR, a termination flag e, E [0, 1]. and the next observation O The dynamics of the MDP are described by the transition probability p($ T ( at). The behaviour of the agent is determined by policy x ;O,; 0), parameterised by 9, where π O + A maps the observation space to the action space. The goal of this policy is to maximise the expected sum of discounted rewards IE "Tr, given that is a predefined discount factor.
Unlike model-free RL, model-based RL does not rely directly on real experiences improve the policy x (O,;; 0) (Sutton & Barto 1998). There are various approaches to obtaining a world model, including Monte Carlo tree search (Schrittwieser al.) 2020). offline imitation learning (DeMoss) et al. 2023) and latent sequence models (Hafner et al. 2019). In this work, we focus on learning a world model f(O,, de:w) w) from actual experiences to capture the dynamics of the POMDP in a latent space. The actual experiences are stored in replay buffer, allowing them to be repeatedly sampled for training the world model. The world model consists of a variational autoencoder (VAE) KKingma
According (Tay al. 2021 long sequence is defined as having length of 1,000 more.
2

---

Published as a conference paper at ICLR 2025
Fr
f2
de
d2
d2
Co Aaoll
0
CTAnno
C. C.
0
B2
G Ae
B.
it 1
G Bag
Head
Drama world model architecture. At each sequence index f, the raw game frames are encoded into and combined with the action a, as input to the Mamba blocks. The input channel dimension is divided by the head dimension p to generate the deterministic recurrent state d, This recurrent state d, is used predict the next embedding 21+1, reward ft. and termination flag é which represent the outcomes based on the current frame and action. The decoder reconstructs the original frame from the encoded embeddings 2, rather than from the predicted embeddings 2, The Mamba-2 block employs a semi-separee semi-separable matrix structure, which can be decomposed into x sub-matrices, enabling more efficient computation and processing.
Figure 1:
& Welling 2014 Hafner et al., 2021I, sequence model, and linear heads to predict rewards and termination flags. The details of world model discussed in Section 2.2
After each update to the world model, a batch of experiences is sampled from the replay buffer to initiate process called 'imagination', Starting from actual initial observation and using action generated by the current behaviour policy, the sequence model generates the next latent state. This process is repeated until the agent collects sufficient imagined samples for policy improvement. We explain this process in detail in Section 2.3
2.1 STATE SPACE MODELLING WITH MAMBA
SSMs are mathematical frameworks inspired by control theory represent the complete state of system at given point in time. These models map input sequence to output sequence
E R + R', where / denotes the sequence length. In structured SSMs, hidden state H E R (n..) is used to track the sequence dynamics, as described by the following equations:
H. AH.-1 Bri
(I)
V/ = HTH,
where A E (n,n) B (), C E (*.1)) and H. E (n,1), in which n represents the prede- fined dimension of the hidden state that remains invariant to the sequence length. To efficiently compute the hidden states, it is common to structure A diagonal matrix, as discussed in (Gu
ct al. 2022b Gupta et al. 2022 Smith et al. 2023 Gu Dao 2024). Additionally, selective SSMs, such Mamba, extend the matrices (A, B. C') be time- varying. introducing extra
dimension corresponding to the sequence length. The shapes these time-varying matrices are
E (T.NNN) B E (T.N)) and C (TT,N)
PIn Mamba, the time variation of A is influenced by discretisation parameter A. For more details, please refer (Gu & Dao 2024)

---

Published as a conference paper at ICLR 2025
Dao & Gu (2024) introduced the concept of structured state space duality (SSD), which further restricts the diagonal matrix A to be scalar multiple of the identity matrix, forcing all diagonal elements to be identical. To address the resulting reduced expressive power, Mamba-2 introduces multi-head technique, akin to attention, by treating each input channel as p independent sequences. Unlike Mamba, which computes SSMs as recurrence, Mamba-2 approaches the sequence trans- formation problem through matrix multiplication, which is more GPU-efficient:
yr = =c= H.
yr Σ T. A,, 1=0
where Ans A,A,-1 ... A441. This allows the SSM be formulated matrix transformation:
y SSM(x; A, B, C)
Mr
M₁.1 {o A.
if 2
jf
Mamba-2 reformulates the state state-space equations as a single matrix multiplication using semi- separable matrices (Vandebril et al. 2005 Dao & Gu (024), which is well known in computational linear algebra, as shown by Figure The matrix M can also written as:
M = L o CBT (TT2)
01
L
a201
02
1
a7-1
a1 ................................ A2
dT-1
where E [0, is an input-dependent scalar. The matrix L bridges the SSM mechanism with the causal self-attention mechanism by removing the softmax function and applying a mask matrix L to the 'attention-like' matrix. It is equivalent to causal linear attention when all = 1. As result, Mamba-2 achieves 2-8 times faster training speeds than Mamba, while maintaining linear scaling with sequence length.
2.2 WORLD MODEL LEARNING
Our world model has two main components: an autoencoder and a sequence model. Additionally includes two MLP heads for reward and termination predictions. The architecture of the world model is illustrated in Figure
2.2.1 DISCRETE VARIATIONAL. AUTOENCODER
The autoencoder extends the standard variational autoencoder (VAE) architecture (Kingma & Welling 2014) by incorporating fully-connected layer discretise the latent embeddings, sistent with previous approaches (Hafner et al., 2021 Robine et al. 2023 Zhang et al.) 2023). The raw observation is three-dimensional image, O E [0, th, at time step t. The encoder compresses the observation into a discrete latent vector, denoted z, p(z, O.). The decoder reconstructs the raw image, O, from 2,- Gradients are passed directly from the decoder to the encoder using the straight-through estimator, bypassing the sampling operation during backpropa- gation (Bengio et al. 2013).
2.2.2 SEQUENCE MODEL
The sequence model simulates the environment in the latent variable space, žt, using a deterministic state variable, d, Note that this is distinct from the hidden states typically used by SSMs, like Mamba and Mamba-2, to track dynamics. At each time step t, the next token in the sequence is determined by both the current latent variable, 2, and the current action dz To integrate these,
4

---

Published as a conference paper at ICLR 2025
concatenate them and project the result using fully connected layer before passing it to the sequence model. Given a sequence length I, the deterministic state derived from all previous latent variables and actions. The sequence model can be expressed as:
Seuqnce model:
d: = f( (e-
(5)
Latent variable predictor:
2141 p(2
We implement the sequence model with Mamba-2 (Dao Gu 2024). Specifically, each time a batch of samples, denoted as O [0, 255] is the batch size, / the sequence length, and h, w,, the image height, width, and channel dimension respectively. After encoding, the batch will be compressed to Z E (b...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 X E R (b.i.d) of the Mamba blocks. To fully leverage GPU parallelism, the training process must strictly avoid sequential dependencies. That is, time step t, the sequence model predicts the latent variable 2,+1, and its target ž++1 depends solely on the observation O++1 as shown in Figure 1 Unlike DreamerV3, where 1++1 p(zc+1 1 de), this approach eliminates sequential dependence.
drawn from the experience buffer & where b
Mamba processes the input tensor It into a sequence of hidden states H E R (b.--1n)) which are then mapped back the deterministic state sequence Db, using time-varying parameters. Since the hidden states operate in fixed dimension n (unlike standard attention mechanisms, where the state scales with the sequence length), Mamba achieves linear computational complexity in l.
Mamba-2 applies similar transformation but leverages matrix multiplication. The input tensor X's dimension d is first split into d/p heads, which are processed independently. The transformation matrix is a specially designed semiseparable lower triangular matrix, which can be decomposed into q x q blocks. Specialised blocks handle causal attention over short ranges and hidden state transformations, enabling efficient GPU computation.
2.3 BEHAVIOUR POLICY LEARNING
The behaviour policy trained within the 'imaagnation',, an autoregressive process driven by the sequence model. Specifically, batch of bing trajectories, each of length limg. is sampled from the replay buffer. Leveraging Mamba's i efficiency with long sequences we use real-world transitions to estimate more informative hidden state for the 'imagination' process. Rollouts begin from the last transition of each sequence (at step Lemg) and continue for h steps. Notably, the rollout does not stop when an episode ends, unlike the prior SSM-based meta-RL model (Lu et al. 2023) where the hidden state must be manually reset, as the Mamba- sequence model automatically resets the state episode boundaries (Gu & Dao 2024).
A key difference between Mamba- and transformer- based world models lies in the 'imagination' process: Mamba updates inference parameters independently of sequence length, accelerating the 'imagination" process, which is a major time-consuming phase in model-based RL. The behaviour policy's 's state concatenates the prior discrete variable &n with the deterministic variable d, exploit the temporal information. While the behaviour policy utilises standard actor-critic architecture, other on-policy algorithms can also be applied. this work, we adopt the recommendations from (Andrychowicz. et al.] (202T) and adjust the loss functions and value normalisation techniques as described in (Hafner et al. 2023).
2.4 DYNAMIC FREQUENCY- BASEDD 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 & Barto 1998). These inaccuracies particularly common carly in training when the model is fitted to limited data. Thus, propose a samplc-efficient method to address this issue, i.c., Dynamic Frequency-
Sampling (DFS).
The primary objective is to sample transitions that the world model has sufficiently learned to ensure reliable 'imagination'. To accomplish this, we maintain two vectors during train- ing, each matching the length of the transition buffer [g]. For the world model, where v e z for i E f1, [ which tracks the number of times a tran-
(v), U....., VIej), sition has been sampled to improve the world model. The resulting sampling probabilities
5

---

Published as a conference paper at ICLR 2025
computed as, (P1, P... = softmax(- similar Robine al. 2023 For 'imagi- nation', b (b1, b2. b(g)), where b, E Et for {1, which counts the times that the transition has been sampled to improve the behaviour policy. The sampling probabilities are so softmax(f(ee) b)), where f(v, b)
denoted as, (P1. P...P....PPP..))
- b max(0, v b)- During training, two cases arise: 1) 3i E (8), 12 2 b;, J(,,, b,) = 0, In this case, the transition has been trained more frequently with the world model than with the behaviour policy, suggest- ing that the world model is likely capable of making accurate predictions from this transition. 2) bj, signaling that the transition is either likely under-trained
Bi E [8|, 02 < bi, f(v,, b,) for the world model rollouts overfitted to the behaviour policy. Consequently, the probability of selecting this transition for behaviour policy training decreases. These two mechanisms ensure that 'imagination' sampling favours transitions learned by the world model, while avoiding excessive determinism.
EXPERIMENTS
In this work, the proposed Drama framework is implemented on top of the STORM infrastructure (Zhang et al. 2023). We evaluate the model using the Atari100k benchmark (Kaiser et al. 2020), which is widely used for assessing the sample efficiency of RL algorithms. Atar100kk limits interac- tions with the environment 100,000 steps (equivalent 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.1 ATARI100R 10K RESULTS
We compare Drama against SOTA MBRL algorithms across 26 Atari games in the Atari100k bench- mark. In Table T the 'Normalised Mean' refers the average normalised score, calculated (evaluated_ score- random caaennsscorr//eeee (Iumannsccree score-randon, _score). For each game, we train secluatdd_coreeeeeee Drama with independent seeds and track training performance using a 5-episode running average, as recommended by Machado et al. (2018). practice also followed in related work (Halner et al., 2023).
Despite utilising an extra-small world model (7M parameters, referred to as the XS model). Drama achieves performance comparable IRIS and TWM. To enable a like-for-like comparison between Drama and DreamerV3 with a similar number of parameters, we evaluate the learning curves of Drama and 12M-parameter variant of DreamerV3 (referred as DreamerV3XS) on the full Atari 100K benchmark. As shown in Figure 4 in the appendix, Drama significantly outperforms Dreamer achieving a normalised mean score of 105 compared to 37 and normalised median score of 27 compared to 7, as presented in Table
Table T demonstrates that Drama, with Mamba-2 as the sequence model, both sample- and parameter-efficien For comparison, SimPLe (Kaiser et al.) 2020) trains video prediction model to optimise a PPO agent (Schulman et al.) 2017). while SPR (Schwarzer et al.) (02T) uses a se- quence model to predict in latent space, enhancing consisteney through data augmentation. TWM (Robine et al.) (202) employs a Transformer-X architecture to capture dependencies among states, actions, and rewards, training a policy-based agent. This method incorporates short-term tempo- ral information into the embeddings to avoid using the sequence model during actual interactions. Similarly, IRIS (iicheli et al. 2023) uses Transformer its sequence model, but generates new samples in image space, allowing lixel-level feature extraction for behaviour policies. DreamerV3 (Hafner et al.) (023), which employs an RNN-based sequence model along with robustness tech- niques, achieves superhuman performance on the Atari 100k benchmark using 200M parameter model-20 times larger than our XS model. STORM (Zhang et al., 2023), which adopts many of DreamerV3's robustness techniques while replacing the sequence model with transformer, reaches similar performance on the Atari100k benchmark as DreamerV3.
Drama excels in games like Boxing and Pong, where the player competes against autonomous opponent in simple, static environments, requiring a less intense autoencoder. This strong perfor- mance indicates that Mamba-2 effectively captures both ball dynamics and the opponent's posi- tion. Similarly, Drama performs well in Asterix, which benefits from its ability to predict ob- ject movements. However, Drama struggles in Breakout, where performance can be improved

---

Published as a conference paper at ICLR 2025
Random
Human
PPO
SimPLe
SPR
TWM
IRIS
STORM
DreamerV3
DramaXS
Alien
228
7128
276
617
842
675
420
984
1118
820
Amidar
1720
74
180
143
205
97
Assault
742
327
527
566
1524
801
Asterix
210
8503
292
1128
962
1117
854
1028
1062
1632
BankHliist
14
753
14
34
345
467
641
398
137
BattleZone
2360
37188
4031
14834
5068
13074
13540
20300
10860
Boxing
12
3
36
78
70
80
82
78
Breakout
30
16
20
ChopperCommand
2222
1642
CrazyClimber
10780
35829
14675
62584
36700
71820
59324
66776
86225
83931
DemonAAtack
152
1971
160
208
518
350
2033
165
577
201
Freeway
Frosthite
65
4335
127
237
1171
1476
259
1316
3377
785
Gopher
258
2412
368
597
661
1675
2236
8240
2160
2757
2657
7254
11044
7946
Jameshond
29
303
100
366
362
463
509
540
372
Kangaroo
52
3035
55
51
3617
1240
838
4208
2643
1384
Krull
1598
2666
2205
6349
8171
9693
KunsteMMstter
258
22736
2090
14862
14783
24535
21760
26183
25900
23920
MsPacman
307
6952
366
1480
1318
1588
999
2673
1521
2270
Pong
19
PrivateEye
25
69571
100
35
87
100
7781
3238
Qhert
164
13455
317
1289
866
3331
746
4522
2921
796
RoadRunner
12
7845
602
5641
12213
9109
9615
17564
19230
14020
Seaquest
68
42055
305
683
558
661
525
962
497
UpNDown
533
11693
1502
3350
10859
15982
3546
7985
46910
7387
Normalised Mean
100
105
127
Normalised Median (%)
100
13
51
49
27
Comparison of game performance metrics for various algorithms across multiple Atari games. For F 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. DreamerV3 notably uses 200M parameter model and achieves good results in series of diverse tasks. STORM does report the number of trainable parameters.
Table 1:
with robust autoencoder in Figure 6 Additionally, Drama excels in games like Krull 1 and MsP acma, which require longer sequence memory, but faces challenges in sparse reward games like Jamesbond and PrivateEye.
3.2 ABLATION EXPERIMENTS
In this section, present three ablation experiments to evaluate key components of Drama. First, we compare dynamic frequency-based sampling performance against uniform sampling on the full Atari 100k benchmark, demonstrating its effectiveness across diverse environments. Secondly, compare Mamba and Mamba-2 on subset of Atari games, including Krull, Boxing, Freeway, and Kangaroo, highlight the differences in their performance when applied to dynamic game- play scenarios. Lastly, we compare the ueng-eeueence processing capabilities of Mamba, Mamba-2, and GRU in custom Grid World environment. This experiment focuses on a prediction task using different sequence models, offering insights into their sequence modelling capabilities, which are crucial for MBRL applications especially long-sequence modelling important.
3.2.1 DYNAMIC FREQUENCY BASED SAMPLING
In this experiment, we compare DFS with the uniform sampling method in the Mamba-2-basse Drama the full Atari100k benchmark. In Table DFS is more effective than the uniform sam- pling overall, achieving a 105% normalised mean score (vs. the uniform's 80%), despite both meth- ods sharing similarly median performance (27% 28%). As shown in Figure DFS shows signif-

---

Published as a conference paper at ICLR 2025
10
Krull
30
Boxing
1.5
Freeway
1.5
Kangaroo
Norm lsed Score
0.5
U.5
2
0.0
202
<3c
6.3
10c
120k
08-
20cc
-00g
601c
Bicc
100c
04
20k
40c
60c
BCC
100K
02
20k
40c
60k
10c
1006
Step
Mamboa-
Narmba 2
Mumba-1 (Man
Hambba
(Min-Mox) Moxt
Figure 2: Mamba Mamba-2. Mamba2 has shown superior performance to Mamba in three of four games. Both Mamba and Mamba-2 use DFS in this experiment.
icant advantages in games requiring adaptation to evolving dynamics, such as Al Alien, Asteri BankHeis it, and Seaquest. Additionally, DFS performs well in opponent-bane ooponent-based games such as Bexing and Pong, where exploiting the weaknesses of the opponent AI is essential. However, DFS performs less effectively in games like Breakout and KungFuMaa uMaster,, likely because the critical game dynamics are accessible early in the gameplay.
3.2.2 MAMBA VS. MAMBA-2
As mentioned in Sec 2.1 Mamba-2 imposes restrictions the diagonal matrix A to improve ef- ficiency. However, whether these restrictions degrade performance of SSMs remains unclear, as prior work lacks conclusive theoretical empirical evidence (Dao & Gul (024). In response this gap, we compare Mamba-2 and Mamba as the backbone the world model in model-based RL. We conducted ablation experiments using DFS, with both architectures configured under identical hyperparameters.
Figure illustrates that Mamba-2 outperforms Mamba in the games Krull, Boxi ng 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 plateau in score. As shown, Mamba experiences a score plateau in Krull, whereas Mamba-2 successfully overcomes this challenge, leading to higher performance. Note that Freeway is sparse reward game requiring high-quality exploration. A positive training effect is achieved only by combining DFS with Mamba-2 without any additional configuration.
3.2.3 SEQUENCE MODELS FOR LONG-SEQUENCE PREDICTABILITY TASKS
To assess the efficiency of Mamba and Mamba-2 2 in long-range modelling compared to Transformers and GRUs, which are widely used in recent MBRL approaches, we present simple yet representa- tive grid world environmen as illustrated in Figure 3a The learning objectives here are twofold: 1) the sequence model must reconstruct (predict) the correct grid- woold geometry over a long se- quence and 2) the sequence model must accurately generate the agent's location within the grid world, reflecting the prior sequence of movements. To achieve this, we represent a trajectory as a long sequence by flattening consecutive frames (row-wise tokenisation of frames) and separating each frame with a movement action a. Let the size of the grid world be lg. Then, each frame can be tokenised into sequence of length 1, = 1,9 + 1, as depicted in Figure 3b) Since >> ly. the task demands strong long-range sequence modelling to ensure geometric and logical consistency in predictions-a core requirement for MBRL sequence models.
We compare GRU, Transformer, Mamba, and Mamba-2 in this grid world environment, where l = 5 and I, 26, considering two sequence lengths: short sequence length / = 8 x ly and a long sequence length 64 1y. Performance measured via training time, memory usage and reconstruction error (lower time consumption and reconstruction error indicate better en- vironment understanding). Results show that, Mamba and Mamba-2 achieve equivalent low error and short training time in both sequence lengths compared to other methods. However, Mamba-2 demonstrates lowest training time over all methods. These findings confirm that the proposed Manba-based architecture presents strong capability to capture essential information, particularly in scenarios involving long sequence lengths.
"mplemmnntttion hased Torres Legueet 2024
8

---

Published as a conference paper at ICLR 2025
Flour
Wall
Agent
Goel
invalid
south
cast
north
Qg
Q22
an
(b)
Figure 3: Illustrations of the grid world environment and reconstruction into sequential format. (a) Sequence of consecutive frames in the grid world environment. The Example presents sequence of consecutive frames, arranged from left to right. Each frame represents 5 x grid, where the outer 16 cells black walls, and the central 3 x 3 grid is the reachable space. The red cell is the controllable agent, which moves according to a random action, and the yellow cell is a fixed goal. The sequence of frames, from left right, illustrates the movement of the agent following the action sequence: east + south + east + north. Once the yellow cell is reached by the agent, the location of the agent and goal will be reset randomly, long sequence. Each grey-shaded box contains 25 flattened grid tokens and one action token.
(b) Reconstructing the grid world into
Method
Training Time (ms)
Memory Usage (%)
Error (%)
208
25
13
15.6 2.6
Mamha-2
1664
214
55
14.2 = 0.3
208
34
14
13.9 0.4
Mamba
1664
299
52
0.4
208
75
66
21.3 = 0.3
GRU
1664
628
68
34.7 25.4
208
45
17
24.7 = 7.4
Transformer
1664
OOM
Table 2: Performance comparison of different methods in the grid world environment. Memory usage is reported as percentage of an 8GB GPU. The error represented as mean + stan- dard deviation. The training time refers the average duration per training step. Notably, the Transformer encounters an out-of-memory (OOM) error during training with long sequences. All experiments are conducted laptop. The definition of Error (%) is provided in Appendix A.6
4 RELATED WORK
4.1 MODEL-BASED RL
The origin of model-based RL can traced back to the Dyna architecture introduced by [Sutton & Barto (1998). although Dyna selects actions through planning rather than learning. Notably, Sutton & Barto (1998) also highlighted the suboptimality that arises when the world model is flawed, espe- cially the environment improves. The concept of learning in 'imagination'
first proposed by Ha Schmidhuber (2018), where a world model predicts the dynamics of the environment. Later, SimPLe (Kaiser al.] 2020) applied MBRL. Atari games, demonstrating improved sample effi-

---

Published as a conference paper at ICLR 2025
ciency compared to SOTA model-frre algorithms. Beginning with Hafner et al. (2019), 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.] 2020: 2021 2023). More recently, inspired by the suc- cess of transformers in NL.P, many MBRL studies have adopted transformer architectures for their sequence models. For instance, IRIS (Micheli et al. 2023) encodes game frames sets of tokens using VQ-VAE (Oord et al.) 2017) and learns sequence dependencies with a transformer. In IRIS, the behaviour policy operates on raw images, requiring an image reconstruction during the 'imagi- nation' process and an additional CNN-LSTM structure to extract information. TWM (Robine et al.) 2023), another transformer- baase world model, uses a different structure. It stacks grayscale frames and does not activate the sequence model during actual interaction phases. However, its behaviour policy only has access to limited frame history, raising questions about whether learning from to- kens that already include this short-term information could be detrimental to the sequence model. STORM (Zhang et al., 2023), closely following DreamerV3, replaces the GRU with a vanilla trans- former. Additionally, it incorporates demonstration technique, populating the buffer with expert knowledge, which has shown to be particularly beneficial in the game F reeway.
4.2 STRUCTURE STATE SPACE MODEL BASED RL
Structured SSMs originally introduced to tackle long-range dependency challenges, comple- menting the transformer architecture (Gu et al. 2022a Gupta et al. (022). However, Mamba and its successor, Mamba-2, have emerged as powerful alternatives, now competing directly with trans- formers (Gu & Dao 2024 Dao & Gu 2024). Deng et al. (2023) implemented an SSM-based world model, comparing against RNN- based and transformer- based models across various pre- diction tasks. Despite this, while SSMs have been applied to world model-based RL (e.g.. Recall to Imagine (R2I) (Samsami et al. 2024)). architectures like Mamba and Mamba-2 remain untested in this framework. Mamba has recently been applied to offline RL, either with a standard Mamba block (Lv et al.) 2024) or Mamba-attention hybrid model (Huang et al.) 2024). Lu al.) (2023) proposed applying modified SSMs to meta-RL, where hidden states are manually reset episode boundaries. Since both Mamba and Mamba-2 are input-dependent, such resets are unnecessary. No- tably, R21 leverages advanced SSMs to enhance long-term memory and credit assignment in MBRL, achieving SOTA performance in memory-intensin tasks, though it exhibits slightly weaker overall performance compared to DreamerV3 (Samsami et al.) 2024).
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, novel sampling method effectively mitigates suboptimality during the early stages of training, contributing to lightweight world model (only 7 million trainable parameters) that is accessible and trainable on standard hardware. Overall, our method achieves normalised score competitive with other SOTA RL algorithms, offering a practical and efficient alternative for model- -baed RL. systems. Although Drama enables longer training and inference sequences, it does not demonstrate decisive advantage that would allow it to dominate other world models on the Atari 100k benchmark. An interesting direction for future work to explore tasks where longer sequences drive superior performance in model-based RL. Additionally. it would be valuable to investigate whether Mamba can help address persistent challenges in model-based RI, such as long-horizon planning, behaviour learning, and informed exploration.
ACKNOWLEDGMENTS
This publication has emanated from research conducted with the financial support of Taighde Éire- ann Research Ireland under Frontiers for the Future grant number 2//957-A/8957 and grant num- ber 18/CRT/6223. For the purpose of Open Access, the author has applied a CC BC public copyright licence any Author Accepted Manuseript version arising from this submission.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Marcin Andrychowiez, Anton Raichuk, Piotr Stańczyk, Manu Orsini, Sertan Girgin, Raphael Marinier, Léonard Hussenot, Matthieu Geist, Olivier Pictquin, Marcin Michalski, Sylvain Gelly, and Olivier Bachem. What Matters In On-Policy Reinforcement Learning? A Large-Scale Em- pirical Study. In International Conference on Learning Representations, 2021.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or Propagating Gradi- ents Through Stochastic Neurons for Conditional Computation, August 2013. URL http: //arxiv.oro/abs/13008 3432 arXiv:1308.35 [cs].
Christopher Bemer, Greg Brockman, Brooke Chan, Vicki Cheung, Przemystaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub Pachocki, Michael Petrov, Henrique P.d.O. Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota 2 with Large Scale Deep Reinforcement Learning, wtt.a////9....6.... arXiv:1912.06680 [cs,
December 2019. URL stat].
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision Transformer: Reinforcement Learning via Sequence Modeling. In Advances in Neural Information Processing Systems, volume 34, pp- 15084-15097, June 2021.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. In Proceedings of the 41st International Conference on Machine Learning. volume 235, pp. 10041-10071, 2024.
Branton DeMoss, Paul Duckworth, Nick Hawes, and Ingmar Posner. DITTO: Offline Imita- tion Learning with World Models, February 2023. URL http:://xxxv...rr/aa/.. 03086 arXiv:2302.03086.
Fei Deng, Junyeong Park, and Sungjin Ahn. Facing Off World Model Backbones: RNNs, Trans- formers, and S4. Advances in Neural Information Processing Systems, volume 36, pp. 72904- 72930, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.
Albert Gu and Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. In First Conference on Language Modeling, 2024.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently Modeling Long Sequences with Structured State Spaces. In International Conference on Learning Representations, 2022a.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Re. On the Parameterization and Initializa- tion of Diagonal State Space Models. In Advances in Neural Information Processing Systems, 2022b.
volume 35, pp.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces effective structured state spaces. In Advances in Neural Information Processing Systems, volume 35, pp. 22982- 22994, 2022.
David Ha and Jürgen Schmidhuber. Recurrent World Models Facilitate Policy Evolution. In Ad- vances Neural Information Processing Systems, volume 31, 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft Actor-Critie Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor. In Proceedings of the 35th International Conference Machine Learning, volume 80, pp. 1861-1870, 2018.
11

---

Published as a conference paper at ICLR 2025
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning Latent Dynamics for Plamning from Pixels. In Proceedings of the 36th International Conference Machine Learning, volume 97, pp. 2555-2565. 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to Control: Learn- ing Behaviors by Latent Imagination. In International Conference on Learning Representations, 2020.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with Discrete World Models. In International Conference on Learning Representations, 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering Diverse Do- mains through World Models, 2023. URL
arXiv:2301.08 1.0104 [cs, stat].
Sili Huang, Jifeng Hu, Zhejian Yang. Liwei Yang, Tao Luo, Hechang Chen, Lichao Sun, and Bo Yang. Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling. In Advances in Neural Information Processing Systems, volume 37, pp. 72688-72709, 2024.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model-Based Reinforcement Learning for Atari. In International Conference on Learning Representations, 2020.
Diederik Kingma and Max Welling. Auto-Encoding Variational Bayes. In International Confer-
ence on Learning Reprexentations, 2014.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured State Space Models for In-Context Reinforcement Learning. In Thirty-eeventh Conference on Neural Information Processing Systems, 2023.
Qi Lv, Xiang Deng, Gongwei Chen, Michael Yu Wang, and Liqiang Nie. Decision Mamba: A Muati-Grained State Space Model with Sell-Evolution Regularization for Offline RL. In Decision Mamba: A Mulri-Grained State Space Model with Self- Evofutvotution Regularization for Offline RI., volume 37, pp. 22827-22849, 2024.
Marlos C. Machado, Mare G. Bellemare, Erik Talvitie, Joel Veness Matthew Hausknecht, and Michael Bowling. Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents. Journal of Artificial Intelligence Research, 61:523-562, March 2018. ISSN 1076-9757. doi: 10.161 10.1613/jair.56..
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are Sample-Efficient World Mod- els. International Conference on Learning Representations, March 2023.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing Atari with Deep Reinforcement Learning, December 2013. URL httt:t//p///xx...t// org abs/ 1at.e.///1...... 5602, arXiv:1312.5602 [cs].
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural Discrete Representation Learn- ing. In Advances in Neural Information Processing Systems, volume 30, 2017.
Jan Robine, Mare Höftmano, Tobias Uelwer, and Stefan Harmeling. Transformer-based World Mod- els Are Happy With 100k Interactions. In International Conference on Learning Representations, March 2023.
Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar. Mastering Memory Tasks with World Models. In International Conference on Learning Representations, 2024.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lil- licrap, and David Silver. Mastering Atari, Go, chess and shogi by planning with learned model. Nature, 588(7839):604-693 December 2020. ISSN 0028-0836, 1476-4687. doi: 10. 10.10388/411866000-03333
12

---

Published as a conference paper at ICLR 2025
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Pol- icy Optimization Algorithms, August 2017. / wttps//.rwgaaab//77.00555555 arXiv:1707.06347 [es].
Max Schwarzer, Ankesh Anand, Rishab Goel, Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-Efficient Reinforcement Learning with Self-Predictivy Representations. In Interna- tional Conference on Learning Representations, 2021.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George Van Den Driess- che, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):4844445 2016. ISSN 0028-0836, 1476-4687. doi: (033/nature16666
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen. Timothy Lillicrap. Fan Hui, Laurent Sifre, George Van Den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of Go without human knowledge. Nature, 550(7676):354 2017. ISSN 0028-0836, 1476-4687. doi: 10. 10.00//auuue.
Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In International Conference Learning Representations. 2023.
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: an introduction. Adaptive com- putation and machine learning. MIT Press, Cambridge, Mass, 1998. ISBN 978-0-262-193388-6
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long Range Arena: A Benchmark for Efficient Transformers. In International Conference Learning Representations, 2021.
Alexandre Legres. mamba.py: A simple, hackable and efficient Mamba implementation in pure PyTorch and MLX., 2024. URL URL/httss: /tgp:///iittaaa com/ a XxxdrTL/mamma. py
Raf Vandebril, M Van Barel, Gene Golub, and Nicola Mastronardi. A bibliography on semiseparable matrices. Calcolo, 42:249-270, 2005. Publisher: Springer.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Infor- mation Processing Systems, volume 30, 2017.
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning. In Advances in Neural Informa- fion Processing Systems, volume 36, pp. 27147-27166, 2023.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vi- sion Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model. In Proceedings of the 41st International Conference on Machine Learning, volume 235, pp. 62429- 62442, 2024.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1
ATART10OK LEARNING CURVES
Alen
Amidar
ASSOuR.
Asterix
2.0
772
1.5
250
200
333
1.5
1.5
23
13.00
202
-20
ChapperCommand
CrstyClimter
Demontcack
3.25
2.5
23.00
1.1
033
20.24
700
mossbbbe
Sopher
Mero
Jemesdonnd
12:32
5000
2.58
4.22
182
Kargurco
SangruNenter
17.00
00.32
3.00
33.01
23.24
2.34
20.24
15.24
10.21
7ong
Prisstetye
22.00
000
11.00
400
12:00
4.5
200
332
5.50
Seaccust
Upp\Dow
20A
1854
225
2006
102
500
Drams
20.00
1002
1004
224
425
14B
70
Figure 4: Atari100k Learming Curve. This figure compares the performance of DramaXS (10 million parameters) and DreamerV3X3 (12 million parameters) on the Atari100k benchmark. DramaXS outperforms DreamerV3XS in most games. Exceptions include PrivateEye and Obert. where Dreamer V3XS performs better.
14

---

Published as a conference paper at ICLR 2025
Game
Random
Human
DramaXS
DreamerV3XS
Alien
228
7128
820
553
Amidar
1720
131
Assault
222
742
539
489
Asterix
210
8503
1632
669
BankHeist
14
753
137
27
BattleZone
2360
37188
10860
5347
Boxing
12
78
Breakout
2
30
7
4
ChopperCommand
811
7388
1642
1032
CrazyClimber
10780
35829
83931
7466
DemonAttack
152
1971
201
64
Freeway
30
15
Frostbite
65
4335
785
144
Gopher
258
2412
2757
287
Hero
1027
30826
7946
3972
Jamesbond
29
303
372
142
Kangaroo
3035
1384
584
Krull
1598
2666
9693
2720
KungFuMaster
258
22736
23920
4282
MsPacman
307
6952
2270
1063
Pong
-21
15
15
-10
PrivateEye
69571
90
207
Qbert
164
13455
796
983
RoadRunner
12
7845
14020
8556
Seaquest
68
42055
497
169
UpNDown
533
11693
7387
6511
Normalised Mean (%)
100
105
37
Normalised Median (%)
100
27
7
Table 3: Atari1008 100K performance table. DramaXS achieves significantly better performance than DreamerV3XS in compact model settings within model-based reinforcement learning, highlighting the parameter efficiency of Mamba- Mamba-based architectures.
15

---

Published as a conference paper at ICLR 2025
A.2
UNIFORM SAMPLING vs. DES LEARNING CURVES
Alen
Anridar
ASSOuP
Astb
144
2.34
332
2.4
1.5
200
492
1.0
500
200
BanWHeist
Balt aZone
Boatig
Broakiout
20.04
100
2.7
4.0
MW
13:04
3.c
200
2
ClopperCammand
Crazyllmmer
Demonttteck
Treasway
4.90
2.80
700
00.02
2.04
300
300
500
20.04
100
Prosthibe
Ocphar
0.0-
3.04
10:04
8.01
000
234
1.0%
4:02
200
Kangarse
Knul
KungFuNaster
MrPacman
0.04
10.04
20.04
0.00
4.02
5.0g
2.00
1.0
2.32
2.02
43c
10.04
1.00
2.08
Pong
Priateetve
Cbert
HeadRuuneer
300
3.58
-100
-000
1.04
Yewgueet
0gAilown
00.04
200
50.04
Unfore (er Mac
400
400
236
$00
2.56
30c
-200
2802
Figure Uniform Sampling vs. DFS Learning Curve. DFS outperforms uniform sampling in 11 games (e.g., Asterix BankHeist. Krull underperforms in 2 games (Breakout. KungFuMaster ), and matches performance 13 games The normalised mean score of DFS (105% surpasses uniform sampling (80% ), while the normalised median is comparable (27% vs. 28% ). DFS demonstrates stronger performance in games requiring exploiting the opponents' strat- egy (e.g., Pong, Boxing) but struggles in environments with early-stage dynamics (Breakout).
16

---

Published as a conference paper at ICLR 2025
Game
Random
Human
DFS
Uniform
Alien
228
7128
820
696
Amidar
1720
131
154
Assault
222
742
539
511
Asterix
210
8503
1632
1045
BankHeist
14
753
137
52
BattleZone
2360
37188
10860
10900
Boxing
12
78
49
Breakout
2
30
7
11
ChopperCommand
811
7388
1642
1083
CrazyClimber
10780
35829
83931
77140
DemonAttack
152
1971
201
151
Freeway
30
15
15
Frostbite
65
4335
785
975
Gopher
258
2412
2757
2289
Hero
1027
30826
7946
7564
Jamesbond
29
303
372
363
Kangaroo
52
3035
1384
620
Krull
1598
2666
9693
7553
KungFuMaster
258
22736
23920
24030
MsPacman
307
6952
2270
2508
Pong
-21
15
15
3
PrivateEye
25
69571
90
76
Qbert
164
13455
796
939
RoadRunner
12
7845
14020
9328
Seaquest
68
42055
497
384
UpNDown
533
11693
7387
5756
Normalised Mean (%)
100
105
80
Normalised Median
100
27
28
Table 4: The Atari100K performance table demonstrates that the Drama XS model, when paired with DFS, achieves higher normalised mean score compared to using the uniform sampling method. This highlights the effectiveness of DFS enhancing performance of Mamba-powered MBRL.
A.3 MORE TRAINABLE PARAMETERS
As model-based RL agents consist of multiple trainable components, hyperparameters tuning for each part can be computationally expensive and not the primary focus of this research. Prior work has demonstrated that increasing the neural network's size often leads to stronger performance on benchmarks (Hafner et al.) 2023). Figure 6, we demonstrate that Drama achieves overall better performance when using more robust autoencoder and larger SSM hidden state dimen- sion n. Notably, the model exhibits significantly improved results in games like Breakout and BankHe st, where pixel-level information plays crucial role.
17

---

Published as a conference paper at ICLR 2025
Breakout
Kangaroo
ChopperCommand
BankHeist
0.4
Nore leed Score
0.6
3.50
0.1
0.2
0.00
201
40x
60x
B0c
100k
70k
402
6912
Bite
100k
06
20c
40cc
60c
BCC
100t
08
306
102
802
800
1006
Step
Dramas
Dramass
amaa (Mix-Max)
Dramals Nox)
Figure 6: S model vs. XS model. We adjusted the game set to emphasise the importance of recog- nising small objects. The model features a more robust autoencoder than the XS model, with additional filters and 3M more trainable parameters. In terms of performance, the model signifi- cantly outperforms the XS model in Break kout and BankHeist, st. However, it underperforms in Kangaroo and shows comparable performance in ChopperConmand.
A.4 LOSS AND HYPERPARAMETERS
A.4.1 VARIATIONAL AUTOENCODER
The hyperparameters shown in Table correspond to the default model, also referred to XS in Figure [] For the model, we simply double the number of filters per layer to obtain a stronger autoencoder.
Hyperparameter
Value
Learning rate
4e-5
Frame shape (h, w. c)
(64, 64, 3)
Layers
5
Filters per layer (Encoder)
(16, 32, 48, 64, 64)
Stride
(1, 2, 2, 2)
Kernel
5
Weight decay
le4
Act
SiLU
Norm
Batch
Table 5: Hyperparameters for the autoencoder.
A.4.2 MAMBA AND MAMBA-2
Similar to the previous section, the values reported in Table 6 correspond to the default model. For the S model, we double the latent state dimension, thereby enabling the recurrent state to retain more task-relevant information. In the Mamba-2 model, the enhanced architecture accommodates larger latent state dimension without substantial increase in training time.
A.4.3 REWARD AND TERMINATION PREDICTION HEADS
Both the reward and termination flag predictors take the deterministic state output from the sequence model to make their predictions. Due to the expressiveness of the temporal information extracted by the sequence model, single fully connected layer is sufficient for accurate predictions.
18

---

Published as a conference paper at ICLR 2025
Hyperparameter
Value
Learning rate
4e-5
Hidden state dimension (d)
512
Layers
2
Latent state dimension (n)
16
Act
SiLU
Norm
RMS
Weight decay
lc-4
Dropout
0.1
Mamba-2: Head dimension (p)
128
Table 6: Hyperparameters for Mamba and Mamba-2. Except the head dimension is only for Mamba- the other hyperparameters are shared. The head number 512/128 4.
Hyperparameter
Value
Hidden units
256
Layers
Act
SiLU
Norm
RMS
Table Hyperparameters for reward and termination prediction heads.
The world model is optimized in end-to-end and self-supervised manner batches of shape (b, 1) drawn from the experience replay.
Σ (0, O.) + Ldgg (w) + 0.1 Crep (w)
((w)
i=1 reconstruction loss
(6)
In p(r((dd:;
In pp(.:/d;:
reward prodiction loss
termination prodiction loss
where
(wy max (1, KL (os(p(244)0441100 g(E (d;
Crep(w)- = max (1, KL [p(z,+1 1ww) (g(((2=-1l:;w)))
and sg(-) represents the stop gradient operation.
A.4.4 ACTOR CRITIC HYPERPARAMETERS
We adopt the behaviour policy learning setup from DreamerV3 (Hafner al. 2023) for simplicity and its demonstrated strong performance, since the behaviour policy model is not central to our primary contribution.
19

---

Published as a conference paper at ICLR 2025
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 (bimg)
1024
Imagine context length (limg)
8
Table 8: Hyperparameters for the behaviour policy.
A.5 PSEUDOCODE DRAMA
Algorithm Training the world model and the behaviour policy
Require: Initialize behavior policy To, world model fus and replay buffer &
1: Loop:
2: Phase 1: Data Collection
3: Collect experience tuple (O,, Ge, Te, ee) using me
4:
Store (O,, de Te. eq) into replay buffer &
5:
Phase 2: World Model Training
6:
Sample trajectories of length from &
7:
Update world model fa using sampled trajectories
8:
Phase 3: Behaviour Model Training
9:
Sample bing trajectories of length 4ing from
10:
Retrieve context from the first ling experiences from the world model fw Generate imagined rollout for h steps using the last experience
12:
Train behavior policy To with imagined rollout
Repeat
A.6
THE GRID WORLD ERROR CALCULATION
The Grid World environment task requires the sequence model to capture two types of sequences. The first, referred as the geometric sequence, involves reconstructing the spatial structure of the map. The environment consists of a grid surrounded by black walls, with single agent cell and goal cell positioned, while all remaining cells are plain floor tiles. Formally, the map M deñindd as a grid where Mii.) ji represents the cell at position (i.j). The geometric sequence requires the sequence model to encode the spatial relationships such that M [i,]]] satisfies the constraints walls (W), floor (F), agent (A), and goal (G), with walls forming the boundary:
W, if (i= 0 or (j - 0 1)1,
F. if (i,j) & A.G,
Miijj]
А. (i, agent position,
G, if (i. = goal position.
The geometric eтror Eg measures violations of the grid's structural constraints. defined as the number of boundary cells incorrectly classified non-wall (M[i.) * W when (i 0 or l 1) or (j 0 or j = I, 1) ). For interior cells, where < < lg - and < j < 1 - 1,
there must be exactly one agent and goal, with all remaining cells being floors.
20

---

Published as a conference paper at ICLR 2025
The second component, referred to the logic sequence, requires predicting the agent's 's next po- sition A, based on the prior action a;-1- This prediction requires the model to retain information about the prior action, reconstruct the geometric sequence, and infer the agent's subsequent position accordingly. The logic error, FG, is defined as a prediction failure, which occurs if: (1) the predicted frame contains invalid configurations (e.g., multiple agents in the interior), (2) the predicted agent position does not match the groudtruth position in the subsequent frame.
The Error (% %) presented in Table represents the average of Eg and E,.
21

---

Published as a conference paper at ICLR 2025
A.7 EXPERIMENT 'IMAGINATION FIGURES
In this section, we analyze reconstructed frames generated by the 'imagination' of the sequence model investigate potential causes of its poor performance in certain games, such as Breakout.
==1
002 001
t= 5
t 9
582
002
002 002
t t 13
== 17
4
Figure 7: Drama XS model's 'imagination' in Breakout. The model exhibits poor performance in Breakout, its autoregressive generation produces reconstructed frames that frequently omit the ball-a key visual element. This systematic omission likely undermines its ability to execute effective policies, contributing suboptimal task performance.
The discrepancies in reconstructed frames (Figure Figure g and the performance gains in Figure 6 collectively suggest that a more robust autoencoder enhances task performance in environments where pixel-level information is critical. This observation is further supported by the Drama XS model's strong performance in Pong (Figure D), game sharing core mechanics with Breakout (e.g., paddles and balls) but with reduced visual complexity due to the absence of multicolored bricks. While systematic analysis is warranted to validate this hypothesis, these results indicate that refining the autoencoder may serve as critical first step in alleviating performance limitations in visually demanding tasks.
22

---

Published as a conference paper at ICLR 2025
t=
003
t=5 5
9
001 003
t 13
= 17
4
Figure 8: Drama S model's 'imaagnation' in Breakout. The Drama S model exhibits significant improvements over the XS variant, with the ball-a critical game cemment-consistenn recon- structed in the majority of autoregressive frames. This enhancement suggests stronger capacity to encode pixel-level details, aligning with its superior task performance.
23

---

Published as a conference paper at ICLR 2025
t=
t= 5
t 9
0
= 13
4
t = 17
4
Figure 9: Drama XS model's 'imagination' in Pong. The Drama XS model exhibits strong perfor- mance in Pong, contrasting sharply with its suboptimal results in Breakout. While both games share core mechanics (c.g.. paddles and balls), Pong's 's absence of multicolored bricks reduces vi- sual complexity, thereby lowering demands on the model's frame-encoding capacity. Consequently, the ball-a a critical element- is consistently reconstructed in the majority of autoregressive frames, supporting effective policy execution.
A.8 WALL-CLOCK TIME COMPARISON OF SEQUENCE MODELS MBRL
As illustrated in Figure T0 we compare the wall-clock time efficiency of sequence models in the Atari100k MBRL task. The results demonstrate that Mamba and Mamba-2 outperform the Trans- former architecture during the imagination phase for all tested sequence lengths. While Mamba-2 exhibits marginal computational overhead compared to Mamba and the Transformer for shorter training sequences, it achieves superior efficiency for longer sequences, making it particularly ad- vantageous for tasks demanding long-range temporal modelling. All models were evaluated under identical experimental conditions, with comparable parameter sizes and training confi configurations.
24

---

Published as a conference paper at ICLR 2025
Training Time by Train Sequence Length
Transformer
300
Mamba2
Training Time Per Step (MS)
Mamba
250
150
100
64
128
256
512
1024
Train Sequence Length
Imagining Time by Imagine Sequence Length
160
Mamba
Transformer
Imagining Time Per Step (MS)
Mamba2
120
100
60
40
20
16
64
Imagine Sequence Length
(b)
Figure 10: Wall-clock time comparison of sequence models in MBRL. Experiments were conducted on consumer-gradd grade laptop with an NVIDIA RTX 2000 Ada Mobile GPU, ensuring practical rele- vance to resource-consee constraind settings. Notably, the Transformer model leveraged key-value (KV) cache to optimise inference speed. Results demonstrate that Mamba-2 achieves superior efficiency for longer sequences in both training and 'imagination' phases. However, it incurs slight com- putational overhead compared to the Transformer and Mamba during training at shorter sequence lengths.
25


Reviews:
[{"review_id": "Thwrl17mCr", "reviewer": "Reviewer_3CFw", "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.\n\nAdditionally, 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.", "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?\n- **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.  \n- **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.\n - **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?\n- **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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n- **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.\n- **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.\n- **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.", "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.\n- **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.\n- **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.\n- **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.\n- **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4oUHUcxPXK", "reviewer": "Reviewer_hT9b", "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.", "questions": "1. At line 222, it is not clear what targets mean.\n2. 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?\n3. What is $h_t$ in behavior policy learning? The deterministic variable is defined as $d_t$ in Eq 5.\n4. 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).\n5. 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.\n6. Was any experiment conducted to see if the behavior model underestimates rewards especially with limited data? \n7. For the ablation presented in Sec 3.2.2, were both variants Mamba-1 and Mamba-2 based WM trained with DFS?\n8. How does the proposed method compare with R2I [1]. Since they propose using SSMs, should it be included as another baseline?\n9. [Typo] Employs is written twice at line 132.\n\n#### References\n[1] Samsami et al., Mastering Memory Tasks with World Models, ICLR’24.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The proposed idea of using SSM for WMs is interesting as they provide crucial benefits over training with Transformers and RNNs. \n- The proposed method achieves good performance with significantly fewer parameters (7M) when compared with baselines.", "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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LDYtCaqvPm", "reviewer": "Reviewer_mwc8", "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.", "questions": "- [Q1] Could the decoder operate based on the output of Mamba-2, such that $d$ rather than $z$ serves as the input?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- [S1] The paper combines established methods in a novel way, effectively addressing an existing gap in world model research.\n- [S2] The proposed model is computationally efficient, requiring only 7 million trainable parameters, making it accessible.", "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.\n- [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.\n- [W3] Several presentation aspects should be improved:\n  - 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.\n  - Many figures contain small text that is difficult to read without zooming.\n  - Minor suggestions:\n    - Highlight the highest scores in Table 1 for easy reference.\n    - Correct notations for all \\hat{} terms (e.g., \\hat{r}_t instead of \\hat{r_t}).\n    - Address minor notation errors: e.g., incorrect $A$ on lines 148, 154, 169 and incorrect $T$ on line 177\n    - Variables on line 218 should be in math mode.\n    - Consider changing \"auto-generative\" to \"autoregressive\" on line 238?\n    - Revise line 264 to read \"tracks the number of *times* the transition has been used.\"\n- [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.\n\nI would consider raising my scores if these issues were addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JpB2WEOk1R", "reviewer": "Reviewer_mCEd", "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.", "questions": "- What is the value of h?\n- \"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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The combination of MAMBA and Sequence-based world models is novel\n- The authors demonstrate comparable results with a smaller model size", "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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)