[Paper Review] Decoupling Representation Learning from Reinforcement Learning

Adam Stooke, Kimin Lee+2 more|arXiv (Cornell University)|Sep 14, 2020

Reinforcement Learning in Robotics39 references65 citations

TL;DR

The paper introduces Augmented Temporal Contrast (ATC), an unsupervised learning task that decouples representation learning from policy learning in RL, and shows ATC-trained encoders can match or exceed end-to-end RL across multiple environments and tasks.

ABSTRACT

In an effort to overcome limitations of reward-driven feature learning in deep reinforcement learning (RL) from images, we propose decoupling representation learning from policy learning. To this end, we introduce a new unsupervised learning (UL) task, called Augmented Temporal Contrast (ATC), which trains a convolutional encoder to associate pairs of observations separated by a short time difference, under image augmentations and using a contrastive loss. In online RL experiments, we show that training the encoder exclusively using ATC matches or outperforms end-to-end RL in most environments. Additionally, we benchmark several leading UL algorithms by pre-training encoders on expert demonstrations and using them, with weights frozen, in RL agents; we find that agents using ATC-trained encoders outperform all others. We also train multi-task encoders on data from multiple environments and show generalization to different downstream RL tasks. Finally, we ablate components of ATC, and introduce a new data augmentation to enable replay of (compressed) latent images from pre-trained encoders when RL requires augmentation. Our experiments span visually diverse RL benchmarks in DeepMind Control, DeepMind Lab, and Atari, and our complete code is available at https://github.com/astooke/rlpyt/tree/master/rlpyt/ul.

Motivation & Objective

• Motivate learning reward-agnostic visual representations for RL to overcome limitations of reward-driven feature learning.
• Propose ATC as an unsupervised task to learn robust encoders from observations independent of RL losses.
• Demonstrate online RL with ATC-encoded representations across diverse environments.
• Benchmark ATC against other unsupervised learning methods and evaluate multi-task generalization.
• Explore ablations and data augmentation effects to understand ATC components.

Proposed method

• Introduce Augmented Temporal Contrast (ATC): contrastive task that associates observations o_t with near-future o_{t+k} within a trajectory.
• Encode augmented observations via a shared CNN encoder; compress to latent code with a linear compressor; predict forward in latent space with a residual predictor.
• Use a momentum encoder for positives and apply InfoNCE loss with negatives from the batch.
• Apply stochastic data augmentation (random shift) to observations; employ a predictor layer to process the anchor code before contrasting.
• Train the encoder with ATC offline (unsupervised) while policy learning occurs solely on latent representations; optionally freeze encoder weights during RL in baselines.
• Optionally introduce a new augmentation (subpixel random shift) to enable latent-image replay and reduce computation.

Experimental results

Research questions

• RQ1 Can representation learning be decoupled from reinforcement learning without sacrificing policy performance?
• RQ2 Is ATC a superior unsupervised task for pretraining encoders for RL compared to existing UL methods?
• RQ3 Do multi-task encoders learned via ATC transfer to new environments or tasks?
• RQ4 What ablations and augmentations are critical for ATC effectiveness?
• RQ5 How does ATC perform across a diverse set of benchmarks (DMControl, DMLab, Atari) and RL algorithms?

Key findings

• ATC-trained encoders, when used online and detached from RL gradients, match or outperform end-to-end RL encoders in most DMControl and DMLab environments and in over half of Atari games tested.
• Pre-trained ATC encoders with frozen weights often surpass other leading UL algorithms in DMControl, DMLab, and Atari benchmarks.
• A single ATC encoder trained on multiple DMControl environments can generalize to new downstream tasks with frozen weights.
• ATC ablations show that data augmentation is necessary for DMControl, and subpixel random shift enables latent-image replay without extra computation.
• Multi-task pretraining across environments can improve transfer for some tasks, though Atari results show limited cross-game transfer, indicating representation capacity and domain differences matter.
• Using ATC as an auxiliary loss or for weight initialization further boosts performance in several Atari games.