arXiv · 2603.02765
Next Embedding Prediction Makes World Models Stronger
Abstract
Capturing temporal dependencies is critical for model-based reinforcement learning (MBRL) in partially observable, high-dimensional domains. We introduce NE-Dreamer, a decoder-free MBRL agent that leverages a temporal transformer to predict next-step encoder embeddings from latent state sequences, directly optimizing temporal predictive alignment in representation space. This approach enables NE-Dreamer to learn coherent, predictive state representations without reconstruction losses or auxiliary supervision. On the DeepMind Control Suite, NE-Dreamer matches or exceeds the performance of DreamerV3 and leading decoder-free agents. On a challenging subset of DMLab tasks involving memory and spatial reasoning, NE-Dreamer achieves substantial gains. These results establish next-embedding prediction with temporal transformers as an effective, scalable framework for MBRL in complex, partially observable environments.
Explore related subjects
Keep this discovery
George Bredis, Nikita Balagansky, Daniil Gavrilov, Ruslan Rakhimov. 2026-03-03. Next Embedding Prediction Makes World Models Stronger. https://arxiv.org/abs/2603.02765
Cite the original work for its findings. Save a collection to share your selection of sources.