PaperlayerView exact revision on arXiv

Source-linked AI summary

CURL: Contrastive Unsupervised Representations for Reinforcement Learning

Aravind Srinivas, Michael Laskin, Pieter Abbeel

arXiv:2004.04136v4cs.LGcs.CVstat.ML

TL;DR

Pixel-based reinforcement learning is sample-inefficient, especially relative to state-based learning, creating a need for better representations from online visual interactions. CURL jointly learns contrastive representations and off-policy control with minimal architectural changes, outperforming prior pixel-based methods by 1.9x on DMControl and 1.2x on Atari at 100k-step benchmarks. On DMControl, it nearly matches state-based SAC’s sample efficiency and establishes substantial contrastive-learning gains across pixel-based control tasks.

  • Problem

    Pixel-based reinforcement learning is sample-inefficient, motivating representation learning that can extract relevant state information from images during online control.

  • Method

    CURL jointly trains a visual encoder with contrastive learning and an off-policy model-free reinforcement-learning objective using minimal architectural changes.

  • Results

    1.9x higher median performance on DMControl and 1.2x higher median performance on Atari are reported over prior pixel-based methods at 100k-step benchmarks.

  • Takeaways & Limitations

    CURL nearly matches state-based SAC’s performance and sample efficiency on most DMControl environments while achieving state-of-the-art data efficiency from pixels.

Abstract

from arXiv · show

We present CURL: Contrastive Unsupervised Representations for Reinforcement Learning. CURL extracts high-level features from raw pixels using contrastive learning and performs off-policy control on top of the extracted features. CURL outperforms prior pixel-based methods, both model-based and model-free, on complex tasks in the DeepMind Control Suite and Atari Games showing 1.9x and 1.2x performance gains at the 100K environment and interaction steps benchmarks respectively. On the DeepMind Control Suite, CURL is the first image-based algorithm to nearly match the sample-efficiency of methods that use state-based features. Our code is open-sourced and available at https://github.com/MishaLaskin/curl.

1. Introduction

CURL addresses the sample inefficiency of pixel-based reinforcement learning by jointly learning contrastive representations and off-policy control from online interactions. Across DMControl and Atari, it improves performance over prior pixel-based methods while nearly matching state-based sample efficiency on DMControl.

  • Motivation: Raw-pixel reinforcement learning is sample-inefficient, motivating representations that extract task-relevant state information from images.The paper emphasizes sample efficiency as important for both real-world robotics and simulation.
  • Approach: CURL applies contrastive learning to augmented observations collected online while reinforcement learning proceeds simultaneously.The approach targets the challenge of learning useful representations from a dynamically changing interaction dataset without offline pretraining.
  • Results: 1.9x median higher performance over Dreamer is achieved by CURL with SAC on DMControl at 100k environment steps.The result is reported across 16 DMControl environments, where CURL also matches state-based SAC on most environments.
  • Results: 1.2x median higher performance over prior methods is achieved by CURL with data-efficient Rainbow on Atari at 100k interaction steps.CURL improves upon Efficient Rainbow on 19 of 26 games and surpasses human efficiency on two games.
  • Contribution: CURL is presented as the first model-free approach to show substantial contrastive-learning data-efficiency gains across pixel-based continuous and discrete control tasks.The paper contrasts these gains with earlier mixed results from contrastive learning for model-free reinforcement learning.
  • Approach: CURL integrates with model-free reinforcement learning using the same latent space and minimal architectural or training-pipeline changes.The authors prioritize a simple, reproducible pipeline with minimal overhead.

2. Related Work

Prior work improves reinforcement-learning sample efficiency through auxiliary self-supervision or world models, while contrastive learning supplies a simpler representation-learning alternative. CURL evaluates this direction on established DMControl and Atari sample-efficiency benchmarks.

  • Self-Supervised Learning: Self-supervised learning aims to learn rich representations from unlabeled high-dimensional data for diverse downstream tasks.Prior examples include BERT, CPC, MoCo, and SimCLR.
  • Contrastive Learning: Contrastive learning learns representations by enforcing similarity constraints between related examples and dissimilarity constraints between unrelated examples.Instance discrimination treats augmented views of the same instance as positive pairs and other instances as negatives.
  • Self-Supervised Learning for RL: Auxiliary-task approaches improve model-free reinforcement learning by adding self-supervised objectives such as future prediction from observations and actions.Future prediction may occur in pixel space or latent space.
  • World Models: World-model approaches learn predictive environment models and use generated rollouts for sampling and planning.This line of work differs from jointly learning an auxiliary unsupervised task with model-free control.
  • Sample-Efficient RL for Image-Based Control: CURL benchmarks pixel-based sample efficiency on the DeepMind Control Suite and Atari Games.The cited benchmarks include the DMControl suite for continuous control and Atari’s 100k interaction-step setting.

3. Background

CURL combines an off-policy RL algorithm with contrastive representation learning over pixel observations. The background introduces SAC, Rainbow DQN, and the contrastive-learning objective underlying this pipeline.

  • RL Algorithms: CURL can use different RL algorithms, with SAC used for continuous control and Rainbow DQN for discrete control.The framework is described as compatible in principle with on-policy or off-policy RL algorithms.
  • Soft Actor Critic: SAC is an off-policy actor-critic algorithm that learns a stochastic policy and two critics by minimizing Bellman error.Its policy and critics are represented by πψ and Qφ1, Qφ2.
  • Soft Actor Critic: SAC trains from replay-buffer transitions, whose target includes observations, actions, rewards, and termination signals.The transition is t = (o, a, o′, r, d), and B denotes the replay buffer.
  • Rainbow DQN: Rainbow DQN combines Q-learning with a convolutional network and multiple improvements for mapping raw pixels to action values.The cited description presents Rainbow as a collection of improvements applied together on top of Nature DQN.
  • Contrastive Learning: Contrastive learning treats augmented views as positive pairs and other samples as negatives, optimizing similarity between queries and keys.The InfoNCE loss can be interpreted as the log-loss of a K-way softmax classifier whose label is the positive key.

4. CURL Implementation

CURL adds a simple instance-discrimination objective to model-free RL using augmented frame stacks, momentum target encoding, and a learned bilinear similarity.

  • Architecture: CURL trains contrastive learning as an auxiliary loss during each batch update while retaining the base RL algorithm.The experiments pair CURL with SAC for DMControl and data-efficient Rainbow DQN for Atari.
  • Discrimination Objective: CURL performs instance discrimination across temporally consecutive frame stacks rather than single images.This setup is combined with momentum encoding for targets and a bilinear InfoNCE score.
  • Discrimination Objective: CURL uses instance discrimination instead of patch discrimination to minimize architectural adjustments and extra design choices.The paper motivates this choice by the brittleness and dynamically generated datasets of RL training.
  • Query-Key Pair Generation: Random augmentations generate query-key views, while applying the same augmentation across each frame stack preserves temporal structure.The implementation uses different random augmentations for x_q and x_k.
  • Similarity Measure: CURL measures query-key agreement with the learned bilinear product sim(q, k) = qT Wk.The paper reports that this measure outperformed the normalized dot product used by MoCo and SimCLR.
  • Target Encoding with Momentum: The key encoder follows the query encoder through an exponential moving average while gradients update only the query encoder.This momentum target procedure is analogous to MoCo.

5. Experiments

The experiments evaluate CURL’s sample efficiency and performance on pixel-based continuous and discrete control benchmarks. DMControl uses SAC-based comparisons, while Atari uses Rainbow-based comparisons under fixed interaction budgets.

  • Evaluation Protocol: DMControl is evaluated at 100k and 500k environment steps, while Atari is evaluated at 100k interaction steps.Atari100k corresponds to 400k environment steps with an action repeat of 4.
  • Evaluation Protocol: Sample efficiency is measured by the steps required for the best baseline to match CURL at a fixed budget, while performance is the return ratio at that budget.The reported fixed budgets are 100k or 500k steps for DMControl and 100k interaction steps for Atari.
  • Benchmarks: The study targets broad pixel-based control, covering continuous DMControl tasks and discrete Atari games.The stated primary goal is sample-efficient control from pixels across environments.
  • DMControl: DMControl comparisons include SAC with state features and five leading pixel-based methods across six selected environments.CURL is also run on sixteen DMControl environments in total.
  • Atari: Atari comparisons include SimPLe, Rainbow variants, random-agent, and human baselines under the Atari100k protocol.The benchmark compares both model-based and model-free methods.
  • Results: CURL achieves state-of-the-art performance on 5 of 6 reported DMControl500k environments and 7 of 26 Atari100k environments.The Atari table also reports improvements over Efficient Rainbow on 19 of 26 games.

6. Results

CURL delivers strong sample-efficiency across DMControl and Atari, outperforming prior pixel-based methods and often approaching state-based performance.

  • DMControl: 1.9x higher median performance than Dreamer on DMControl100k establishes CURL as state-of-the-art on 5 of 6 benchmarked environments.CURL is also reported as 4.5x more data-efficient than Dreamer.
  • DMControl: CURL nearly matches or sometimes surpasses state-based SAC sample efficiency across most of 16 DMControl environments.At DMControl500k, it matches the median state-based score.
  • DMControl: CURL converges near the optimal score of 1000 on most of 16 DMControl experiments within 500k steps.It also matches the state-based median across six extensively benchmarked environments.
  • Atari: 17.5% median human-normalized score places CURL above SimPLe at 14.4% and Efficient Rainbow DQN at 16.1% on Atari100k.The corresponding mean HNS values are 38.1%, 44.3%, and 28.5%, respectively.
  • Atari: CURL improves over Efficient Rainbow on 19 of 26 Atari games and achieves a 1.3x average improvement across those games.Median improvements over SimPLe and Efficient Rainbow are 1.2x and 1.1x, respectively.
  • Atari: CURL surpasses human performance on JamesBond and Krull, reaching 1.6 and 2.5 human-normalized scores, respectively.

7. Ablation Studies

The paper uses ablations to examine what CURL representations capture, how much RL contributes, and why performance varies across DMControl environments.

  • Ablation Studies: The ablations test whether CURL learns visual features alone or also captures temporal dynamics.
  • Ablation Studies: The studies evaluate policy performance when representations are learned solely with contrastive learning and without an RL signal.
  • Ablation Studies: The ablations investigate why CURL matches state-based RL performance on some DMControl environments but not others.

8. Conclusion

CURL is presented as a contrastive representation-learning method that achieves state-of-the-art data efficiency on pixel-based reinforcement-learning benchmarks.

  • Conclusion: CURL achieves state-of-the-art data efficiency across a diverse set of pixel-based reinforcement-learning benchmark environments.
  • Conclusion: The method uses contrastive learning with minimal architectural changes in a model-free RL pipeline.
  • Conclusion: The authors identify real-world robotics as a potential application where reinforcement-learning data efficiency is important.

A. Implementation Details

CURL couples SAC or Efficient Rainbow with contrastive representation learning using shared encoders, data augmentation, momentum keys, and a bilinear similarity objective. Its implementation emphasizes minimal architectural changes and reproducible training across continuous and discrete control.

  • Architecture and control: CURL couples SAC with a shared encoder for actor and critic representations, while Efficient Rainbow provides the Atari control baseline.The DMControl implementation builds on SAC; Atari experiments use Efficient Rainbow with grayscale observations and random-crop augmentation.
  • Contrastive objective: CURL uses momentum encoding for keys and a learned bilinear inner product to measure query-key agreement.The bilinear similarity is q^T Wk, and the paper reports it outperforming normalized dot-product similarity in an ablation.
  • Optimization and sampling: CURL constructs contrastive batches from replay-buffer observations and can generate arbitrarily many keys from augmented views without additional sampling.Unlike a memory-bank approach, negatives are constructed on the fly for each minibatch.
  • Optimization and sampling: The contrastive and reinforcement-learning objectives use equal weighting and learning rate, avoiding an additional balancing hyperparameter.The implementation describes this as a simplification relative to methods that jointly train auxiliary models such as VAEs.
  • Data augmentation: Random crops are applied consistently across the four-frame stack so the model can correlate spatiotemporal patterns.DMControl observations are rendered at 100 × 100, randomly cropped to 84 × 84 for training, and center-cropped for evaluation.

D. Further Investigation of Data-Efficiency in Contrastive RL

Additional investigations compare CURL with state-based SAC, test temporal discrimination and representation detachment, and examine compute-aware update ratios. CURL often matches state-based data efficiency, while harder settings expose limitations and trade-offs.

  • State-based comparison: CURL matches state-based SAC data efficiency on most of 16 DMControl environments but lags on more challenging tasks.The comparison uses three seeds per selected environment, with more seeds for six environments in the main comparison.
  • State-based comparison: 4.5x is CURL’s average data-efficiency advantage over Dreamer when comparing the steps required to reach CURL’s 100k-step score.The comparison is clipped at 1M steps and uses environments for which Dreamer reports results.
  • Compute trade-offs: A 3:1 gradient-update ratio improves performance for CURL and SLAC relative to 1:1 updates, but adds significant compute and wall-clock overhead.CURL achieves state-of-the-art performance on three of four DMControl500k environments under this ratio.
  • Temporal representations: Temporal discrimination generally outperforms visual-only discrimination, except in reacher and ball-in-cup environments.Those exceptions suggest that learning dynamics is unnecessary in those two environments; walker uses a different action repeat in this ablation.
  • Task-agnostic representations: Detached CNN representations support near-optimal learning on most environments but significantly underperform on cheetah.The authors leave further exploration of task-agnostic representations for future work.

E.4. Removing Data Augmentation for the Actor Critic

Removing augmentation from the actor-critic inputs tests whether contrastively learned features alone improve control. Contrastive features improve pixel SAC substantially, but state-information limits remain on difficult tasks.

  • Contrastive features alone: The no-policy-augmentation ablation updates the contrastive loss on augmented views while computing actor and critic losses on the original observation.The combined loss is the sum of the contrastive, critic, and actor losses.
  • Contrastive features alone: 2x is the average improvement over pixel SAC when augmentations are used only for the contrastive objective.The Figure 10 caption reports a 2.0x mean and 1.7x median improvement on DMControl500k; augmenting SAC inputs further benefits most environments.
  • State information: CURL environments that match state-based SAC have low pixel-to-proprioceptive-state prediction error, while harder environments have higher error.The regression uses stacks of three consecutive frames as inputs and simulator-extracted proprioceptive states as targets.
  • State information: The investigation suggests that degraded performance on challenging tasks may reflect insufficient underlying-state information in the pixels.The authors leave further investigation of this explanation for future work.

G. Connection to work on data augmentations

CURL is positioned alongside augmentation-based methods that optimize control directly rather than adding a contrastive objective. Its reward-independent representations offer a different trade-off for task-agnostic learning and unsupervised pretraining.

  • Comparison with augmentation methods: RAD and DrQ directly feed augmented observations to policy or value networks, whereas CURL adds an auxiliary contrastive consistency objective.The cited comparison describes RAD and DrQ as versions of CURL without the auxiliary contrastive loss.
  • Comparison with augmentation methods: When rich rewards are available and average reward is the objective, RAD and DrQ are likely to work better than CURL.The stated reason is that they directly optimize the objective, while CURL introduces an additional auxiliary consistency objective.
  • Task-agnostic representations: CURL’s contrastive objective is reward-independent, supporting task-agnostic latent spaces and data-efficient controllers across multiple tasks.The detached-encoder ablation provides evidence that simple MLPs can use CURL features without fine-tuning the encoder on many DMControl tasks.
  • Task-agnostic representations: Unsupervised pretraining without rewards followed by RL fine-tuning across tasks is identified as a possible future use of CURL.The passage frames this as a potential scenario for future data-efficient reinforcement learning, especially in real-world settings.
  • Scope of comparison: The paper concludes that the preferable algorithm depends on the researcher’s or practitioner’s objective rather than admitting a single universal winner.The comparison distinguishes reward-focused optimization from reward-independent representation learning.
Loading 2004.04136v4…