Source-linked AI summary
Next Embedding Prediction Makes World Models Stronger
George Bredis, Nikita Balagansky, Daniil Gavrilov, Ruslan Rakhimov
TL;DR
Model-based reinforcement learning needs latent states that remain predictive over time under partial observability, while pixel reconstruction can add burden and task-irrelevant visual capacity. NE-Dreamer uses a causal temporal transformer to predict and align next encoder embeddings within a decoder-free Dreamer-style agent. It improves long-horizon memory and navigation in DMLab Rooms, matches strong baselines on DMC, and attributes gains to predictive sequence modeling.
Problem
Partially observable, high-dimensional MBRL requires compact latent states that integrate information over time, while pixel reconstruction can add modeling burden and emphasize task-irrelevant visual detail.
Method
NE-Dreamer replaces pixel reconstruction with causal-transformer prediction and alignment of the next encoder embedding from latent state sequences.
Results
NE-Dreamer improves long-horizon memory and navigation in DMLab Rooms while matching strong baselines on the DeepMind Control Suite.
Takeaways & Limitations
Next-embedding prediction with a causal transformer is presented as a practical, scalable foundation for representation learning in model-based reinforcement learning.
Takeaways & Limitations
The experiments focus on environments where long-term structure, rather than fine visual detail, is the primary challenge; performance on high-fidelity tasks remains open.
Abstract
from arXiv · showhide
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.
1. Introduction
NE-Dreamer addresses the need for temporally predictive latent states in partially observable, high-dimensional control by replacing pixel reconstruction with next-embedding prediction. Its predictive sequence modeling improves memory- and navigation-heavy performance while retaining a Dreamer-style pipeline.
- Partial observability requires latent states to integrate information over time for long-horizon prediction and control.
- Pixel reconstruction provides rich features but adds modeling burden and can spend capacity on visually detailed, task-irrelevant factors.
- Same-timestep agreement is insufficient for partial observability because representations must predict future events across time.
- NE-Dreamer predicts each next encoder embedding with a causal temporal transformer and aligns it using a redundancy-reduction objective.
- NE-Dreamer achieves substantially higher performance than same-size prior methods on partially observable DMLab environments under matched compute and model capacity.The benchmark uses 50M environment steps, 5 seeds, and 12M parameters.
- Ablations attribute the gains to predictive sequence modeling through the causal transformer and next-step target shift rather than reconstruction or auxiliary tricks.
2. Related Work
Prior work uses latent dynamics, reconstruction, or task-oriented decoder-free objectives, but same-step representation targets do not explicitly enforce temporal predictiveness. NE-Dreamer extends representation prediction to future embeddings by aligning a causal forecast with a stop-gradient target.
- Latent world models learn compact states for long-horizon prediction and decision-making from high-dimensional observations.
- Dreamer-style agents use RSSM-based latent imagination for actor–critic learning, while NE-Dreamer changes how the latent representation is learned.
- Reconstruction-based methods provide dense supervision but may allocate capacity to visual details weakly coupled to reward.
- Decoder-free methods remove pixel reconstruction and instead choose latent anchors and learning-signal time indices, including task-oriented reward or value supervision.
- NE-Dreamer makes temporal prediction explicit by forecasting the next encoder embedding from history and aligning it to a stop-gradient target.
- Existing collapse-prevention methods often apply redundancy reduction or bootstrapping to paired views at the same timestep.
- NE-Dreamer extends redundancy reduction to future prediction, using a forecasted embedding aligned with a stop-gradient target to enforce temporal coherence.
3. Method
NE-Dreamer replaces pixel reconstruction with causal next-embedding prediction, while retaining an RSSM-based latent world model and Dreamer-style imagined actor–critic learning.
- Next-embedding predictive alignment: NE-Dreamer predicts the next encoder embedding from history using a causal temporal transformer and aligns it with a stop-gradient target.The objective uses Barlow Twins redundancy reduction between predicted and target embeddings.
- Latent world model: The RSSM uses deterministic state h_t and stochastic latent z_t, with posterior sampling during training and prior sampling during imagination.The encoder maps observations to embeddings before recurrent latent inference.
- Latent world model: The latent world model combines reward and continuation likelihoods, prior–posterior regularization, and the next-embedding loss.This objective replaces the pixel decoder used by standard Dreamer.
- Next-embedding predictive alignment: The Barlow Twins objective encourages large diagonal correlations and small off-diagonal correlations for next-step predictions.Embeddings are normalized per dimension over valid minibatch transitions before computing cross-correlations.
- Actor-Critic Learning: NE-Dreamer trains its actor and critic on imagined latent trajectories, using horizon H = 15 steps for batch actor–critic updates.The imagined state is s_t = (h_t, ẑ_t), with actions sampled from the policy and values estimated by the critic.
- Actor-Critic Learning: The learning procedure and hyperparameters match DreamerV3, so the reported gains are attributed to the representation-learning objective.Policy gradients are backpropagated through the world model for continuous actions.
4. Experiments
The experiments evaluate NE-Dreamer across DMLab Rooms and DMC against decoder-based, decoder-free, and model-free baselines, then isolate the contributions of temporal prediction. NE-Dreamer improves long-horizon DMLab performance while matching DreamerV3 and competitive decoder-free methods on DMC under matched evaluation conditions.
- DMLab Rooms: NE-Dreamer improves memory- and navigation-heavy DMLab Rooms performance across four tasks, learning reliably and achieving substantially higher final returns than all baseline methods.The benchmark targets partial observability, long-horizon credit assignment, memory, and spatial-layout reasoning.
- Mechanism ablations: Removing the causal temporal transformer causes performance to collapse on all DMLab Rooms tasks, indicating that causal sequence modeling is indispensable in this partially observable setting.The ablation removes the transformer while keeping the rest of the pipeline unchanged.
- Mechanism ablations: Removing the next-step target shift nearly eliminates the full method’s gains, showing that explicit temporal prediction is needed beyond current-step matching or reconstruction.The ablation retains the temporal transformer but trains against the current-step embedding instead of predicting the next-step target.
- Mechanism ablations: Removing the lightweight projector has only a minor effect on asymptotic performance, although it can affect optimization speed and stability.The projector may smooth alignment or improve conditioning, but it is not fundamentally responsible for the observed gains.
- Benchmark comparison: Under the unified protocol, NE-Dreamer matches DreamerV3 and competitive decoder-free baselines on the DeepMind Control Suite.World-model methods use matched capacity, and all agents follow the same training protocol across five random seeds.
- Representation diagnostics: Post-hoc reconstructions show that NE-Dreamer preserves task-relevant objects, spatial layout, and features over time more consistently than Dreamer and R2-Dreamer.The comparison uses a decoder trained after the fact on frozen latents; same-timestep methods show transient attributes that fade across subsequent latents.
5. Discussion
NE-Dreamer’s temporal prediction mechanism supports gains in memory-intensive environments while preserving standard-control performance, though its evaluation scope is concentrated on long-term structure rather than fine visual detail.
- Mechanism: NE-Dreamer predicts next encoder embeddings from history using a causal temporal transformer and aligns them with stop-gradient targets.Barlow Twins is used for alignment, though other non-degenerate alignment losses could substitute.
- Mechanism: A causal temporal transformer enables multi-step prediction of long-horizon dependencies without additional rollout cost.The sequence model compresses history into latent features predictive of future states.
- Results: NE-Dreamer substantially outperforms decoder-free and decoder-based baselines on memory- and planning-intensive DMLab Rooms tasks at equal model size and compute.On standard DMC benchmarks, it matches prior methods.
- Limitations: The experiments focus on environments where long-term structure, rather than fine visual detail, is the primary challenge.Whether prediction-based objectives match reconstruction on high-fidelity tasks remains open.
6. Conclusion
The paper presents NE-Dreamer, a decoder-free Dreamer-style agent that predicts and aligns next encoder embeddings with a causal temporal transformer.
- NE-Dreamer learns world-model representations by predicting and aligning the next encoder embedding using a causal temporal transformer.
A. Technical details
The study’s hyperparameters are primarily inherited from DreamerV3, with minimal modifications for the proposed representation-learning objective.
- The primary hyperparameters are based on DreamerV3, with minimal modifications related to the proposed representation-learning objective.
B. DMC detailed results
The DMC detailed-results section provides per-task learning curves across all 20 tasks and reports the main hyperparameter settings used for comparison.
- Per-task learning curves are provided for all 20 tasks in the DMC benchmark.
- The main hyperparameters are identical to DreamerV3 unless otherwise noted.