# Cross-Trajectory Representation Learning for Zero-Shot Generalization in RL

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

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 52 references that are in the map.
- First posted on arXiv 4 Jun 2021.

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

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [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)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Learning deep representations by mutual information estimation and maximization](https://mlsanity.com/paper/1808.06670/) (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)
- [Diversity is All You Need: Learning Skills without a Reward Function](https://mlsanity.com/paper/1802.06070/) (2018)
- [Learning Representations by Maximizing Mutual Information Across Views](https://mlsanity.com/paper/1906.00910/) (2019)
- [Self-labelling via simultaneous clustering and representation learning](https://mlsanity.com/paper/1911.05371/) (2019)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [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)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Learning to Drive in a Day](https://mlsanity.com/paper/1807.00412/) (2018)
- [Decoupling Representation Learning from Reinforcement Learning](https://mlsanity.com/paper/2009.08319/) (2020)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2102.11271/) (2021)
- [DARLA: Improving Zero-Shot Transfer in Reinforcement Learning](https://mlsanity.com/paper/1707.08475/) (2017)
- [Unsupervised State Representation Learning in Atari](https://mlsanity.com/paper/1906.08226/) (2019)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)
- [Phasic Policy Gradient](https://mlsanity.com/paper/2009.04416/) (2020)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Learning One Representation to Optimize All Rewards](https://mlsanity.com/paper/2103.07945/) (2021)

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

- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning](https://mlsanity.com/paper/2209.09203/) (2022)
- [Procedural Generalization by Planning with Self-Supervised World Models](https://mlsanity.com/paper/2111.01587/) (2021)
- [Improving Zero-shot Generalization in Offline Reinforcement Learning using Generalized Similarity Functions](https://mlsanity.com/paper/2111.14629/) (2021)
- [Understanding and Addressing the Pitfalls of Bisimulation-based Representations in Offline Reinforcement Learning](https://mlsanity.com/paper/2310.17139/) (2023)
- [Rethinking Value Function Learning for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2210.09960/) (2022)
- [Constrained Ensemble Exploration for Unsupervised Skill Discovery](https://mlsanity.com/paper/2405.16030/) (2024)
- [Accelerating exploration and representation learning with offline pre-training](https://mlsanity.com/paper/2304.00046/) (2023)
- [Learning Action-based Representations Using Invariance](https://mlsanity.com/paper/2403.16369/) (2024)
- [Discovering Hierarchical Achievements in Reinforcement Learning via Contrastive Learning](https://mlsanity.com/paper/2307.03486/) (2023)
- [CompilerDream: Learning a Compiler World Model for General Code Optimization](https://mlsanity.com/paper/2404.16077/) (2024)
- [Generalization in Visual Reinforcement Learning with the Reward Sequence Distribution](https://mlsanity.com/paper/2302.09601/) (2023)
- [Learning Generalizable Agents via Saliency-Guided Features Decorrelation](https://mlsanity.com/paper/2310.05086/) (2023)
- [Prototypical context-aware dynamics generalization for high-dimensional model-based reinforcement learning](https://mlsanity.com/paper/2211.12774/) (2022)
- [CROP: Towards Distributional-Shift Robust Reinforcement Learning using Compact Reshaped Observation Processing](https://mlsanity.com/paper/2304.13616/) (2023)
- [Sequential Action-Induced Invariant Representation for Reinforcement Learning](https://mlsanity.com/paper/2309.12628/) (2023)
- [Style-Agnostic Reinforcement Learning](https://mlsanity.com/paper/2208.14863/) (2022)
- [The Sandbox Environment for Generalizable Agent Research (SEGAR)](https://mlsanity.com/paper/2203.10351/) (2022)
- [KAGE-Bench: Fast Known-Axis Visual Generalization Evaluation for Reinforcement Learning](https://mlsanity.com/paper/2601.14232/) (2026)
- [Sparse Masked Attention Policies for Reliable Generalization](https://mlsanity.com/paper/2602.19956/) (2026)

## Abstract (arXiv)

A highly desirable property of a reinforcement learning (RL) agent -- and a major difficulty for deep RL approaches -- is the ability to generalize policies learned on a few tasks over a high-dimensional observation space to similar tasks not seen during training. Many promising approaches to this challenge consider RL as a process of training two functions simultaneously: a complex nonlinear encoder that maps high-dimensional observations to a latent representation space, and a simple linear policy over this space. We posit that a superior encoder for zero-shot generalization in RL can be trained by using solely an auxiliary SSL objective if the training process encourages the encoder to map behaviorally similar observations to similar representations, as reward-based signal can cause overfitting in the encoder (Raileanu et al., 2021). We propose Cross-Trajectory Representation Learning (CTRL), a method that runs within an RL agent and conditions its encoder to recognize behavioral similarity in observations by applying a novel SSL objective to pairs of trajectories from the agent's policies. CTRL can be viewed as having the same effect as inducing a pseudo-bisimulation metric but, crucially, avoids the use of rewards and associated overfitting risks. Our experiments ablate various components of CTRL and demonstrate that in combination with PPO it achieves better generalization performance on the challenging Procgen benchmark suite (Cobbe et al., 2020).