Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning

Abstract

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.

ICLR 2021 (Spotlight). Website: https://agarwl.github.io/pse

References in corpus (23)

• UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction
• A Simple Framework for Contrastive Learning of Visual Representations
• Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks
• Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor
• Reinforcement Learning with Augmented Data
• A Study on Overfitting in Deep Reinforcement Learning
• dm_control: Software and Tasks for Continuous Control
• Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels
• Illuminating Generalization in Deep Reinforcement Learning through Procedural Level Generation
• Zero-Shot Task Generalization with Multi-Task Deep Reinforcement Learning
• Assessing Generalization in Deep Reinforcement Learning
• How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks
• Generalization and Regularization in DQN
• A Dissection of Overfitting and Generalization in Continuous Reinforcement Learning
• Decoupling Representation Learning from Reinforcement Learning
• Generalization in Reinforcement Learning with Selective Noise Injection and Information Bottleneck
• Automatic Data Augmentation for Generalization in Deep Reinforcement Learning
• Data-Efficient Reinforcement Learning with Self-Predictive Representations
• Natural Environment Benchmarks for Reinforcement Learning
• Predictive Information Accelerates Learning in RL
• The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels
• Invariant Policy Optimization: Towards Stronger Generalization in Reinforcement Learning
• Rotation, Translation, and Cropping for Zero-Shot Generalization

Cited by in corpus (10)

• A Survey of Zero-shot Generalisation in Deep Reinforcement Learning
• Automatic Data Augmentation for Generalization in Deep Reinforcement Learning
• Representation Matters: Offline Pretraining for Sequential Decision Making
• Structure in Deep Reinforcement Learning: A Survey and Open Problems
• Provable Representation Learning for Imitation with Contrastive Fourier Features
• Decoupling Value and Policy for Generalization in Reinforcement Learning
• Towards Deeper Deep Reinforcement Learning with Spectral Normalization
• Learning Temporally-Consistent Representations for Data-Efficient Reinforcement Learning
• Learning Representations for Pixel-based Control: What Matters and Why?
• Towards an Understanding of Default Policies in Multitask Policy Optimization