Decoupling Representation Learning from Reinforcement Learning

2020/09/14 by Adam Stooke, Stooke, Adam, Kimin Lee +5 · 65 citations
Computer Science · Mathematics · #Adversarial Robustness in Machine Learning #Artificial Intelligence (cs.AI) #Artificial intelligence #Benchmark (surveying) #Computer Vision and Pattern Recognition (cs.CV) #Computer science #Decoupling (probability) #Domain Adaptation and Few-Shot Learning #Encoder #FOS: Computer and information sciences #Feature learning #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Machine learning #Pattern recognition (psychology) #Reinforcement Learning in Robotics #Reinforcement learning #Task (project management) #cs.AI #cs.CV #cs.LG #stat.ML

paper · pdf · doi:10.48550/arxiv.2009.08319

published in arXiv (Cornell University), 9870-9879 (Cornell University) · Improved related works and fixed code hyperlink

openalex publication_date 2020/09/14 · arxiv created 2021/05/16 · arxiv updated 2021/05/18 · openalex created_date 2025/10/10 · openalex updated_date 2026/08/04

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.

Citations

• Reinforcement Learning with Unsupervised Auxiliary Tasks
• DeepMind Control Suite
• QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation
• Data-Efficient Image Recognition with Contrastive Predictive Coding
• Leveraging Procedural Generation to Benchmark Reinforcement Learning
• A Simple Framework for Contrastive Learning of Visual Representations
• Bootstrap Latent-Predictive Representations for Multitask Reinforcement Learning
• Momentum Contrast for Unsupervised Visual Representation Learning
• Data-Efficient Reinforcement Learning with Self-Predictive\n Representations

Cited by

• Videos are Sample-Efficient Supervisions: Behavior Cloning from Videos via Latent Representations
• Spectral Representation-based Reinforcement Learning
• PvP: Data-Efficient Humanoid Robot Learning with Proprioceptive-Privileged Contrastive Representations
• World Models for Autonomous Navigation of Terrestrial Robots from LIDAR Observations
• Enhancing Tactile-based Reinforcement Learning for Robotic Control
• Oracle-Guided Masked Contrastive Reinforcement Learning for Visuomotor Policies
• Generalization in Reinforcement Learning by Soft Data Augmentation
• Training Larger Networks for Deep Reinforcement Learning
• Mutual Information Tracks Policy Coherence in Reinforcement Learning
• Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation
• Autonomous Learning From Success and Failure: Goal-Conditioned Supervised Learning with Negative Feedback
• Learning Markov State Abstractions for Deep Reinforcement Learning
• SECANT: Self-Expert Cloning for Zero-Shot Generalization of Visual Policies
• Behavior From the Void: Unsupervised Active Pre-Training
• Stackelberg Coupling of Online Representation Learning and Reinforcement Learning
• In-Context Reinforcement Learning via Communicative World Models
• Revisiting Bisimulation Metric for Robust Representations in Reinforcement Learning
• Continual Learning in Deep Networks: an Analysis of the Last Layer
• Towards Robust Bisimulation Metric Learning
• Combining Pre-Trained Models for Enhanced Feature Representation in Reinforcement Learning
• TD-MPC-Opt: Distilling Model-Based Multi-Task Reinforcement Learning Agents
• A Survey of State Representation Learning for Deep Reinforcement Learning
• Pretraining Representations for Data-Efficient Reinforcement Learning
• Automatic Data Augmentation for Generalization in Deep Reinforcement Learning
• Action-Adaptive Continual Learning: Enabling Policy Generalization under Dynamic Action Spaces
• Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning
• Policy Search, Retrieval, and Composition via Task Similarity in Collaborative Agentic Systems
• Merging and Disentangling Views in Visual Reinforcement Learning for Robotic Manipulation
• Learning Visual Robotic Control Efficiently with Contrastive Pre-training and Data Augmentation
• Return-based Scaling: Yet Another Normalisation Trick for Deep RL
• Normalizing Flows are Capable Models for Continuous Control
• Imagine Beyond! Distributionally Robust Auto-Encoding for State Space Coverage in Online Reinforcement Learning
• Which Mutual-Information Representation Learning Objectives are Sufficient for Control?
• Accelerating Visual-Policy Learning through Parallel Differentiable Simulation
• DSADF: Thinking Fast and Slow for Decision Making
• Reinforcement Learning with Latent Flow
• MInCo: Mitigating Information Conflicts in Distracted Visual Model-based Reinforcement Learning

Related

• CRC-RL: A Novel Visual Feature Representation Architecture for Unsupervised Reinforcement Learning
• Temporal Disentanglement of Representations for Improved Generalisation in Reinforcement Learning
• TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning
• DRIBO: Robust Deep Reinforcement Learning via Multi-View Information Bottleneck
• Masked Contrastive Representation Learning for Reinforcement Learning