PaperlayerView exact revision on arXiv

Source-linked AI summary

Contrastive World Models

Bonnie Li

arXiv:2609.22175v1cs.LGcs.AIcs.CV

TL;DR

Pixel-reconstruction objectives are common in world models, while the robustness of recent latent-space alternatives to complex visual nuisance factors remains unshown. Contrastive World Models replaces Dreamer’s reconstruction objective with a Deep InfoMax-like objective over future observation patches. It matches both baselines in the simple setting and outperforms them with distractors or natural video backgrounds.

  • Problem

    Reconstruction or pixel-prediction objectives remain common in world models, while recent latent-space alternatives have not been shown robust to complex visual nuisance factors.

    Source: §1 · p. 2

  • Method

    Contrastive World Models replaces Dreamer’s observation reconstruction with a Deep InfoMax-like objective maximizing mutual information between state-action sequences and local patches of future observations.

    Source: §4 · p. 7

  • Results

    The method matches reconstruction-based and momentum-based baselines in the simple setting and significantly outperforms them when simple distractors or natural video backgrounds are introduced.

    Source: §4 · p. 7

  • Takeaways & Limitations

    Removing the pixel decoder reduces per-step compute and yields substantially faster wall-clock training steps across all three tested settings.

    Source: §3.2 · p. 7

    Experiments are small-scale, use the older RSSM architecture, and are limited to the DeepMind Control Suite benchmark.

    Source: §4 · p. 7

Abstract

from arXiv · show

World models trained via pixel reconstruction can struggle in visually complex environments, where irrelevant information dominates the objective and distract the model from information relevant to planning and control. We present Contrastive World Models, an approach for learning latent dynamics models without pixel reconstruction. Building on Dreamer, we replace observation reconstruction in the standard world model objective with a Deep InfoMax-like lower bound that maximizes the mutual information between state-action sequences and local patch features of future observations, encouraging state representations to retain information that is predictive of the future without requiring the model to reconstruct visually irrelevant details. We evaluate our approach in small-scale experiments across three settings of increasing visual complexity. Our method matches Dreamer and a momentum prediction baseline in the default setting, and substantially outperforms both once distractors or natural video backgrounds are introduced, while also training more efficiently by removing the pixel decoder entirely. Our approach is general and makes minimal assumptions beyond access to state-action sequences and future observations. These results suggest that contrastive, infomax-based objectives are a principled and promising direction for building world models that are robust to visual nuisance factors, a property particularly relevant for transferring model-based RL agents to the real world.

1 Introduction

World models use learned dynamics to support planning and efficient action, but reconstruction-based objectives remain common. This work combines contrastive representation learning with Dreamer as an alternative aimed at visually complex environments.

1 Introduction

  • Dreamer and Genie advances use reconstruction or pixel-prediction objectives to learn latent state representations.The passage describes Dreamer’s RSSM-based progression and Genie’s action-controllable video generation.Source: §1 · p. 2
  • Contrastive objectives distinguish related views or timesteps from unrelated ones, preserving information that discriminates the correct future without reconstructing pixels.Examples include InfoNCE-based methods and BYOL’s momentum-target prediction objective.Source: §1 · p. 2
  • An earlier contrastive Dreamer variant underperformed pixel reconstruction on most tasks, while CSWM was evaluated only on latent-ranking metrics in simple deterministic domains without rewards or downstream control.CSWM models transitions over object slots with a graph neural network and a TransE-style energy-based hinge loss.Source: §1 · p. 2
  • LeWorldModel has been evaluated at small scale in low-distraction environments with stationary backgrounds, and its robustness to complex visual nuisance factors has not been shown.The passage describes LeWM as a roughly 15M-parameter JEPA-based model evaluated on four low-distraction environments.Source: §1 · p. 2
  • Contrastive World Models replaces Dreamer’s observation reconstruction with a Deep InfoMax-like mutual-information objective, removing the pixel decoder.The approach combines contrastive representation learning with a Dreamer world model.Source: §1 · p. 3

2 Contrastive World Models

The approach builds on Dreamer’s latent-dynamics model for POMDPs and replaces pixel reconstruction with an InfoMax objective. It trains representations using mutual information between state-action sequences and local features of future observations.

2 Contrastive World Models

  • The method replaces Dreamer’s reconstruction-based observation model with an InfoMax objective for learning state representations in a POMDP.The objective is intended to learn robust representations without an explicit pixel decoder.Source: §2 · p. 3, §2.1 · p. 4, §2.2 · p. 4

2.1 Latent Dynamics Model

  • The RSSM encodes each observation into a latent vector, updates posterior state features, predicts the next prior from state and action, and predicts rewards from the posterior.The prior prediction is Gaussian, and the RSSM jointly optimizes its components.Source: §2.1 · pp. 3–4
  • The RSSM lower bound combines observation and reward likelihoods with a KL regularizer between posterior and prior transitions.The bound makes marginal-likelihood learning tractable by introducing a variational posterior and applying Jensen’s inequality.Source: §2.1 · pp. 3–4

2.2 InfoMax Representations

  • The InfoMax objective maximizes mutual information between state-action sequences and local patches of future state encodings, using an InfoNCE lower bound.The state posterior is concatenated with the action sequence leading to the future state.Source: §2.2 · pp. 4–5
  • The training objective combines Deep InfoMax representation loss I with KL transition objective D scaled by λ1 and reward prediction objective R scaled by λ2, respectively.The InfoMax reformulation suggests mutual information maximization for learning state representations without requiring an explicit pixel decoder.Source: §2.2 · pp. 4–5

3 Experiments

The experiments compare InfoMax with Dreamer and Momentum Prediction under stationary, synthetic-distractor, and natural-video backgrounds, finding comparable default performance but stronger InfoMax results as visual complexity increases, alongside faster training steps.

3 Experiments

  • The experiments test stationary default DMC backgrounds, moving colored-ball distractors, and full-color Kinetics video backgrounds, averaging results over three seeds.The full-color natural-video condition differs from grayscale backgrounds used in DBC and could be more challenging for agents.Source: §3 · p. 5, Figure 3 · p. 5

3.1 Baselines and Experimental Setup

  • InfoMax, Dreamer, and Momentum Prediction share an actor-critic planner and differ in their representation-learning objective; each trains for 1 × 10^6 steps on three DMC tasks.Momentum Prediction uses a momentum-encoder target, while Dreamer uses pixel reconstruction.Source: §3.1 · pp. 5–6

3.2 Results

  • In default DMC, all three methods eventually reach comparable asymptotic performance.The uncluttered observations provide a sufficient learning signal for pixel reconstruction and near-equivalent InfoMax performance.Source: Figure 4 · p. 6, §3.2 · p. 7
  • With moving distractors, InfoMax matches or exceeds Dreamer on two tasks, while Momentum Prediction collapses across all three.The cheetah-run result is especially notable because Dreamer plateaus at substantially lower reward.Source: Figure 5 · p. 6, §3.2 · p. 7
  • In the natural-video setting, InfoMax substantially outperforms Dreamer and Momentum Prediction on every task.Dreamer makes little task progress and Momentum Prediction plateaus near random performance.Source: Figure 6 · p. 6, §3.2 · p. 7
  • InfoMax is consistently substantially faster per training step than Dreamer across all three settings, without reducing encoder or dynamics-model parameter counts.Removing the pixel decoder eliminates the need to train a pixel-decoding network.Source: §3.2 · p. 7

4 Conclusion

Contrastive World Models replace pixel reconstruction with a mutual-information objective and match baseline performance in visually simple settings while outperforming them with visual distractions. The work is limited to small-scale RSSM experiments on DeepMind Control Suite, with larger domains and newer architectures left for further study.

4 Conclusion

  • The method matches reconstruction-based and momentum-based baselines without distractions, then significantly outperforms them with simple distractors or natural video backgrounds.Removing the pixel decoder also yields more efficient training.Source: §4 · p. 7
  • Contrastive World Models replace RSSM observation reconstruction with a Deep InfoMax-like lower bound on mutual information between state-action sequences and future-observation patch features.Source: §4 · p. 7
  • The evidence is limited to small-scale experiments using the older RSSM architecture on the established DeepMind Control Suite benchmark.Source: §4 · p. 7
  • The objective assumes access to state-action sequences and future observations, and future work could scale it, use action-free video for pretraining, or add language-specified goals.It does not require action-conditioned reconstruction.Source: §4 · p. 7

Related papers

Loading 2609.22175v1…