EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling
Authors: Jia-Hua Lee
Bor-Jiun Lin
Wei-Fang Sun
Chun-Yi Lee
Venue: NeurIPS 2025 spotlight ()
Keywords: Model-based Reinforcement Learning
Atari 100k
Doom
Crafter
MAMBA
Diffusion
World Model
Primary area: reinforcement_learning
TL;DR: EDELINE combines diffusion models with state space models to create a world model for reinforcement learning that overcomes memory limitations in previous approaches.

Abstract: World models represent a promising approach for training reinforcement learning agents with significantly improved sample efficiency. While most world model methods primarily rely on sequences of discrete latent variables to model environment dynamics, this compression often neglects critical visual details essential for reinforcement learning. Recent diffusion-based world models condition generation on a fixed context length of frames to predict the next observation, using separate recurrent neural networks to model rewards and termination signals. Although this architecture effectively enhances visual fidelity, the fixed context length approach inherently limits memory capacity.
In this paper, we introduce EDELINE, a unified world model architecture that integrates state space models with diffusion models. Our approach outperforms existing baselines across visually challenging Atari 100k tasks, memory-demanding Crafter benchmark, and 3D first-person ViZDoom environments, demonstrating superior performance in all these diverse challenges. Code is available at https://github.com/LJH-coding/EDELINE.

EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling
Jia-Hua Leel Bor-Jiun Lin? Wei-Fang Sun Chun-Yi Lee²
Department of Computer Science, National Tsing Hua University, Hsinchu, Taiwan Department of Computer Science, National Taiwan University, Taipei, Taiwan 3 NVIDIA AI Technology Center (NVAITC), NVIDIA Corporation, Santa Clara, CA, USA
Abstract
World models represent promising approach for training reinforcement learning agents with significantly improved sample efficiency. While most world model methods primarily rely on sequences of discrete latent variables model environ- ment dynamics, this compression often neglects critical visual details essential for reinforcement learning. Recent diffusion-based world models condition generation on a fixed context length of frames to predict the next observation, using separate recurrent neural networks model rewards and termination signals. Although this architecture effectively enhances visual fidelity, the fixed context length approach inherently limits memory capacity. In this paper, we introduce EDELINE, unified world model architecture that integrates state space models with diffusion models. Our approach outperforms existing baselines across visually challenging Atari 100k tasks, memory deennddg Crafter benchmark, and 3D first-person ViZDoom environments, demonstrating superior performance all these diverse challenges. LJH- Code is available at htttt:://ggiiiii www.o///www.o.oon/ooo/ooooooool/
Introduction
World models ! constitute foundational element of modern reinforcement learning (RL) by simu- lating environment dynamics for agent planning and reasoning. The capacity to learn environment representations [2]3 facilitates policy optimization through imagined trajectories, which substantially enhances sample efficiency 4 relative to conventional RL approaches. This capability is especially valuable for real-world applications in robotics and autonomous systems. Recent advances have shown exceptional performance across diverse environments, including real-world tasks (
Existing world models fall into two principal paradigms: latent-space models and generative models. employ recurrent neural networks (RNNs) or variants to predict future states within a compressed latent space for efficient policy optimization. This compression, however, introduces information loss that compromises generality and reconstruction quality. Genera- tive models, particularly diffusion- based approaches [8] have transformed world modeling through high-fidelity visual predictions via noise-reversal processes. Nevertheless, prior generative models depend on fixed-length observation-actio extended temporal dependencies. This limitation presents a challenge especially in partially observ- able environments where agents must retain and reason over prolonged observation sequences for informed decisions. Moreover, the architectural segregation of reward prediction, termination signals, and observation modeling in existing frameworks can potentially lead to suboptimal representation sharing and optimization conflicts that further impair performance.
Latent-space approaches 6
7 2
windows that truncate historical context and fail to capture
In order to mitigate long sequence dependency issues, recent state space models (SSMs) 9 10 12 [3 provide a complementary advantage through their capacity to model long-term dependencies efficiently. With linear-time complexity and selective state updates [12, SSMs can process theo- retically unbounded sequences while preserving critical historical information. This capability is particularly valuable for world modeling, where accurate trajectory prediction often necessitates retention and reasoning across extended observation-action histories. the world modeling domain,
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

recent work such as R2I [14] has addressed fundamental challenges in long- sing-temmmnn memory and credit assignment, demonstrating the importance of enhanced temporal reasoning capabilities.
Based on these considerations, we introduce EDELINE (Enhancing Diffusion- World Models via LINear-Time Sequence Modeling), unified framework that integrates the advantages of diffusion models and SSMs. EDELINE advances the state-of-the-a (SOTA) through three key innovations: (I) Memory Enhancement: A recurrent embedding module (REM) based on Mamba SSMs that processes unbounded observation-action sequences to enable adaptive memory retention beyond fixed-context limitations, (2) Unified Framework: Direct conditioning of reward and termination prediction on REM hidden states that eliminates separate networks for efficient representation sharing, and (3) Dynamic Loss Harmonization: Adaptive weighting of observation and reward losses that addresses scale disparities in multi-task optimization. To validate EDELINE's s effectiveness, we conduct extensive evaluations. achieves 1.87 human normalized scores on the sample-efficiency challenging Atari 100k [15] benchmark and surpasses all model-based methods that do not use look-ahead search. The ablation studies confirm the effectiveness of each architectural component for world modeling performance. To substantiate EDELINE's capacity to preserve long temporal information for consistent predictions, we evaluate performance on environments that require long- term memory capability: MiniGrid-Memory 16. Crafter 17.. and ViZDoom [8] Both qualitative and quantitative results show superior temporal consistency in modeling and imaginary quality across 2D and 3D environments. Furthermore, EDELINE shows significant performance improvements compared to prior diffusion- based world models. Our contributions can be summarized as follows:
We introduce a unified architecture EDELINE that integrates a Next-Frame Predictor for future observation imaginary, Recurrent Embedding Module for temporal sequence processing, and Reward/Termination Predictor to address the long-term memory limitations in existing models.
EDELINE utilizes an SSMs-based embedding module that overcomes the fixed context limitations of prior diffusion-basse methods and enhances performance in memory demoory-demanding environment
Our experimental insights across various benchmarks validate EDELINE's precise prediction of reward-critical elements where prior diffusion based world models exhibit structural inaccuracies.
2 Related Work
2.1 Diffusion Models
Diffusion models have revolutionized high-resolution image generation through their noise-reversal process. Foundational works including DDPM 19 and DDIM (20) established core prineiples for
subsequent developments. Score-based models 21 22 enhanced sampling efficiency through gradi- ent estimation of data distributions, while energy- based models (33 introduced robust optimization properties via probabilistic state modeling. The application of diffusion models in RL has expanded significantly. These models serve as policy networks for efficient offline learning (24 25 22),, en- able diverse strategy generation in planning tasks 127 28 and provide novel approaches to reward modeling [29]. MetaDiffuser [30] showed effectiveness as conditional planners in offline meta- RL, while others have adopted diffusion models for trajectory modeling and synthetic experience generation B1..
2.2
World Models for Sample-Efficient RL
World models serve as fundamental component in model-based RL, enabling sample-efficient and safe learning through simulated environments. SimPLe [15) established the groundwork by introducing world models to the Atari domain and proposing the Atari 100k benchmark. Dreamer (6]
advanced this field through RL from latent space predictions, which DreamerV2 (7] further refined with discrete latents to mitigate compounding errors. DreamerV3 [2 achieved a sigmificant milestone by demonstrating human-level performance across multiple domains with fixed hyperparameters.
Recent architectural innovations have expanded world model capabilities. TWM 32 and STORM 33 incorporated Transformer architectures 34 for enhanced sequence modeling, while IRIS (35] developed a discrete image token language for structured learning. R2I [14] addressed fundamental challenges in long-term memory and credit assignment. The integration of generative modeling approaches has further advanced world model capabilities. DIAMOND breakthrough by incorporating diffusion models, which achieved superior visual fidelity and state-of- the-art (SOTA) performance on the Atari 100k benchmark. This success inspired broader applications of generative approaches in interactive environments.
marked a significant

---

2.3 Generative Game Engines
Another line of research explores training generative world engines using pre-collected datasets instead of RL-in-the-loop learning. GameGAN 36 pioneered this direction through GAN-based environment modeling, while Genie (37 advanced these capabilitiess by generating complex plat- former environments from image prompts. GameNGen [88] established new standards for visual fidelity and scalability through diffusion-based environment simulation. Recently, GAIA [39] 40) and Pandora [4 have further expanded this field by developing generative world models that leverage video, text, and action inputs to produce realistic scenarios. It is important to note that these generative world methodologies do not involve RL-in-the-loop training and are therefore orthogonal to our work.
2.4 State Space Models (SSMs) in Reinforcement Learning
SSMs have demonstrated remarkable effectiveness in modeling long-term dependencies and struc- tured dynamics. Foundational works [42 9
introduced efficient sequence modeling approaches. Recent applications in RL include structured SSMs for in-context learning [45], DecisionMamba [46]ss adaptation of Decision Transformer gration of SSMs in world model learning.
43 IO,
IT
44
[2] [3]
47 architecture, and Drama [48] inte-
Our work advances the of-the-arrr by synergistically combining diffusion baasdd world models with SSMs. While previous works have explored these approaches separately, EDELINE demonstrates that their integration enables superior temporal consistency and extended imagination horizons in world model learning. This architectural innovation enhances the visual fidelity of diffusion models through the integration of SSMs' efficient temporal processing capabilities, which addresses the key limitations of existing world model approaches.
3 Background
In this section, we focus on the essential concepts necessary for understanding our EDELINE framework. We provide additional background material on score-bases core-based diffusion models and multi- task world model learning in Appendix C
3.1 Reinforcement Learning and World Models
The problem considered in this study focuses on image-hased reinforcement learning (RL.), for- mulated Partially Observable Markov Decision Process (POMDP) [49 defined by tuple (S, O, P, RR, O, 7). Our formulation specifically considers high-dimensional image observations as inputs, as described in Section 1. The state space S comprises states s, E S, while the action space A can be either discrete or continuous with actions a E A. The observation space O contains image observations o E xx A transition function P the environment dynamics p($ St, a.), while the reward function R S x A x S > R maps
A x S + [0, 1] characterizes
transitions scalar rewards Ft E R. The observation function O S x O + [o, 1] establishes observation probabilities p(o, (s The objective centers on learning policy π that maximizes the expected discounted return E. E220 7"1,). with discount factor 7 E [0. 1]. Model-based Reinforce- ment Learning (MBRL) (50] achieves this objective by learning world model that encapsulates the environment dynamics p(0 reloge, age). MBRL enables learning in imagination through three systematic stages: (1) collecting real environment interactions, (2) updating the world model, and (3) training the policy through world model interactions.
3.2 Linear-Time Sequence Modeling with Mamba
SSMs provide an alternative paradigm to attention- based architectures for sequence modeling. The Mamba architecture [12] introduces a selective state space model that offers linear time complexity and efficient parallel processing. which employs auriable-depender projection matrices to implement its selective mechanism, thus overcoming the inherent limitations of computational inefficiency and quadratic complexity in conventional SSMs [42 12]. The foundational mechanism of
Mamba is characterized by a linear continuous-time state space formulation via first-order differential equations as follows:
Ox(()) = Az(t) + B(u(t))u(t)
Ot
(1)
y(t) = C(u(t))=(t).
where st() represents the latent state, v(()) denotes the input, and y(() indicates the output. The matrix adheres to specifications from (10) The primary innovation compared to traditional SSMs lies in B(u(t)) and C(u(t)), which function state-dependent linear operators enable
3

---

selective state updates based on input content. For discretization, the system employs the zero- order-hold (ZOH) rule [3]] to transform the A and B matrices into A = exp( AA) and B = (cx)- - T) A B, where the step size A serves as a variable-dependent parameter. This transformation enables SSMs to process continuous inputs as discrete signals and converts the original Linear Time-Invariant (LTT) equation into a recurrence format.
3.3 Diffusion-based World Model Learning
To adapt diffusion models for world modeling, which offers superior sample quality and tractable likelihood estimation, key requirement is modeling the conditional distribution p 1 logt: age). where On and d₃ represent observations and actions at time step t. The denoising process incorporates both the noised next observation and the conditioning context as input: Do(o T, Oge, ase)). While diffusion- based world models [8] have shown promise, the Mfate-of-the-ar approach DIAMOND g exhibits limitations although it achieves superior performance on the Atari 100k benchmark [TS] These models face two critical limitations. The first limitation stems from their constrained condition- ing context, which typically considers only the most recent observations and actions. For instance, DIAMOND restricts its context to the last four observations and actions in the sequence. This constraint impairs the model's capacity to capture long-term dependencies and leads to inaccurate predictions in scenarios that require extensive historical context. The second limitation in current diffusion-basee world models lies in their architectural separation of predictive tasks. For example, DIAMOND implements separate recurrent neural network for reward and termination prediction. This separation prevents the sharing of learned representations between the diffusion model and these predictive tasks and results in reduced overall learning efficiency of the system.
4 Motivational Experiments
T=0
4...8
To substantiate the memory limitations of the DIAMOND model, we conducted experiments using the MiniGrid MemoryS7 and MemoryS9 environments [6]. These experiments evaluate memory consistency and temporal prediction capabilities in world models. Fig. [ presents qualitative comparisons among the DIAMOND model, our proposed EDELINE method, and the Oracle world model. The world models were trained on MiniGrid MemoryS7, and their state predictions were evaluated. The upper half of Fig. (a) shows an initial action sequence of four consecutive left turns followed by four no- op actions. DIAMOND and EDELINE then autoregressively predicted the subsequent states. The lower half of Fig. (a) reveals the predic- tion outcomes. At timestep 9, DIAMOND gen- erated incorrect predictions with extra key object in the state. The model's predictions at timesteps 10 11 exhibited inaccuracies in the
The
Tran
Tow
True
EDELINE DIAMOND
Mow
(a) Qualitative results MemoryS7
Memory57
MemoryS9
0.900
0.500
0.750
0.750
0.500
0.600
outer wall representations. These prediction er-
0.450
0.450
0.300
0.300
rors result from the environment's partial ob-
0.150
0.150
servability and DIAMOND's four-frame con-
0.000
0.000
text limitation. The input of four idle frames at timestep 9 led to loss of earlier state context. In contrast, EDELINE maintained accurate predic- tions throughout the sequence. This improved performance stems from EDELINE's architec- tural design, which addresses the memory con- straints inherent in DIAMOND. The accurate prediction at timestep 12 by DIAMOND can be attributed to its access to the 8th state, which enabled accurate predictions at timestep 13 through the correct 12th state context. Fig. [ (b) illustrates DIAMOND's performance limitations in both Me and M9 environments due insufficient memory capacity partially observable scenarios. In contrast, EDELINE shows near-optimal performance under these conditions.
EDEUNE
DIAMOND
(b) Quantitative results in MemoryS7/59
Figure 1: Motivational examples for both qualita- tive and quantitative evidences demonstrate that DIAMOND faces difficulties in imagining accu- rate future under memorization tasks.

---

Parameter
Add notse d(c)
Encoder Layer
60 #+1
Embedding Layer
Conv Attention
Frams Black Next-Frame Predictor Fi
Layer
Adapltec Group Normaization
Cross Attention
SSLU
P
Action Emboddding
Observation Encoder
SILU
Recurrent Embedding Module d
Reward Termination Predictor Pa
Figure 2: The EDELINE world model includes three principal components: (1) A U-Net-like Next- Frame Predictor enhanced by adaptive group normalization and cross-attention mechanisms, (2) A Recurrent Embedding Module built on Mamba architecture for temporal sequence processing, and (3) A gewaad//ermination Predictor implemented through linear layers. The EDELINE framework uses shared hidden representations across the components for efficient world model learning.
5 Methodology
Conventional diffusion- based world models [8] demonstrate promise in learning environment dy- namics yet face fundamental limitations in memory capacity and horizon prediction consistency. To address these challenges, this paper presents EDELINE, illustrated in Fig. 2 unified architecture that integrates state space models (SSMs) with diffusion-based world models. EDELINE's core innovation lies in its integration of SSMs for encoding sequential observations and actions into hidden embeddings, which a diffusion model then processes for future frame prediction. This hybrid design maintains temporal consistency while generating high-quality visual predictions. A Convolutional Neural Network based actor processes these predicted frames to determine actions, thus enabling autoregressive generation of imagined trajectories for policy optimization.
5.1 World Model Learning
The core architecture of EDELINE consists of Recurrent Embedding Module (REM) f that processes the history of observations and actions (oo, dg. 01, 01 ..., Ot, az) to generate a hidden embedding ht through recursive computation. This embedding enables the Next-Frame Predictor Ps to generate predictions of the subsequent observation ôt+1- The architecture further incorporates dedicated Reward and Termination Predictors to estimate the reward A and episode termination signal d, respectively. The trainable components of EDELINE! world model are formalized as:
Recurrent Embedding Module: h4 S fo 0t)
Next-Frame Predictor: ôt+1 ~ (p(00411h4)
Reward Predictor: i, pp((f(lh))
Termination Predictor: d, ~ (de(be))
5.1.1 Recurrent Embedding Module
While DIAMOND, the current state-of-the-art in diffusion- world models, relies on a fixed context window of four previous observations and actions sequence, the proposed EDELINE ar- chitecture advances beyond this limitation through recurrent architecture for extended temporal sequence processing. At each timestep t, the Recurrent Embedding Module processes the current observation-action pair (0t, ) to update a hidden embedding h4 = (ht ot, at). The implemen- tation of REM utilizes Mamba [12], an SSM architecture that offers distinct advantages for world modeling. This architectural selection is motivated by the limitations of current sequence process- ing methods deep learning. Self-attention-base Transformer architectures, despite their strong modeling capabilities, suffer from quadratic computational complexity which impairs efficiency.

---

Traditional recurrent architectures including Long Short-Term Memory (LSTM) [52] and Gated Recurrent Unit (GRU) [53] experience gradient instability issues that affect dependency learning. In contrast, SSMs provide an effective alternative through linear-time sequence processing coupled with robust memorization capahilities via their state-space formulation. The adoption of Mamba emerges as promising choice due to its demonstrated effectiveness in modeling temporal patterns across various sequence modeling tasks. Appendix presents comprehensive ablation study that evaluates different architectural choices for the REM.
5.1.2 Next-Frame Predictor
While motivated by DIAMOND's success in diffusion-based world modeling, EDELINE introduces significant architectural innovations in its Next-Frame Predictor to enhance temporal consistency and feature integration. At time step i, the model conditions on both the last frames and the hidden embedding he from the Recurrent Embedding Module predict the next frame ô The predictive distribution ps po(of+al ( is implemented through a denoising diffusion process, where D& functions as the denoising network. Let V = (r, of 1......... og, he) represent the conditioning information, where represents the diffusion time. The denoising process can be formulated as 0. = (o + ). To effectively integrate both visual and hidden information, D& employs two complementary conditioning mechanisms. First, the architecture incorporates Adaptive Group Normalization (A (AGN) [54) layers within each residual block to condition normalization parameters on the hidden embedding It, and diffusion time т, which establishes context-aware feature normalization [54]. This design significantly extends DIAMOND's s implementation, which limits AGN conditioning to 7 and action embeddings only. The second key innovation introduces cross-attention blocks inspired by Latent Diffusion Models (LDMs) [55], which utilize fú2 and as context vectors. The UNet's feature maps generate the query, while h2 and T project keys and values. This novel attention mechanism, which is absent in DIAMOND, facilitates the fusion of spatial-temporal features with abstract dynamics encoded in h2 The observation modeling loss Lab(()) is defined based on Eq. and can be formulated as follows:
Lob (0) = T.. (1 - 1..
5.1.3 Reward/ / Termination Predictor
EDELINE advances beyond DIAMOND's architectural limitations through an integrated approach to reward and termination prediction. Rather than employing separate neural networks, EDELINE leverages the rich representations from its REM. The reward and termination predictors are imple- mented as multilayer perceptrons (MLPs) that utilize the deterministic hidden embedding h4 as their conditioning input. This architectural unification enables efficient representation sharing across all predictive tasks. EDELINE processes both reward and termination signals as probability distributions conditioned on the hidden embedding: (f (he) and po po (de (he) respectively. The predictors are optimized via negative log- eogllikelihood losses, expressed as:
Lrew (Φ) = - In (r (0) = ---In Lnppo(e (dl (ha)-
This unified architectural design represents an improvement over DIAMOND's separate network approach, where reward and termination predictions require independent representation learning from the world model. The integration of these predictive tasks with shared representations enables REM to learn dynamics that encompass all relevant aspects of the environment. The architectural efficiency facilitates enhanced learning effectiveness and better performance.
5.1.4 EDELINE World Model Training
To implement the loss functions defined in Eq. (2) and Eq. (3) we adopt a training strategy tailored the architecture of the EDELINE world model. The Fext-Frame Predictor trained to reconstruct future observations from hidden embeddings sampled from arbitrary timesteps. In contrast to DIAMOND, which operates over a fixed four-frame window without memory. EDELINE maintains memory across the full trajectory. Directly applying the observation loss at every timestep, as done DIAMOND, would ineur significantly higher computational costs due to EDELINE'S recurrent embedding module. To mitigate this, we leverage the Mamba parallel scan algorithm to efficiently compute hidden embeddings across all timesteps, and approximate the full loss by randomly sampling single target timestep for reconstruction. Specifically, we sample trajectory segment of fixed

---

.......,........ length T from the replay buffer, indexed by I = ft, ,t.....t + T - 1}. Given (of, az)sez, we initialize the hidden state h and compute the hidden embeddings {h, hez in parallel using Mamba. We then randomly select target timestep E (1 ..... + T - }), where B denotes the burn-in period, and compute the observation reconstruction loss as:
(o) = |00 - -yll", where j Uniform + B...., ++ T - 1}. and 0j p(o)) hy-1). (4)
This approach enables EDELINE achieve computational efficiency comparable to DIAMOND, while retaining the advantages of long-term memory. For reward and termination prediction, EDE- LINE utilizes cross-entropy losses averaged over the sampled trajectory segment, which can be formulated as follows: Crew (6) T LiEZ CrossEnt(fr, ri), (6) Liez CrossEnt(d,, d:).
The detailed algorithm is presented in Appendix D
This unified training approach, combining random sampling strategies with dynamic loss harmoniza- tion, demonstrates superior efficiency compared to DIAMOND's separate network methodology, as validated in our results presented in Section 6 and Appendix D.5 Moreover, the quantitative analysis presented in Appendix D.7 reveals substantial reductions in world model training duration.
The world model integrates an innovative end-to-end training strategy with self-supervised approach. EDELINE extends the harmonization technique from HarmonyDream [56) through the adoption of harmonizers Wo and w,, which dynamically balance the observation modeling loss Lot (ø) and reward modeling loss ( This adaptive mechanism results in the total loss function L()):
((Φ) wo Cor + Lrewro (d) + Lend )) + log(w,)) + log(w,").
(5)
5.2 Agent Behavior Learning
To enable fair comparison and demonstrate the effectiveness of EDELINE's world model architecture, the agent architecture adopts the same optimization framework as DIAMOND. Specifically, the agent integrates policy тe and value Ve networks with REINFORCE value baseline and Bellman error optimization using X-returns [8]. The training framework executes a procedure with three key phases: experience collection, world model updates, and policy optimization. This method, as formalized in Algorithm [] follows the established paradigms model-based RL literature [LS] 6 35 g]. To
ensure reproducibility. we provide extensive details in the Appendix, with documentation of objective functions, and the hyperparameter configurations Appendices D.11 respectively.
Experiments
This section presents our experimental results of EDELINE, Additionally, we include ablation studies for each component of our approach in Appendix D.8 We also validate the benefits of our unified architecture through representation analysis Appendix D.5 and evaluate the quality of our model's imagination capabilities in Appendix D.6
6.1 Experimental Setup
We evaluate EDELINE on the Atari 100k benchmark (15), which serves as the standard evaluation protocol in recent model-based RL literature for fair comparison. In addition, our experimental validation extends to ViZ.oom [8] MiniGrid [16], and Crafter (17) environments demonstrate
broader applicability. To ensure statistical significance, all reported results represent averages across three independent runs. The Atari 100k benchmark [15] encompasses 26 diverse Atari games that evaluate various aspects of agent capabilities. Each agent receives strict limitation of 100k environment interactions for learning, in contrast conventional Atari agents that typically require 50 million steps. EDELINE's performance is evaluated against current state-of -fee--ff-heeeeeee the-art world model- based approaches, including DIAMOND [8 STORM [3). DreamerV3 2]. IRIS (35), TWM [32), and Drama 48]. For evaluating 3D scene understanding capabilities, we employ ViZDoom scenarios that demand sophisticated 3D spatial reasoning in first-person environments. This provides a crucial testing ground beyond the third-person perspective of Atari environments. Furthermore, the Crafter and MiniGrid memory scenarios evaluate memorization capabilities through tasks that require information retention across extended time horizons.
7

---

Table 1: Game scores and overall human-normalized scores on the 26 games in the Atari 100k benchmark. Results are averaged over seeds, with bold numbers indicating the best performing method for each metric.
Game
Random
Human
SimPLe
TWM
IRIS
STORM
Dreamer
Drama
DIAMOND
EDELINE (ours)
Alien
227.8
7127.7
616.9
674.6
420.0
983.6
959.4
820
744.1
974.6
Amidar
Assauit
222.4
742.0
527.2
682.6
1524.4
801.0
705.6
539
1526.4
1225.8
Asterix
210.0
8503.3
1128.3
1116.6
853.6
1028.0
932.5
1632
3698.5
4224.5
BankHHeist
753.1
34.2
466.7
648.7
19.7
BatteZonee
2360.0
57187,5
4031.2
5068,D
13074,0
13540.0
12250.0
10860
4702.0
5683
Boxing
0.1
12.1
7.8
77.5
70.1
79.7
78.0
78
86.9
88,1
83.7
15.9
ChogperCommand
811.0
7387.8
979.4
1697.
1565.0
1888.0
410.0
1642
1369.8
2047.
CrazyClimber
10780.5
35829,4
62583.6
71820.4
59324,2
66776.00
97190.0
52242
99167.8
1017810
DemonAttack
350.2
1016.1
Freeway
0.0
29.6
16.7
24.3
31.1
33.5
0.0
15
33.3
33,8
Froscbite
65.2
4334,7
236.9
1475.6
259.1
1316.0
909.4
785
274.1
286.8
Gopher
596.8
1674.8
2357
58979
Hero
1027.0
50826.4
2656.6
7254.0
7057.4
11044.3
11160.5
7946
$621.8
12780.8
Jameshond
29.0
302.8
100.5
362.4
462.7
509.0
444.6
372
427.4
784.3
Kangaroo
52.0
51.2
1384
5270.0
Krull
1598.0
2665.5 5
2204.8
6349.2
6616.4
8412.66
7781.
9693
8610.1
9748.8
Master
258.5
22736.3
14862.5
2455A.6
81759.88
26182.0
21420.0
18713.6
3144800
MsPacman
307.3
6951.6
1480.0
1588.
999.1
2673.5
1526.9
2270
1958.2
1849.
Pong
.20.7
14.6
12.8
18.8
14.6
11.3
18.4
15
20.4
20.5
PrivaleFiye
24.9
69571.3
35.0
100.0
7781,0
881.6
90
114.3
99.5
Obert
13455.0
1288.8
3330.8
745.7
4522.5
3405. 1
4499.3
6776.2
ReadRunner
11.5
7845.0
5640.6
9109.0
9614.6
17564.0
15565.0
14020
20673.2
32020.0
Senques:
68.4
42054.7
683.3
774.4
661.3
$25.2
618.0
497
581.2
2140.1
UpNDown
533.4
11693.2
3350.3
15981.7
3546.2
7983.0
7567.1
7387
3856.3
5650.3
#Superhuman (t)
N/A
8
10
10
9
11
13
0.966
1.666
Modian (?)
0.000
1.000
0.134
0.505
0.289
0.580
0-485
0.270
0.373
0.817
IQM
0.000
1.000
0.130
0.459
0.501
0.636
0.487
0.641
0,940
Optimality Gap (J)
0.000
0.513
0.510
Mean
Median
Interquartile Mean
EDELINE (ours)
DIAMOND
STORM
DreamerV3
IRIS
TWM
SimPLe
2.0
0.75
0.75
Human Normalined
Figure 3: Mean, Median and IQM HNS.
6.2 Atari 100k Experiments
Following standard evaluation paradigms for world models, we evaluate EDELINE on the Atari 100k benchmark. For performance quantification, we adopt the human-normalized score (HNS) (57)
which measures agent performance relative to human and random baselines:
agent score random score
HNS
(6)
human score random score
Fig. 3 presents stratified bootstrap confidence intervals following 58 recommendations for point estimate limitations. can be observed that EDELINE exhibits exceptional performance across this benchmark. Our approach surpasses human players 1.87, median of 0.82, and IQM of 0.94. These metrics demonstrate superior performance compared existing model-based reinforcement learning baselines without look-ahead search techniques. For detailed quantitative analysis, Table provides comprehensive scores for all games. The superior performance EDELINE stems from ability to preserve the visual fidelity advantages of diffusion- based world models. Its architecture demonstrates significant improvements over DIAMOND in environments such as BankHeist and Hero. The enhanced performance arises from EDELINE'S memorization capabilities, combined with harmonizer-enablee precise visual detail selection for reward-relevant features. Appendices [D2 D.3 and and provide additional training curves, qualitative analyses, and performance metrics, respectively.
games with superhuman mean HNS of

---

Table 2: Comparison of different methods on Crafter in terms of average return and world model parameter count.
Method
Avg Return
#World Model Params
EDELINE
11.5 + 0.9
11M
DreamerV3 XL
9.2 + 0.3
200M
A-IRIS
37M 7.7 + 0.5
25M
DreamerV3 M
6.2
IRIS
5.7
48M
DIAMOND
2.8 0.5
10.4M
6.3 Crafter Experiments
To further evaluate EDELINE's memory enhancement capabilities, we conducted experiments on Crafter [17]a procedurally generated survival environment that presents complex memory challenges. Crafter specifically designed to assess 'wide and deep exploration, long-term reasoning and credit assignment, and generalization (2).
Crafter requires substantial memory capabilities due to its demand for agents to retain information about previously collected resources, crafted items, and explored territories for optimal decision- making, which establishes it as ideal testbed for the evaluation of our memory-enhinn memory-enhanced architec- ture. Within IM environment step budget, EDELINE achieves superior performance compared to state-of-the-a art baselines including DIAMOND 8
DreamerV3 [2]. A-IRIS [59], and IRIS [35, de-
spite its relatively modest parameter count IIM. These results highlight the significant advantages resulting from the integration of Mamba's 's memory capabilities with diffusion's 's generative abilities.
Table 2 presents our experimental results with a 1M environment step budget. The results reveal that EDELINE significantly outperforms all baselines with 25% higher returns than Dreamer XL despite the utilization of 18 fewer parameters. Most importantl, EDELINE delivers 4. 4.1% im- provement over DIAMOND with comparable parameter count, which demonstrates the substantial benefits of our enhanced memory mechanism. In order to provide additional insight into EDELINE's memory advantages, we conducted a qualitative analysis of model predictions over extended imagi- nation. This analysis, as presented Appendix [D.10] visually demonstrates EDELINE'S superior ability to remember environmental features and maintain consistency when revisiting previously explored areas, critical capability for effective planning in complex environments.
Desic
Dead CCrnider
DelendCenter
PredictPos cam
205
94
73C
Figure 4: Training curves comparing EDELINE (blue) and DIAMOND (green) across five ViZDoom scenarios. The solid lines represent the mean performance over three seeds, with shaded regions indicating standard deviation.
6.4 ViZDoom Experiments
To evaluate EDELINE's memory retention and imagination consistency in more challenging scenar- ios, we conduct experiments in the first-person 3D environment ViZDoom, which presents increased complexity compared Atari 100k's vvo-dimensional perspective. Our evaluation encompasses five default scenarios (i.e., Basic, DeadlyCorridor, DefendCenter, HealthGathering, PredictPosition), with detailed scenario descriptions and reward configurations available Appendix [D. The experi- mental results presented in Fig. demonstrate EDELINE's superior performance over DIAMOND in scenarios that demand sophisticated 3D scene modeling (DeadlyCorrido) and spatiotemporal prediction (PredictPosition).

---

DIAMOND EDELINE Tiue
Figure 5: Qualitative comparison on the DeadlyCorridor scenario. Each row shows ground truth, EDELINE's predictions, and DIAMOND's predictions, respectively. Colored boxes highlight suc- cessful (green) and failed (red) predictions of task-critical visual elements including enemies, particle effects from hits, and the armor. EDELINE accurately predicts the reward-relevane visual details. DIAMOND captures only basic environment structure.
The DeadlyCorridor scenario provides particularly compelling demonstration of EDELINE's memory advantages. This environment requires the agent to navigate corridor, eliminate enemies on both sides, and ultimately reach armor the corridor's end. The task necessitates sustained spatial awareness throughout relatively long trajectory. As Fig. (o.55))l illlusttsss EDELINE's s SSM effectively integrates information from the entire history, whereas DIAMOND remains constrained to only the past four frames. This memory limitation significantly impairs DIAMOND' 's ability to maintain consistent predictions in 3D environments where current observations provide only partial information about the scene. The visual comparison reveals that EDELINE accurately models both the game environment and particle effects while correctly predicts the corridor's 's end and armor object in the final steps. This substantiates EDELINE's comprehension of the agent's 's position relative to the corridor's end through its extended memory capabilities, while DIAMOND's predictions become increasingly blurry and inconsistent due to its memory constraints. It also reveals that the enhanced spatial awareness directly contributes EDELINE's superior performance in the environment.
Conclusions
In this work, we addressed the limitations of current diffusion-based world models in handling long-term dependencies and maintaining prediction consistency. Through the integration Mamba SSMs, EDELINE effectively processed extended intervattionaattiinn sequences through its recurrent embedding module, which enabled adaptive memory retention beyond fixed-context approaches. The unified framework eliminated architectural separation between observation, reward, and termination prediction, which fostered efficient representation sharing. Dynamic loss harmonization further mitigated optimization conflicts arising from multi-task learning. Extensive experiments Atari 100K, Crafter, MiniGrid, and ViZDoom validated EDELINE's state-of-the- performance, with significant improvements in quantitative metrics and qualitative prediction fidelity.
Acknowledgements
The authors gratefully acknowledge the support from the National Science and Technology Council (NSTC) Taiwan under grant numbers NSTC 144-2221-E0002-069-M 212-MY3, and NSTC 114-2216 8-E-A49-0026 We also express our sincere appreciation to NVIDIA Corporation and the NVIDIA AI Technology Center (NVAITC) for the donation of GPUs and access to the T supercomputer. Furthermore, the authors extend their gratitude to the National Center for High-Peerformacce Computing for providing the necessary computational and storage resources.
NSTC 113-2221-E-00
10

---

References
[1] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2018.
[2] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, pages 1-7,2022 2025.
[3] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, and al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 2020.
[4] Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2021.
[5] Philipp Wu, Alejandro Escontrela, Danijar Hafner, Picter Abbeel, and Ken Goldberg. Day- dreamer: World models for physical robot learning. In Proc. Conf. on Annual Conference on Robot Learning (CoRL), 2022.
[6] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representa- fions (ICLR), 2020.
[7] Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations (ICLR), 2021.
[8] Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, and et al. Diffu- sion for world modeling: Visual details matter in atari. In Proc. Conf. Neural Information Processing Systems (NeurIPS), 2024.
[9] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In Int. Conf. on Learning Representations (ICLR), 2022.
[10] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2022.
[1]] Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In Int. Conf. on Learning Representations (ICLR), 2023.
[12] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, In Proc. Int. Conf. on Language Modeling (CoLM), 2024.
[13] Albert Gu and Tri Dao. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Proc. Int. Conf. on Machine Learning (ICML), 2024.
[14] Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar. Mas- tering memory tasks with world models. The Twelfth International Conference on Learning Representations (ICLR), 2024.
[15] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Mitos, Blažej Osiński, Roy H Campbell, and al. Model based reinforcement learning for atari. International Conference on Learning Representations (ICLR), 2020.
[16] Maxime Chevalier- Boisverr, Bolun Dai, Mark Towers, Rodrigo De Lazcano Vicent,, Lucas Willems, and et al. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2023.
[17] Danijar Hafner. Benchmarking the spectrum of agent capabilities. In International Conference on Learning Representations (ICLR), 2022.
[18] Michal Kempka, Marek Wydmuch, Grzegorz Rune, Jakub Toczek, and Wojciech Jaskowski. Viz- doom: doom-based ai research platform for visual reinforcement learning. MiXi:1600500099 2016.
11

---

[19] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[20] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations (ICLR), 2021.
[21] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2019.
[22] Yang Song, Jascha Sohl- Sobl--iccsteen, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and et al. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021.
[23] Yilun Du, Conor Durkan, Robin Strudel, Joshua B. Tenenbaum, Sander Dieleman, and et al. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and meme. In Proc. Int. Conf. on Machine Learning (ICML), 2023.
[24] Wang Zhendong, Hunt Jonathan J, and Zhou Mingyuan. Diffusion policies expressive policy class for offline reinforcement learning. In Int. Conf. on Learning Representations (ICLR), 2023.
[25] Ajay Anurag, Du Yilun, Gupta Abhi, B. Tenenbaum Joshua, S. Jaakkola Tommi, and et al. Is conditional generative modeling all you need for decision making? In Int. Conf. on Learning Representations (ICLR), 2023.
[26] Pearce Tim, Rashid Tabish, Kanervisto Anssi, Bignell Dave, Sun Mingfei, and et al. Imitating human behaviour with diffusion models. In Int. Conf. on Leurning Representations (ICLR), 2023.
[27] Janner Michael, Du Yilun, B. Tenenbaum Joshua, and Levine Sergey. Planning with diffusion for flexible behavior synthesis. Proc. Int. Conf. on Machine Learning (ICML), 2023.
[28] Liang Zhixuan, Mu Yao, Ding Mingyu, Ni Fei, Tomizuka Masayoshi, and Luo Ping. Adaptd- iffuser: Diffusion models as adaptive self-evolving planners. In Proc. Int. Conf. on Machine Learning (ICML), 2023.
[29] Nuti Felipe Pinto Coelho, Franzmeyer Tim, and Henriques Joao Extracting reward functions from diffusion models. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2023.
[30] Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, and al. Metadiffuser: Diffusion model as conditional planner for offline meta-rl. In Proc. Int. Conf. on Machine Learning (ICML), 2023.
[31] Lu Cong, Ball Philip J., Teh Yee Whye, and Parker-Holder Jack. Synthetic experience replay. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2023.
[32] Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-bras based world models are happy with 100k interactions. In The Eleventh International Conference on Learning Representations (ICLR), 2023.
[33] Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. STORM: Efficient stochas- tic transformer based world models for reinforcement learning. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2023.
[34] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett,, editors, Advances in Neural Information Processing Systems (NeurIPS), volume 30. Curran Associates, Inc., 2017.
[35] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference Learning Representations (ICLR), 2023.
[36] Seung Wook Kim, Yuhao Zhou, Jonah Philion, Antonio Torralba, and Sanja Fidler. Learning to Simulate Dynamic Environments with GameGAN. IEFE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
12

---

[37] Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, Yusuf Aytar, Sarah Maria Elis- abeth Bechtle, Feryal Behbahani, Stephanie C.Y. Chan, Nicolas Heess, Lucy Gonzalez, Simon Osindero, Sherjil Ozair, Scott Reed, Jingwei Zhang, Konrad Zolna, Jeff Clune, Nando de Freitas, Satinder Singh, and Tim Rocktäschel. Genie: Generative interactive environments. In Forty-first International Conference on Machine Learning (ICML), 2024.
[38] Dani Valevski, Yaniv Leviathan, Moab Arar, and Shlomi Fruchter. Diffusion models are real- time game engines, In The Thirteenth International Conference on Learning Representations (ICLR), 2025.
[39] Anthony Hu, Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie Shotton, and Gianluca Corrado. Gaia-l: A generative world model for autonomous driving, 2023.
[40] Lloyd Russell, Anthony Hu, Lorenzo Bertoni, George Fedoseev, Jamie Shotton, Elahe Arani, and Gianluca Corrado. Gaia-2: A controllable multi-view generative world model for autonomous driving, 2025.
[41] Jiannan Xiang, Guangyi Liu, Yi Gu, Qiyue Gao, Yuting Ning, Yuheng Zha, Zeyu Feng, Tianhua Tao, Shibo Hao, Yemin Shi, Zhengzhong Liu, Eric P. Xing. and Zhiting Hu. Pandora: Towards general world model with natural language actions and video states, 2024.
[42] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2020.
[43] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2022.
[44] Hasani Ramin, Lechner Mathias, Wang Tsun-Hsuan, Chahine Makram, Amini Alexander, and et al. Liquid structural state-space models. In Int. Conf. on Learning Representations (ICLR), 2023.
[45] Lu Chris, Schroecker Yannick, Gu Albert, Parisotto Emilio, Nicolaus Foerster Jakob, and et al. Structured state space models for in-context reinforcement learning. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2023.
[46] Sili Huang, Jifeng Hu, Zhejian Yang, Liwei Yang. Tao Luo, and et al. Decision mamba: Reinforcement learning via hybrid selective sequence modeling. In Proc. Conf. Neural Information Processing Systems (NeurIPS), 2024.
[47] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, and et al. Decision trans- former: Reinforcement learning via sequence modeling. In Proc. Conf. on Neural Information Processing Systems (NeurIPS), 2021.
[48] Wenlong Wang. Ivana Dusparie, Yucheng Shi, Ke Zhang, and Vinny Cahill. Drama: Mamba- enabled model- reinforcement learning is sample and parameter efficient. In The Thirleenth International Conference on Learning Representations (ICLR), 2025.
[49] Karl Johan Åström. Optimal control of markov processes with incomplete state information. Journal of Mathematical Analysis and Applications, 10:174-205, 1965.
[50] Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine Learning, 3:9.44, 1988.
[51] Yang Chifu, Gao Shuang, and Xue Zhu. Improving the closed-loop tracking performance using the first-order hold sensing technique with experiments. MxXi::1801.01263 2018.
[52] S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural Computation, 9:1735-1780, 1997.
[53] J. Chung, C. Gulcehre, K. Cho, and Y. Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. Proc. Conf. on Neural Information Processing Systems Workshop (NeurIPSW), 2014.
13

---

[54] Heliang Zheng, Jianlong Fu, Yanhong Zeng, Jiebo Luo, and Zheng-Jun Zha. Learning semantic- aware normalization for generative adversarial networks. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[55] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models, 2021.
[56] Haoyu Ma, Jialong Wu, Ningya Feng, Chenjun Xiao, Dong Li, Jianye HAO, Jianmin Wang, and Mingsheng Long. Harmonydream: Task harmonization inside world models. In Forty-first International Conference on Machine Learning (ICML), 2024.
[57] Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In Proceedings of The 33rd International Conference on Machine Learning (ICML), 2016.
(58) Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances Neural Information Processing Systems (NeurIPS), 2021.
[59] Vincent Micheli, Eloi Alonso, and François Fleuret. Efficient world models with context-aware tokenization. In Forty-first International Conference on Machine Learning (ICML), 2024.
[60] Jaseha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning (ICML), pages 2256-2265. PMLR, 2015.
[61] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances neural information processing systems (NeurIPS), 34:8780-878 8794,, 2021.
[62] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems (NeurIPS), 13:12438-12448 2020.
[63] Chen-Hao Chao. Wei-Fang Sun, Bo-Wun Cheng, Yi-Chen Lo, Chia-Che Chang, Yu-Lun Liu, Yu-Lin Chang, Chia-Ping Chen, and Chun-Yi Yi Lee. Denoising likelihood score matching for con- ditional bosee data generation. In International Conference on Learning Representations (ICLR), 2022.
[64] Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems (NeurIPS), 35:26565-26577, 2022.
[65] Rich Caruana. Multitask learning. Machine Learning, 28:41-75, 1997.
[66] Lior Cohen, Kaixin Wang, Bingyi Kang, and Shie Mannor. Improving token-based world models with parallel observation prediction. In Forty-first International Conference on Machine Learning (ICML), 2024.
[67] Maxime Burchi and Radu Timofte. Learning transformer-basse world models with contrastive predictive coding. In The Thirteenth International Conference on Learning Representations (ICLR), 2025.
[68] ZiRui Wang, DENG Yue, Junfeng Long, and Yin Zhang. Parallelizing model-based reinforce- ment learning over the sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS), 2024.
[69] Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, 2018
14

---

A Limitation
Although EDELINE maintains computational efficiency equivalent that of DIAMOND and si- multaneously addresses its memory limitations, the fundamental computational requirements of diffusion- based world models remain substantial relative to alternative approaches. The iterative characteristics of the diffusion process necessitate multiple forward passes throughout both training and inference phases, which consequently elevates computational overhead. The enhancement of computational efficiency within diffusion processes for world modeling applications constitutes promising avenue for future research endeavors.
B Broader Impact
EDELINE addresses fundamental limitations inherent in iiffusion-based world models through the enhancement of memory capabilities and imagination consistency. Our research contributes three pivotal advances data-efficient reinforcement learning: (1) overcome the fixed context window constraints of prior methodologies such as DIAMOND through the implementation of selective state space models, which substantially enhances memory capabilities, (2) we establish end-to-end training methodology that effectively leverages hidden embeddings and incorporates Harmony Dream augment world modeling performance, and (3) demonstrate superior performance across visually rich 2D environments (i.e., Atari), challenging 3D first-person perspectives (i.e., ViZDoom), and memory-intente
seeeduurallygeeeerate survival scenarios (i.e., Crafter), which illustrates the potential of our approach to address increasingly complex domains.
World models represent a promising direction for the improvement of sample efficiency and safety in reinforcement learning through the reduction of direct environment interaction requirements. While EDELINE advances this research domain, acknowledge that imperfections in world models may continue to result in suboptimal agent behaviors. We believe our contributions toward more accurate and memory-capable world models will serve mitigate such risks future applications.
15

---

C Additional Background Material Section
C.1 Score-based Diffusion Generative Models
Diffusion probabilistic modeling [60 19 61] and score-based generative modeling 21 62 63 can
be unified through a forward stochastic differential equation (SDE) formulation [22] The forward
diffusion process (x"] } with continuous time variable T transforms the data distribution pd p data to
prior distribution p" -prror expressed as:
dx f(x, r)dT + g(r)dw,
where f(x, T) represents the drift coefficient, g(T) denotes the diffusion coefficient, and w is the Wiener process. The corresponding reverse-time SDE can then be formulated as:
dx [f(x, -) - "(7)² log [x)] d + g(r)dW;
(8)
where w is the reverse-time Wiener process. Eq. (8) enables sampling from Dº when the (Stein) score function Vx log (x) is available. A common approach to estimate the score function is through the introduction of a denoiser Da, which is trained to minimize the following objective:
E Bongin Sppan
[x + o) -x x
(9)
where is Gaussian noise with zero mean and variance determined by variance scheduler (T) that follows a noise distribution ptrai, and (x" + n) corresponds to the perturbed data x The score function can then be estimated through: Vx log (x) = (Do(x; a) x). In practice, modeling the denoiser Dg directly can be challenging due to the wide range of noise scales. To address this, EDM (64) introduces design space that isolates key design choices, including preconditioning functions (Caiip Cout, Cin. Croise to modulate the unconditioned neural network F' to represent Dg. which can be formulated as:
Dx(x; = Cakip(o) + Cour Fo(c((())) Croise (o)).
(10)
The preconditioners serve distinct purposes: (a) and Coul (a) maintain unit variance for network inputs and outputs across noise levels, Groise provides transformed noise level conditioning, and (a) adaptively balances signal mixing. This principled framework improves the robustness and efficiency of diffusion models, enabling state-of-the-ard performance across various generative tasks.
C.2 Multi-task Essence of World Model Learning
Modern world models 15 7 typically address two fundamental prediction tasks: the modeling of environment dynamies through observations and the prediction of reward signals. The learning of these tasks requires distinct considerations based on the complexity of the environment. In simple low- dimensional settings, separate learning approaches suffice for each task. However, the introduction of high-dimensional visual inputs fundamentally alters this paradigm, as partial observability creates an inherent coupling between state estimation and reward prediction. This coupling necessitates joint learning through shared representations, an approach that aligns with established multi-task learning principles [65]. The implementation of such joint learning through shared representations introduces several technical challenges. The integration of multiple learning objectives requires careful consideration of their relative importance and interactions. A fundamental difficulty stems from the inherent scale disparity between nigh-dimensional visual observations and scalar reward signals. This disparity manifests in the world model learning objective, which combines observation modeling (0), reward modeling (0), and dynamics modeling Ca(0) losses with weights w., Wy. Wa to control relative contributions:
( = w. (.) + w. (a))
(11)
HarmonyDream 56 demonstrated that observation modeling tends to dominate this objective due to visual inputs' high dimensionality compared to scalar rewards. Their work introduced a variational formulation:
[(0, Wor Wy, wa) =
Σ H(C,(0))
- Σ w,L,(0) + log( W
(12)
wi (E(o.r,d)
where (O) w)) = LW,C + log(1/w;) dynamically balances the losses by maintaining Ejw". L] = 1. This harmonization technique can substantially enhance sample efficiency and performance. Our work extends these insights through the integration of dynamic task balancing mechanisms into EDELINE world model architecture
16

---

D Additional Experiments and Experiment Details
In this section, provide extensive additional experiments and detailed experimental setups that complement the primary results presented in the main manuscript. These supplementary materials include in-depth analyses, ablation studies, and technical details that further validate and illuminate EDELINE's s advantages.
D.1 EDM Network Preconditioners and Training
Following DIAMOND 8 we the EDM preconditioners from 64 for normalization and rescaling to improve network training:
on Vo(T)² odata
(13)
(
Eot vo(T)² + Jauta
(14)
Enoise log(
(15)
Odato
Cakip
(16)
data + (T)
where data 0.5.
The noise parameter (T) sampled maximize the effectiveness of training:
log(a(T)) (Pmmean: Piaa),
(17)
where Pmeen -04, Pad 1.2.
17

---

D.2 Atari 100K Training Curves
Fig. presents the detailed training curves for all individual games in the Atari 100k benchmark.
Ainn
Am Amidar
Asssul:
Adteri
Buncheest
1.2
4.00
1.45
13K
My
200
630
600
133
420
-00
420
1.000
560
320
330
2.000
23R
L00K
0.000
506
185K
0.002
23C
736
0.00
33X
120K
0030
23K
73K
BattieZone
Box ng
Breacout
ChooperCommand
Crazylimber
100
50
50-
10-
240
1058
301
10C
60-
300
50
40-
10
160
120
120
610
2.00
3,000
110
3.60
25R
106
100k
0.500
256
10C
MK
100c
0.663
7M
75c
190c
0.003
9.0
75C
130C
20030
256
5.36
PMC
1900
freeway
Gopphrr
Mero
300
340
1.55
20
210
225
900
M
240
3.00
348
20
210
7.58
601
4.30
400
10
120
2.40
120
150
2,000
3.000
25k
S0R
756
160K
3,000
256
506
756.
1836
0.000
33C
504
756.
1306
0.000
236
536
756
1006
0.000
236
53K
736
1006
Jamosbond
1anganoo
KungFuMaster
106
60
228
20k
1.98
208
248
3.25
66
208
300
69
330
150
3.0
0.09
0.000
3.000
230
3.000
25K
50K
15k
100K
0.000
25K
56c
756.
183k
0.603
33K
524
756
1006
0.000
25K
53X
75c
1000
0.00
236
53x
73K
1006
140
140
418
409
2.48
6.
6.00
1.20
3.00
350
3.00
0.000
2.000
25R
50C
756
185k
2,000
256
50c
756
1854
0.007
25C
35C
756
1916
0.500
25K
536
75C
1704
0.70
25R
53K
75
1006
1.8
750
3.50
1.0
320
2.00
1300
1.0
300
2,400
25K
75K
0.000
50c
226
756
SEILING
Figure 6: Training curves of EDELINE on the Atari1008 benchmark for individual games (400K environment steps). The solid lines represent the average scores over 3 seeds, and the filled areas indicate the standard deviation across these seeds.
18

---

D.3 Atari 100k Qualitative Analysis
To provide deeper insights into EDELINE's superior performance, we conduct qualitative analysis on three Atari games where our approach demonstrates the most significant improvements over DIAMOND: BankHeist, DemonAttack, and Hero. Fig. presents temporal sequences comparing ground truth gameplay with predictions from both EDELINE and DIAMOND world models.
In BankHeist, agents maximize scores through repeated map traversal to encounter enemies. EDELINE's SSM-enhanced world model maintains consistent tracking of the player character posi- tion throughout prediction sequences, while DIAMOND's model shows progressive degradation with the character eventually disappearing from predictions. For DemonAtrack, EDELINE successfully captures the relationship between enemy hits and score updates in its predictions. DIAMOND preserves basic visual structure but fails to reflect these crucial state transitions. The Hero envi- romment showeases EDELINE's long-range prediction capabilities, accurately capturing sequences of the player breaking obstacles and navigating new areas. These qualitative results highlight how EDELINE's architectural innovations SSM-based memory, unified training, and diffusion model- ing enable robust state tracking. action-consequence modeling, and temporal consistency. These capabilities directly contribute EDELINE's superior performance on the Atari 100k benchmark.
T=o
10
20
30
35
True
Bankfeeis
EDELINE
DIAMOND
True
DemonAttack
EDELINE
DIAMOND
True
DIMOND EDELINE
Nero
Figure 7: Qualitative comparison of world model predictions on three Atari games. Each panel shows temporal sequences comparing ground truth with EDELINE and DIAMOND predictions. In BankHeist (top), EDELINE successfully tracks the player character while DIAMOND loses this information. DemonAttack (middle) demonstrates EDELINE'S accurate modeling of score updates upon successful hits. Hero (bottom) showcases EDELINE'S ability to maintain consistent predictions of complex successful (green) and failed (red) predictions of key game elements.
interactions across extended sequences. Colored boxes highlight
19

---

D.4 Extended Atari 100k Benchmark Performance Analysis
1.00
SimPLe TWM
Dreamerva STORM
DIMMNND EDELINE (ounS)
EDELINE (ours)
Fraztion of nuns with scDe
DIAMOND
STORM
DreamerV8
0:50
IRIS
0.28
TWM
SimPLe
0.00
0.45
0.60
0.75
Human Normalized Score (T)
Human Normalized Score
Performance profiles
Optimality gaps
Figure 8: Additional performance analyses. (a) Performance profiles showing fraction of runs achiev- ing scores above different human normalized score thresholds. (b) Optimality gaps demonstrating relative distance theoretical optimal performance across different methods.
To provide additional performance insights, we analyze EDELINE using performance profiles and optimality gaps [58]. Fig. [[a] shows empirical cumulative distribution of human- coman-normalized scores (HNS) across all 26 Atari games. The y-axis indicates the fraction of games achieving scores above each HNS threshold т. EDELINE consistently maintains higher fraction of games across different thresholds compared to baseline methods, particularly in the middle range (T between 1 and 4). We further analyze model performance through optimality gaps shown in Fig. 8(b) The optimality gap measures the distance between model performance and theoretical optimal behavior. EDELINE achieves the smallest optimality gap among all compared methods, demonstrating its effectiveness in approaching optimal performance across the Atari 100k benchmark suite. These detailed analyses complement the aggregate metrics presented in Section 6.2 by providing a more granular view of performance distribution and theoretical efficiency.
(Aaon 100N)
(loor 100%)
270
13%
50%
110%
575
44%
Lowe)
Pletatt
TROTTER
lood
wolk
aperthemn ieac
110t mord mode
Dheer
Figure 9: Comparison of ofate-of-the-arre model- RL methods without using look-ahead search techniques on the Atari 100k benchmark. EDELINE outperforms all existing model-based approaches on the Atari 100k benchmark. Previous methods can be categorized by their world model architectures: Transformer-basse
models (TWM (32,, IRIS (35] REM (66] STORM (3]
A-IRIS [59], and TWISTER [67]), RNN-based models (DreamerV3 2. PaMoRL. (68], and Har- monyDream (56)). SSM-based model (Drama [48)), and diffusion-based model (DIAMOND (8)).
Our proposed EDELINE advances the state-of-the- thrt--rr-t by integrating diffusion modeling with state space models, combining their respective strengths in visual generation and temporal modeling.
20

---

D.5 Atari 100k Linear Probing Analysis
To quantitatively evaluate the information content captured in EDELINE'S hidden representations, we conducted comprehensive línear probing analysis, which assesses how effectively the model's hidden states encode information critical for both reward prediction and observation prediction tasks.
Table 3: Linear probing loss of reward and observation prediction from EDELINE and DIAMOND hidden states. Liew (reward loss) was measured using cross- roos-entropy loss on reward prediction, while Lots (observation loss) was measured using MSE loss on predicting encoded observation features. EDELINE outperforms DIAMOND in observation prediction with an average 57.3% reduction in loss, while both models achieve comparable performance reward prediction (i.e., average feew of 0.0246 vs. 0.0258). These results demonstrate that EDELINE'S unified hidden representation successfully encodes information required for both tasks.
Environment
EDELINE Laew
DIAMOND Lrew
EDELINE Lats
DIAMOND Labs
Alien
0.0434
0.0663
0.6161
1.1798
Amidar
0.0186
0.0187
0.6430
0.7520
Assault
0.0779
0.0270
0.4894
0,4199
Asterix
0.0114
0.0018
0.8247
1,4411
BankHeist
0.0023
0.0039
0.0312
2.5398
BattleZone
0.0129
0.0220
0.2265
0.1413
Boxing
0.0482
0.0011
0.3089
0.5328
Breakout
0.0615
0.0264
0.6224
0,4871
ChopperCommand
0.0261
0.0113
0.6815
1.3819
CrazyClimber
0.0029
0.0034
0.2844
0,5513
DemonAttack
0.0609
0.0726
0,8588
1.0039
Freeway
0.0000
0.0000
0.2163
0.3911
Frostbite
0.0001
0.0002
0.1992
0.7435
Gopher
0.0090
0.0091
0.3139
0.9558
Hero
0.0065
0.0088
0.0708
0.1013
Jamesbond
0.0054
0.0058
0.5931
1,1080
Kangaroo
0.0009
0.0046
0.0992
0.5054
Krull
0.0728
0.0455
0.3216
0.8615
KungFuMaster
0.0048
0.0043
0.2674
0,8735
MsPacman
0.0044
0.0669
0,2510
0,8769
Pong
0.0000
0.0000
0.1770
1.3162
PrivateFye
0.0014
0.0144
0.1188
0,4931
Qbert
0.0240
0.0025
0.1813
0,2530
RoadRunner
0.0658
0.0310
0.2032
0.7896
Seaquest
0.0008
0.0299
0,2159
0.7852
UpNDown
0.1012
0.2110
0.5967
1,0783
ViZDoom-Dead DeadlyCorrido
0.0008
0.0094
0.3354
1.2848
Average Loss
0.0246
0.0258
0.3610
0.8462
We trained linear probes on the hidden states extracted from both EDELINE and DIAMOND models. For our experimental setup, we collected 50 trajectories as training data and 10 trajectories as validation data for each environment, with five independent random seeds employed to ensure statistical robustness. For reward prediction, we trained a single linear layer to predict rewards for 50 epochs. For observation feature prediction, we trained multi-layer perceptron (MLP) to predict encoded features of the subsequent observation (extracted from the pre-trained actor-critic network encoder) for 100 epochs. For EDELINE, utilized the Mamba hidden states, while for DIAMOND, we employed the LSTM hidden states from its reward-termination model. We conducted these experiments across 26 Atari environments and the challenging ViZDoom-Deade Seaonn-DDealyyorriidd environment to ensure comprehensive evaluation.
Table presents the results of our linear probing analysis. EDELINE and DIAMOND demonstrate comparable performance on reward prediction, with average losses of 0.0246 and 0.0258 respectively, EDELINE outperforms DIAMOND in 15 out of 27 environments on this task. More significantly, EDELINE substantially outperforms DIAMOND in observation prediction, with an average 57.3% reduction in loss. EDELINE achieves lower observation prediction loss in 24 out of 27 environments. The substantial improvement in observation prediction while maintaining comparable reward predic- tion performance demonstrates that EDELINE'S unified hidden representation successfully captures information required for both tasks.
21

---

D.6 Generation Quality Evaluation
To quantitatively assess the generation quality of our proposed world model, we conducted a compre- hensive evaluation measuring pixel-wise Mean Squared Error (MSE) across 26 Atari environments and the challenging ViZDoom- ViZDoom-DeadlyComm DeadllCorridorr environment. This analysis provides direct evidence of EDELINE's improved predictive accuracy compared to DIAMOND.
Table 4: Pixel-wise MSE comparison between EDELINE and DIAMOND across 27 environments. Lower values (in bold) indicate better performance. The rightmost column shows the imagination horizon length for each environment. The bottom reports the average normalized score (EDE- LINE MSE DIAMOND MSE), with values below 1.0 indicating EDELINE'S overall superior performance.
Environment
EDELINE
DIAMOND
Imagine Horizon Length
Alien
0.0063
0.0064
635
Amidar
0.0177
0,0235
1029
Assault
0.1011
0.2528
566
Asterix
0,0190
0.0177
1493
BankHeist
0,1510
0.1733
2019
BattleZone
0.0306
0.0292
1407
Boxing
0.0068
0,0174
1371
Breakout
0.0420
0.0415
1796
ChopperCommand
0.0084
0,0082
3006
CrazyClimher
0.0666
0.0686
3219
DemonAttack
0.0312
0.0318
1474
Freeway
0.0015
0.0030
1986
Frostbite
0.0338
0.0370
Gopher
0.0152
0.0172
2689
Hero
0.1559
0,2005
2103
Jamesbond
0.2967
0.3023
2212
Kangaroo
0.0061
0.0063
3241
Krull
0,1577
0.1575
1326
KungFuMaster
0.0210
0.0210
1864
MsPacman
0.0179
0.0197
892
Pong
0,0035
0.0034
1532
PrivateEye
0,1043
0.0638
2454
Qbert
0,0435
0.0460
1434
RoadRunner
0.0526
0.0532
960
Scaquest
0.0107
0.0136
1810
UpNDown
0.1223
0.1234
1279
ViZDoom- Baaoooo-DDeallyCCC DeadlyCorridor
0.0179
0.0184
Mean Normalized Score
0.918
1.000
Table Jableeetttttttt the comprehensive results of our pixel-wise MSE comparison. The analysis reveals that EDELINE achieves lower MSE than DIAMOND in 20 out of 27 environments. For meaningful comparison across environments with varying visual complexity, we calculated the average normal- ized score (EDELINE MSE/ DIAMOND MSE) across all environments. Our analysis demonstrates that EDELINE achieves an average normalized score of 0.918, which represents an overall 8.2% reduction in pixel-level error compared to DIAMOND.
The superior generation quality of EDELINE be attributed to its utilization of Mamba to incor- porate longer observation history beyond the four frames employed by DIAMOND. This capability enables EDELINE to maintain more consistent predictions extended horizons, which proves particularly advantageous in environments that necessitate memory (such as ViZDoom) or involve complex dynamics.
22

---

D.7 Training Time Profile
We performed detailed training time profiling of EDELINE to evaluate its computational efficiency compared to DIAMOND. Sable/5provicce comprehensive breakdown of training time components across different scales, while Table 6 direetly compares EDELINE with DIAMOND. Our profiling reveals that EDELINE achieves notable efficiency improvements in world model training. For world model updates, EDELINE is approximately 26.8% faster than DIAMOND. This efficiency gain stems from our unified architecture approach. While DIAMOND employs a two-stage training process (dif- fusion model for observations plus separate CNN-LSTM network for reward/termination prediction), EDELINE's unified architecture enables joint learning of observations, rewards, and terminations through shared representations. In addition, Mamba's parallel scan algorithm contributes to this computational advantage during training. For actor-critic updates, EDELINE requires approximately 17.7% more computation time than DIAMOND. This difference occurs because while both methods use identical actor-critic architectures, EDELINE's world model involves more complex inference during imagined rollouts due to SSM processing and cross-attention mechanisms. Specifically, as shown in Table each imagination step requires 2m.4ms, with 17.2ms dedicated to observation prediction and 6.5ms reward/termination prediction. Overall, these timing differences balance out, resulting in comparable total training time between the two approaches. This demonstrates that EDELINE successfully maintains computational efficiency comparable to DIAMOND while delivering significantly enhanced memory capabilities and performance.
Table 5:
Detailed breakdown of training time components across different scales. Profiling performed using a Nvidia RTX 4090 with default hyperparameters. Measurements are representative, exact durations depend on specific hardware, environment, and training stage.
Single update
Time (ms)
Detail (ms)
Total
548.6
148.6 + 400
World model update
148.6
Actor-Critic model update
400
15 x 24.4 34
Imagination step (x 15)
24.4
17.2 + 6.5 0.7
Next observation prediction
17.2
Denoising step (x 3)
5.7
Reward/Terninatio prediction
6.5
Action prediction
0.7
Loss computation and backward
34
Epoch
Time (s)
Detail (s)
Total
219.4
59.4 160
World model
59,4
400 x 148.6 x x10
Actor-Critic model
160
400 x 400 x10-3
Run
Time (days)
Detail (days)
Total
2.9
2.5 + 0.4
Training time
2.5
1000
219.4 (24 3600)
Other (collection, evaluation, checkpointing)
0.4
Table 6: Computational efficiency comparison between DIAMOND and EDELINE across different training stages, showing per-stage timing and relative differences.
Module
DIAMOND (ms/s/days)
EDELINE (ms/s/days)
Difference ( (%)
Single Update
Total
543 ms
548.6 ms
+1.03%
World Model Update Actor-Critic Model Update
203 ms
148.6 ms
-26.80%
340 ms
400 ms
+17.65%
Epoch
Total
217
219.4 s
+1.11% %
World Model
81
59.4
-26.67%
Actor-Critt Model
136
160 8
+17.65%
23

---

D.8 Ablation Studies
To systematically validate the effectiveness of EDELINE'S key architectural components, we perform comprehensive ablation studies across multiple environments. These studies isolate the contribution of each design decision and demonstrate their impact model performance. The ablation studies in Appendix D.8. and D.8.2 focus on five representative environments for validating the proposed key components, These include four Atari games where EDELINE demonstrates significant improvements over DIAMOND (BankHeist, DemonAttack, Hero, Seaquest), and MuniGrid-Memory capability evaluation. This selection provides comprehensive validation across visual prediction quality and memorization requirements. To compare Transformer and Mamba architectures for the Recurrent Embedding Module, we conducted experiments on memory tasks such as MiniGrid MemoryS7/S9 and Crafter. These results are presented in Appendix D.8.4
for memory
D.8.1 Choice of REM architecture
To validate the selection of Mamba for REM, we compare its performance against traditional linear- time sequence models GRU and LSTM across five environments. Fig. [OJillustrates that although all models achieve reasonable performance, Mamba demonstrates more stable learning curves and superior final performance, particularly in memory-intensive tasks such as BankHeist and MiniGrid- MemoryS9. GRU and LSTM models exhibit increased training variance with lower final scores, which validates Mamba's s effectiveness.
Banthaeat
DemonAttack
Hero
Mamory50
Se Seaquaatt
2
215
6.6
30C
756
0.452
0.500
902
002
280
372
256
22k
15K
25R
756
756
756. 2008
CRU
MAMBA
Figure 10: Performance comparison of different linear-time sequence models as REM architecture across five environments. Training curves show mean and standard deviation over three seeds. Mamba (red) shows more stable training progression and superior final performance compared to GRU (blue) and LSTM (green).
D.8.2 Cross-Attention in Next-Frame Predictor
To evaluate whether cross-attention improves the Next-Frame Predictor's ability process information-rich hidden embeddings, we examine EDELINE with and without this mechanism. As depicted in Fig. [1 the original EDELINE demonstrates superior performance in BankHeist, MemoryS9, and Seaquest where rich contextual information processing proves essential. The Memo- ryS9 environment validates this necessity, models must integrate historical information for complete representation reconstruction. Cross-attention enables effective fusion of hidden embeddings with visual features for temporal context integration.
BankMtubi
Hemory59
Scarwestt
2.4k
1.3
300
118
1.2K
104
752
250
300
502
220
500
3150
153
255
S0K
358
1200
256
13C
116
122
2-6
12c
746
20C
41C
420
626
1064
25c
6
156
100
EDCLINE
Figure 11: Ablation study comparing EDELINE with cross-attention blocks (blue) and without (green) across five test environments. Training curves depict mean and standard deviation over three seeds. EDELINE's cross-attention mechanism provides advantages in environments requiring rich contextual information processing.
24

---

D.8.3 Effect of Harmonizers
To validate the effectiveness of harmonizers, we evaluate performance across the full Atari 100k benchmark.
Trus
EDELINE
EDELINE
(a) BankHeist
Trus
EDELINE
wil Hemonntees
EDELINE
(b) RoadRunner
Figure 12: Effect of harmonizers world model predictions. (a) Without harmonizers, EDELINE
fails to maintain consistent character modeling BankHeist. (b) RoadRunner, removal of harmonizers leads to loss of seward-relevant visual details. Colored boxes highlight successful (green) and failed (red) predictions of key game elements.
Table presents quantitative performance comparison between EDELINE with and without har- monizers. The ablation results demonstrate significant performance degradation across multiple environments when harmonizers removed. with particularly notable drops in games requiring precise reward related visual detail retention (BankHeist and RoadRunner). Qualitative analysis on these environments, which showed the largest performance improvements with harmonizers, reveals how harmonizers contribute to world model performance. In BankHeist (Fig. [2(a)), removing harmonizers causes the world model lose track of game characters, failing maintain consistent agent representation across prediction sequences. Similarly, in RoadRunner (Fig. [2(b). the model without harmonizers fails to capture reward-relevant visual details, degrading its ability to model critical game state information.
These results demonstrate how harmonizers help achieve crucial balance between observation modeling and reward prediction. Without harmonizers, the world model struggles with fine- fine-grained task-relevant observations in both environments tracking small character sprites in BankHeist and capturing critical game state details in RoadRunner. By maintaining this dynamic equilibrium between observation and reward modeling, harmonizers enable EDELINE to effectively leam compact lask-centric dynamics while preserving essential visual details for sample-efficient learning.
25

---

Table To evaluate the benefits of Harmonizers in detail, we compare DreamerV3 and its variant with Harmonizers CaaronnyDraayy alongside EDELINE without Harmonizers and the complete EDELINE on the 26 games the Atari 100k benchmark. Bold numbers indicate the highest scores.
EDELINE w/o Harmonizers
Game
Random
Human
DreamerV3
HammonyDream
EDELINE (ours)
Alien
227.8
7127.7
959,4
889.7
1086.2
974.6
Amidar
5.8
1719.5
139.1
141.1
212.5
299.5
Assuult
222.4
742.0
705.6 6
1002.7
1180.1
1225.8
Asterix
210.0
8503.3
932.5
4612.8
4224.5
BankHeist
14.2
753.1
648.7
1068.6
139.8
854.0
BattleZone
2360.0
37187.5
12250.0
16456.0
2685.0
5683.3
Boxing
0.1
78.0
79.6
88.3
88.1
Breakout
1.7
30.5
31.1
52.6
255.7
250.5
ChopperCommand
811.0
7387.8
410.0
1509.6
2616.5
2047.3
CrazyClimber
10780.5
35829,4
97190.0
82739.0
96889.5
101781.0
DemonAttack
152.1
1971.0
303,3
202.6
924.2
1016.1
Freeway
0.0
29.6
0.0
33.8
33.8
Frostbite
65.2
4334.7
909,4
678.7
289.0
286.8
Gopher
257.6
2412.5
3730.0
13042.8
7982.7
6102.3
1027.0
30826.4
11160.5
13378.0
9366.2
12780.8
Jameshond
29.0
302.8
444.6
317.1
527.5
784_3
Kangaroo
$2.0
3035.0
4098.3
5117.6
3970.0
$270.0
Kruli
2665.5
7753.6
9748.8
Kunsterastter
258.5
22736.3
21420.0
22274.0
14088.5
31448.0
MsPacman
307.3
6951.6
1326.9
1680.7
1773.3
1849.3
Pong
-20.7
18.4
18.6
18.4
20.5
PrivatelEye
24.9
69571.3
881.6
2932.2
73.0
99.5
Qbert
163.9
13455.0
3405.1
3932.5
5062.3
6776.2
RoadRunner
11.5
7845.0
15565.0
14646.4
23272.5
32020.0
Seaquest
68.4
42054.7
618.0
665.3
1277.2
2140.1
UpNDown
533.4
11693.2
7567.1
10873.6
2844.6
5650.3
#Superhuman (1)
Mean (1)
0.000
1,000
1.124
1.364
1.674
1.866
Table 8: Performance comparison between Transformer and Mamba-base REM architectures across memory-demandeee environments. Values for MiniGrid environments represent success rates, while Crafter values show average return over three independent seeds after 1M environment steps.
REM Architecture
MiniGrid- Mentrry5-
MiniGrid- Memorry-
Crafter
REM Inference FLOPs
Transformer
0.980
0.978
10.2 + 0.8 11.5 = 0.9
263.484M
Mamba
0.981
0.982
213.307M
D.8.4 Transformer-baset Recurrent Embedding Module
To investigate the efficacy of Mamba compared Transformer architecture, we implemented Transformer-based variant of the Recurrent Embedding Module (REM). This ablation study aims assess both performance and computational efficiency across memory-demanding environments, which provides critical insights into architecture selection for world modeling. We evaluated both SeniGrid-MMmoryy
architectures the SuniGrid-MMmorrSSS specifically challenge a model's ability to retain and utilize historical information.
and Crafter environments, which
As demonstrated in Table 8 both Transformer and Mamba architectures achieve comparable perfor- mance across the evaluated memory-demanding environments. On the MiniGrid-Ma Memmory tasks, both architectures attain success rates approaching optimal performance (>0.97). In the more complex Crafter environment, which requires nuanced long-term memory and planning capabilities, Trans- former achieves average return of 10.2 0.8 while Mamba achieves 11.5 + 0.9, demonstrating that both architectures can effectively model long-term dependencies in world modeling contexts.
From computational perspective, the REM inference FLOPs comparison reveals that Mamba requires 213.307M FLOPs compared to Transformer's 's 263.484144 48344 FLOPs, representing approximately 19% lower computational cost attributed to its linear time complexity. These results substantiate that both Transformer and Mamba represent viable architectural alternatives for the REM component.
Our primary objective is to addresses the memory limitations of existing approaches. Our selection of Mamba motivated by its simplicity and linear-time complexity, which achieves satisfactory performance while maintaining computational efficiency. Further exploration of superior architectural designs that enhance both efficiency and accuracy for diffusion-based world modeling constitutes an important direction for future research.
26

---

D.8.5 Frame Stacking Ablation
To validate the necessity of EDELINE's frame-stacking design, we investigate whether both SSM hidden embeddings and explicit frame stacking are essential for achieving optimal performance. Our architectural design maintains both components based on the hypothesis that they serve complemen- tary roles the diffusion process: the hidden embedding captures long-term temporal dependencies and global context across the entire sequence history, while the explicit frame stacking provides pixel-based spatial details from recent observations that are crucial for pixel-level prediction accuracy. The diffusion model benefits from both the abstract temporal representations (via cross attention with hidden embedding) and direct access to recent visual features (via frame stacking) for generating high-fidelity predictions.
To empirically validate this design choice, we conducted an ablation study comparing three configu- rations: (1) EDELINE with both components, (2) EDELINE with hidden embedding only (without frame stacking), and (3) DIAMOND as a baseline. We evaluated these variants on two environments with contrasting characteristics: Atari Breakout, which requires minimal long-term memory while de- manding high visual detail processing, and MiniGrid- Memorrrr, which requires substantial memory capabilities while having simpler visual complexity.
Table 9: Frame Stacking Ablation Study. Performance comparison demonstrating the necessity of both SSM hidden states and explicit frame stacking across environments with different memory and visual complexity requirements.
Environment
DIAMOND
EDELINE w/o frame-stack
EDELINE
MiniGrid- MenGorddMemmoyy99
0.380
0.985
0.982
Atari Breakout
132.5
145.6
250.5
As shown in Table 9 the results demonstrate distinct patterns across the two environments. In Atari Breakout, EDELINE without frame stacking achieves score of 145.6, which represents an improvement DIAMOND's 's 132.5 but remains substantially below the full EDELINE's 250.5. This significant performance gap in a visually complex environment highlights the critical importance of explicit frame stacking for preserving fine-grained visual details necessary for accurate pixel-level prediction. In contrast, MiniGrid- MMmmdrySS9 presents different pattern: EDELINE without frame stacking achieves success rate of 0.985, which is comparable to the full EDELINE's 0.982 and substantially outperforms DIAMOND's ' 0.380. This near-optimal performance a memory- -nttoryiittesiie but visually simpler environment validates that the hidden embedding alone can effectively capture the temporal dependencies required for success when visual complexity is limited.
These results substantiate our architectural design by demonstrating that the two components serve complementary and essential roles. The hidden embedding provides the long-term memory capacity necessary for temporal reasoning, evidenced by the strong performance in MiniGrid-Menn even without frame stacking. Conversely, the explicit frame stacking preserves the Memoryy9 visual information crucial for environments with rich visual complexity, as demonstrated by the substantial performance improvement in Atari Breakout when both components are present. This ablation study validates that our dual input design is not redundant but rather represents a deliberate architectural choice that enables EDELINE to excel across diverse environments with varying memory and visual complexity requirements.
27

---

D.9 ViZDoom Environment Specifications
This appendix provides detailed specifications for the ViZDoom scenarios used in our experiments. We evaluate EDELINE on five key scenarios that test different agent capabilities in first-person 3D environments.
D.9.1 DeadlyCorridor
Objective: Navigate through enemy fire to acquire armor at corridor end
Reward Mechanism:
+1 for each enemy killed
+1 for armor acquisition
-0.01 for each damage instance received
-0.01 move backward
Evaluation Metric: Binary success if armor acquired before episode termination
Episode Termination: Agent death 512 environment interaction steps.
D.9.2 HealthGathering
Objective: Survive by collecting medkits in toxic environment
Reward Mechanism:
+1 for each medkit collected
Evaluation Metric: Final health percentage (0-100%)
Episode Termination: Agent death or 512 environment interaction steps.
D.9.3 PredictPosition
Objective: Anticipate enemy movement land delayed rocket hit
Reward Mechanism:
- +1 for successful enemy kill
Evaluation Metric: Binary success if enemy eliminated
Episode Termination: Successful kill 75 environment interaction steps.
D.9.4 Basic
Objective: Eliminate stationary enemy with limited ammunition
Reward Mechanism:
+1 for successful kill
Evaluation Metric: Binary success if enemy eliminated
Episode Termination: Successful kill or 75 environment interaction steps.
D.9.5 DefendCenter
Objective: Survive against infinite enemy waves
Reward Mechanism:
- +1 per enemy killed
-- for agent death
Evaluation Metric: Total enemies killed per episode
Episode Termination: Agent death 512 environment interaction steps.
All scenarios use the following common configuration parameters unless otherwise specified:
Observation space: (64,6433)
Action space: Discrete movement and attack actions
28

---

D.10 Crafter Experiments Qualitative Results
GT
EDELINE
5.39
5397
EDELINE
DIAMOND
GT
EDELINE
EDELINE
DIAMOND
GT
EDELINE
EDELINE
DIAMOND
Figure 13: Qualitative comparison of world model predictions in Crafter. We compare GT (Ground Truth), EDELINE, EDELINE" (Transformer-based variant), and DIAMOND. The red boxes indicate prediction errors, the pink boxes show acceptable errors, and the green boxes highlight correct predictions of key elements.
In Fig. 13 we can observe comparative results between the different model architectures. timesteps 1-10 (top row), all three methods demonstrate reasonable consistency in their predictions. Starting from timesteps 11 11-20 (middle row), notable differences emerge: DIAMOND fails to correctly represent crafted materials, and both DIAMOND and EDELINE" incorrectly predict outcomes when the player chops trees and crafts new materials. EDELINE maintains higher prediction accuracy throughout these interactions. In timesteps 21-30 (bottom row), all methods generate acceptable
29

---

predictions for previously unseen regions. While these differ from Ground Truth, they remain acceptable as they preserve gameplay-critical elements. During sleep actions, both EDELINE and DIAMOND fail to predict the sleeping state, which we classify acceptable error given that sleep random event triggered independently of player actions. Most importantly, when the player returns previously visited areas, only EDELINE successfully remembers and reproduces critical environmental features such as trees. EDELINE" and DIAMOND both fail to recall the crafting table and trees outside the map border. This evidence demonstrates EDELINE's superior memorization capability for maintaining environmental consistency over long horizons.
30

---

D.11 Hyperparameters
Table 10: Hyperparameters for EDELINE.
Hyperparameter
Symbol
Value
General
Number of epochs
1000
Training steps per epoch
400
Environment steps per epoch
100
Batch size
32
Sequence Length
T
19
Epsilon (greedy) for collection
0.01
Observation shape
(h,w.c)
(64,64.3)
Actor Critic
Imagination horizon
15
Discount factor
0.985
Entropy weight
0.001
X-returs coefficient
X
0.95
LSTM dimension
512
Residual blocks layers
[1,1.1,1]
Residual blocks channels
[32,32.64.64]
World Model
Number of conditioning observations
L
4
Burn-in length
B
4
Hidden embedding dimension
mmode
512
Reward Termination Model Hidden Units
Mamba
Mamba layers
n_ layers
3
Mamba state dimension
d_ state
16
Expand factor
2
ID Convolution Dimension
_conv
4
Residual blocks layers
[1,1.1,]]
Residual blocks channels
[64.64.64.64]
Action embedding dimenstion
128
Diffusion
Sampling method
Euler
Number of denoising steps
3
Condition embedding dimension
Residual blocks layers
[2,2,2,2]
Residual blocks channels
[64.64.64.64]
Optimization
Optimizer
AdamW
Learning rate
lc-4
Epsilon
1e-8
Weight decay (World Model)
le-2
Weight decay (Actor-Critic)
0
Max grad norm (World Model)
1.0
Max grad norm (Actor-Critie)
100.0
31

---

E EDELINE Algorithm
We summarize the overall training procedure of EDELINE in Algorithm Srrrithmmbbbaeee bellow, which is modified from Algorithm I in g. We denote D the replay dataset where the agent stores data collected from the real environment, and other notations are introduced previous sections are self collanxalanatory
Algorithm 1: EDELINE
Procedure training.loop(f for epochs do
collect_experiee ence (steps_collect)
for steps_world__mon do
=pdate__oord__ode10 for steps_ aactrrcccitttt do update__ctoorcritiic
Procedure cowlect_eepeeeeeeeeeeeeeeeeeeeee collectt_xperieecee)))
of f- env reset ()
for ! = 0 to
do
Sample a, To(a; (a, of))
0.+1, rt, de + env www..tteeeee.e..... step(a,)
DU (ot. 0r )
if d, 1 then
0f f env reset()
Procedure update_war sorrtte_worlldddo 1eclid_moee@0) ():
Sample indexes := (t, T - 1
// sequence length
Sample sequence (of, at Tos ds)EEI D DD
Initialize h4-1
// MAMBA hidden state
Parallel for I do
h4 - fo as, h1-1) processed in parallel via MAMBA parallel scan P: (Fili)h
d, Appooddlii (di(be))
Compute Crew (Ø) Eiez CE(fi, r:)
CE: cross cross-entropy loss
Compute Lend (Ø) = EiET CE(d,, di)
// CE: cross-entropy loss // burn-in steps // log-normal sampling from EDM 11 identity schedule from EDM
Sample index ~ (n BB,
-
Sample log(a) ([Pmean, p2
Define = Sample of Nof. 0²I)
// add Gaussian noise
Compute 06 = (o T, 00 o .0 1
Compute observation modeling loss Loba (0) = 1109 - 09112
Update @ according to Eq. 5
Procedure ceattt__acto_crrrtiic
Sample initial buffer (of-8+1: at-B+1, of) ~
// Burn LSTM states mo, Vo and MAMBA states So with buffer for i= i f 1
Sample 0; ~ (a; [o?
Compute h; folo;, ai h4-1)
Sample reward r., next observation of and termination d, via Ps
Compute (o()) for i ... H
Compute RI. losses Lv (0) and (0)
Update To and Vo
32

---

Actor-Critic Learning Objectives
We follow DIAMOND g in the design of our agent behavior learning. Let or, rt, and d, denote the observations, rewards, and boolean episode terminations predicted by our world model. We denote H as the imagination horizon, Vo as the value network, To as the policy network, and dr as the actions taken by the policy within the world model.
X-returns balance bias and variance in the regression target. Given an imagined trajectory of length H, we define the A-return recursively:
For value network training,
1r+ d)) A)Vo(ou-1) + AA++1] ift H Vo(oH)
A,
(18)
ift
The value network V/ is trained to minimize Lv()), the expected squared difference with A-returns over imagined trajectories:
H
Cv (8) = Exe
Σ (Vo(o)) - sg(A)))²
(19)
where sg(-) denotes the gradient stopping operation, following standard practice 35
For policy training. we leverage the ability to generate large amounts of on-policy trajectories in imagination using REINFORCE objective [69]. The policy trained minimize:
H-1
(=) -B log(nolaglo:)))g((4 Vo(o:)) + +)((To(a)oo)))
(20)
where Vo(o:) serves as a baseline to reduce gradient variance, and the entropy term H with weight encourages sufficient exploration.
33

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The main claims presented in the abstract and introduction accurately reflect the paper's contributions and scope. The methodology in SectionnSthooughhh develops the proposed approach, while the experiments in Sections 4 and provide comprehensive empirical validation for these claims, demonstrating the effectiveness of EDELINE across various environments and benchmarks.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contribu- tions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
It is fine include aspirational goals motivation long it clear that these goals not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations of this work are summarized in the Appendix
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to vi- olations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs, In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-ted system might not be used reliably to provide closed captions for online lectures because fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by review- ers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
34

---

Answer: [NA]
This paper does not present novel theoretical results that would require formal proofs. Instead, it builds upon existing theoretical frameworks of diffusion models and state space models, applying them to world modeling for reinforcement learning. Appendix C.T and Section praetionnprovideeee provides the necessary theoretical background on score-based diffusion models, reinforcement learning formulations, and sequence modeling with Mamba, with appropriate references to established work. The main contribution is an architectural innovation that combines these existing theoretical frameworks rather than developing new theoretical results requiring formal proofs.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
This paper discloses the information needed to reproduce the experimental results. The experimental configurations, detailed hyperparameter setups, and hardware requirements for the experiments are elaborated in Appendix D
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access hosted model (e.g., the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g., large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g..
35

---

to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described supplemental material?
Answer: [Yes]
Justification: We provide open access anonymized codebase at 4open. /ciienn///EDELLIIE which contains detailed installation instructions, hyperpa- rameter configurations, and implementation details necessary reproduce our experimental results.
Guidelines:
The answer NA means that paper does include experiments requiring code.
Please see the NeurIPS code and data submission guidelines
for more details.
While we encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer: Papers cannot be rejected simply for not including code, unless this is central to the contribution (e-g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurlPS code and data submission guidelines (https cc/ppblic/guiids//ooeeeeee
cemmiissionii ssisnPolicyl for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Our full training procedure is detailed in Algorithm I in Appendix We provide comprehensive hyperparameters, including environment configurations, number of epochs, batch sizes, data collection schedules, optimizer settings, and architectural details in Appendix()) The paper clearly describes the evaluation protocols for all benchmarks with appropriate statistical reporting across multiple seeds. All experimental settings necessary to understand and reproduce results are available in our anonymized codebase: httph: htttps://anonmmouts 4open. science/ r/ /EDELLINE-BBB
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The paper consistently reports statistical significance across all experiments. For
the Atari 100k benchmark, we present stratified bootstrap confidence intervals following (58] 's
recommendations (Fig. 3 and report means across seeds for all game scores (Table I For
36

---

ViZDoom experiments, Fig. 4 shows training curves with shaded regions indicating standard deviation across three seeds. Similarly. for Crafter experiments, Table average returns with standard deviation. This consistent approach to reporting statistical significance across different environments enables proper assessment of EDELINE's performance improvements over baselines and ensures the reliability of our experimental findings.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main
claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e-g.. Normally distributed errors). It should be clear whether the error bar the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates). If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experi- ments?
Answer: [Yes]
Justification: We provide detailed information on the computational resources required for our experiments in Appendix D.7 Table 5 the appendix present comprehensive breakdowns of compute requirements, including hardware specifications (Nvidia RTX 4090 GPU) and execution times across different scales (single updates, epochs, and complete runs).
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experi- mental runs well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips. cc/publicc/thiccSuuieeeee
Answer: [Yes]
Justification: All authors have reviewed the NeurIPS Code of Ethics and confirmed that the research conducted in this paper complies with it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
37

---

The authors should make sure to preserve anonymity (e.g., if there is special consideration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The paper discusses potential societal impacts in the impact statement provided in Appendix B
Guidelines:
The answer NA means that there is no societal impact the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it needed to point that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have made our best effort to include relevant citations and licenses in our paper and codebase.
38

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paaerrsw thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license
dataset. For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to asset's creators.
13. New assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification:
We provide detailed documentation alongside our codebase.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: This paper does not involve research with human subjects
Guidelines:
39

---

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is important, original, non-standard component the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs are only used for editing grammar.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM ottp:p///eeerrp..cccc//oeeeeeeee should or should not be described.
for what
40


Reviews:
[{"review_id": "E6lst3TjYq", "reviewer": "Reviewer_dLo2", "summary": "This paper introduces EDELINE, a unified world model architecture that integrates State Space Models (Mamba) with diffusion models to model longer action, observation sequences.  It aims to overcome the fixed context length limitation of prior methods (e.g., 4 frames for DIAMOND) and unify the training architecture end-to-end. EDELINE reports improved performance, temporal consistency, and imaginary quality across several environments.", "questions": "- While EDELINE theoretically supports \"unbounded\" sequences, the current evaluations are primarily within standard game episode lengths. What are practically sufficient context lengths needed here? Do the authors have insights on true long-horizon tasks in the real world?\n- For reproducibility and to fully understand the model's scale, could the authors provide detailed specifications of the Mamba architecture used as well as the training sequence lengths utilized across the different environments?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n - The paper identified a context length problem in the prior work and proposes a REM module to model it.\n - The framework proposes a unified architecture which integrates reward and termination prediction within shared hidden states.\n\nWeakness:\n  - The choice of mamba, though discussed by the author, is not sufficiently supported by experiments. A few improvements would be helpful\n      1. The comparison between the other REM modules supports this key design choice and should be placed in the main body instead of the appendix.\n      2. It is stated that \"Self-attention-based Transformer architectures, despite their strong modeling capabilities, suffer from quadratic computational complexity which impairs efficiency\". This needs to be empirically justified. Transformers do suffer from quadratic complexity, but would that be a problem in this setting? For example, the examples provided discussed action step 12, which is well within Transformer's \"comfort zone\". And empirically, the real computation time of a Transformer might not be worse than a Mamba in many settings. Given this, this particular statement is not sufficient to rule out the choice of Transformers.  \n            - on a side note, if we would like to discuss how important context lengths are, I would suggest sharing stats of the context lengths used during training, and also practically useful context length in inference. \n            - I am not saying Mamba is a wrong choice, but I disagree with ruling out a potential candidate architecture with a statement like this, which may misguide future explorations.\n     3. It appears the training details are not sufficiently presented. I can find some hyperparameters in the appendix, but I have a hard time finding the Mamba layer details, etc.", "was_revised": true, "final_justification": "I have a few main comments about this work initially:\n\n-  I disagree with ruling out a potential candidate architecture (transformer) with a statement like this, which may misguide future explorations.\n   - The authors have argued thoroughly here. Though I still would think there are no strong justification to claim Mamba saves more compute when the context is short, I also find it impractical to have the authors try all architectures when the focus of this paper is not this comparison. I think this is acceptable if the authors rephrase the statement and provide a more candid overview of the architecture choice.\n- It appears the training details are not sufficiently presented. I can find some hyperparameters in the appendix, but I have a hard time finding the Mamba layer details, etc.\n  - This is probably an overlook on my side while going through the appendix\n\n\nIn general, I still think this paper is borderline but I think an \"borderline accept\" is reasonable.", "rebuttal": "We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.\n\n> Q1. The comparison between the other REM modules should be placed in the main body.\n\nWe thank the reviewer for this valuable suggestion. We acknowledge the reviewer's recommendation and will incorporate the comparison between different REM modules into the main body of the revised manuscript to better support our architectural design choices.\n\n> Q2. This particular statement is not sufficient to rule out the choice of Transformers.\n\nWe thank the reviewer for this important observation. We agree with the reviewer that Transformer represents a viable and potentially promising architectural choice for our framework. Our initial experiments employed a window size of 16 to maintain computational feasibility and fair comparison with existing baselines during our preliminary investigations. However, we acknowledge this may have inadvertently disadvantaged Transformer architectures. When we removed this limitation and extended the window size to match Crafter's maximum episode steps (10,000), we found that Transformer performance became comparable to Mamba, as shown in Table R1.\n\nOur experimental results demonstrate that both Transformer and Mamba architectures achieve strong performance in our framework. We selected Mamba not due to its superiority over Transformer, but rather as we found that the simpler linear-time complexity solution already achieves satisfactory performance for our objectives. Given that both architectures perform well, and considering that the primary theme and most important aspect of this work is efficiency, we opted for the more computationally straightforward approach while acknowledging that Transformer remains an equally valid choice. We will update our discussion to incorporate the perspective of Transformer's capabilities in the revised manuscript.\n\n|REM Architecture|Avg Return|REM Inference FLOPs|#REM Parameters|\n|-|-|-|-|\n|Transformer|8.9 +/- 0.6|238.318M|7.251M|\n|Mamba|**11.5 +/- 0.9**|**213.307M** (-10.4%)|**5.670M** (-21.8%)|\n|Transformer w/o window limit|**10.2 +/- 0.8**|263.484M (+10.6%)|7.251M (0 %)|\n\n**Table R1: Recurrent Embedding Module (REM) Architecture Comparison on Crafter Environment.** We conducted additional comparisons between Transformer and Mamba architectures as the REM component in EDELINE. The evaluation includes performance (average return), computational efficiency (i.e., REM single-step inference FLOPs measured at context length = 32), and model complexity (i.e., number of REM parameters). The results demonstrate that both architectures are viable choices, with Transformer achieving competitive performance when the window size limitation is removed.\n\n> Q3. I would suggest sharing stats of the context lengths used during training, and also practically useful context length in inference.\n\nWe thank the reviewer for the suggestion. During world model training, we employ different sequence lengths depending on the environment: 19 steps for Atari, MiniGrid, and ViZDoom, and 32 steps for Crafter. For the imagination phase when training the Actor-Critic component, we use four steps of world model burn-in followed by 15 steps of imagination horizon for Atari, MiniGrid, and ViZDoom, while Crafter employs 32 steps of burn-in with a 15-step imagination horizon. Please note that due to Mamba's characteristics, there are no specific context length limitations.\n\n> Q4. The training details are not sufficiently presented. \n\nWe thank the reviewer for pointing this out. We have provided comprehensive implementation details in our supplementary materials. Algorithm 1 in Appendix E presents our overall training procedure, while Appendix D.12 contains hyperparameters used during training. In the revised version, we will include a detailed table of contents and reorganize the supplementary materials to avoid potential confusion and improve accessibility.\n\n> Q5. What are practically sufficient context lengths needed here? Do the authors have insights on true long-horizon tasks in the real world?\n\nWe thank the reviewer for the question. As detailed in our response to Q3 above, the sequence lengths vary by environment and training phase. While we agree with the reviewer regarding real-world applicability opportunities, please note that the primary theme, focus, and contributions of this work concern addressing the memory limitations of existing diffusion-based world models. Our contributions enhance the potential for deploying these models in real applications, especially when long-term memory capabilities are essential for effective decision-making.\n\n> Q6. Could the authors provide detailed specifications of the Mamba architecture used as well as the training sequence lengths utilized across the different environments?\n\nWe thank the reviewer for the question regarding reproducibility. We have provided Mamba-related configurations in Section D.12, including the number of Mamba layers, Mamba's latent state size, and expansion ratio. Additional settings follow standard Mamba defaults as established in the original implementation. The training sequence lengths across different environments are as described in our response to Q3 above. To further improve clarity and accessibility, we will reorganize these technical specifications for better accessibility in the revised manuscript."}, {"review_id": "RLH28b7cse", "reviewer": "Reviewer_asDt", "summary": "This paper proposes EDELINE, a unified world model architecture that enhances long-term memory capabilities by integrating state space models (SSMs) with diffusion models. The proposed method addresses two limitations of previous diffusion-based world models like DIAMOND: (1) the inability to capture long-range temporal dependencies due to fixed-length context windows, and (2) the architectural separation of observation, reward, and termination modeling, which hinders representation sharing. EDELINE introduces three key innovations: a Recurrent Embedding Module (REM) built on the Mamba SSM for unbounded sequence modeling; a unified architecture that shares representations across all prediction heads; and a dynamic loss harmonization mechanism to balance training objectives. Experiments show that EDELINE achieves state-of-the-art performance with improved memory utilization and predictive consistency.", "questions": "1.\t‌How does EDELINE uniquely advance beyond existing Mamba-diffusion works?\n\n2.\tWhy is the input $o_{t+1}$ required for predicting $\\hat{o}^0_{t+1}$ in the next-frame predictor module of Figure 2? \n\n3.\tCan the author provide the complete mathematical formulation showing how $h_t$ is derived from $h_{t-1}$ using matrices $A$, $B$, $C$, and $\\Delta$ in Figure 2?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n1.\tThe paper conducts extensive experiments on multiple benchmarks, including Atari 100k, Crafter, and ViZDoom, demonstrating EDELINE's superior performance across diverse environments.\n\n2.\tThe authors ablate key components across multiple environments and visualize predicted observations to quantify representation quality and generation fidelity.\n\n\nWeaknesses:\n\n1.\tIt is unclear how EDELINE differs from and improves upon existing works that also integrate SSMs with diffusion models [1] or design world models to predict multiple environment elements [2]. The authors may want to clarify the key novelty of their method.\n\n2.\tThere exists symbol abuse and inconsistency in the paper. For example, the use of 'O' for both the observation space and the observation function, and the overloading of the symbol 'B' for both the matrix defined for SSMs and the burn-in period. This lack of clarity could confuse readers.\n\n3.\tIn Figure 2, it is unclear why the prediction of the next frame $ \\hat{o}^0_{t+1} $ requires the input $o_{t+1}$. This seems contradictory as $o_{t+1}$ is supposed to be the prediction target. The authors should clarify the data flow and the exact meanings of notations.\n\n\n[1] Teng, Yao, et al. \"Dim: Diffusion mamba for efficient high-resolution image synthesis.\" arXiv preprint arXiv:2405.14224 (2024).\n\n[2] Lin, Jessy, et al. \"Learning to model the world with language.\" arXiv preprint arXiv:2308.01399 (2023).", "was_revised": true, "final_justification": "The authors have adequately resolved my earlier concerns regarding architectural novelty relative to prior SSM-with-diffusion works and clarified the notation. Accordingly, I raise my score to Borderline accept.", "rebuttal": "We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.\n\n> Q1. It is unclear how EDELINE differs from and improves upon existing works that also integrate SSMs with diffusion models or design world models to predict multiple environment elements. The authors may want to clarify the key novelty of their method.\n\nWe thank the reviewer for the question. Our primary contribution lies in designing a unified architecture that leverages SSMs to address the memory limitations of existing diffusion-based world models within the Model-based Reinforcement Learning domain, while maintaining computational efficiency through Mamba's parallel scan algorithm during training.\n\nWe would like to clarify that most current works combining diffusion models with Mamba, such as DiM [1], DiMSUM [2], Matten [3], and other recent approaches, primarily utilize Mamba's linear time complexity and unbounded context advantages as a replacement for traditional attention-based architectures in image generation and video synthesis tasks. These works, however, concentrate primarily on architectural fusion rather than addressing the specific challenges of diffusion-based world models. Regarding world models that predict multiple environment elements like Dynalang [4], those world models focus on grounding diverse language types through GRU-based sequence modeling, while our proposed EDELINE specifically targets the memory constraints of diffusion-based world models like DIAMOND (e.g., fixed 4-frame context windows).\n\nOur key novelty lies in integrating Mamba SSMs to enable **unbounded temporal consistency** in diffusion world modeling while maintaining superior visual fidelity. This approach addresses a fundamentally different challenge from existing approaches.\n\n[1] Teng, Yao, et al. \"DiM: Diffusion mamba for efficient high-resolution image synthesis.\" arXiv preprint arXiv:2405.14224 (2024).\n\n[2] Phung, Hao, et al. \"DiMSUM: Diffusion Mamba--A Scalable and Unified Spatial-Frequency Method for Image Generation.\" NeurIPS (2024).\n\n[3] Gao, Yu, et al. \"Matten: Video generation with mamba-attention.\" arXiv preprint arXiv:2405.03025 (2024).\n\n[4] Lin, Jessy, et al. \"Learning to model the world with language.\" ICML (2024).\n\n> Q2. There exists symbol abuse and inconsistency in the paper. \n\nWe appreciate the reviewer's attention to this detail. We will address these notation issues in our revised manuscript by replacing the observation space symbol $O$ with $\\mathcal{O}$ and changing the burn-in length symbol $B$ to $\\mathcal{B}$ to eliminate the overloading of symbols.\n\n> Q3. Why is the input $o_{t+1}$ required for predicting $\\hat{o}^0_{t+1}$ in the next-frame predictor module of Figure 2?\n\nWe thank the reviewer for this question. Figure 2 illustrates the **training process** of EDELINE's world model. During training, the noised version of $o_{t+1}$ ​ serves as the input to the diffusion process, where the model performs denoising based on the current conditioning information and noise level. This follows the standard training procedure for conditional diffusion models. During **inference**, no ground truth frame is available. Instead, the model generates the next frame $\\hat{o}\\_{t+1}$ by starting from pure Gaussian noise and progressively denoising through multiple steps without requiring $o_{t+1}$ as input. This approach is consistent with previous diffusion-based world models such as DIAMOND and follows established practices in conditional diffusion literature. We will incorporate this clarification in the revised manuscript.\n\n> Q4. Can the author provide the complete mathematical formulation showing how $h_t$ is derived from $h_{t-1}$ using matrices $A, B, C,$ and $\\Delta$ in Figure 2?\n\nWe thank the reviewer for this question. Our implementation follows the established Mamba literature. We concatenate the current timestep's observation $o_t$ and action $a_t$ ​as input, denoted as $u_t = (o_t, a_t)$. We project $u_t$ to initialize the input-dependent matrices $B, C$, and $\\Delta$, while matrix $A$ remains input-independent. The discretization is then performed using $\\Delta$ to obtain the discrete matrices $\\tilde{A}$ and $\\tilde{B}$. Let $x_t$ denote the Mamba latent state at timestep $t$. The Mamba architecture updates the latent state through the following recurrence:\n\n$x_t = \\tilde{A} x_{t-1} + \\tilde{B} u_t$\n\nThe hidden embedding $h_t$ in the EDELINE architecture is obtained by applying matrix $C$ to the mamba latent state $x_t$:\n\n$h_t = C x_t$\n\nThis architecture enables Mamba to store **unbounded contextual information** within the latent state and compute the current timestep's hidden embedding $h_t$ through the selective mechanism. As a result, $x_{t-1}$ contains the information from $h_{t-1}$ ​since $h_{t-1} = C x_{t-1}$, establishing the complete derivation chain from $h_{t-1}$​ to $h_t$​ through the latent state recurrence. To enhance the clarity of Figure 2, we will include the Mamba latent states $x_{t-1}$ and $x_t$​ in the revised version of the diagram."}, {"review_id": "hIgP42rZfN", "reviewer": "Reviewer_Jqhz", "summary": "This paper introduces EDELINE for enhancing the memory capacity of the diffusion-based world model, which is limited by the fixed context length. Specifically, EDELINE integrates the state space model with diffusion models, presenting a unified architecture of the world model. Experimental results show that EDELINE can achieve superior performance across various challenging benchmarks.", "questions": "### Question\n\n1. What's the inference time of EDELINE compared to other baselines?\n2. Could the authors provide additional experiments on challenging visual RL? (e.g., MineCraft, CS or dm_control)?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths\n1. The writing is clear. The research problem is well-illustrated, and the proposed method intuitively makes sense.\n2. The performance of the proposed method is superior to existing SOTAs.\n\n### Weakness\n1. To further investigate the proposed world model architecture, it would be great to have the inference and training speed comparison between different world models.\n2. More challenging visual RL benchmarks used in previous works can be considered (e.g., Minecraft and FPS games).", "was_revised": true, "final_justification": "I have read the response and other reviews, and my concerns have been addressed. \nThis paper is well-written with convincing motivation and comprehensive experiments.\nThus, I tend to keep my positive score (Boardline Accept).", "rebuttal": "We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.\n\n> Q1. What's the training time and inference time of EDELINE compared to other baselines?\n\nWe thank the reviewer for this important question. We provide the total training time for DreamerV3 [1], IRIS [2], DIAMOND [3], and EDELINE on a single RTX 4090 GPU for Atari 100k benchmark training. This time encompasses both world model training and policy training using the world model and data collection.\n\nIn addition, we provide the single-step imagination time (i.e., world model inference) for these baselines. While EDELINE's inference time is slightly higher than DIAMOND's, the adoption of Mamba's parallel scan algorithm and the random target timestep selection technique during world model training makes EDELINE more efficient in world model training compared to DIAMOND. This enables EDELINE to maintain computational efficiency comparable to DIAMOND while addressing DIAMOND's memory capacity limitations. The detailed comparison is presented in Table R2.\n\n||DreamerV3 [1]|IRIS [2]|DIAMOND [3]|EDELINE (Ours)|\n|-|-|-|-|-|\n|Training days|**< 1**|4.1|2.9|2.9|\n|Inference Time (ms)|**13.9**|89.6|20.4|24.4|\n|Atari 100k HNS|1.124|1.046|1.459|**1.866**|\n\n**Table R2: Training and Inference Time Comparison.** Comparison of training efficiency, inference speed, and performance across different world model architectures on Atari 100k benchmark, including Atari 100k Human Normalized Score (HNS).\n\n[1] Hafner, Danijar, et al. \"Mastering diverse domains through world models.\" Nature (2025).\n\n[2] Micheli, Vincent, Eloi Alonso, and François Fleuret. \"Transformers are sample-efficient world models.\" ICLR (2023).\n\n[3] Alonso, Eloi, et al. \"Diffusion for world modeling: Visual details matter in atari.\" NeurIPS (2024).\n\n> Q2. Could the authors provide additional experiments on challenging visual RL?\n\nWe thank the reviewer for this suggestion. We conducted additional evaluations on the Battle scenario in ViZDoom, a 3D FPS game scenario proposed in [1]. This scenario features a 3D enclosed maze with multiple ammunition pickups, health packs, and continuously moving monsters that attack the player. The evaluation metric is the number of monsters killed by the player. We employed a sparse reward function where the agent receives a reward of 1 only upon successfully killing a monster, with the agent limited to 1M environment interactions. This setting simultaneously challenges the agent's capabilities in handling complex visual inputs, sparse rewards, 3D FPS gameplay, long-term memory, and sample efficiency. As demonstrated in Table R3 below, EDELINE significantly outperforms DIAMOND [2] due to its memory advantages and unified architecture.\n\n|Method|Avg Return|\n|-|-|\n|EDELINE (Ours)|**7.270**|\n|DIAMOND [2]|2.395|\n\n**Table R3: ViZDoom Battle Scenario Results.** Performance comparison on a challenging 3D FPS environment requiring long-term memory and sparse reward handling.\n\n[1] Dosovitskiy, Alexey, and Vladlen Koltun. \"Learning to act by predicting the future.\" ICLR (2017).\n\n[2] Alonso, Eloi, et al. \"Diffusion for world modeling: Visual details matter in atari.\" NeurIPS (2024)."}, {"review_id": "kVDyZLnKsA", "reviewer": "Reviewer_t2dk", "summary": "This paper proposes EDELINE, a diffusion-based world model that incorporates long-term memory through State Space Models (SSMs). The authors address a key limitation of existing diffusion-based world models—their lack of long-term memory—by integrating SSM architectures to enable better temporal modeling over extended sequences.\n\nThe key contributions include: (1) a novel architecture that combines diffusion models with SSMs for world modeling, (2) a memory module trained on reward and termination signal prediction to enhance decision-making capabilities, and (3) a dynamic loss harmonization technique to balance multiple training objectives. The authors evaluate EDELINE on standard benchmarks including Atari 100k, Crafter, and selected ViZDoom tasks, demonstrating superior performance compared to baseline methods.", "questions": "- Line 18, \"efficiencyy\" looks as a typo.\n- In Figure 3, the labels are too small, can you increase this?\n- In line 328, it looks like there is a typo in \"It architecture demonstrates ....\".", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "## Strengths\n\n**Well-Motivated Research Direction.** The paper addresses a clear limitation of existing diffusion-based world models, particularly DIAMOND's inability to effectively handle long-term dependencies. Integrating memory mechanisms through SSMs represents a natural and well-justified advancement in this research area.\n\n**Clear Problem Presentation.** Section 4 effectively demonstrates the limitations of prior work, providing concrete evidence of memory-related failures that motivate the proposed approach. This structure helps readers understand the specific problem being addressed and appreciate the necessity of the solution.\n\n**Consistent Performance Improvements.** EDELINE demonstrates substantial and consistent performance gains over baseline methods across multiple diverse benchmarks (Atari 100k, Crafter, ViZDoom). The results provide strong empirical support for the authors' claims about the benefits of incorporating long-term memory.\n\n**Comprehensive Experimental Analysis.** The paper provides thorough experimental validation beyond the main results, including qualitative visualizations on Atari 100k, linear probing analysis to understand learned representations, architectural ablation studies, and computational efficiency analysis. This comprehensive evaluation strengthens the paper's contributions and provides valuable insights for future work.\n\n## Weaknesses\n\n**Limited Architectural Innovation and Potential Redundancy.** While EDELINE achieves strong empirical results, the integration of SSMs with diffusion models appears relatively straightforward and potentially suboptimal. The architecture maintains both SSM hidden states (containing full historical context) and explicit frame stacking (last L frames) as inputs to the diffusion model, creating apparent redundancy. While this dual approach might provide benefits—such as giving additional emphasis to recent observations that are crucial for next-frame prediction—the authors provide no empirical justification for this design choice.", "was_revised": true, "final_justification": "The authors properly addressed my raised concerns, especially for the architectural innovation with their newly introduced ablation study. Hence, I maintain my positive score.", "rebuttal": "We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.\n\n> Q1. Limited Architectural Innovation and Potential Redundancy. \n\nWe appreciate the reviewer's observation. We would like to bring to the reviewer's attention that this represents significant architectural innovation, as evidenced by the superior performance compared to prior works in this domain. Our design choice to maintain both components stems from the complementary roles they serve in the diffusion process. The hidden embedding captures long-term temporal dependencies and global context across the entire sequence history, while the explicit frame stacking provides pixel-based spatial details from recent observations that are crucial for pixel-level prediction accuracy. The diffusion model benefits from both the abstract temporal representations (via cross-attention with hidden embedding) and direct access to recent visual features (via frame stacking) for generating high-fidelity predictions.\n\nFor empirical validation to address this concern, we conducted an ablation study comparing: (1) EDELINE with both components, (2) Hidden embedding only (without frame stacking), and (3) DIAMOND. As shown in Table R4, the results demonstrate that in Atari Breakout, which requires minimal long-term memory while demanding high visual detail processing, EDELINE without frame stacking performs slightly higher than DIAMOND, though still significantly worse than full EDELINE. In MiniGrid-MemoryS9, which requires substantial memory capabilities while having simpler visual complexity, EDELINE without frame stacking can still successfully learn to solve the task and outperforms DIAMOND. Note that our analysis reveals that removing either component leads to degradation in either temporal consistency or visual fidelity.\n\nWe appreciate the reviewer's suggestion and consider new architectural design or more elegant integration mechanisms as interesting avenues for future research directions. However, we would like to respectfully emphasize that our current framework achieves the primary objective of combining long-term memory with high-quality visual prediction while maintaining computational efficiency, which is the main theme of this manuscript. The experimental results faithfully validate our contributions, and we have experimental evidence demonstrating that these components are complementary in achieving superior performance.\n\n||DIAMOND|EDELINE w/o frame-stack|EDELINE|\n|-|-|-|-|\n|MemoryS9|0.380|**0.985**|**0.982**|\n|Breakout|132.5|145.6|**250.5**|\n\n**Table R4: Frame Stacking Ablation Study.** Performance comparison demonstrating the necessity of both SSM hidden states and explicit frame stacking across environments with different memory and visual complexity requirements.\n\n> Q2. Typo in Line 18 and Line 328.\n\nWe thank the reviewer for catching these typographical errors. We will correct these issues in the revised manuscript.\n\n> Q3. In Figure 3, the labels are too small.\n\nWe appreciate the reviewer's feedback regarding figure readability. We will increase the label sizes in Figure 3 to enhance clarity in the revised version."}]

Author rebuttal:

Meta-review: This paper proposes a model that combines State Space Models (SSMs) with diffusion models for model-based RL.
The reviewers unanimously agreed that the paper addresses an important and well-motivated problem, a limited memory in diffusion-based world models, and that the empirical results are strong and consistent across a diverse set of challenging benchmarks.

During the discussion period, reviewers raised several questions regarding architectural choices, novelty in comparison to other diffusion works, and requests for additional computational analysis and more challenging benchmarks. The authors addressed the concerns with new experiments during the rebuttal period. 

Given the importance of the problem and the comprehensive empirical evidence, I believe that this paper makes a solid contribution to the field. Thus, I recommend acceptance.
Decision: Accept (spotlight)