Decoupling Representation Learning from Reinforcement Learning

By Adam Stooke, Kimin Lee, Pieter Abbeel, Michael Laskin. 2020

402 citations recorded by the source corpus.

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 this https URL.

Subfields: Artificial Intelligence

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

Citations in the corpus

Hyperparameter Selection for Imitation Learning (2021, methodology): “While this is not directly applicable to vision-based observations, there are plenty of techniques which can be used to compute vector embedding for vision observations in RL/IL, e.g. Stooke et al. (2020); Lee et al. (2020).”
Representation Matters: Offline Pretraining for Sequential Decision Making (2021, methodology): “…is used, we apply the contrastive loss between f ( φ ( s t )) and φ target ( s t + i ), where f is a learned function and φ target denotes a non-trainable version of φ , with weights corresponding to a slowly moving average of the weights of φ , as in Stooke et al. (2020); He et al. (2020).”
Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning (2021, methodology): “The use of such data augmentations now forms an essential component of several recent visual RL algorithms [Srinivas et al., 2020, Raileanu et al., 2020, Yarats et al., 2021a, Stooke et al., 2020, Hansen and Wang, 2021, Schwarzer et al., 2020b].”
Contrastive Learning as Goal-Conditioned Reinforcement Learning (2022, methodology): “While such representations might emerge from end-to-end training [7, 79, 118, 125], prior work has found it necessary to equip RL algorithms with perception-specific loss functions [31, 43, 70, 74, 89, 91, 101, 140] or data augmentations [72, 74, 117, 136], effectively decoupling the representation ”
Learning Markov State Abstractions for Deep Reinforcement Learning (2021, background): “…loss objectives typically aim to distinguish either sequential states from non-sequential ones (Shelhamer et al., 2016; Anand et al., 2019; Stooke et al., 2020), real states from predicted ones (Van den Oord et al., 2018), or determine whether two augmented views came from the same or…”
Model-augmented Prioritized Experience Replay (2022, methodology): “…et al., 2017; Shelhamer et al., 2016) or to enhance representation learning for pixel-based observations (Lin et al., 2019; Agarwal et al., 2021; Stooke et al., 2021; Lee et al., 2020b; Schrittwieser et al., 2020; Hessel et al., 2021; Hafner et al., 2020b; Srinivas et al., 2020; Guo et al.,…”
Bayesian Robust Optimization for Imitation Learning (2020, background): “Note that the assumption of a linear reward function is not necessarily restrictive as these features can be arbitrarily complex nonlinear functions of the state and could be obtained via unsupervised learning from raw state observations [12, 13, 55].”
PI-QT-Opt: Predictive Information Improves Multi-Task Robotic Reinforcement Learning at Scale (2022, background): “While existing studies of learning predictive information representations in RL [2, 3, 4, 8, 9] have largely been limited to learning single-task specialist agents in simulated environments such as DM-Control [10] and Atari [11], our hypotheses can be seen as extending the generalization results in ”
Contrastive Representation Learning: A Framework and Review (2020, neutral): “[94] proposed the Augmented Temporal Contrast (ACT) for pre-training representations that are transferable to multiple RL tasks.”
Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning (2021, background): “…learning has been applied for imposing state self-consistency (Laskin et al., 2020b), capturing predictive information (Oord et al., 2018; Mazoure et al., 2020; Lee et al., 2020b) or encoding transition dynamics (van der Pol et al., 2020; Stooke et al., 2020; Schwarzer et al., 2020) within an MDP.”
Adapt to Environment Sudden Changes by Learning a Context Sensitive Policy (2022, methodology): “Self-Supervised Representation Learning has been used in several RL domains in recent work (Laskin, Srinivas, and Abbeel 2020; Zhang et al. 2021; Stooke et al. 2021).”
Curious Representation Learning for Embodied Intelligence (2021, methodology): “Subsequently, we find that weights from either PointNav weights or ATC perform better than training a policy from scratch.”
Representation Learning for Online and Offline RL in Low-rank MDPs (2021, background): “…several empirical works have shown that leveraging representation learning techniques developed in supervised or unsupervised learning settings can accelerate the search for good decision-making strategies (Silver et al., 2018; Stooke et al., 2021; Srinivas et al., 2020; Yang and Nachum, 2021).”
A Framework for Efficient Robotic Manipulation (2020, background): “R O ] 14 D ec 2020 progress in unsupervised representation learning [31], [32] and data augmentation [10], [33] has signiﬁcantly improved the efﬁciency of learning with RL in simulated robotic [21] and video game [34] environments.”
IKEA Furniture Assembly Environment for Long-Horizon Complex Manipulation Tasks (2019, background): “Moreover, recent progress in vision-based RL methods [31]–[33] shows comparable sample efficiency to statebased RL.”
Provably efficient representation selection in Low-rank Markov Decision Processes: from online to offline RL (2021, methodology): “In deep RL, deep neural networks have been used to learn good representations in different settings [Diuk et al., 2008, Stooke et al., 2020, Yang et al., 2020].”
Continual Learning in Deep Networks: an Analysis of the Last Layer (2021, background): “Decoupling the feature extractor from a sub-network has shown to be useful, for example, in reinforcement learning [29, 37, 48].”