A central challenge in image-based Model-Based Reinforcement Learning (MBRL) is to learn representations that distill essential information from irrelevant visual details. While promising, reconstruction-based methods often waste capacity on large task-irrelevant regions. Decoder-free methods instead learn robust representations by leveraging Data Augmentation (DA), but reliance on such external regularizers limits versatility. We propose R2-Dreamer, a decoder-free MBRL framework with a self-supervised objective that serves as an internal regularizer, preventing representation collapse without resorting to DA. The core of our method is a redundancy-reduction objective inspired by Barlow Twins, which can be easily integrated into existing frameworks. On DeepMind Control Suite and Meta-World, R2-Dreamer is competitive with strong baselines such as DreamerV3 and TD-MPC2 while training 1.59x faster than DreamerV3, and yields substantial gains on DMC-Subtle with tiny task-relevant objects. These results suggest that an effective internal regularizer can enable versatile, high-performance decoder-free MBRL. Code is available at https://github.com/NM512/r2dreamer.
Figures & tables
Figure 1: Comparison of representation learning mechanisms in world models. (a) R2-Dreamer learns representations without a decoder or DA. It uses an internal redundancy reduction objective LBT that aligns the latent state st (via a projector) with the embedding of the observation ot . (b) Dreamer relies on a decoder to learn representations by reconstructing the observation o^t from the latent state st , guided by a reconstruction loss Lrecon . (c) DreamerPro removes the decoder but depends on DA. It enforces consistency between augmented views of the observation aug(ot) using a spatial loss LSwAV and a temporal loss LTemp that leverages an Exponential Moving Average (EMA) of the encoder weights.
Figure 2: Examples drawn from benchmark tasks. Left: Meta-World Assemble. Center: DMC Reacher (hard). Right: DMC-Subtle Reacher with a significantly smaller target.
Figure 3: Mean and median performance over 20 DMC tasks, with standard deviation across seeds. R2-Dreamer is competitive with the baselines on average without requiring a decoder or DA.
Figure 4: Aggregated performance on Meta-World 50 tasks using the mean and the median across tasks. R2-Dreamer achieves strong results even on contact-rich manipulation tasks, remaining competitive with the baselines on average.
Figure 5: Performance on five challenging DMC-Subtle tasks. R2-Dreamer substantially outperforms the baselines, demonstrating its robustness to subtle but critical visual information.
Figure 6: Policy saliency maps on the DMC-Subtle Reacher task. For clarity, the target location is marked with a yellow dot. The two rows show results from two different environment seeds, which are identical across all methods.
Figure 7: Ablation results on 20 DMC tasks using the mean and the median across tasks. Our internal redundancy reduction objective proves more effective and robust than reliance on heuristic DA.
Figure 8: Comparison of R2-Dreamer with and without DA on the DMC-Subtle benchmark. The results highlight that DA can be detrimental in tasks requiring high precision, as it may distort subtle but critical information, underscoring the importance of a DA-free approach for such environments.
Method
Training Time (hours)
R2-Dreamer
4.4
Dreamer
7.0
DreamerPro
10.4
Dreamer (Author’s JAX impl.)
6.6
Table 1: Computational efficiency comparison on the DMC Walker Walk task. Wall-clock time is measured for 1 million environment steps on a single NVIDIA GeForce RTX 3080 Ti GPU.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: DMC-Subtle benchmark. Top: original versions of the five tasks (left to right: Ball in Cup Catch, Cartpole Swingup, Finger Turn, Point Mass, Reacher). Bottom: modified versions with downscaled task-critical objects in the same order.
Figure 10: Per-task learning curves for all 20 DMC tasks.
Figure 11: Per-task learning curves for all 50 Meta-World tasks.
Figure 12: Per-task learning curves for the ablation study across all 20 DMC tasks.
Name
Symbol
Value
General
Replay capacity
—
5×106
Batch size
B
16
Batch length
T
64
Activation
—
RMSNorm+SiLU
Learning rate
—
4×10−5
Appendix
Table 2: Main hyperparameters. Our settings are identical to DreamerV3, with key changes to the representation learning loss.
Model-based reinforcement learning (MBRL) agents operating in high-dimensional observation spaces, such as Dreamer, rely on learning abstract representations for effective planning and control. Existing approaches typically employ reconstruction-based objectives in the observation space, which can render representations sensitive to task-irrelevant details. Recent alternatives trade reconstruction for auxiliary action prediction heads or view augmentation strategies, but perform worse in the Crafter environment than reconstruction-based methods. We close this gap between Dreamer and reconstruction-free models by introducing a JEPA-style predictor defined on continuous, deterministic representations. Our method matches Dreamer's performance on Crafter, demonstrating effective world model learning on this benchmark without reconstruction objectives.
Michael Hauri, Friedemann Zenke
Friedrich Miescher Instiute for Biomedical Research, Fabrikstrasse 24, 4056 Basel, Switzerland
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.
George Bredis, Nikita Balagansky, Daniil Gavrilov +1
Model-based reinforcement learning (MBRL) achieves strong sample efficiency by planning within learned latent dynamics, yet its performance degrades substantially under unseen visual distractions such as background variations, lighting changes, or camera shifts. Unlike model-free RL, where encoder perturbations affect only single-step predictions, MBRL suffers from a two-level vulnerability: visual distractions first push encoder outputs out of distribution, and these errors then compound through recursive latent rollouts over the planning horizon. We propose visual generalization via latent-space consistency in model-based RL (VIGOR), a framework that enables zero-shot generalization to unseen visual distractions while retaining the sample efficiency of its MBRL backbone. VIGOR integrates three interdependent components: (i) asymmetric weak-to-strong augmentation, which pairs weak-only and weak-to-strong latent views within a single batch; (ii) dynamics-level consistency, which enforces augmentation-invariant transition predictions through direct latent regression; and (iii) encoder-level stabilization, which prevents encoder drift under the cross-augmentation supervision imposed by dynamics-level consistency. Evaluations on the DeepMind Control Suite (DMC) and Robosuite show that VIGOR outperforms state-of-the-art model-free and model-based baselines, surpassing the second-best baseline by 3.4% on DMC and 43.6% on Robosuite. Ablations further show that VIGOR's robustness is augmentation-agnostic: replacing the default augmentation with alternatives from distinct perturbation families preserves strong generalization, confirming that latent-space consistency, not the augmentation choice, drives robustness.
Mingyu Park, Samyeul Noh, Hyun Myung +1
Robotics Program KAIST, ETRI · Electrical Engineering KAIST