PaperlayerView exact revision on arXiv

Source-linked AI summary

Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning

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

arXiv:2101.05265v2cs.LGcs.AIstat.ML

TL;DR

RL policies trained on few environments often fail on unseen environments, motivating representations that exploit sequential behavioral structure. The paper introduces PSM and a contrastive metric-embedding procedure, whose PSM-based instantiation, PSEs, improves generalization across three diverse benchmarks. The jumping-task analysis also identifies a narrow-grid extrapolation boundary for learning-based methods.

  • Problem

    RL methods trained on few environments often struggle to generalize to unseen environments, including semantically equivalent jumping tasks with unseen obstacle locations.

  • Method

    The paper defines PSM from similarity between local and future optimal policies, then uses a contrastive procedure to learn PSEs that encode this state similarity.

  • Results

    PSEs improve generalization across jumping from pixels, LQR with spurious correlations, and Distracting DM Control Suite benchmarks.

  • Takeaways & Limitations

    The paper shows that exploiting reinforcement learning’s sequential structure can produce effective representations for cross-environment generalization.

  • Takeaways & Limitations

    On the narrow jumping grid, learning-based methods generalize mainly near training tasks even with augmentation, reflecting limited extrapolation.

Abstract

from arXiv · show

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.

1 INTRODUCTION

RL agents trained on few environments often fail to generalize to unseen but semantically equivalent environments. The paper addresses this by incorporating sequential behavioral structure into representation learning through PSM and PSEs.

  • Deep RL agents trained on a few jumping tasks struggle when test obstacles appear at unseen locations.
  • Recent generalization methods largely adapt supervised-learning techniques and rarely exploit similarity in actions across temporal observations.
  • The approach embeds states closely when their optimal policies in those states and future states are similar.
  • PSM measures cross-environment state similarity through long-term optimal behavior, is reward-agnostic, and upper-bounds suboptimality of transferred policies.
  • PSEs instantiate a general contrastive procedure with PSM to encode task-relevant behavioral invariances explicitly.
  • PSEs improve generalization across jumping from pixels, LQR with spurious correlations, and Distracting DM Control Suite benchmarks.

2 PRELIMINARIES

The paper formalizes zero-shot generalization across related MDPs with shared action spaces and disjoint state spaces. It introduces contrastive embeddings that represent a supplied state similarity metric.

  • The setting uses a collection of related MDPs sharing an action space while having disjoint state spaces.
  • The union state space S contains configurations from all environments, enabling pseudometric comparisons across environments.
  • Training environments are sampled from a distribution, and the learned policy is evaluated zero-shot on unseen environments from that distribution.
  • The π-bisimulation metric recursively combines expected-reward differences with Wasserstein distances between policy-induced transition distributions.
  • The contrastive setup maps states to embeddings using cosine similarity and a SimCLR-style loss that separates positive from other pairs.

3 POLICY SIMILARITY METRIC

PSM measures state similarity through local and future optimal behavior rather than rewards alone. Its theory connects nearest-neighbor state matching to bounded policy-transfer suboptimality.

  • Reward-based state similarity can be too restrictive or permissive when rewards and policies disagree.
  • PSM replaces absolute reward differences with a probability pseudometric between optimal policies and compares their behavior across future states.
  • The DIST term captures local optimal-behavior differences, while the Wasserstein term captures long-term optimal-behavior differences.
  • PSM supports nearest-neighbor transfer by matching each target state to the closest source state under the metric.
  • Theorem 1 gives a non-vacuous transfer bound when d*(x̃_y, y) < 1/(1+γ), and zero distance implies an optimal transferred policy.
  • PSM provides a principled route from generalization across inputs to generalization across environments.

4 LEARNING CONTRASTIVE METRIC EMBEDDINGS

The paper learns contrastive metric embeddings by using a state similarity metric to define weighted positive and negative pairs. PSEs are CMEs learned with PSM and trained jointly with the RL agent.

  • The general CME procedure uses a state metric to define positive and negative pairs and assign their contrastive-loss importance weights.
  • Algorithm 1 samples pairs of training MDPs and updates the embedding parameters to minimize the expected contrastive loss.
  • A Gaussian kernel converts metric distances into soft similarities, with scale β controlling sensitivity.
  • For each target anchor, the nearest source state forms the positive pair and remaining source states form negative pairs.
  • The soft SimCLR loss trains the state-to-embedding function using metric-weighted similarity between each positive pair.
  • PSEs are CMEs learned with PSM and can be trained jointly with the RL agent as an auxiliary objective.

5 JUMPING TASK FROM PIXELS: A CASE STUDY

The jumping-task case study evaluates whether PSEs generalize from 18 training tasks to 268 unseen tasks. Across configurations, PSEs improve generalization, complement data augmentation, and depend on combining PSM with contrastive metric embeddings.

  • Task and setup: The task tests pixel-based invariance learning by requiring precise jumps over obstacles across varying obstacle positions and floor heights.The evaluation splits the problem into 18 seen training tasks and 268 unseen test tasks, creating a low-data, highly unbalanced classification problem.
  • Generalization results: PSEs achieve better performance than standard regularization with only 18 training examples, without incorporating additional domain knowledge or data augmentation.The comparison covers multiple train/test splits and grid configurations.
  • Generalization results: PSEs outperform bisimulation transfer in the wide and random grids because PSM is reward-agnostic, whereas expected jump returns vary with obstacle position.Bisimulation transfer additionally assumes oracle access to unseen-environment dynamics and rewards plus tabular state spaces.
  • Data augmentation: PSEs are complementary to RandConv: combining them substantially improves RandConv across all grid configurations, while bisimulation-preserving representations reduce generalization by 30 −50% relative to PSEs.The study also reports that PSEs remain effective when data augmentation hurts performance.
  • Limitations: Learning-based methods are ineffective on the narrow grid without data augmentation, and even with RandConv they generalize mainly near training tasks.The authors attribute this boundary to poor neural-network extrapolation, which is more apparent without data-augmentation inductive bias.
  • Ablations and representations: PSEs generalize significantly better than alternatives because combining PSM with contrastive metric embeddings aligns states sharing long-term behaviorally relevant features.PSEs align states at equal distance from the obstacle, while alternative embeddings leave behaviorally equivalent states distant or impose incorrect invariances.

6 ADDITIONAL EMPIRICAL EVALUATION

The empirical evaluation tests PSM and PSEs on spurious-feature and visual-distraction benchmarks, finding near-optimal or improved generalization despite limited training environments and approximate policies.

  • LQR with spurious correlations: PSM-based representations learn task-relevant variation while ignoring distractors that are spuriously correlated with optimal actions during training.The proposed conjecture states that, with zero PSM aggregation error, gradient-descent policy learning is independent of distractors.
  • LQR with spurious correlations: PSM state aggregation achieves near-optimal LQR generalization, while competing methods generalize abysmally despite optimally solving training environments.The LQR agent trains on two environments with fixed distractors and is evaluated on unseen distractors.
  • Distracting DM Control Suite: PSEs scale to Distracting DM Control Suite by using learned policies as proxies for unavailable optimal policies when computing PSM and collecting representation-learning data.The suite evaluates whether agents can ignore high-dimensional visual distractors irrelevant to the task.
  • Distracting DM Control Suite: PSEs outperform state-of-the-art data augmentation on Distracting DM Control Suite, even when PSM is approximated from pretrained policies.The comparison uses SAC with DrQ and unseen visual distractions, including 30 unseen videos after training on two videos.
  • Distracting DM Control Suite: PSEs produce substantial gains over DrQ and also outperform DrQ initialized with the same pretrained policies.The latter comparison isolates how the methods use the auxiliary information supplied by pretrained policies.

7 RELATED WORK

The paper relates PSM and PSEs to bisimulation metrics, contrastive representation learning, and data augmentation, while positioning them as complementary approaches to generalization.

  • Bisimulation metrics: PSM builds on π-bisimulation but is reward-agnostic and is theoretically linked to an upper bound on the suboptimality of transferred policies.Traditional bisimulation relies on reward information, whereas PSM measures long-term behavioral proximity across environments.
  • Contrastive representation learning: PSEs use contrastive learning to encode behavioral similarity across MDPs and can be combined with methods encoding self-consistency, predictive information, or transition dynamics.These related contrastive objectives can provide additional invariances alongside PSEs.
  • Data augmentation: PSEs complement data augmentation and improve performance when combined with random convolutions on the jumping task and DrQ on Distracting Control Suite.The paper frames PSEs as explicitly encoding task-relevant invariances rather than relying on them to emerge indirectly.

8 CONCLUSION

The paper advances RL generalization by combining a behavior-based state similarity metric with contrastive metric embeddings, showing the value of exploiting sequential decision-making structure.

  • Contributions: PSM defines state similarity through proximity of optimal behavior, while contrastive metric embeddings learn representations from a specified similarity metric.PSEs are the combination of PSM with this general contrastive embedding procedure.
  • Conclusion: Overall, the paper shows benefits from exploiting the inherent structure of reinforcement learning when learning representations for generalization.The conclusion identifies PSEs as the resulting representation approach.

A PROOFS

The appendix develops technical results for Wasserstein-based pseudometrics, contraction and fixed-point properties, policy-transfer bounds, and computational assumptions for PSM.

  • Technical preliminaries: The proofs use pseudometrics and the 1-Wasserstein distance to compare probability distributions over states.Pseudometrics allow distinct states to have zero distance, and Wasserstein distance measures minimal transport cost under a pseudometric.
  • Fixed-point analysis: The operator F is shown to be a contraction for γ < 1 and therefore to have a unique fixed point d∗.The contraction follows from ||F(d) − F(d′)|| ≤ γ||d − d′|| and the Banach Fixed Point Theorem.
  • Policy-transfer bound: Theorem 1 provides a bound on policy transfer by relating the total-variation deviation of a transferred policy to the policy-similarity distance and future behavior.The proof proceeds by induction over the horizon, with the base case bounded by d∗.
  • Scope and limitations: Exact bisimulation is impractical for large state spaces because it is expensive and requires tabular representations, while stochastic PSM approximation remains open.The appendix also identifies investigation of alternative long-term behavior distances as future work.
  • Computing PSM: Computing PSM in the presented setting uses dynamic programming, with deterministic-environment recursions resembling Q-learning updates.When optimal policies are unavailable, an approximate policy produces an approximate PSM whose error depends on policy suboptimality.

C.2 PSM CONNECTIONS TO DATA AUGMENTATION AND BISIMULATION

PSM is designed to capture behavioral similarity across environments by comparing optimal policies and their long-term behavior, avoiding reward-dependent mismatches that can misidentify equivalent states. Its generalized form supports arbitrary policy comparisons, and approximation error decreases with policy suboptimality.

  • PSM versus bisimulation: PSM can identify behaviorally similar states even when bisimulation assigns them arbitrarily large distance because their rewards differ.The paper states that PSM-similar states have similar optimal policies, whereas bisimulation can fail under reward variation.
  • PSM versus bisimulation: When ry > (1 + 1/γ) rx, bisimulation incorrectly makes x0 more similar to y1 than y0 despite the behavioral correspondence.The stated inequality yields d∼(x0, y1) < d∼(x0, y0), illustrating reward-driven misalignment.
  • PSM and data augmentation: PSM can quantify whether data augmentations preserve optimal behavior, avoiding aliases between states whose optimal actions differ after transformations.The paper highlights image flips in left-right goal-reaching tasks as an augmentation that can reverse the optimal action.
  • Generalized PSM: A generalized PSM compares arbitrary policies through a recursive metric construction, and its approximation error is bounded by the policy’s suboptimality.As policy suboptimality decreases, the PSM approximation becomes more accurate, improving the resulting PSEs.

D L2 METRIC EMBEDDINGS

The paper distinguishes behavior-based representation learning from prior generalization methods adapted from supervised learning or transfer settings. It motivates policy similarity embeddings as a way to exploit the sequential structure of reinforcement learning rather than relying only on generic regularization or domain-specific information.

  • Metric embeddings: A squared-loss metric embedding objective matches representation distances exactly to metric distances, but the paper argues this may be overly restrictive.The paper contrasts this objective with its alternative contrastive embedding procedure.
  • Prior approaches: Prior generalization methods commonly use regularization, stochasticity, noise injection, diverse training conditions, or self-attention architectures adapted from supervised learning.These approaches target the learning process but do not explicitly encode behavioral similarity across temporal observations.
  • Behavior-based representations: PSEs instead exploit behavior similarity, a property tied to the sequential structure of reinforcement learning.The paper presents this as a contrast with approaches that do not explicitly use sequential decision-making properties.
  • Meta-learning: Meta-learning seeks parameterizations that require few gradient steps to perform well on new tasks, including methods targeting zero-shot generalization.The described methods typically optimize training-environment losses together with an auxiliary loss.
  • Domain-specific methods: Other zero-shot generalization approaches may depend on domain-specific knowledge about entity equivalences or agent control, while causality-based methods do not scale to high-dimensional observations.These constraints distinguish such methods from the paper’s behavior-similarity approach.

F JUMPING TASK WITH PIXELS

The jumping benchmark tests pixel-based generalization across obstacle and floor configurations, using few training tasks and multiple evaluation regimes. The supplied materials also describe PSE implementations and related LQR experiments targeting invariance to distractors.

  • Task behavior: Optimal trajectories consist of right actions, one jump at a task-dependent distance from the obstacle, and subsequent right actions.This trajectory structure makes behaviorally relevant timing central to the task.
  • Task and setup: The jumping task requires precise timing of a jump over an obstacle using pixel observations and only right and jump actions.Different tasks shift floor height and obstacle position, so successful generalization requires the appropriate invariances.
  • Architecture: The jumping network uses three convolutional layers followed by a 256-unit fully connected layer and a policy-output layer.The PSE projection uses a single ReLU layer with k = 64 units from the penultimate network layer.
  • Training objective: The total jumping-task loss combines imitation learning with an auxiliary PSE loss weighted by α.The auxiliary term is denoted LCME in the supplied description.
  • LQR distractor invariance: In LQR, PSE aggregation targets distractor invariance, and under zero state-aggregation error the learned policy matrix is independent of distractors.The accompanying argument derives Kd(Wdx − Wdy) = 0 for states visited by the optimal policy.

H.4 TRAINING AND EVALUATION SETUP

The evaluation setup measures generalization on unseen environments using repeated seeds, episode returns, and fixed checkpoint protocols. Distracting DM Control experiments use DrQ settings and compare standard DrQ with DrQ augmented by PSEs.

  • Evaluation protocol: Experiments use five random seeds per task to compute means and standard deviations or errors of evaluation results.The reported setup follows the stated DrQ configuration with K = 2 and M = 2.
  • Implementation details: The Distracting DM Control setup uses task-specific action-repeat hyperparameters and three consecutive frames as the observational input.The action-repeat values are summarized in the referenced hyperparameter table.
  • Random initialization: DrQ augmented with PSEs outperforms DrQ during training from random initialization on unseen environments.Curves report average episode return across five seeds, with shaded standard-deviation regions.
  • Pretrained initialization: With pretrained initialization, DrQ augmented with PSEs outperforms DrQ on most environments.The figure reports average return across five seeds and evaluates each checkpoint using 100 episodes with unseen backgrounds.
  • Contrastive implementation: The contrastive implementation computes embedding similarity using cosine similarity and optimizes a contrastive loss over positive and negative weights.The supplied implementation excerpts identify cosine similarity and the returned loss corresponding to Equation 4.
Loading 2101.05265v2…