# Return-Based Contrastive Representation Learning for Reinforcement Learning

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

## Key facts

- Cited by 36 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 55 references that are in the map.
- First posted on arXiv 22 Feb 2021.

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

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (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)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [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)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model](https://mlsanity.com/paper/1907.00953/) (2019)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [Learning to Navigate in Complex Environments](https://mlsanity.com/paper/1611.03673/) (2016)
- [Unsupervised State Representation Learning in Atari](https://mlsanity.com/paper/1906.08226/) (2019)
- [MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning](https://mlsanity.com/paper/2006.16908/) (2020)
- [Multi-task Deep Reinforcement Learning with PopArt](https://mlsanity.com/paper/1809.04474/) (2018)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)
- [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)
- [Universal Successor Features Approximators](https://mlsanity.com/paper/1812.07626/) (2018)
- [Bootstrap Latent-Predictive Representations for Multitask Reinforcement Learning](https://mlsanity.com/paper/2004.14646/) (2020)
- [Playing hard exploration games by watching YouTube](https://mlsanity.com/paper/1805.11592/) (2018)
- [Deep Reinforcement and InfoMax Learning](https://mlsanity.com/paper/2006.07217/) (2020)
- [Loss is its own Reward: Self-Supervision for Reinforcement Learning](https://mlsanity.com/paper/1612.07307/) (2016)
- [The Value-Improvement Path: Towards Better Representations for Reinforcement Learning](https://mlsanity.com/paper/2006.02243/) (2020)
- [Predictive Information Accelerates Learning in RL](https://mlsanity.com/paper/2007.12401/) (2020)
- [A Geometric Perspective on Optimal Representations for Reinforcement Learning](https://mlsanity.com/paper/1901.11530/) (2019)

## Built on by (most cited first, 30 of 36)

- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [For SALE: State-Action Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2306.02451/) (2023)
- [RLCD: Reinforcement Learning from Contrastive Distillation for Language Model Alignment](https://mlsanity.com/paper/2307.12950/) (2023)
- [PlayVirtual: Augmenting Cycle-Consistent Virtual Trajectories for Reinforcement Learning](https://mlsanity.com/paper/2106.04152/) (2021)
- [Mask-based Latent Reconstruction for Reinforcement Learning](https://mlsanity.com/paper/2201.12096/) (2022)
- [Deep Learning Methods for Abstract Visual Reasoning: A Survey on Raven's Progressive Matrices](https://mlsanity.com/paper/2201.12382/) (2022)
- [Contrastive UCB: Provably Efficient Contrastive Self-Supervised Learning in Online Reinforcement Learning](https://mlsanity.com/paper/2207.14800/) (2022)
- [Behavior Contrastive Learning for Unsupervised Skill Discovery](https://mlsanity.com/paper/2305.04477/) (2023)
- [Cross-Domain Policy Adaptation by Capturing Representation Mismatch](https://mlsanity.com/paper/2405.15369/) (2024)
- [CoBERL: Contrastive BERT for Reinforcement Learning](https://mlsanity.com/paper/2107.05431/) (2021)
- [Towards Deeper Deep Reinforcement Learning with Spectral Normalization](https://mlsanity.com/paper/2106.01151/) (2021)
- [Unsupervised Representation Learning in Deep Reinforcement Learning: A Review](https://mlsanity.com/paper/2208.14226/) (2022)
- [Learn Goal-Conditioned Policy with Intrinsic Motivation for Deep Reinforcement Learning](https://mlsanity.com/paper/2104.05043/) (2021)
- [Value-Consistent Representation Learning for Data-Efficient Reinforcement Learning](https://mlsanity.com/paper/2206.12542/) (2022)
- [Contrastive Learning for Enhancing Robust Scene Transfer in Vision-based Agile Flight](https://mlsanity.com/paper/2309.09865/) (2023)
- [MOBODY: Model Based Off-Dynamics Offline Reinforcement Learning](https://mlsanity.com/paper/2506.08460/) (2025)
- [RePreM: Representation Pre-training with Masked Model for Reinforcement Learning](https://mlsanity.com/paper/2303.01668/) (2023)
- [CR-VAE: Contrastive Regularization on Variational Autoencoders for Preventing Posterior Collapse](https://mlsanity.com/paper/2309.02968/) (2023)
- [Application of Deep Reinforcement Learning to Payment Fraud](https://mlsanity.com/paper/2112.04236/) (2021)
- [Learning List-wise Representation in Reinforcement Learning for Ads Allocation with Multiple Auxiliary Tasks](https://mlsanity.com/paper/2204.00888/) (2022)
- [Sequential Action-Induced Invariant Representation for Reinforcement Learning](https://mlsanity.com/paper/2309.12628/) (2023)
- [Feasibility Consistent Representation Learning for Safe Reinforcement Learning](https://mlsanity.com/paper/2405.11718/) (2024)
- [MOOSS: Mask-Enhanced Temporal Contrastive Learning for Smooth State Evolution in Visual Reinforcement Learning](https://mlsanity.com/paper/2409.02714/) (2024)
- [Pretrained Encoders are All You Need](https://mlsanity.com/paper/2106.05139/) (2021)
- [Learning Goal-Conditioned Representations for Language Reward Models](https://mlsanity.com/paper/2407.13887/) (2024)
- [Revisiting Estimation Bias in Policy Gradients for Deep Reinforcement Learning](https://mlsanity.com/paper/2301.08442/) (2023)
- [Domain Adaptation of Reinforcement Learning Agents based on Network Service Proximity](https://mlsanity.com/paper/2303.01013/) (2023)
- [Stackelberg Coupling of Online Representation Learning and Reinforcement Learning](https://mlsanity.com/paper/2508.07452/) (2025)
- [Gaze on the Prize: Shaping Visual Attention with Return-Guided Contrastive Learning](https://mlsanity.com/paper/2510.08442/) (2025)
- [Policy-Induced Self-Supervision Improves Representation Finetuning in Visual RL](https://mlsanity.com/paper/2302.06009/) (2023)

## Abstract (arXiv)

Recently, various auxiliary tasks have been proposed to accelerate representation learning and improve sample efficiency in deep reinforcement learning (RL). However, existing auxiliary tasks do not take the characteristics of RL problems into consideration and are unsupervised. By leveraging returns, the most important feedback signals in RL, we propose a novel auxiliary task that forces the learnt representations to discriminate state-action pairs with different returns. Our auxiliary loss is theoretically justified to learn representations that capture the structure of a new form of state-action abstraction, under which state-action pairs with similar return distributions are aggregated together. In low data regime, our algorithm outperforms strong baselines on complex tasks in Atari games and DeepMind Control suite, and achieves even better performance when combined with existing auxiliary tasks.