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.

• The paper introduces next embedding prediction using a causal transformer to learn temporally coherent latent representations without pixel reconstruction.
• It demonstrates significant improvements on memory-demanding tasks, notably in DMLab Rooms, by reducing computational overhead and ensuring long-horizon planning.
• Experimental ablations confirm that predictive sequence modeling is critical for robust performance in partially observable reinforcement learning environments.

NE-Dreamer: Temporal Predictive Alignment for Decoder-Free World Models

Motivation and Background

Model-based reinforcement learning (MBRL) is highly dependent on constructing latent state representations that enable long-horizon prediction and control, especially in high-dimensional, partially observable environments. Traditional pixel reconstruction objectives, exemplified by Dreamer variants, facilitate representation learning but introduce significant generative complexity and incentivize modeling task-irrelevant factors. Decoder-free methods alleviate computational overhead but typically enforce only instantaneous agreement, which is insufficient for environments where success requires temporal integration. NE-Dreamer reframes latent representation learning by directly optimizing temporal predictive alignment in embedding space—predicting the next encoder embedding from history via a causal temporal transformer and aligning it with the target using a redundancy reduction metric (Barlow Twins).

Figure 1: The NE-Dreamer pipeline augments Dreamer's RSSM with a causal temporal transformer for next-embedding prediction, substituting reconstruction losses with temporal predictive alignment.

Methodology

NE-Dreamer retains the recurrent state-space model (RSSM) as its latent backbone, with deterministic and stochastic states. The model learns to predict the next-step encoder embedding $\hat{e}_{t+1}$ using a temporal transformer $T_\theta$, conditioned on history up to time $t$ ($h_{\leq t}, z_{\leq t}, a_{\leq t}$). The predicted embedding is aligned to a stop-gradient target $e^\star_{t+1}$ (computed from the next observation $x_{t+1}$) using a Barlow Twins loss, enforcing invariance (high diagonal correlation) and redundancy reduction (low off-diagonal correlation) in the cross-correlation matrix between $\hat{e}_{t+1}$ and $e^\star_{t+1}$.

This objective, termed next-embedding prediction, establishes causality and prevents representational collapse commonly encountered in bootstrapped or self-supervised predictive frameworks. The actor-critic learning pipeline remains identical to DreamerV3 (policy/value learning in latent space by imagined rollouts), ensuring observed effects stem from the representation learning mechanism rather than augmentations in the policy or dynamics model.

Experimental Evaluation

The experiments are structured around three core claims:

• NE-Dreamer delivers substantial improvements in memory and navigation on DMLab Rooms tasks, outperforming both decoder-based and decoder-free world-model baselines under matched compute/model size.
• Gains are attributed to predictive sequence modeling—removal of the temporal transformer or next-step target shift sharply degrades performance, while omitting the lightweight projection head impacts only optimization speed/stability.
• NE-Dreamer matches or slightly exceeds baseline performance on near-saturated DeepMind Control Suite (DMC) tasks, demonstrating no regression when pixel reconstruction is omitted.

  Figure 2: Summary of DMLab benchmark—NE-Dreamer consistently outperforms DreamerV3 and strong decoder-free baselines under identical conditions.

  

  Figure 3: NE-Dreamer achieves the highest scores on all DMLab Rooms tasks, especially where temporal memory and navigation dominate.

  

  Figure 4: Ablation results: removing the temporal transformer or next-step target causes significant performance drops, isolating predictive sequence modeling as the key mechanism.

NE-Dreamer’s superiority is particularly pronounced in environments requiring the retention and recall of spatial information over extended time horizons, with object and layout consistency persisting over time (Figure 5).

Figure 5: Post-hoc decoder reconstructions from frozen latents—NE-Dreamer maintains temporal consistency, unlike alternatives which lose task-relevant information.

Performance parity in DMC with DreamerV3 and decoder-free baselines (Figure 6) confirms that temporal predictive alignment does not negatively affect traditional continuous control settings.

Figure 6: DMC results—removal of reconstruction maintains or marginally improves returns across standard continuous-control benchmarks.

Representation Diagnostics

A notable diagnostic involves training a post-hoc pixel decoder on frozen NE-Dreamer latents. Reconstructions demonstrate that NE-Dreamer’s representations preserve decision-critical features and spatial layout consistently across trajectories, validating that predictive alignment enforces temporal stability and prevents latent drift toward transient or visually detailed spurious features. By contrast, same-timestep objectives enable representational collapse with task-specific attributes appearing intermittently or vanishing across time.

Implications and Future Directions

The findings establish next-embedding prediction with a causal transformer as a scalable, domain-agnostic foundation for robust representation learning in MBRL, particularly suited for demanding partially observable settings. The method offers a principled path away from generative reconstruction, mitigating the computational and optimization burdens, and focuses on persistent, decision-relevant representations. The role of predictive sequence modeling is empirically and ablation-wise isolated as the principal contributor to long-horizon performance gains.

Current experiments emphasize tasks where the core challenge is temporal structure rather than fine visual fidelity. An open question is whether decoder-free objectives—centered on next-embedding prediction—can match or surpass pixel-level reconstruction in high-fidelity, visually complex domains. Exploration of alternative alignment losses and broader validation across diverse tasks will further elucidate the theoretical and practical boundaries of this approach. The results potentially inform future architectures for world models, encouraging designs with explicit temporal causality and predictive alignment.

Conclusion

NE-Dreamer pioneers a decoder-free world model that leverages causal temporal transformers for next-embedding prediction, achieving state-of-the-art performance in memory-intensive navigation benchmarks and demonstrating robust performance in standard continuous-control tasks. The improvements are attributed to predictive sequence modeling, as validated by ablations and representation diagnostics. This framework advances scalable, efficient, and temporally coherent latent representation learning for model-based reinforcement learning and lays the groundwork for future research targeting environments with complex temporal dependencies and partial observability.

Figure 7: NE-Dreamer's per-task learning curves for all DMC tasks—no tradeoff between performance in hard and standard RL domains.