# Contrastive Unsupervised Learning of World Model with Invariant Causal Features

arXiv: https://arxiv.org/abs/2209.14932 · HTML: https://mlsanity.com/paper/2209.14932/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 21 of its 34 references that are in the map.
- First posted on arXiv 29 Sep 2022.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)

## Builds on (most cited first, 21 of 21)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://mlsanity.com/paper/1511.07289/) (2015)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Gibson Env: Real-World Perception for Embodied Agents](https://mlsanity.com/paper/1808.10654/) (2018)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Discriminative Unsupervised Feature Learning with Exemplar Convolutional Neural Networks](https://mlsanity.com/paper/1406.6909/) (2014)
- [Representation Learning via Invariant Causal Mechanisms](https://mlsanity.com/paper/2010.07922/) (2020)
- [iGibson 1.0: a Simulation Environment for Interactive Tasks in Large Realistic Scenes](https://mlsanity.com/paper/2012.02924/) (2020)
- [Unsupervised State Representation Learning in Atari](https://mlsanity.com/paper/1906.08226/) (2019)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)

## Built on by (most cited first, 5 of 5)

- [Four Principles for Physically Interpretable World Models](https://mlsanity.com/paper/2503.02143/) (2025)
- [Hieros: Hierarchical Imagination on Structured State Space Sequence World Models](https://mlsanity.com/paper/2310.05167/) (2023)
- [LanGWM: Language Grounded World Model](https://mlsanity.com/paper/2311.17593/) (2023)
- [Transferable Delay-Aware Reinforcement Learning via Implicit Causal Graph Modeling](https://mlsanity.com/paper/2605.12312/) (2026)
- [The Loss Is Not Enough: Sampling Conditions and Inductive Bias in Contrastive Representation Learning](https://mlsanity.com/paper/2606.04280/) (2026)

## Abstract (arXiv)

In this paper we present a world model, which learns causal features using the invariance principle. In particular, we use contrastive unsupervised learning to learn the invariant causal features, which enforces invariance across augmentations of irrelevant parts or styles of the observation. The world-model-based reinforcement learning methods independently optimize representation learning and the policy. Thus naive contrastive loss implementation collapses due to a lack of supervisory signals to the representation learning module. We propose an intervention invariant auxiliary task to mitigate this issue. Specifically, we utilize depth prediction to explicitly enforce the invariance and use data augmentation as style intervention on the RGB observation space. Our design leverages unsupervised representation learning to learn the world model with invariant causal features. Our proposed method significantly outperforms current state-of-the-art model-based and model-free reinforcement learning methods on out-of-distribution point navigation tasks on the iGibson dataset. Moreover, our proposed model excels at the sim-to-real transfer of our perception learning module. Finally, we evaluate our approach on the DeepMind control suite and enforce invariance only implicitly since depth is not available. Nevertheless, our proposed model performs on par with the state-of-the-art counterpart.