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 1 ). We demonstrate that PSEs improve generalization on diverse benchmarks, including LQR with spurious correlations, a jumping task from pixels, and Distracting DM Control Suite. Source code would be made available at agarwl.github.io/pse.

Efficient Embedding of Semantic Similarity in Control Policies via Entangled Bisimulation

“…Contrastive learning (Srinivas et al, 2020;Chen et al, 2020;Hjelm et al, 2018) allows the algorithm designer to specify positive and negative matches in representation space, and embed a similarity measure based on this by maximizing agreement between positive matches, and minimizing it w.r.t negative matches; contrastive learning has been shown to improve sample efficiency. Observation prediction and reconstruction (Hafner et al, 2019;Sekar et al, 2020) also provides a rich auxiliary training signal, but forces the agent to model and reconstruct task-irrelevant distractors, which may be a significant disadvantage on natural scenarios (Zhang et al, 2020;Agarwal et al, 2021).…”

“…The works of (Agarwal et al, 2021;Zhang et al, 2020) (DBC, PSM) tackle the issue of approximating policy bisimulation (or related notions) over partially observable Markov decision processes (POMDP). Both approximations unfortunately suffer from drawbacks where their proposed bisimulation optimization objectives have nontrivial biases for estimating Wasserstein distances, and provably lose important metric properties like state self similarity.…”

“…Two recent approaches to embed π−bisimulation metrics are DBC (Zhang et al, 2020) and PSM (Agarwal et al, 2021). We first analyze these approaches before presenting our solution.…”

“…Lemma 4.7. Let E π be the largest π-bisimulation relation, and let E P SM likewise be the largest average-action bisimulation relation in (Agarwal et al, 2021). The entangled bisimulation relation E satisfies…”

“…Correspondence to: Martin Bertran <martin.bertran@duke.edu>. lation metrics (Ferns et al, 2011;Ferns & Precup, 2014;Castro, 2020), where states that are indistinguishable with respect to future reward trajectories are grouped together (Zhang et al, 2020) (deep bisimulation for control, DBC), or states which produce the same action sequence under the optimal policy are similarly grouped (Agarwal et al, 2021) (policy similarity metrics, PSM). Bisimulation tackles generalization by incorporating and measuring important invariances of the policy and environment dynamics into the representation learning process.…”