# Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning

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

## Key facts

- Cited by 123 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 67 references that are in the map.
- First posted on arXiv 13 Jan 2021.

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

- [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)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [The 2017 DAVIS Challenge on Video Object Segmentation](https://mlsanity.com/paper/1704.00675/) (2017)
- [Learning to Generalize: Meta-Learning for Domain Generalization](https://mlsanity.com/paper/1710.03463/) (2017)
- [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)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [Implicit Regularization in Matrix Factorization](https://mlsanity.com/paper/1705.09280/) (2017)
- [Decoupling Representation Learning from Reinforcement Learning](https://mlsanity.com/paper/2009.08319/) (2020)
- [How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks](https://mlsanity.com/paper/2009.11848/) (2020)
- [A Study on Overfitting in Deep Reinforcement Learning](https://mlsanity.com/paper/1804.06893/) (2018)
- [A Tour of Reinforcement Learning: The View from Continuous Control](https://mlsanity.com/paper/1806.09460/) (2018)
- [Assessing Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/1810.12282/) (2018)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)
- [Generalization and Regularization in DQN](https://mlsanity.com/paper/1810.00123/) (2018)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Network Randomization: A Simple Technique for Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/1910.05396/) (2019)
- [Generalization in Reinforcement Learning with Selective Noise Injection and Information Bottleneck](https://mlsanity.com/paper/1910.12911/) (2019)
- [The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels](https://mlsanity.com/paper/2101.02722/) (2021)
- [Zero-Shot Task Generalization with Multi-Task Deep Reinforcement Learning](https://mlsanity.com/paper/1706.05064/) (2017)
- [Observational Overfitting in Reinforcement Learning](https://mlsanity.com/paper/1912.02975/) (2019)

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

- [Evolving Curricula with Regret-Based Environment Design](https://mlsanity.com/paper/2203.01302/) (2022)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Foundation Models for Decision Making: Problems, Methods, and Opportunities](https://mlsanity.com/paper/2303.04129/) (2023)
- [Bigger, Better, Faster: Human-level Atari with human-level efficiency](https://mlsanity.com/paper/2305.19452/) (2023)
- [Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability](https://mlsanity.com/paper/2107.06277/) (2021)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [Representation Matters: Offline Pretraining for Sequential Decision Making](https://mlsanity.com/paper/2102.05815/) (2021)
- [Decoupling Value and Policy for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2102.10330/) (2021)
- [DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2110.14565/) (2021)
- [Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning](https://mlsanity.com/paper/2212.08860/) (2022)
- [Learning Task Informed Abstractions](https://mlsanity.com/paper/2106.15612/) (2021)
- [Towards Robust Bisimulation Metric Learning](https://mlsanity.com/paper/2110.14096/) (2021)
- [AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning](https://mlsanity.com/paper/2107.02729/) (2021)
- [MICo: Improved representations via sampling-based state similarity for Markov decision processes](https://mlsanity.com/paper/2106.08229/) (2021)
- [Proto-Value Networks: Scaling Representation Learning with Auxiliary Tasks](https://mlsanity.com/paper/2304.12567/) (2023)
- [What Makes Good Contrastive Learning on Small-Scale Wearable-based Tasks?](https://mlsanity.com/paper/2202.05998/) (2022)
- [Bisimulation Makes Analogies in Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2204.13060/) (2022)
- [Provable Representation Learning for Imitation with Contrastive Fourier Features](https://mlsanity.com/paper/2105.12272/) (2021)
- [Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning](https://mlsanity.com/paper/2209.09203/) (2022)
- [Proper Value Equivalence](https://mlsanity.com/paper/2106.10316/) (2021)
- [Run-time Observation Interventions Make Vision-Language-Action Models More Visually Robust](https://mlsanity.com/paper/2410.01971/) (2024)
- [DRIBO: Robust Deep Reinforcement Learning via Multi-View Information Bottleneck](https://mlsanity.com/paper/2102.13268/) (2021)
- [On the Importance of Exploration for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2306.05483/) (2023)
- [On the Generalization of Representations in Reinforcement Learning](https://mlsanity.com/paper/2203.00543/) (2022)
- [Continuous MDP Homomorphisms and Homomorphic Policy Gradient](https://mlsanity.com/paper/2209.07364/) (2022)
- [Conditional Mutual Information for Disentangled Representations in Reinforcement Learning](https://mlsanity.com/paper/2305.14133/) (2023)
- [Value Function Spaces: Skill-Centric State Abstractions for Long-Horizon Reasoning](https://mlsanity.com/paper/2111.03189/) (2021)
- [The Generalization Gap in Offline Reinforcement Learning](https://mlsanity.com/paper/2312.05742/) (2023)
- [Learning Representations for Pixel-based Control: What Matters and Why?](https://mlsanity.com/paper/2111.07775/) (2021)
- [Investigating the Properties of Neural Network Representations in Reinforcement Learning](https://mlsanity.com/paper/2203.15955/) (2022)

## Abstract (arXiv)

Reinforcement learning methods trained on few environments rarely learn policies that generalize to unseen environments. To improve generalization, we incorporate the inherent sequential structure in reinforcement learning into the representation learning process. This approach is orthogonal to recent approaches, which rarely exploit this structure explicitly. Specifically, we introduce a theoretically motivated policy similarity metric (PSM) for measuring behavioral similarity between states. PSM assigns high similarity to states for which the optimal policies in those states as well as in future states are similar. We also present a contrastive representation learning procedure to embed any state similarity metric, which we instantiate with PSM to obtain policy similarity embeddings (PSEs). We demonstrate that PSEs improve generalization on diverse benchmarks, including LQR with spurious correlations, a jumping task from pixels, and Distracting DM Control Suite.