# Masked Contrastive Representation Learning for Reinforcement Learning

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

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 15 Oct 2020.

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

- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Investigating the Role of Negatives in Contrastive Representation Learning](https://mlsanity.com/paper/2106.09943/) (2021)
- [PlayVirtual: Augmenting Cycle-Consistent Virtual Trajectories for Reinforcement Learning](https://mlsanity.com/paper/2106.04152/) (2021)
- [Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learning](https://mlsanity.com/paper/2307.01849/) (2023)
- [Mask-based Latent Reconstruction for Reinforcement Learning](https://mlsanity.com/paper/2201.12096/) (2022)
- [Masked Image Modeling: A Survey](https://mlsanity.com/paper/2408.06687/) (2024)
- [Revisiting Plasticity in Visual Reinforcement Learning: Data, Modules and Training Stages](https://mlsanity.com/paper/2310.07418/) (2023)
- [A Comprehensive Survey of Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2210.04561/) (2022)
- [Cross-Domain Policy Adaptation by Capturing Representation Mismatch](https://mlsanity.com/paper/2405.15369/) (2024)
- [Does Self-supervised Learning Really Improve Reinforcement Learning from Pixels?](https://mlsanity.com/paper/2206.05266/) (2022)
- [CoBERL: Contrastive BERT for Reinforcement Learning](https://mlsanity.com/paper/2107.05431/) (2021)
- [An Empirical Investigation of Representation Learning for Imitation](https://mlsanity.com/paper/2205.07886/) (2022)
- [Light-weight probing of unsupervised representations for Reinforcement Learning](https://mlsanity.com/paper/2208.12345/) (2022)
- [Robust Representation Learning by Clustering with Bisimulation Metrics for Visual Reinforcement Learning with Distractions](https://mlsanity.com/paper/2302.12003/) (2023)
- [VideoMiner: Iteratively Grounding Key Frames of Hour-Long Videos via Tree-based Group Relative Policy Optimization](https://mlsanity.com/paper/2510.06040/) (2025)
- [Uncertainty Estimation Using Riemannian Model Dynamics for Offline Reinforcement Learning](https://mlsanity.com/paper/2102.11327/) (2021)
- [MOBODY: Model Based Off-Dynamics Offline Reinforcement Learning](https://mlsanity.com/paper/2506.08460/) (2025)
- [Learning Better with Less: Effective Augmentation for Sample-Efficient Visual Reinforcement Learning](https://mlsanity.com/paper/2305.16379/) (2023)
- [Learning to Discover Skills through Guidance](https://mlsanity.com/paper/2310.20178/) (2023)
- [MA2CL:Masked Attentive Contrastive Learning for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2306.02006/) (2023)
- [Offline Imitation from Observation via Primal Wasserstein State Occupancy Matching](https://mlsanity.com/paper/2311.01331/) (2023)
- [Representation Learning For Efficient Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2406.02890/) (2024)
- [MOOSS: Mask-Enhanced Temporal Contrastive Learning for Smooth State Evolution in Visual Reinforcement Learning](https://mlsanity.com/paper/2409.02714/) (2024)
- [Representation Learning for Sequential Volumetric Design Tasks](https://mlsanity.com/paper/2309.02583/) (2023)
- [Contrastive learning-based agent modeling for deep reinforcement learning](https://mlsanity.com/paper/2401.00132/) (2023)
- [Scaling DRL for Decision Making: A Survey on Data, Network, and Training Budget Strategies](https://mlsanity.com/paper/2508.03194/) (2025)
- [Gaze on the Prize: Shaping Visual Attention with Return-Guided Contrastive Learning](https://mlsanity.com/paper/2510.08442/) (2025)
- [Toward Active Object Detection for UAVs in the Wild: A Large-Scale Dataset, Benchmark and Method](https://mlsanity.com/paper/2607.09078/) (2026)
- [Rollout Total Correlation for Deep Reinforcement Learning](https://mlsanity.com/paper/2209.05333/) (2022)
- [Intrinsic Dynamics-Driven Generalizable Scene Representations for Vision-Oriented Decision-Making Applications](https://mlsanity.com/paper/2405.19736/) (2024)

## Abstract (arXiv)

Improving sample efficiency is a key research problem in reinforcement learning (RL), and CURL, which uses contrastive learning to extract high-level features from raw pixels of individual video frames, is an efficient algorithm~citepsrinivas2020curl. We observe that consecutive video frames in a game are highly correlated but CURL deals with them independently. To further improve data efficiency, we propose a new algorithm, masked contrastive representation learning for RL, that takes the correlation among consecutive inputs into consideration. In addition to the CNN encoder and the policy network in CURL, our method introduces an auxiliary Transformer module to leverage the correlations among video frames. During training, we randomly mask the features of several frames, and use the CNN encoder and Transformer to reconstruct them based on the context frames. The CNN encoder and Transformer are jointly trained via contrastive learning where the reconstructed features should be similar to the ground-truth ones while dissimilar to others. During inference, the CNN encoder and the policy network are used to take actions, and the Transformer module is discarded. Our method achieves consistent improvements over CURL on 14 out of 16 environments from DMControl suite and 21 out of 26 environments from Atari 2600 Games. The code is available at https://github.com/teslacool/m-curl.