Source-linked AI summary
Decoupling Representation Learning from Reinforcement Learning
Adam Stooke, Kimin Lee, Pieter Abbeel, Michael Laskin
TL;DR
Reward-driven feature learning in image-based RL can be limited, motivating separation of representation learning from policy learning. The paper introduces ATC, which learns from augmented temporally adjacent observations with contrastive loss, and finds that ATC-trained encoders generally match or exceed end-to-end RL and other tested UL methods. The results support reward-free representation learning across several RL settings, while transfer across diverse Atari games remains limited and some games still benefit from RL encoder training.
Problem
Reward-driven visual feature learning in deep RL can be limited by sparse rewards and single-task utility, motivating reward-free representation learning.
Method
ATC trains a convolutional encoder to contrast augmented observations with nearby future observations, decoupling encoder training from policy learning.
Results
ATC matches or outperforms end-to-end RL in most tested environments and matches or outperforms prior UL algorithms across all tested domains with frozen pretrained encoders.
Takeaways & Limitations
Reward-free ATC representation learning can replace RL encoder training while maintaining or improving online performance across a broad range of environments.
Takeaways & Limitations
Transfer across diverse Atari games is mixed, and a few Atari environments still perform better when the RL loss trains the encoder.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper decouples visual representation learning from policy learning through reward-free unsupervised learning, introducing ATC for reinforcement learning. Across diverse benchmarks, frozen or UL-trained encoders generally match or outperform end-to-end RL, though some Atari environments still benefit from RL encoder updates.
- Reward-driven visual feature learning can be limited by sparse rewards and by utility restricted to a single task.
- ATC learns visual representations by contrasting augmented observations with nearby future observations, without requiring rewards or actions.
- Online RL with UL: ATC-trained encoders can replace end-to-end RL encoders without degrading policy performance across the tested DMControl and DMLab environments.
- Online RL with UL: ATC achieves nearly equal or greater performance in 5 of 8 Atari games, while auxiliary-loss or initialization use improves the remaining three.
- Encoder Pre-Training Benchmarks: ATC matches or outperforms prior unsupervised-learning algorithms when frozen encoders are pretrained on expert demonstrations across all tested domains.
- Multi-Task Encoders: A single encoder trained on four DMControl environments generalizes successfully to four held-out environments, whereas transfer across eight Atari games is mixed.
- Ablations and Encoder Analysis: Ablations examine ATC components, and subpixel random shift preserves performance while enabling latent-image augmentation for replay benefits.
2. Related Work
Prior RL representation methods commonly use unsupervised objectives jointly with policy learning, while other approaches pretrain representations before control. ATC combines temporal association, augmentation, and contrastive learning in a distinct architecture.
- UNREAL, CPC, and related methods add unsupervised or self-supervised auxiliary losses to jointly train RL agents.
- ST-DIM and DRIML use temporal contrastive objectives, while CURL combines augmentation and a momentum encoder without temporal contrast.
- Other work pretrains visual representations using reconstruction, object-centric, or contrastive forward-model objectives before downstream control.
- ATC associates augmented observations through an encoder, compressor, residual predictor, and contrastive loss with a future observation processed by a momentum encoder.
3. Augmented Temporal Contrast
ATC trains an encoder to associate an augmented observation with one from a nearby future timestep using contrastive learning. Its architecture predicts and contrasts compact codes, while only the convolutional encoder is shared with the RL agent.
- ATC pairs an observation with one from a specified near-future timestep and applies stochastic random-shift augmentation before encoding.
- The convolutional encoder maps the anchor to a latent image, which a compressor converts into a compact code.
- A residual predictor advances the anchor code as an implicit forward model, and a transformation matrix supports contrastive comparison.
- Target encoder parameters are updated by an exponential moving average of the online encoder parameters.
- The convolutional encoder alone is shared with the RL agent, separating representation-learning components from policy learning.
- InfoNCE uses bilinear logits between predicted anchor codes and momentum-encoded positive codes, with other batch positives serving as negatives.
4. Experiments
Across DMControl, DMLab, and Atari, ATC supports detached encoder training, frozen-encoder pre-training, multitask transfer, and augmentation-based efficiency, with performance depending on domain and task sparsity.
- Online RL with ATC: ATC-trained encoders matched or exceeded end-to-end RL across all six DMControl environments and both DMLab environments tested.Sparse-reward DMControl and DMLab tasks particularly benefited from detached representation learning.
- Online RL with ATC: In Atari, detached ATC training worked well in 5 of 8 games, while auxiliary ATC or weight initialization improved performance in the remaining difficult games.Breakout and Space Invaders particularly suffered with fully detached training.
- Encoder Pre-Training Benchmarks: ATC matched or outperformed competing unsupervised methods across all four DMControl test environments and was the only method matching or exceeding end-to-end RL in every case.Comparisons included Augmented Contrast and a time-delayed VAE.
- Encoder Pre-Training Benchmarks: ATC performed significantly better than Pixel Control and CPC in DMLab LASERTAG, although the methods performed similarly in EXPLORE.ATC also avoids actions, deconvolution, and recurrence used by the comparison methods.
- Multi-Task Encoders: A single encoder pretrained on four DMControl environments supported efficient RL in held-out domains, whereas multitask Atari pre-training showed diminished transfer in most games.The Atari limitation was partly alleviated by widening the encoder, indicating representation capacity constraints.
- Ablations and Encoder Analysis: Subpixel random shift restored raw-image augmentation performance on latent images, allowing the pretrained encoder to be bypassed during policy updates; ATC accuracy itself did not strongly predict downstream RL performance.Trajectory-segment batches helped Breakout, while individual transitions were equal or better elsewhere; random shift could improve RL despite reducing contrastive accuracy.
5. Conclusion
The paper presents reward-free representation learning with ATC as a way to maintain or improve online RL performance while decoupling encoder training from policy learning. A few Atari environments remain cases where including the RL loss performs better, leaving opportunities for further UL improvements.
- ATC can fully replace RL for training convolutional encoders while maintaining or improving online performance.
- A few Atari games remain environments where including the RL loss surpasses the UL-only method.
- The authors identify further opportunities to improve unsupervised learning for reinforcement learning.
Algorithm 1
Algorithm 1 alternates environment interaction, policy updates with gradients stopped at the encoder, and ATC updates using sampled anchor-positive observations and a momentum encoder.
- The agent samples environment transitions through the encoder and stores observations in a replay buffer.
- A momentum encoder is initialized and updated using an exponential moving average of the ATC encoder parameters.
- Policy parameters are updated with the chosen on-policy or off-policy RL formula while stopping gradients into the encoder.
- The encoder and contrastive model are updated with ATC gradients from sampled anchor and positive observations.
A.2. Additional Figures
Multi-task encoders show mixed transfer across eight Atari games, while LASERTAG attention maps indicate that some ATC encoders focus on task-relevant objects similarly to RL-trained encoders.
- Multi-Task Atari Encoders: Multi-task encoders with frozen weights produce mixed performance across eight Atari games, with partial improvement from increased network capacity.
- Multi-Task Atari Encoders: Training on seven Atari games and testing on a held-out game yields diminished but non-zero performance, indicating limited feature transfer.
- LASERTAG Attention Maps: The LASERTAG pixel-control encoder focuses on the score, whereas the ATC encoder focuses on the enemy similarly to an RL-trained encoder.
- LASERTAG Attention Maps: The LASERTAG UL encoder focuses properly on the coin, similar to the RL-trained encoder.
A.3. RL Settings
The RL settings specify hyperparameters for DMControl RAD-SAC, Atari PPO, and DMLab PPO experiments, including optimizers, learning rates, convolutional architectures, batch sizes, and discounting.
- The settings include separate hyperparameter tables for DMControl RAD-SAC, Atari PPO, and DMLab PPO.
- DMControl: The DMControl configuration uses Adam optimization, stacked frames, action repeats, and environment-specific learning rates.
- Network Architecture: The encoder uses four convolutional layers with filters [32, 32, 32, 32], strides [2, 2, 2, 1], filter size 3, and latent dimension 50.
- Training Hyperparameters: The configuration specifies batch sizes of 512 for Cheetah and Pendulum and 256 for Rest, with discount γ .99.
A.4. Online ATC Settings
Online ATC comparisons used coarse hyperparameter sweeps, with environment- and algorithm-specific settings selected as the best tested configurations.
- Coarse hyperparameter sweeps tuned each competing unsupervised-learning algorithm, and comparisons used each algorithm’s best tested setting.
- DMControl: In DMControl, VAE used 1-step temporal reconstruction and 5 × 10^4 updates, while ATC used 1-step temporal contrast, 5 × 10^4 updates, and batch size 256.
- DMLab: In DMLab, Pixel Control followed prior settings with offline empirical returns, while CPC selected rollout length 16 after testing batch-time shapes.
- Atari: In Atari, the best tested temporal differences were 3 steps for VAE and ST-DIM and 1 step for the inverse model, with random-shift augmentation helping the latter.