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: The paper proposes a method to simultaneously learn effective representations and efficient exploration in a reward-free context. The algorithm iterates between minimizing a contrastive loss and maximizing an intrinsic reward derived from a k-NN entropy estimation of the state distribution. Then, authors empirically evaluate the method over a set of visual Mujoco tasks and Atari games.