# Contrastive UCB: Provably Efficient Contrastive Self-Supervised Learning in Online Reinforcement Learning

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

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 58 references that are in the map.
- First posted on arXiv 29 Jul 2022.

## Builds on (most cited first, 30 of 42)

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Deep Reinforcement Learning for Autonomous Driving: A Survey](https://mlsanity.com/paper/2002.00444/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Provably Efficient Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/1907.05388/) (2019)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Time-Contrastive Networks: Self-Supervised Learning from Video](https://mlsanity.com/paper/1704.06888/) (2017)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Reinforcement Learning in Healthcare: A Survey](https://mlsanity.com/paper/1908.08796/) (2019)
- [Contextual Decision Processes with Low Bellman Rank are PAC-Learnable](https://mlsanity.com/paper/1610.09512/) (2016)
- [Decoupling Representation Learning from Reinforcement Learning](https://mlsanity.com/paper/2009.08319/) (2020)
- [Model-Based Reinforcement Learning with Value-Targeted Regression](https://mlsanity.com/paper/2006.01107/) (2020)
- [FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs](https://mlsanity.com/paper/2006.10814/) (2020)
- [Deep Reinforcement Learning framework for Autonomous Driving](https://mlsanity.com/paper/1704.02532/) (2017)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [Provably Efficient Exploration in Policy Optimization](https://mlsanity.com/paper/1912.05830/) (2019)
- [Reinforcement Learning in Feature Space: Matrix Bandit, Kernels, and Regret Bound](https://mlsanity.com/paper/1905.10389/) (2019)
- [Nearly Minimax Optimal Reinforcement Learning for Linear Mixture Markov Decision Processes](https://mlsanity.com/paper/2012.08507/) (2020)
- [Provably efficient RL with Rich Observations via Latent State Decoding](https://mlsanity.com/paper/1901.09018/) (2019)
- [Pessimistic Model-based Offline Reinforcement Learning under Partial Coverage](https://mlsanity.com/paper/2107.06226/) (2021)
- [Unsupervised State Representation Learning in Atari](https://mlsanity.com/paper/1906.08226/) (2019)
- [Is a Good Representation Sufficient for Sample Efficient Reinforcement Learning?](https://mlsanity.com/paper/1910.03016/) (2019)
- [Representation Learning for Online and Offline RL in Low-rank MDPs](https://mlsanity.com/paper/2110.04652/) (2021)
- [When to use parametric models in reinforcement learning?](https://mlsanity.com/paper/1906.05243/) (2019)
- [Kinematic State Abstraction and Provably Efficient Rich-Observation Reinforcement Learning](https://mlsanity.com/paper/1911.05815/) (2019)

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

- [Hybrid RL: Using Both Offline and Online Data Can Make RL Efficient](https://mlsanity.com/paper/2210.06718/) (2022)
- [Feedback Efficient Online Fine-Tuning of Diffusion Models](https://mlsanity.com/paper/2402.16359/) (2024)
- [Spectral Decomposition Representation for Reinforcement Learning](https://mlsanity.com/paper/2208.09515/) (2022)
- [Behavior Contrastive Learning for Unsupervised Skill Discovery](https://mlsanity.com/paper/2305.04477/) (2023)
- [Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models](https://mlsanity.com/paper/2405.19673/) (2024)
- [From Words to Actions: Unveiling the Theoretical Underpinnings of LLM-Driven Autonomous Systems](https://mlsanity.com/paper/2405.19883/) (2024)
- [Representation Learning for General-sum Low-rank Markov Games](https://mlsanity.com/paper/2210.16976/) (2022)
- [Latent Variable Representation for Reinforcement Learning](https://mlsanity.com/paper/2212.08765/) (2022)
- [Contrastive Learning for Enhancing Robust Scene Transfer in Vision-based Agile Flight](https://mlsanity.com/paper/2309.09865/) (2023)
- [Contrastive Imitation Learning for Language-guided Multi-Task Robotic Manipulation](https://mlsanity.com/paper/2406.09738/) (2024)
- [Spectral Representation Learning for Conditional Moment Models](https://mlsanity.com/paper/2210.16525/) (2022)
- [Spectral Representation for Causal Estimation with Hidden Confounders](https://mlsanity.com/paper/2407.10448/) (2024)
- [Bridging the Sim-to-Real Gap from the Information Bottleneck Perspective](https://mlsanity.com/paper/2305.18464/) (2023)
- [Provable Representation with Efficient Planning for Partial Observable Reinforcement Learning](https://mlsanity.com/paper/2311.12244/) (2023)
- [Constrained Ensemble Exploration for Unsupervised Skill Discovery](https://mlsanity.com/paper/2405.16030/) (2024)
- [OVD-Explorer: Optimism Should Not Be the Sole Pursuit of Exploration in Noisy Environments](https://mlsanity.com/paper/2312.12145/) (2023)
- [Pessimistic Value Iteration for Multi-Task Data Sharing in Offline Reinforcement Learning](https://mlsanity.com/paper/2404.19346/) (2024)
- [Pessimism Meets Risk: Risk-Sensitive Offline Reinforcement Learning](https://mlsanity.com/paper/2407.07631/) (2024)
- [Uncertainty-Aware Reward-Free Exploration with General Function Approximation](https://mlsanity.com/paper/2406.16255/) (2024)
- [FAR: Failure-Aware Retry for Test-Time Recovery and Continual Policy Improvement](https://mlsanity.com/paper/2607.01111/) (2026)
- [An Empirical Study on the Power of Future Prediction in Partially Observable Environments](https://mlsanity.com/paper/2402.07102/) (2024)
- [Primal-Dual Spectral Representation for Off-policy Evaluation](https://mlsanity.com/paper/2410.17538/) (2024)
- [Contrastive Diffuser: Planning Towards High Return States via Contrastive Learning](https://mlsanity.com/paper/2402.02772/) (2024)
- [Ensemble Successor Representations for Task Generalization in Offline-to-Online Reinforcement Learning](https://mlsanity.com/paper/2405.07223/) (2024)
- [CURLing the Dream: Contrastive Representations for World Modeling in Reinforcement Learning](https://mlsanity.com/paper/2408.05781/) (2024)
- [Robust Scene Transfer for PointGoal Navigation via Privileged Sensor Guided Contrastive Learning](https://mlsanity.com/paper/2606.05506/) (2026)

## Abstract (arXiv)

In view of its power in extracting feature representation, contrastive self-supervised learning has been successfully integrated into the practice of (deep) reinforcement learning (RL), leading to efficient policy learning in various applications. Despite its tremendous empirical successes, the understanding of contrastive learning for RL remains elusive. To narrow such a gap, we study how RL can be empowered by contrastive learning in a class of Markov decision processes (MDPs) and Markov games (MGs) with low-rank transitions. For both models, we propose to extract the correct feature representations of the low-rank model by minimizing a contrastive loss. Moreover, under the online setting, we propose novel upper confidence bound (UCB)-type algorithms that incorporate such a contrastive loss with online RL algorithms for MDPs or MGs. We further theoretically prove that our algorithm recovers the true representations and simultaneously achieves sample efficiency in learning the optimal policy and Nash equilibrium in MDPs and MGs. We also provide empirical studies to demonstrate the efficacy of the UCB-based contrastive learning method for RL. To the best of our knowledge, we provide the first provably efficient online RL algorithm that incorporates contrastive learning for representation learning. Our codes are available at https://github.com/Baichenjia/Contrastive-UCB.