Decoupling Representation Learning from Reinforcement Learning

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 hiddenurl . 1 INTRODUCTION . Ever since the first fully-learned approach succeeded at playing Atari games from screen images ( Mnih et al. , 2015 ) , standard practice in deep reinforcement learning ( RL ) has been to learn visual features and a control policy jointly , end-to-end . Several such deep RL algorithms have matured ( Hessel et al. , 2018 ; Schulman et al. , 2017 ; Mnih et al. , 2016 ; Haarnoja et al. , 2018 ) and have been successfully applied to domains ranging from real-world ( Levine et al. , 2016 ; Kalashnikov et al. , 2018 ) and simulated robotics ( Lee et al. , 2019 ; Laskin et al. , 2020a ; Hafner et al. , 2020 ) to sophisticated video games ( Berner et al. , 2019 ; Jaderberg et al. , 2019 ) , and even high-fidelity driving simulators ( Dosovitskiy et al. , 2017 ) . While the simplicity of end-to-end methods is appealing , relying on the reward function to learn visual features can be severely limiting . For example , it leaves features difficult to acquire under sparse rewards , and it can narrow their utility to a single task . Although our intent is broader than to focus on either sparse-reward or multi-task settings , they arise naturally in our studies . We investigate how to learn visual representations which are agnostic to rewards , without degrading the control policy . A number of recent works have significantly improved RL performance by introducing auxiliary losses , which are unsupervised tasks that provide feature-learning signal to the convolution neural network ( CNN ) encoder , additionally to the RL loss ( Jaderberg et al. , 2017 ; van den Oord et al. , 2018 ; Laskin et al. , 2020b ; Guo et al. , 2020 ; Schwarzer et al. , 2020 ) . Meanwhile , in the field of computer vision , recent efforts in unsupervised and self-supervised learning ( Chen et al. , 2020 ; Grill et al. , 2020 ; He et al. , 2019 ) have demonstrated that powerful feature extractors can be learned without labels , as evidenced by their usefulness for downstream tasks such as ImageNet classification . Together , these advances suggest that visual features for RL could possibly be learned entirely without rewards , which would grant greater flexibility to improve overall learning performance . To our knowledge , however , no single unsupervised learning ( UL ) task has been shown adequate for this purpose in general vision-based environments . In this paper , we demonstrate the first decoupling of representation learning from reinforcement learning that performs as well as or better than end-to-end RL . We update the encoder weights using only UL and train a control policy independently , on the ( compressed ) latent images . This capability stands in contrast to previous state-of-the-art methods , which have trained the UL and RL objectives jointly , or Laskin et al . ( 2020b ) , which observed diminished performance with decoupled encoders . Our main enabling contribution is a new unsupervised task tailored to reinforcement learning , which we call Augmented Temporal Contrast ( ATC ) . ATC requires a model to associate observations from nearby time steps within the same trajectory ( Anand et al. , 2019 ) . Observations are encoded via a convolutional neural network ( shared with the RL agent ) into a small latent space , where the InfoNCE loss is applied ( van den Oord et al. , 2018 ) . Within each randomly sampled training batch , the positive observation , ot+k , for every anchor , ot , serves as negative for all other anchors . For regularization , observations undergo stochastic data augmentation ( Laskin et al. , 2020b ) prior to encoding , namely random shift ( Kostrikov et al. , 2020 ) , and a momentum encoder ( He et al. , 2020 ; Laskin et al. , 2020b ) is used to process the positives . A learned predictor layer further processes the anchor code ( Grill et al. , 2020 ; Chen et al. , 2020 ) prior to contrasting . In summary , our algorithm is a novel combination of elements that enables generic learning of the structure of observations and transitions in MDPs without requiring rewards or actions as input . We include extensive experimental studies establishing the effectiveness of our algorithm in a visually diverse range of common RL environments : DeepMind Control Suite ( DMControl ; Tassa et al . 2018 ) , DeepMind Lab ( DMLab ; Beattie et al . 2016 ) , and Atari ( Bellemare et al. , 2013 ) . Our experiments span discrete and continuous control , 2D and 3D visuals , and both on-policy and off policy RL algorithms . Complete code for all of our experiments is available at hiddenurl . Our empirical contributions are summarized as follows : Online RL with UL : We find that the convolutional encoder trained solely with the unsupervised ATC objective can fully replace the end-to-end RL encoder without degrading policy performance . ATC achieves nearly equal or greater performance in all DMControl and DMLab environments tested and in 5 of the 8 Atari games tested . In the other 3 Atari games , using ATC as an auxiliary loss or for weight initialization still brings improvements over end-to-end RL . Encoder Pre-Training Benchmarks : We pre-train the convolutional encoder to convergence on expert demonstrations , and evaluate it by training an RL agent using the encoder with weights frozen . We find that ATC matches or outperforms all prior UL algorithms as tested across all domains , demonstrating that ATC is a state-of-the-art UL algorithm for RL . Multi-Task Encoders : An encoder is trained on demonstrations from multiple environments , and is evaluated , with weights frozen , in separate downstream RL agents . A single encoder trained on four DMControl environments generalizes successfully , performing equal or better than end-to-end RL in four held-out environments . Similar attempts to generalize across eight diverse Atari games result in mixed performance , confirming some limited feature sharing among games . Ablations and Encoder Analysis : Components of ATC are ablated , showing their individual effects . Additionally , data augmentation is shown to be necessary in DMControl during RL even when using a frozen encoder . We introduce a new augmentation , subpixel random shift , which matches performance while augmenting the latent images , unlocking computation and memory benefits . 2 RELATED WORK . Several recent works have used unsupervised/self-supervised representation learning methods to improve performance in RL . The UNREAL agent ( Jaderberg et al. , 2017 ) introduced unsupervised auxiliary tasks to deep RL , including the Pixel Control task , a Q-learning method requiring predictions of screen changes in discrete control environments , which has become a standard in DMLab ( Hessel et al. , 2019 ) . CPC ( van den Oord et al. , 2018 ) applied contrastive losses over multiple time steps as an auxiliary task for the convolutional and recurrent layers of RL agents , and it has been extended with future action-conditioning ( Guo et al. , 2018 ) . Recently , PBL ( Guo et al. , 2020 ) surpassed these methods with an auxiliary loss of forward and backward predictions in the recurrent latent space using partial agent histories . Where the trend is of increasing sophistication in auxiliary recurrent architectures , our algorithm is markedly simpler , requiring only observations , and yet it proves sufficient in partially observed settings ( POMDPs ) . ST-DIM ( Anand et al. , 2019 ) introduced various temporal , contrastive losses , including ones that operate on “ local ” features from an intermediate layer within the encoder , without data augmentation . CURL ( Laskin et al. , 2020b ) introduced an augmented , contrastive auxiliary task similar to ours , including a momentum encoder but without temporal contrast . Mazoure et al . ( 2020 ) provided extensive analysis pertaining to InfoNCE losses on functions of successive time steps in MDPs , including local features in their auxiliary loss ( DRIML ) similar to ST-DIM , and finally conducted experiments using global temporal contrast of augmented observations in the Procgen ( Cobbe et al. , 2019 ) environment . Most recently , MPR ( Schwarzer et al. , 2020 ) combined data augmentation with multi-step , convolutional forward modeling and a similarity loss to improve DQN agents in the Atari 100k benchmark . Hafner et al . ( 2019 ; 2020 ) ; Lee et al . ( 2019 ) proposed to leverage world-modeling in a latent-space for continuous control . A small number of model-free methods have attempted to decouple encoder training from the RL loss as ablations , but have met reduced performance relative to end-to-end RL ( Laskin et al. , 2020b ; Lee et al. , 2020 ) . None have previously been shown effective in as diverse a collection of RL environments as ours ( Bellemare et al. , 2013 ; Tassa et al. , 2018 ; Beattie et al. , 2016 ) . Finn et al . ( 2016 ) ; Ha & Schmidhuber ( 2018 ) are example works which pretrained encoder features in advance using image reconstruction losses such as the VAE ( Kingma & Welling , 2013 ) . Devin et al . ( 2018 ) ; Kipf et al . ( 2019 ) pretrained object-centric representations , the latter learning a forward model by way of contrastive losses ; Yan et al . ( 2020 ) introduced a similar technique to learn encoders supporting manipulation of deformable objects by traditional control methods . MERLIN ( Wayne et al. , 2018 ) trained a convolutional encoder and sophisticated memory module online , detached from the RL agent , which learned read-only accesses to memory . It used reconstruction and one-step latent-prediction losses and achieved high performance in DMLab-like environments with extreme partial observability . Our loss function may benefit those settings , as it outperforms similar reconstruction losses in our experiments . Decoupling unsupervised pretraining from downstream tasks is common in computer vision ( Hénaff et al. , 2019 ; He et al. , 2019 ; Chen et al. , 2020 ) and has favorable properties of providing task agnostic features which can be used for training smaller taskspecific networks , yielding significant gains in computational efficiency over end-to-end methods . 3 AUGMENTED TEMPORAL CONTRAST . Our unsupervised learning task , Augmented Temporal Contrast ( ATC ) , requires a model to associate an observation , ot , with one from a specified , near-future time step , ot+k . Within each training batch , we apply stochastic data augmentation to the observations ( Laskin et al. , 2020b ) , namely random shift ( Kostrikov et al. , 2020 ) , which is simple to implement and provides highly effective regularization in most cases . The augmented observations are encoded into a small latent space where a contrastive loss is applied . This task encourages the learned encoder to extract meaningful elements of the structure of the MDP from observations . Our architecture for ATC consists of four learned components - ( i ) a convolutional encoder , fθ , which processes the anchor observation , ot , into the latent image zt = fθ ( AUG ( ot ) ) , ( ii ) a linear global compressor , gφ to produce a small latent code vector ct = gφ ( zt ) , ( iii ) a residual predictor MLP , hψ , which acts as an implicit forward model to advance the code pt = hψ ( ct ) + ct , and ( iv ) a contrastive transformation matrix , W . To process the positive observation , ot+k into the target code c̄t+k = gφ̄ ( fθ̄ ( AUG ( ot+k ) ) , we use a momentum encoder ( He et al. , 2019 ) parameterized as a slowly moving average of the weights from the learned encoder and compressor layer : θ̄ ← ( 1− τ ) θ̄ + τθ ; φ̄← ( 1− τ ) φ̄+ τφ . ( 1 ) The complete architecture is shown in Figure 1 . The convolutional encoder , fθ , alone is shared with the RL agent . We employ the InfoNCE loss ( Gutmann & Hyvärinen , 2010 ; van den Oord et al. , 2018 ) using logits computed bilinearly , as l = ptWc̄t+k . In our implementation , every anchor in the training batch utilizes the positives corresponding to all other anchors as its negative examples . Denoting an observation indexed from dataset O as oi , and its positive as oi+ , the logits can be written as li , j+ = piWc̄j+ ; our loss function in practice is : LATC = −EO [ log exp li , i+∑ oj∈O exp li , j+ ] . ( 2 )

Summary: This paper presents a new unsupervised learning method for learning latent representations for visual RL control domains. The method, Augmented Temporal Contrast (ATC), can be used alone to learn a representation to be combined with an RL algorithm, or as an auxiliary task in an end-to-end system. ATC matches or outperforms comparable end-to-end systems in several environments. The paper provides an extensive experimental study to support its claims.