[Paper Review] Return-Based Contrastive Representation Learning for Reinforcement Learning

Guoqing Liu, Chuheng Zhang+6 more|arXiv (Cornell University)|Feb 22, 2021

Reinforcement Learning in RoboticsComputer Science46 references18 citations

TL;DR

This paper proposes Return-Based Contrastive Representation Learning (RCRL), a novel contrastive auxiliary task in deep reinforcement learning that uses return signals to train representations to distinguish state-action pairs with similar versus different returns. RCRL improves sample efficiency and performance in low-data regimes on Atari and DMControl benchmarks, outperforming strong baselines and synergizing with existing methods.

ABSTRACT

Recently, various auxiliary tasks have been proposed to accelerate representation learning and improve sample efficiency in deep reinforcement learning (RL). However, existing auxiliary tasks do not take the characteristics of RL problems into consideration and are unsupervised. By leveraging returns, the most important feedback signals in RL, we propose a novel auxiliary task that forces the learnt representations to discriminate state-action pairs with different returns. Our auxiliary loss is theoretically justified to learn representations that capture the structure of a new form of state-action abstraction, under which state-action pairs with similar return distributions are aggregated together. In low data regime, our algorithm outperforms strong baselines on complex tasks in Atari games and DeepMind Control suite, and achieves even better performance when combined with existing auxiliary tasks.

Motivation & Objective

• To address the sample inefficiency of deep reinforcement learning in low-data regimes by introducing a return-aware auxiliary task.
• To develop a theoretically grounded representation learning method that captures the structure of a new state-action abstraction based on return distributions.
• To improve policy learning by training representations to be invariant to return-irrelevant features while preserving return-relevant distinctions.
• To demonstrate that return-based contrastive learning can be combined effectively with existing unsupervised contrastive methods for further performance gains.

Proposed method

• RCRL constructs positive and negative samples based on return similarity: positive pairs share the same or similar return, negative pairs have dissimilar returns.
• It trains a contrastive loss using a momentum encoder and a predictor head to distinguish between positive and negative state-action representations.
• The method leverages Z^π-irrelevance abstraction, a novel state-action abstraction that groups state-action pairs with similar return distributions under a policy π.
• Z-learning is proposed to approximate Z^π-irrelevance using sampled returns rather than full return distributions, enabling practical training.
• The auxiliary loss is integrated with main RL algorithms (e.g., Rainbow, SAC) using equal weighting and shared learning rates, reducing hyperparameter sensitivity.
• Positive pairs are collected from consecutive segments in trajectories to ensure return similarity, and negative pairs are sampled randomly from replay memory.

Experimental results

Research questions

• RQ1 Can return-based contrastive learning improve representation quality and sample efficiency in deep reinforcement learning?
• RQ2 Does the proposed Z^π-irrelevance abstraction reduce the effective size of the state-action space while preserving Q-value accuracy?
• RQ3 How does RCRL perform in low-data regimes compared to strong baselines on Atari and DMControl environments?
• RQ4 Can RCRL be effectively combined with existing unsupervised contrastive methods to further improve performance?

Key findings

• RCRL outperforms strong baselines on 53 out of 55 Atari games in the low-data regime (100k interactions), achieving a median human-normalized score of 182.4%.
• On the DMControl suite, RCRL achieves superior performance compared to baselines, with significant gains in complex continuous control tasks.
• In high-data regimes (1.5M interactions), RCRL still outperforms the baseline on 4 out of 5 Atari games, indicating broad applicability.
• Representation analysis shows that RCRL learns representations where similar return distributions yield similar embeddings, and dissimilar returns yield dissimilar ones, accelerating generalization.
• The method is robust to hyperparameter tuning, using equal loss weights and learning rates for both auxiliary and main RL tasks without performance degradation.
• When combined with CURL, RCRL achieves even better performance, demonstrating compatibility and synergy with existing unsupervised contrastive methods.