EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling

TL;DR

EDELINE introduces a novel world model architecture combining state space and diffusion models, significantly improving memory capacity and performance in reinforcement learning environments with complex visual and memory demands.

Contribution

This paper presents EDELINE, a unified model that enhances memory in diffusion-based world models by integrating state space models, surpassing existing methods in diverse challenging environments.

Findings

Outperforms baselines on Atari 100k tasks

Achieves superior results on Crafter benchmark

Excels in 3D ViZDoom environments

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.