# Integrating Contrastive Learning with Dynamic Models for Reinforcement Learning from Images

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 41 references that are in the map.
- First posted on arXiv 2 Mar 2022.

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

- [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)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Learning deep representations by mutual information estimation and maximization](https://mlsanity.com/paper/1808.06670/) (2018)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [End-to-End Training of Deep Visuomotor Policies](https://mlsanity.com/paper/1504.00702/) (2015)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [MINE: Mutual Information Neural Estimation](https://mlsanity.com/paper/1801.04062/) (2018)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Target-driven Visual Navigation in Indoor Scenes using Deep Reinforcement Learning](https://mlsanity.com/paper/1609.05143/) (2016)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Unsupervised Learning for Physical Interaction through Video Prediction](https://mlsanity.com/paper/1605.07157/) (2016)
- [Learning to Poke by Poking: Experiential Learning of Intuitive Physics](https://mlsanity.com/paper/1606.07419/) (2016)
- [Deep Spatial Autoencoders for Visuomotor Learning](https://mlsanity.com/paper/1509.06113/) (2015)
- [Unsupervised State Representation Learning in Atari](https://mlsanity.com/paper/1906.08226/) (2019)
- [The Conditional Entropy Bottleneck](https://mlsanity.com/paper/2002.05379/) (2020)
- [Predictive Information Accelerates Learning in RL](https://mlsanity.com/paper/2007.12401/) (2020)
- [PVEs: Position-Velocity Encoders for Unsupervised Learning of Structured State Representations](https://mlsanity.com/paper/1705.09805/) (2017)
- [Independently Controllable Features](https://mlsanity.com/paper/1703.07718/) (2017)

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

- [Reinforcement Learning with Neural Radiance Fields](https://mlsanity.com/paper/2206.01634/) (2022)
- [A Comprehensive Survey of Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2210.04561/) (2022)
- [Unsupervised Representation Learning in Deep Reinforcement Learning: A Review](https://mlsanity.com/paper/2208.14226/) (2022)
- [SNeRL: Semantic-aware Neural Radiance Fields for Reinforcement Learning](https://mlsanity.com/paper/2301.11520/) (2023)
- [Combining Reconstruction and Contrastive Methods for Multimodal Representations in RL](https://mlsanity.com/paper/2302.05342/) (2023)
- [Efficient Policy Adaptation with Contrastive Prompt Ensemble for Embodied Agents](https://mlsanity.com/paper/2412.11484/) (2024)
- [Multimodal Information Bottleneck for Deep Reinforcement Learning with Multiple Sensors](https://mlsanity.com/paper/2410.17551/) (2024)
- [Sample-efficient Real-time Planning with Curiosity Cross-Entropy Method and Contrastive Learning](https://mlsanity.com/paper/2303.03787/) (2023)
- [Maximum Likelihood Learning of Latent Dynamics Without Reconstruction](https://mlsanity.com/paper/2505.23569/) (2025)
- [Learning State-Space Models of Dynamic Systems from Arbitrary Data using Joint Embedding Predictive Architectures](https://mlsanity.com/paper/2508.10489/) (2025)
- [Rollout Total Correlation for Deep Reinforcement Learning](https://mlsanity.com/paper/2209.05333/) (2022)
- [Trajectory Entropy Reinforcement Learning for Robust Robot Motor Skill Learning](https://mlsanity.com/paper/2505.04193/) (2025)

## Abstract (arXiv)

Recent methods for reinforcement learning from images use auxiliary tasks to learn image features that are used by the agent's policy or Q-function. In particular, methods based on contrastive learning that induce linearity of the latent dynamics or invariance to data augmentation have been shown to greatly improve the sample efficiency of the reinforcement learning algorithm and the generalizability of the learned embedding. We further argue, that explicitly improving Markovianity of the learned embedding is desirable and propose a self-supervised representation learning method which integrates contrastive learning with dynamic models to synergistically combine these three objectives: (1) We maximize the InfoNCE bound on the mutual information between the state- and action-embedding and the embedding of the next state to induce a linearly predictive embedding without explicitly learning a linear transition model, (2) we further improve Markovianity of the learned embedding by explicitly learning a non-linear transition model using regression, and (3) we maximize the mutual information between the two nonlinear predictions of the next embeddings based on the current action and two independent augmentations of the current state, which naturally induces transformation invariance not only for the state embedding, but also for the nonlinear transition model. Experimental evaluation on the Deepmind control suite shows that our proposed method achieves higher sample efficiency and better generalization than state-of-art methods based on contrastive learning or reconstruction.