Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning

By Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, Marc G. Bellemare. 2021

197 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computational Theory and Mathematics

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Theory and Mathematics

Citations in the corpus

Provable Representation Learning for Imitation with Contrastive Fourier Features (2021, background): “Of these approaches, bisimulation is the most theoretically mature [17, 13], and several recent works apply bisimulation principles to derive practical representation learning objectives [20, 41, 6].”
Representation Matters: Offline Pretraining for Sequential Decision Making (2021, background): “While these methods originally required explicit knowledge of the reward and dynamics functions of the environment, a number of recent works have translated these ideas to stochastic representation learning objectives using deep neural networks (Gelada et al., 2019; Zhang et al., 2020; Agarwal et al”
Value Function Spaces: Skill-Centric State Abstractions for Long-Horizon Reasoning (2021, methodology): “, 2019) and representing state similarity using the bisimulation metric (Castro, 2020; Zhang et al., 2021a; Agarwal et al., 2021).”
Learning Domain Invariant Representations in Goal-conditioned Block MDPs (2021, background): “These assumptions do not hold in GBMDPs. Contrastive Metric Embeddings (CME) [51] instead uses π∗-bisimulation metric but it also requires extra information of the optimal policy.”
Bisimulation Makes Analogies in Goal-Conditioned Reinforcement Learning (2022, background): “Gelada et al. (2019) present results that their learned representation bounds the bisimulation metric and Zhang et al. (2021); Agarwal et al. (2021); Castro et al. (2021) use on-policy versions of the bisimulation metric (Castro, 2020) to learn an embedding space.”
Deep Reinforcement Learning amidst Continual Structured Non-Stationarity (2021, background): “Deep RL algorithms generally struggle to generalize to out-of-distribution environment conditions (Kumar et al., 2020; Mendonca et al., 2020; Agarwal et al., 2021).”
Model-augmented Prioritized Experience Replay (2022, methodology): “…the exploration (Pathak et al., 2017; Shelhamer et al., 2016) or to enhance representation learning for pixel-based observations (Lin et al., 2019; Agarwal et al., 2021; Stooke et al., 2021; Lee et al., 2020b; Schrittwieser et al., 2020; Hessel et al., 2021; Hafner et al., 2020b; Srinivas et al.,…”
AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning (2021, methodology): “Following the recent published work on reliable evaluation for RL (Agarwal et al., 2021b), we utilize 4 evaluation metrics: median, interquartile mean (IQM), mean, MAML (Finnetal.,2017) CAVIA (Zintgrafetal., and optimality gap.”