PaperlayerView exact revision on arXiv

Source-linked AI summary

Contrastive Learning as Goal-Conditioned Reinforcement Learning

Benjamin Eysenbach, Tianjun Zhang, Ruslan Salakhutdinov, Sergey Levine

arXiv:2206.07568v2cs.LGcs.AI

TL;DR

Deep RL often struggles to learn useful representations end to end, motivating separate losses or data augmentation. The paper instead casts contrastive representation learning as goal-conditioned RL by making representation inner products correspond to value functions. Across goal-conditioned, image-based, and offline tasks, contrastive RL outperforms prior methods, while remaining simpler in some variants.

  • Problem

    End-to-end representation learning in deep RL is unstable, so the paper asks how RL algorithms can acquire good representations without auxiliary perception losses or data augmentation.

  • Method

    Contrastive learning is applied to action-labeled trajectories so learned representation inner products correspond to a goal-conditioned value function.

  • Results

    Contrastive RL outperforms prior methods across goal-conditioned tasks, including image-based and offline settings; it beats all baselines on five of six benchmark tasks.

  • Takeaways & Limitations

    RL algorithms that resemble representation learning can provide simpler methods and better overall performance for goal-conditioned RL.

  • Takeaways & Limitations

    The work studies only goal-conditioned RL, leaving application to arbitrary RL problems open.

Abstract

from arXiv · show

In reinforcement learning (RL), it is easier to solve a task if given a good representation. While deep RL should automatically acquire such good representations, prior work often finds that learning representations in an end-to-end fashion is unstable and instead equip RL algorithms with additional representation learning parts (e.g., auxiliary losses, data augmentation). How can we design RL algorithms that directly acquire good representations? In this paper, instead of adding representation learning parts to an existing RL algorithm, we show (contrastive) representation learning methods can be cast as RL algorithms in their own right. To do this, we build upon prior work and apply contrastive representation learning to action-labeled trajectories, in such a way that the (inner product of) learned representations exactly corresponds to a goal-conditioned value function. We use this idea to reinterpret a prior RL method as performing contrastive learning, and then use the idea to propose a much simpler method that achieves similar performance. Across a range of goal-conditioned RL tasks, we demonstrate that contrastive RL methods achieve higher success rates than prior non-contrastive methods, including in the offline RL setting. We also show that contrastive RL outperforms prior methods on image-based tasks, without using data augmentation or auxiliary objectives.

1 Introduction

The paper asks how RL can learn useful representations without separate perception losses or data augmentation. It casts contrastive representation learning as goal-conditioned RL, linking learned representation inner products to value functions and yielding competitive methods.

  • Motivation: End-to-end representation learning in deep RL is fragile because representation and RL learning depend on each other.Prior work therefore often adds perception-specific losses or data augmentation.
  • Approach: The approach applies contrastive learning to action-labeled trajectories for goal-conditioned RL.Positive and negative state relationships are used to learn representations without manually specified reward functions.
  • Approach: Sampling positive pairs from the discounted state occupancy measure makes representation inner products exactly correspond to a value function.This establishes a formal connection between contrastive learning and goal-conditioned value estimation.
  • Contributions: Contrastive RL generalizes C-learning and produces a simpler method with similar performance and another method that consistently outperforms prior methods.The framework also supports image-based and offline goal-reaching tasks.
  • Contributions: On image-based tasks, contrastive RL outperforms methods using data augmentation and auxiliary objectives without using either.The comparison includes prior representation-learning methods.

2 Related Work

The paper positions contrastive RL as a direct alternative to separating representation learning from RL. It connects hindsight relabeling, similarity-based value functions, and contrastive learning while focusing on goal-conditioned problems.

  • Representation Learning in RL: Prior deep RL methods often add separate reconstruction or contrastive objectives because end-to-end representation learning is challenging.This paper instead avoids a separate representation-learning objective.
  • Contrastive Representation Learning: Contrastive learning normally makes positive representations similar and negative representations dissimilar, whereas this paper uses it to acquire a value function directly.The resulting value function can be used for action selection without additional RL.
  • Goal-Conditioned RL: Goal-conditioned RL has been addressed with temporal-difference, imitation, model-based, hierarchical, and planning-based methods.Automatic goal sampling and exploration are treated as orthogonal problems here.
  • Goal-Conditioned RL: Hindsight relabeling is reinterpreted as a standalone RL algorithm rather than merely a technique added to another algorithm.The analysis connects this view to data-driven goal-conditioned RL.
  • Goal-Conditioned RL: The paper shows that C-learning already performs contrastive learning and that alternative contrastive RL methods can be simpler and achieve higher performance.This links prior similarity-based value functions to contrastive critics.
  • Related Representation Methods: Unlike universal value functions and successor features, the proposed approach does not require manually specified reward functions or features to drive representation learning.Its analysis is closer to methods estimating discounted state occupancy as an inner product.

3 Preliminaries

Goal-conditioned RL models reaching sampled goal states as a multi-task problem with goal-dependent rewards and dynamics. The preliminaries define discounted state occupancy sampling and contrastive learning as the ingredients for the later value-function construction.

  • Goal-Conditioned Reinforcement Learning: Goal-conditioned RL is defined by states, actions, initial-state distribution, dynamics, goal distribution, and a reward function for each goal.It is equivalent to multi-task RL in which tasks correspond to reaching goal states.
  • Goal-Conditioned Reinforcement Learning: The goal-dependent reward is the probability density of reaching the goal at the next time step.This avoids requiring a human-specified distance metric.
  • Goal-Conditioned Reinforcement Learning: The objective samples a goal and optimizes the policy to reach and remain at that goal.The formulation uses a goal-conditioned policy over trajectories.
  • Discounted State Occupancy: The discounted state occupancy measure samples a time offset from a geometric distribution and then the state visited after that many steps.The paper denotes such future states by s_t+ and combines experience from multiple policies using an average stationary distribution.
  • Contrastive Representation Learning: Contrastive representation learning uses positive pairs from a joint distribution and negative examples to make similar pairs close and dissimilar pairs separate.This supplies the representation-learning framework used by contrastive RL.

4 Contrastive Learning as an RL Algorithm

Contrastive learning can directly implement goal-conditioned RL by learning a critic whose representation inner product corresponds to a Q-function. This framework supports policy improvement, generalizes C-learning, and yields a simpler algorithm without common auxiliary components.

  • 4.2 Contrastive Learning Estimates a Q-Function: Contrastive learning estimates a Q-function when applied to state-action pairs and future states sampled from discounted occupancy measures.Positive pairs use future states from the same trajectory, while negative pairs use future states from random state-action pairs.
  • 4.2 Contrastive Learning Estimates a Q-Function: The Bayes-optimal critic’s exponentiated value equals the goal-conditioned Q-function up to the multiplicative constant 1/p(sf).The critic can therefore be interpreted as an unnormalized density model whose partition function can be ignored for action selection.
  • 4.3 Policy Learning: The learned critic can guide actions because it measures how state-action pairs affect the likelihood of reaching a goal.The actor is trained using the critic, while the complete algorithm alternates critic fitting, policy updates, and data collection.
  • 4.5 Convergence Guarantees: Under tabular-state and Bayes-optimal-critic assumptions, one contrastive RL update achieves higher rewards than the initial goal-conditioned policy.Applying the method repeatedly after recollecting data corresponds to approximate policy improvement.
  • 4.5 Convergence Guarantees: Contrastive RL avoids multiple Q-values, target Q networks, data augmentation, and auxiliary objectives.The main-text experiments omit a filtering step required for the convergence proof because ablations show it can hurt practical performance.
  • 4.6 Contrastive RL Algorithms: Contrastive RL generalizes C-learning, while NCE is much simpler and can achieve comparable or higher performance.C-learning uses temporal-difference learning to distinguish future from random goals; NCE uses contrastive learning directly.

5 Experiments

The experiments evaluate contrastive RL against prior goal-conditioned, representation-learning, and offline RL methods across state-, image-, partially observed, and offline tasks. Contrastive RL is competitive or superior across these settings, with simpler variants matching prior methods and combined variants often performing best.

  • Comparing to prior goal-conditioned RL methods: Contrastive RL (NCE) solves the challenging sawyer bin task and outperforms prior methods on both state-based pushing tasks.On image-based tasks, it outperforms baselines on the three more challenging tasks, including sawyer push and sawyer bin, where baselines make no progress.
  • Comparing to prior representation learning methods: Contrastive RL (NCE) outperforms prior representation-learning methods on all image-based tasks without auxiliary objectives or image-specific data augmentation.Autoencoder and DrQ improve TD3+HER on selected tasks, but contrastive RL performs better across the evaluated tasks.
  • Probing the dimensions of contrastive RL: Contrastive RL (NCE) is simpler than prior contrastive methods, while NCE + C-learning consistently ranks among the best methods.CPC outperforms NCE on three tasks, and C-learning outperforms NCE on three different tasks but performs worse elsewhere.
  • Partial Observability and Moving Cameras: 35% success with a moving camera and partial observability compares with 75% using a fixed static camera.The moving-camera experiment occludes the puck location at episode start, yet contrastive RL still successfully handles the task.
  • Contrastive RL for Offline RL: Contrastive RL (NCE) outperforms all baselines on five of six offline D4RL AntMaze tasks.It achieves a 7% to 9% absolute improvement over IQL on the most challenging large tasks and a median absolute improvement of 15% over GCBC.

6 Conclusion

The paper frames contrastive learning as a direct route to goal-conditioned RL, yielding simpler algorithms and better overall performance. Its scope remains limited to goal-conditioned problems, leaving broader RL applications open.

  • Contrastive representation learning can be used to perform goal-conditioned RL and reinterpret prior RL methods.
  • The framework suggests simpler contrastive RL algorithms and algorithms with better overall performance.
  • The work examines only goal-conditioned RL problems, so applying these methods to arbitrary RL problems remains open.

A Additional Related Work

The related-work discussion positions this paper as a connection between representation learning, unsupervised skills, and goal-conditioned RL. Its key contribution is identifying the missing link between representation learning and goal-conditioned RL.

  • Unsupervised skill discovery and contrastive RL both learn without manually specified reward functions and optimize a lower bound on mutual information.
  • The central challenge in connecting unsupervised skills to goal-conditioned policies is grounding each skill in a specific goal.
  • The paper’s key contribution is connecting representation learning to goal-conditioned RL.

B Proofs

The proofs establish how contrastive objectives relate to occupancy measures, Q-functions, and policy improvement. They also explain the regularization needed for infoNCE critics to support goal-reaching actions.

  • The proof begins from the discounted state occupancy measure and rewrites trajectory probabilities using the Markov property and linearity of expectation.
  • The NCE policy-improvement proof relates Q-values for the average policy to those for the goal-conditioned policy, with filtered relabeling parameter ϵ.
  • Using approximate Q-values yields approximate policy improvement, according to the cited policy-improvement lemma.
  • The infoNCE objective requires regularizing the arbitrary action-dependent term c(s, a) so critic-maximizing actions correspond to reaching the desired goal.

D Contrastive RL (NCE + C-learning)

The combined NCE + C-learning method sums the two objectives and improves estimation efficiency by sharing samples across related loss terms. This sampling construction increases the effective sample size substantially.

  • NCE + C-learning uses a simple, unweighted sum of the C-learning and NCE objectives.
  • The method combines positive samples from next-state and future-state distributions through a mixture distribution.
  • 96%: the sampling trick increases effective sample size from 130 to 256.
  • The outer-product estimation trick is not used because it cannot estimate one loss term directly and was found to hurt performance for another.

E Experimental Details

The experiments implement contrastive RL within SAC, evaluate it across state- and image-based goal-conditioned tasks, and test representation quality, goal dimensionality, goal sampling, and transfer.

  • Implementation: The implementation modifies SAC to sample trajectory triplets, parameterize the critic as an inner product, and use the contrastive critic loss.The actor loss remains unchanged except for removing the entropy term on state-based tasks.
  • Implementation: State-action and goal encoders use two-layer 256-unit ReLU networks, while image tasks use 64 × 64 × 3 observations with a CNN encoder.Representation normalization, final activations, and learnable temperature are omitted because they hurt performance.
  • Representation evaluation: Contrastive RL learns representations that achieve lower test error than TD3+HER and a random CNN encoder when predicting shortest-path distance.The probe uses image-based nine-room navigation, where shortest-path distance accounts for walls rather than Euclidean distance.
  • Goal dimensionality: Contrastive learning outperforms a forward model starting with a 4-dimensional goal in the sawyer push experiment.The experiment varies the goal from 1-dimensional gripper position to 7-dimensional object-and-gripper poses and measures Euclidean distance to the goal.
  • Goal sampling: Training on only random goals performs best in the actor-loss goal-sampling ablation.The theory permits any full-support goal distribution, while the experiment compares future, random, and mixed goals.
  • Transfer: After 1M pretraining steps, transferred representations sometimes accelerate learning on new tasks, including transfer between the sawyer bin and sawyer push tasks.The benefit is small on fetch push but more visible on the challenging sawyer bin task.

F.5 Robustness to Environment Perturbations

The additional experiments examine robustness to visual and viewpoint changes and visualize the learned image representations and their gradients.

  • Environment perturbations: Changing the table color caused failure, reducing success to 0 ± 0%, whereas changing object color reduced success from 78 ± 5% to 73 ± 10%.The initial-arm-position change remained robust at 87 ± 6%; each setting used 100 total trials.
  • Representation visualization: The learned state-action representations separate reaching, picking, and placing phases in a TSNE embedding of bin-picking observations.The visualization uses five observations, a goal image, and image representations learned by Contrastive RL (NCE).
  • Representation visualization: TSNE embeddings show that the untrained encoder misses maze structure, while Contrastive RL and TD3 + HER capture it.The comparison uses image observations generated across 270 locations in the Spiral11x11 maze.
  • Gradient analysis: Figure 17 compares cosine similarity between normalized critic gradients with respect to goal images for an untrained network and the learned network.The caption frames lower gradient similarity as reducing interference between updates at different states or tasks.

G Failed Experiments

Several representation-learning modifications failed to consistently improve contrastive RL performance.

  • Failed experiments: Representation normalization, momentum encoders, frame stacking, and loss scaling did not consistently improve performance.Frame stacking tended to decrease performance slightly, while equalizing the total weight of positive and negative examples had no effect.
Loading 2206.07568v2…