Unsupervised Active Pre-Training for Reinforcement Learning

1 INTRODUCTION . Deep reinforcement learning ( RL ) provides a general framework for solving challenging sequential decision-making problems , it has achieved remarkable success in advancing the frontier of AI technologies thanks to scalable and efficient learning algorithms ( Mnih et al. , 2015 ; Lillicrap et al. , 2015 ; Schulman et al. , 2015 ; 2017 ) . These landmarks include outperforming humans in board ( Silver et al. , 2016 ; 2018 ; Schrittwieser et al. , 2019 ) and computer games ( Mnih et al. , 2015 ; Berner et al. , 2019 ; Schrittwieser et al. , 2019 ; Vinyals et al. , 2019 ; Badia et al. , 2020a ) , and solving complex robotic control tasks ( Andrychowicz et al. , 2017 ; Akkaya et al. , 2019 ) . Despite these successes , a key challenge with Deep RL is that it requires a huge amount of interactions with the environment before it learns effective policies , and needs to do so for each encountered task . Environments are required to have carefully designed task-specific reward functions to guide the RL algorithms ( Andrychowicz et al. , 2017 ; Ng et al. , 1999 ) , which further limits its wide applications of Deep RL . This is in contrast to how intelligent creatures learn in the absence of external supervisory signals , acquiring abilities in a task-agnostic manner by exploring the environment . Unsupervised pre-training is a framework that trains models without expert supervision has obtained promising results in computer vision ( Oord et al. , 2018 ; He et al. , 2019 ; Chen et al. , 2020b ; Caron et al. , 2020 ; Grill et al. , 2020 ) and natural language modeling ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ; Peters et al. , 2018 ; Brown et al. , 2020 ) . The key insight of unsupervised pre-training techniques is learning a good representation or initialization from a massive amount of unlabeled data such as ImageNet ( Deng et al. , 2009 ) , Instagram image set ( He et al. , 2019 ) , Wikipedia , and WebText ( Radford et al. , 2019 ) which are easier to collect and scales to millions or trillions of data points . As a result , The learned representation when fine-tuned on the downstream tasks can solve them efficiently without needing any supervision or in a few-shot manner . Driven by the significance of the massive abundance of unlabeled data relative to labeled data , we pose the following question : is enabling efficient unsupervised pretraining for deep RL as easy as increasing the amount of unlabeled data ? Unlike the computer vision or language domains , in reinforcement learning it ’ s not obvious where to extract large pools of unlabeled data . A natural choice is pretraining on ImageNet and transfer the encoder to reinforcement learning tasks . We experimented with using ImageNet data for unsupervised representation learning as initialization of the encoder in deep RL agent , specifically , we used the momentum contrast ( He et al. , 2019 ; Chen et al. , 2020c ) method which is one of the state-of-the-art methods for representation learning . We used DrQ ( Kostrikov et al. , 2020 ) as the RL optimization algorithm . The results on DMControl are shown in Figure 1 . We can see that using ImageNet pre-trained representations does not lead to any significant improvement over training from scratch . We also experimented with using supervised pre-trained ResNet features as initialization similar to Levine et al . ( 2016 ) ( details in Appendix ) but the results are no different . This seems in contrast to the preeminent successes of ImageNet pre-trained models in various computer vision downstream tasks ( see e.g . Krizhevsky et al. , 2012 ; Zeiler & Fergus , 2014 ; Hendrycks et al. , 2019 ; Chen et al. , 2020a ) . On the other hand , previous research in robotics also found that ImageNet pre-training did not help ( Julian et al. , 2020 ) . We hypothesize that the reason for the discrepancy is that the ImageNet data distribution is far from the induced sample distribution encountered during RL training . It is therefore necessary to collect data from the RL agent induced distribution . To investigate this hypothesis , we also experimented with training RL agents by ‘ exhaustively ’ collecting data during the reward-free interaction . Specifically , during pre-training phase , the only reward signal is defined by the count-based exploration ( Bellemare et al. , 2016 ; Ostrovski et al. , 2017 ) which is one of the state-of-the-art methods for exploration ( Taïga et al. , 2019 ) , and the density estimation model is PixelCNN ( Van den Oord et al. , 2016 ) . The results of using the resulting pre-trained policy as initialization are shown in Figure 1 . We can see that pre-trained initialization in Cheetah environment does not improve significantly over random initialization on Cheetah Run task . Similarly , pre-trained initialization in Hopper environment only leads to a small improvement over baseline . The reason for this ineffectiveness is that density modeling at the pixel level is difficult especially in the low-data and non-stationary regime . The results on DMControl demonstrate that simply increasing the amount of unlabeled data does not work well , therefore we need a more systematical strategy that caters to RL . In this paper , we address the issue by proposing to actively collect novel data by exploring unknown areas in the task agnostic environment . Our means is maximizing the entropy of visited state distribution subject to some prior constraints . The entropy maximization principle ( Jaynes , 1957 ) originated in statistical mechanics , where Jaynes showed that entropy in statistical mechanics and information theory were equivalent . Our motivation is that the resulting representation and initialization will encode both prior information while being as agnostic as possible , and can be adapted to various downstream tasks . While the entropy maximization principle seems simple , it is practically difficult to calculate the Shannon entropy ( Shannon , 2001 ) as a density model is needed . To remedy this , we resort to the particle-based entropy estimator ( Singh et al. , 2003 ; Beirlant , 1997 ) which has wide applications in various machine learning areas ( Sricharan et al. , 2013 ; Pál et al. , 2010 ; Jiao et al. , 2018 ) . The particle-based entropy estimator is known to be asymptotically unbiased and consistent . Specifically , it computes the average of the Euclidean distance of each sample to its nearest neighbors . We compute the entropy in the latent representation space , for this we adapt the idea of contrastive learning ( Hadsell et al. , 2006 ; Gutmann & Hyvärinen , 2010 ; Mnih & Kavukcuoglu , 2013 ; He et al. , 2019 ; Chen et al. , 2020b ) to encode image observations to representation space . Our approach alternates between training the encoder via contrastive learning and RL style optimization of maximizing expected reward where reward is defined by the particle-based entropy . After the pre-training phase , we can either fine-tune the encoder representation for test tasks that have different action space dimension or fine-tune the policy initialization for tasks with the same action space dimension . Since our method actively collects data during the pre-training phase , the method is named as Active Pre-Training ( APT ) . We empirically evaluate APT on the Atari game suite and DMControl suite by exposing task-specific reward to the agent after a long unsupervised pre-training phase . On the full suite of Atari games , fine-tuning APT pre-trained models achieves human-level performance on 12 games . On the Atari 100k benchmark ( Kaiser et al. , 2019 ) , our fine-tuning achieves 1.3× higher human median scores than state-of-the-art training from scratch and 4× higher scores than state-of-the-art pre-training RL . On DMControl suite , fine-tuning APT pre-trained models beating all baselines in terms of asymptotic performance and data efficiency and solving tasks that are extremely difficult for training from scratch . The contributions of our paper can be summarized as : ( i ) We propose a new approach for pretraining in RL . ( ii ) We show that our pre-training method significantly improves data efficiency of solving downstream tasks on DMControl and Atari suite . ( iii ) We demonstrate that pre-training with particle-based entropy maximization in contrastive representation space significantly outperforms prior count-based approaches that rely on density modeling . 2 RELATED WORK . 2.1 INTRINSIC MOTIVATION AND EXPLORATION . The learning process of RL agents becomes highly inefficient in sparse supervision tasks when relying on standard exploration techniques . This issue can be alleviated by introducing intrinsic motivation , i.e. , denser reward signals that can be automatically computed . These rewards are generally taskagnostic and might come from state visitation count bonus ( Bellemare et al. , 2016 ; Tang et al. , 2017 ; Ostrovski et al. , 2017 ; Zhao & Tresp , 2019 ) , learning to predict environment dynamics ( Meyer & Wilson , 1991 ; Pathak et al. , 2017 ; Burda et al. , 2018a ; Sekar et al. , 2020 ) , distilling random neural networks ( Burda et al. , 2018b ; Choi et al. , 2018 ) , hindsight relabeling ( Andrychowicz et al. , 2017 ) , learning options ( Sutton et al. , 1999 ) through mutual information ( Jung et al. , 2011 ; Mohamed & Rezende , 2015 ) , information gain ( Lindley , 1956 ; Sun et al. , 2011 ; Houthooft et al. , 2016 ) , successor features ( Kulkarni et al. , 2016 ; Machado et al. , 2018 ) , maximizing mutual information between behaviors and some aspect of the corresponding trajectory ( Gregor et al. , 2016 ; Florensa et al. , 2017 ; Warde-Farley et al. , 2018 ; Hausman et al. , 2018 ; Shyam et al. , 2019 ) , using imitation learning to return to the furthest discovered states ( Ecoffet et al. , 2019 ) , self-play curriculum ( Schmidhuber , 2013 ; Sukhbaatar et al. , 2017 ; Liu et al. , 2019 ) , exploration in latent space ( Vezzani et al. , 2019 ) , injecting noise in parameter space ( Fortunato et al. , 2017 ; Plappert et al. , 2017 ) , learning to imitate self ( Oh et al. , 2018 ) , predicting improvement measure ( Schmidhuber , 1991 ; Oudeyer et al. , 2007 ; Lopes et al. , 2012 ; Achiam & Sastry , 2017 ) , and unsupervised auxiliary task ( Jaderberg et al. , 2016 ) . The work by Badia et al . ( 2020b ) also considers k-nearest neighbor based intrinsic reward to incentive exploration , and shows improved exploration in sparse reward games . Our work differs in that we consider reward-free settings and the objective of our intrinsic reward is based on particle-based entropy instead of count bonus . The work closest to ours is Hazan et al . ( 2019 ) which presents provably efficient exploration algorithms under certain conditions . However , their method directly estimates state visitations through a density model which is difficult to scale . In contrast , our work turns to particle based entropy maximization in contrastive representation space . Concurrent work by Mutti et al . ( 2020 ) shows maximizing particle-based entropy can improve data efficiency in solving downstream continuous control tasks . However , their method relies on importance sampling and on-policy RL which suffers from high variance and is difficult to scale . In contrast , our work resorts to a biased but lower variance entropy estimator which is scalable for high dimensional observations and suitable for off-policy RL optimization .

Summary: This submission presents a technique for unsupervised pre-training of representations and policies for RL. Unsupervised representation learning has obtained impressive results in supervised scenarios, and adapting these methods to RL is an important research direction. One of the main challenges in the RL setting is that of defining the distribution of data to learn from, as well as sampling from it. The learned representations are unlikely to be useful for observations that are out of the pre-training distribution, so it is desirable to perform representation learning on data that is representative of the full state space. Previous works (Hazan et al., Lee et al.) proposed strategies to train agents that induce maximally entropic state visitation distributions, but they involve density estimation whose underlying assumptions are not well suited for pixel observations. The authors propose to overcome these limitations by using a particle based entropy estimate in the learned representation space. The pre-trained representations and policies can be used for RL from pixels, obtaining faster convergence and higher end scores than the considered baselines in both DMControl and Atari.