Title: EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling
Venue: NeurIPS 2025
Accepted: yes
Keywords: Model-based Reinforcement Learning

Atari 100k

Doom

Crafter

MAMBA

Diffusion

World Model
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\n- 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 ....\".\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-13"}

comments: 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.

STRENGTHS AND WEAKNESSES:
### Strengths
1. The writing is clear. The research problem is well-illustrated, and the proposed method intuitively makes sense.
2. The performance of the proposed method is superior to existing SOTAs.

### Weakness
1. To further investigate the proposed world model architecture, it would be great to have the inference and training speed comparison between different world models.
2. More challenging visual RL benchmarks used in previous works can be considered (e.g., Minecraft and FPS games).

QUESTIONS:
### Question

1. What's the inference time of EDELINE compared to other baselines?
2. Could the authors provide additional experiments on challenging visual RL? (e.g., MineCraft, CS or dm_control)?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-06-26

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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).\n\nQUESTIONS:\n1.\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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-05"}