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.
Figures & tables
Figure 1: DMLab Benchmark Summary. Under matched compute and model capacity (50M environment steps; 5 seeds; 12M parameters), NE-Dreamer outperforms strong decoder-based (DreamerV3) and decoder-free world-model baselines (R2-Dreamer, DreamerPro) on the DMLab Rooms memory/navigation tasks.
Figure 2: Method overview. NE-Dreamer keeps Dreamer’s RSSM dynamics and imagination-based actor–critic, but replaces same-step pixel reconstruction with next-embedding prediction using a causal temporal transformer, improving long-horizon performance under partial observability.
Figure 3: DMLab Rooms: improved long-horizon memory/navigation. Under matched compute and model capacity ( 50 M environment steps; 5 seeds; 12M parameters), NE-Dreamer outperforms strong decoder-based (DreamerV3) and decoder-free world-model baselines (R2-Dreamer, DreamerPro) on four Rooms tasks. The largest gains occur when success depends on maintaining state over long horizons rather than reacting to short-lived visual cues.
Figure 4: Mechanism on DMLab Rooms: predictive sequence modeling is the key. Under matched compute and model capacity ( 50 M environment steps; 5 seeds; mean ± std), removing the causal temporal transformer ( w/o transformer ) or removing the next-step target shift ( w/o shift ) substantially reduces performance. Removing the lightweight projector ( w/o projector ) mainly affects optimization speed/stability, with smaller impact on final returns.
Figure 5: Post-hoc decoder reconstruction reveals temporal consistency. Rows show ground-truth observations (GT) and reconstructions from a post-hoc decoder trained on frozen latents. NE-Dreamer preserves task-relevant objects and spatial layout consistently over time (marked green circles), while same-timestep methods (Dreamer, R2-Dreamer) exhibit temporal inconsistency, where task-specific attributes appear transiently and then fade (marked red circles).
Figure 6: DMC: removing reconstruction does not hurt standard control. On near-saturated pixel-based continuous-control benchmarks, NE-Dreamer matches or slightly exceeds strong decoder-based (DreamerV3) and decoder-free world-model baselines (R2-Dreamer, DreamerPro) under a unified protocol (1M environment steps; 5 seeds; 12M parameters). Per-task learning curves can be found in Appendix B
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
Symbol
Setting
General
Replay Buffer Capacity
—
5×106
Batch Size
B
16
Batch Length
T
64
Optimizer
—
Adam
Activation
—
RMSNorm + SiLU
Appendix
Table 1: Main hyperparameters. Our settings are identical to DreamerV3 unless otherwise noted. All method-based hyperparameters identical to original implementations too.
Figure 7: Per-task learning curves for all 20 DMC tasks