Next Embedding Prediction Makes World Models Stronger

2026-03-03

AI Summary

This work addresses the challenge of modeling temporal dependencies in partially observable, high-dimensional environments, where existing model-based reinforcement learning approaches often struggle. The authors propose NE-Dreamer, a decoder-free model-based agent that leverages a temporal Transformer to directly predict the next-step encoder embedding in latent space. By eliminating reconstruction losses and auxiliary supervision, NE-Dreamer learns coherent and predictive state representations solely through temporal prediction alignment over embedding sequences. The method matches or exceeds the performance of DreamerV3 on the DeepMind Control Suite and demonstrates significant improvements on challenging DMLab tasks that require memory and spatial reasoning.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision
• Multiagent Systems: Multiagent Learning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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.

Problem

• temporal dependencies
• model-based reinforcement learning
• partially observable environments
• state representation
• high-dimensional domains

Innovation

• next-embedding prediction
• temporal transformer
• decoder-free MBRL
• predictive state representations
• model-based reinforcement learning