Reinforcement Learning

Next Embedding Prediction Makes World Models Stronger

GGeorge BredisNNikita BalaganskyDDaniil GavrilovRRuslan Rakhimov
Published
March 3, 2026
Authors
4
Word Count
4,743

Next-embedding prediction enables stronger world models without pixel reconstruction overhead.

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.

Key Takeaways

  • 1

    NE-Dreamer predicts next-step embeddings instead of reconstructing pixels, improving temporal coherence in world models.

  • 2

    The method outperforms DreamerV3 on memory-intensive navigation tasks without pixel reconstruction or data augmentation.

  • 3

    Decoder-free approaches require explicit temporal constraints to maintain predictive representations under partial observability.

Limitations

  • Strong performance shown primarily on DMLab memory/navigation tasks; generalization to other domains unclear.

  • Approach evaluated against limited baselines; comparison with other recent decoder-free methods incomplete.

Keywords

temporal transformernext-step encoder embeddingstemporal predictive alignmentrepresentation spacemodel-based reinforcement learningDeepMind Control SuiteDMLabcoherent state representationspredictive state representations

More in Reinforcement Learning

View all
Next Embedding Prediction Makes World Models Stronger | Paperchime