Self-Supervised Structured Representations for Deep Reinforcement Learning

Recent reinforcement learning ( RL ) methods have found extracting high-level features from raw pixels with self-supervised learning to be effective in learning policies . However , these methods focus on learning global representations of images , and disregard local spatial structures present in the consecutively stacked frames . In this paper , we propose a novel approach that learns self-supervised spatial representations ( S3R ) for effectively encoding such spatial structures in an unsupervised manner . Given the input frames , the spatial latent volumes are first generated individually using an encoder , and they are used to capture the change in terms of spatial structures , i.e. , flow maps among multiple frames . To be specific , the proposed method establishes flow vectors between two latent volumes via a supervision by the image reconstruction loss.This enables for providing plenty of local samples for training the encoder of deep RL . We further attempt to leverage the spatial representations in the self-predictive representations ( SPR ) method that predicts future representations using the action-conditioned transition model . The proposed method imposes similarity constraints on the three latent volumes ; warped query representations by estimated flows , predicted target representations from the transition model , and target representations of future state . Experimental results on complex tasks in Atari Games and DeepMind Control Suite demonstrate that the RL methods are significantly boosted by the proposed self-supervised learning of spatial representations . The code is available at https : //sites . google.com/view/iclr2022-s3r . 1 INTRODUCTION . Deep reinforcement learning ( RL ) has been an appealing tool for training agents to solve various tasks including complex control and video games ( François-Lavet et al. , 2018 ) . While most approaches have focused on training deep RL agent under the assumption that compact state representations are readily available , this assumption does not hold in the cases where raw visual observations ( e.g . images ) are used as inputs for training the deep RL agent . Without designing an effective algorithm that uses model-based policy and value improvement operators ( Schrittwieser et al. , 2021 ) , or without attempting to use additional image augmentation ( Laskin et al. , 2020a ; Kostrikov et al. , 2020 ) , learning visual features from raw pixels only using a reward function might fail to learn good features in terms of the performance and sample efficiency . To address this challenge , a number of deep RL approaches ( Sermanet et al. , 2018 ; Dwibedi et al. , 2018 ; Anand et al. , 2019 ; Laskin et al. , 2020b ; Mazoure et al. , 2020 ; Stooke et al. , 2020 ; Schwarzer et al. , 2021 ) leverage the recent advance of self-supervised learning which effectively extracts highlevel features from raw pixels in an unsupervised fashion . In Laskin et al . ( 2020b ) ; Stooke et al . ( 2020 ) , they propose to train the convolutional encoder for pairs of images using a contrastive loss ( van den Oord et al. , 2018 ) . For training the RL agent , given a query and a set of keys consisting of positive and negative samples , they minimize the contrastive loss such that the query matches with the positive sample more than any of the negative samples ( Laskin et al. , 2020b ; Stooke et al. , 2020 ) . While the parameters of the query encoder are updated through back-propagation using the contrastive loss ( van den Oord et al. , 2018 ) , the parameters of the key encoder are computed with an exponential moving average ( EMA ) of the query encoder parameters . The output representations of the query encoder are passed to the RL algorithm for training the agent . Schwarzer et al . ( 2021 ) proposes the self-predictive representations ( SPR ) method that trains the RL agent by leveraging the query and positive sample only , following Grill et al . ( 2020 ) that achieves state-of-the-arts performance without the negative samples in the self-supervised learning . Especially , it extends a model-free RL agent by adding the transition model that predicts its own latent representation multiple steps into the future ( Schwarzer et al. , 2021 ) . The predicted future representations and the representations for future states computed using a target encoder serve as the query and positive sample , respectively . These approaches have shown compelling performance and high sample efficiency on the complex control tasks when compared to existing image-based RL approaches ( Kaiser et al. , 2019 ; van Hasselt et al. , 2019 ; Kielak , 2020 ) . While these approaches can effectively encode the global representations of images with the selfsupervised representation learning , there has been no attention on the local spatial structures present in the consecutively stacked images . Our key observation is that spatial deformation , i.e. , the change in terms of the spatial structures across the consecutive frames , can provide plenty of local samples for training the RL agent . A two-frame flow estimation ( Dosovitskiy et al. , 2015 ; Ilg et al. , 2017 ; Jonschkowski et al. , 2020 ) , which has been widely used for video processing and recognition in computer vision , can be an appropriate tool in modeling the spatial deformation . In this work , we propose a novel approach , termed self-supervised spatial representations ( S3R ) , that learns spatial representations for effectively encoding the spatial structures in a self-supervised fashion . Note that we use the term ‘ spatial representations ’ to indicate a set of feature maps extracted over frames for inferring locally-varying flow maps . The spatial representations generated from an encoder are used to predict the flow maps among the input frames by minimizing an image reconstruction loss in a self-supervised manner . A flow-based warping is then applied to generate future representations . We further extend our framework by leveraging SPR method ( Schwarzer et al. , 2021 ) . As depicted in Figure 1 , we impose similarity constraints on the three representations ; query representations warped by the estimated flow maps , self-predicted target representations from the transition model ( Schwarzer et al. , 2021 ) , and target representations of future state . Note that Shang et al . ( 2021 ) encodes the temporal information by concatenating latent differences of input frames , but the simple substraction operation has certain limitations in effectively capturing the spatial deformation . Contrary to Amiranashvili et al . ( 2018 ) in which the flow map is directly fed into RL algorithms with a stack of images , our method leverages the flow maps to effectively encode the spatial representations of the images and warp spatial representations at the future state . Our contribution is summarized as follows . • Our method learns the spatial representations using the self-supervised flow model for encoding local spatial structures from the consecutive frames used in RL algorithms . • We propose to impose the similarity constraint on the spatial representations as well as the global representations , providing plenty of supervision for training the encoder of deep RL . • We compute the future representations through the flow-based warping for imposing the similarity constraint with target representations . 2 RELATED WORK . Self-supervised Representation Learning : The self-supervised representation learning aims to learn general features from large-scale unlabeled images or videos without expensive data annotations . The contrastive methods have achieved state-of-the-art performance in the self-supervised representation learning . The contrastive learning aims to bring positive samples closer while sep- arating negative samples from each other ( Hadsell et al. , 2006 ) . Wu et al . ( 2018 ) formulate the contrastive learning as a non-parametric classification problem at the instance level , and propose to learn visual features with the memory bank and noise contrastive estimation ( NCE ) ( Gutmann & Hyvärinen , 2010 ; Mnih & Kavukcuoglu , 2013 ) . The method in van den Oord et al . ( 2018 ) proposes a probabilistic contrastive loss , called InfoNCE , for inducing representations by leveraging positive and negative samples . The InfoNCE loss has widely been adopted in Chen et al . ( 2020 ) ; He et al . ( 2020 ) ; Hénaff et al . ( 2020 ) ; Tian et al . ( 2020 ) . Chen et al . ( 2020 ) present a simple framework for contrastive self-supervised learning without specialized architecture ( Bachman et al. , 2019 ; Hénaff et al. , 2020 ) or memory bank ( Wu et al. , 2018 ) , but it requires a large batch size for using enough negative samples when computing the InfoNCE loss ( van den Oord et al. , 2018 ) . He et al . ( 2020 ) propose to build a dynamic dictionary with a queue to avoid the use of large batches when collecting negative samples , and also uses the moving averaged ( momentum ) encoder for target data ( positive and negative samples of query data ) . Grill et al . ( 2020 ) use the momentum encoder to produce representations of the targets as a means of stabilizing the bootstrap step . This enables for learning the representations with only positive samples , which are generated by data augmentation , for a given query without the need to carefully set up negative samples . The method in Chen & He ( 2020 ) further extends this idea by using only stop-gradient operation without using the momentum update . While these approaches focuses on learning global representations of a single image , our method proposes to learn spatial representations for effectively encoding the spatial structures ( i.e. , flow map ) in the consecutive images . Self-supervised Representation Learning in Deep RL : Representation learning is crucial for RL algorithms to learn policies with high-dimensional visual observations . The future prediction conditioned on the past observations and actions serves as auxiliary tasks to improve the sample efficiency of model free RL algorithms . Gelada et al . ( 2019 ) train a transition model to predict representations of future states together with a reward prediction loss . Guo et al . ( 2020 ) present Predictions of Bootstrapped Latents ( PBL ) that builds on multi-step predictive representations of future observations for deep RL . The method in Schwarzer et al . ( 2021 ) propose Self-Predictive Representations ( SPR ) based on an action-conditioned transition model that can predict future representations computed using a target ( momentum ) encoder . Our method attempts to predict the future representations through a warping operation using flow maps computed from the spatial representations of the consecutive frames . Contrastive learning has been used to extract desired latent representations of visual observations used in the RL algorithms . For training robot agents , Sermanet et al . ( 2018 ) present the time-contrastive networks ( TCN ) that train viewpoint-invariant representations using a metric learning . This work was extended in Dwibedi et al . ( 2018 ) by embedding multiple frames at each timestep for learning task-agnostic representations such as position and velocity attributes in continuous control tasks . In Anand et al . ( 2019 ) , the representations for RL algorithms are learned by maximizing mutual information ( Hjelm et al. , 2019 ) across spatially and temporally distinct features of an encoder of visual observations . Schwarzer et al . ( 2021 ) leverage the self-supervised learning ( Grill et al. , 2020 ) for imposing the similarity constraint between self-predictive and target representations . Laskin et al . ( 2020b ) introduce Contrastive Unsupervised representations for Reinforcement Learning ( CURL ) that learns the representations from visual inputs using the InfoNCE loss ( van den Oord et al. , 2018 ) . Stooke et al . ( 2020 ) present Augmented Temporal Contrast ( ATC ) using image augmentations and InfoNCE loss ( van den Oord et al. , 2018 ) for representation learning , and decouples it from policy learning . From a different perspective , Hansen et al . ( 2020 ) propose to adapt the policy network through self-supervised representation learning in unseen environments where it is difficult to predict changed rewards . Our method imposes the similarity constraint on the spatial representations as well as the global representations , thus providing plenty of supervision for training the encoder of deep RL . Visual Correspondence Learning : Visual correspondence estimation is a long-standing research in the computer vision community . It aims to establish a pair of corresponding pixels between two ( or more ) views taken under different locations ( stereo matching ) or timestep ( optical flow ) . Recent methods for stereo matching ( Žbontar & LeCun , 2016 ; Chang & Chen , 2018 ; Zhang et al. , 2019 ) and optical flow estimation ( Dosovitskiy et al. , 2015 ; Ilg et al. , 2017 ; Sun et al. , 2018 ) have been advanced largely thanks to the expressive power of deep networks . Though both approaches share a similar objective of finding corresponding pixels across views , the optical flow is known to be effective for encoding temporal motion trajectories , while the stereo matching is tailored to predicting 3D depth map in the scene . The commonly used architecture for two-frame optical flow estimation involves the feature map extraction of two frames , correlation volume computation , a series of convolutions for refinement , and flow regression . While state-of-the-arts flow estimation methods require using ground truth flow maps as an explicit supervision ( Dosovitskiy et al. , 2015 ; Ilg et al. , 2017 ; Sun et al. , 2018 ) , some unsupervised learning approaches have attempted to infer flow maps with an image reconstruction loss for imposing the constraint that corresponding pixels should have similar intensities ( Ren et al. , 2017 ; Meister et al. , 2018 ; Wang et al. , 2018 ; Jonschkowski et al. , 2020 ) . In our work , we present the self-supervised flow network that learns spatial representations from the consecutive frames used in the RL algorithms .

Summary: This paper proposes a representation learning method that leverages unsupervised signals like flow and forward models and the constraints between them and apply it to the state representation problem of RL.   The method is adding a number of auxiliary rewards to the torso which flow and a latent transition model that are constrained to agree with a contrastive loss. The flow has an architecture inspired by FlowNet and trained with warping reconstruction and a spatial smoothness term. A latent transition model is trained contrastively. Then constraints are added to make the latent of the next step agree with both of these predictions.   The method is run on 13 Atari tasks and compared there with competitive baselines and it largely outperforms them. An ablation without baselines is performed a few dm control suite tasks