Behavior From the Void: Unsupervised Active Pre-Training

1 Introduction . Reinforcement learning ( RL ) provides a general framework for solving challenging sequential decision-making problems . When combined with function approximation , it has achieved remarkable success in advancing the frontier of AI technologies . These landmarks include outperforming humans in computer games [ 40 , 51 , 64 , 5 ] and solving complex robotic control tasks [ 3 , 1 ] . Despite these successes , they have to train from scratch to maximize extrinsic reward for every encountered task . This is in sharp contrast with how intelligent creatures quickly adapt to new tasks by leveraging previously acquired behaviors . Unsupervised pre-training , a framework that trains models without expert supervision , has obtained promising results in computer vision [ 43 , 23 , 14 ] and natural language modeling [ 63 , 16 , 11 ] . The learned representation , when fine-tuned on the downstream tasks , can solve them efficiently in a few-shot manner . With the models and datasets growing , performance continues to improve predictably according to scaling laws . Driven by the significance of massive unlabeled data , we consider an analogy setting of unsupervised pre-training in computer vision where labels are removed during training . The goal of pre-training is to have data efficient adaptation for some downstream task defined in the form of rewards . In RL with unsupervised pre-training , the agent is allowed to train for a long period without access to environment reward , and then only gets exposed to the reward during testing . We first test an array of existing methods for unsupervised pre-training to identity which gaps and challenges exist , we evaluate count-based bonus [ 10 ] , which encourages the agent to visit novel states . We apply count-based bonus to DrQ [ 33 ] which is current state-of-the-art RL for training from pixels . We also evaluate ImageNet pre-trained representations . The results are shown in Figure 1 . We can see that count-based bonus fails to outperform train DrQ from scratch . We hypothesize that the ineffectiveness stems from density modeling at the pixel level being difficult . ImageNet pre-training does not outperform training from scratch either , which has also been shown in previous research 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . in real world robotics [ 29 ] . We believe the reason is that neither of existing methods can provide enough diverse data . Count-based exploration faces the difficult of estimating high dimensional data density while ImageNet dataset is out-of-distribution for DMControl . To address the issue of obtaining diverse data for RL with unsupervised pre-training , we propose to actively collect novel data by exploring unknown areas in the task-agnostic environment . The underlying intuition is that a general exploration strategy has to visit , with high probability , any state where the agent might be rewarded in a subsequent RL task . Concretely , our approach relies on the entropy maximization principle [ 27 , 53 ] . Our hope is that by doing so , the learned behavior and representation can be trained on the whole environment while being as task agnostic as possible . Since entropy maximization in high dimensional state space is intractable as an oracle density model is not available , we resort to the particle-based entropy estimator [ 55 , 8 ] . This estimator is nonparametric and asymptotically unbiased . The key idea is computing the average of the Euclidean distance of each particle to its nearest neighbors for a set of samples . We consider an abstract representation space in order to make the distance meaningful . To learn such a representation space , we adapt the idea of contrastive representation learning [ 14 ] to encode image observations to a lower dimensional space . Building upon this insight , we propose Unsupervised Active Pre-Training ( APT ) since the agent is encouraged to actively explore and leverage the experience to learn behavior . Our approach can be applied to a wide-range of existing RL algorithms . In this paper we consider applying our approach to DrQ [ 33 ] which is a state-of-the-art visual RL algorithm . On the Atari 26 games subset , APT significantly improves DrQ ’ s data-efficiency , achieving 54 % relative improvement . On the full suite of Atari 57 games [ 40 ] , APT significantly outperforms prior state-of-the-art , achieving a median human-normalized score 3× higher than the highest score achieved by prior unsupervised RL methods and DQN . On DeepMind control suite , APT beats DrQ and unsupervised RL in terms of asymptotic performance and data efficiency and solving tasks that are extremely difficult to train from scratch . The contributions of our paper can be summarized as : ( i ) We propose a new approach for unsupervised pre-training for visual RL based a nonparametric particle-based entropy maximization . ( ii ) We show that our pre-training method significantly improves data efficiency of solving downstream tasks on DMControl and Atari suite . 2 Problem Setting . Reinforcement Learning ( RL ) An agent interacts with its uncertain environment over discrete timesteps and collects reward per action , modeled as a Markov Decision Process ( MDP ) [ 48 ] , defined by 〈S , A , T , ρ0 , r , γ〉 where S ⊆ R nS is a set of nS-dimensional states , A ⊆ R nA is a set of nAdimensional actions , T : S ×A×S → [ 0 , 1 ] is the state transition probability distribution . ρ0 : S → [ 0 , 1 ] is the distribution over initial states , r : S × A → R is the reward function , and γ ∈ [ 0 , 1 ) is the discount factor . At environment state s ∈ S , the agent take actions a ∈ A , in the ( unknown ) environment dynamics defined by the transition probability T ( s′|s , a ) , and the reward function yields a reward immediately following the action at performed in state st. We define the discounted return G ( st , at ) = ∑∞ l=0 γ lr ( st+l , at+l ) as the discounted sum of future rewards collected by the agent . In value-based reinforcement learning , the agent learns an estimate of the expected discounted return , a.k.a , state-action value function Qπ ( st , at ) = Est+1 , at+1 , ... [ ∑∞ l=0 γ lr ( st+l , at+l ) ] . A common way of deriving a new policy from a state-action value function is to act ǫ-greedily with respect to the action values ( discrete ) or to use policy gradient to maximize the value function ( continuous ) . Unsupervised Pre-Training RL In pretrained RL , the agent is trained in a reward-free MDP 〈S , S0 , A , T , G〉 for a long period followed by a short testing period with environment rewards R provided . The goal is to learn a pretrained agent that can quickly adapt to testing tasks defined by rewards to maximize the sum of expected future rewards in a zero-shot or few-shot manner . This is also known as the two phases learning in unsupervised pretraining RL [ 20 ] . The current state-of-the-art methods maximize the mutual information ( I ) between policy-conditioning variable ( w ) and the behavior induced by the policy in terms of state visitation ( s ) . max I ( s ; w ) = maxH ( w ) −H ( w|s ) , where w is sampled from a fixed distribution in practice as in DIAYN [ 17 ] and VISR [ 20 ] . The objective can then be simplified as max−H ( w|s ) . Due to it being intractable to directly maximize this negative conditional entropy , prior work propose to maximize the variational lower bound of the negative conditional entropy instead [ 7 ] . The training then amounts to learning a posterior of task variable conditioning on states q ( w|s ) . −H ( w|s ) ≥ Es , w [ log q ( w|s ) ] . Despite successful results in learning meaningful behaviors from reward-free interactions [ e.g . 41 , 18 , 26 , 17 , 20 ] , these methods suffer from insufficient exploration because they contain no explicit exploration . Another category considers the alternative direction of maximizing the mutual information [ 12 ] . max I ( s ; w ) = maxH ( s ) −H ( s|w ) . This intractable quantity can be similarly lowered bound by a variational approximation [ 7 ] . I ( s ; w ) ≥ Es , w [ qθ ( s|w ) ] − Es [ log p ( s ) ] , where Es [ log p ( s ) ] can then be approximated by a posterior of state given task variables Es [ log p ( s ) ] ≈ Es , w [ log q ( s|w ) ] . Despite their successes , this category of methods do not explore sufficiently since the agent receives larger rewards for visiting known states than discovering new ones as theoretically and empirically evidenced by Campos et al . [ 12 ] . In addition , they have only been shown to work from explicit state-representations and it remains unclear how to modify to learning from pixels . In the next section , we introduce a new nonparametric unsupervised pre-training method for RL which addresses these issues and outperforms prior state-of-the-arts on challenging visual-domain RL benchmarks . 3 Unsupervised Active Pre-Training for RL . We want to incentivize the agent with a reward rt to maximize entropy in an abstract representation space . Prior work on maximizing entropy relies on estimating density of states which is challenging and non-trivial , instead , we take a two-step approach . First , we learn a mapping fθ : R nS → RnZ that maps state space to an abstract representation space first . Then , we propose a particle-based nonparametric approach to maximize the entropy by deploying state-of-the-art RL algorithms . We introduce how to maximize entropy via particle-based approximation in Section 3.1 , and describe how to learn representation from states in Section 3.2 3.1 Particle-Based Entropy Maximization . Our entropy maximization objective is built upon the nonparametric particle-based entropy estimator proposed by Singh et al . [ 55 ] and Beirlant [ 8 ] and has has been widely studied in statistics [ 28 ] . Its key idea is to measure the sparsity of the distribution by considering the distance between each sampled data point and its k nearest neighbors . Concretely , assuming we have number of n data points { zi } n i=1 from some unknown distribution , the particle-based approximation can be written as Hparticle ( z ) = − 1 n n ∑ i=1 log k nvki + b ( k ) ∝ n ∑ i=1 log vki , ( 1 ) where b ( k ) is a bias correction term that only depends on the hyperparameter k , and vki is the volume of the hypersphere of radius ‖zi − z ( k ) i ‖ between zi and its k-th nearest neighbor z ( k ) i . ‖ · ‖ is the Euclidean distance . vki = ‖zi − z ( k ) i ‖ nZ · πnZ/2 Γ ( nZ/2 + 1 ) , ( 2 ) Encoder Particle-based Entropy Maximation observations Projectionrepresentation Expected Reward Normalization Contrastive Loss representation representation Minibatch representations K-nearest neighbors distance Reward Figure 2 : Diagram of the proposed method APT . On the left shows the objective of APT , which is to maximize the expected reward and minimize the contrastive loss . The contrastive loss learns an abstract representation from observations induced by the policy . We propose a particle-based entropy maximization based reward function such that we can deploy state-of-the-art RL methods to maximize entropy in an abstraction space of the induced by the policy . On the right shows the idea of our particle-based entropy , which measures the distance between each data point and its k nearest neighbors . where Γ is the gamma function . Intuitively , vki reflects the sparsity around each particle and equation ( 1 ) is proportional to the average of the volumes around each particle . By substituting equation ( 2 ) into equation ( 1 ) , we can simplify the particle-based entropy estimation as a sum of the log of the distance between each particle and its k-th nearest neighbor . Hparticle ( z ) ∝ n ∑ i=1 log ‖zi − z ( k ) i ‖ nZ . ( 3 ) Rather than using equation ( 3 ) as the entropy estimation , we find averaging the distance over all k nearest neighbors leads to a more robust and stable result , yielding our estimation of the entropy . Hparticle ( z ) : = n ∑ i=1 log c+ 1 k ∑ z ( j ) i ∈Nk ( zi ) ‖zi − z ( j ) i ‖ nZ , ( 4 ) where Nk ( · ) denotes the k nearest neighbors around a particle , c is a constant for numerical stability ( fixed to 1 in all our experiments ) . We can view the particle-based entropy in equation ( 4 ) as an expected reward with the reward function being r ( zi ) = log ( c+ 1k ∑ z ( j ) i ∈Nk ( zi ) ‖zi − z ( j ) i ‖ nZ ) for each particle zi . This makes it possible to deploy RL algorithms to maximize entropy , concretely , for a batch of transitions { ( s , a , s′ ) } sampled from the replay buffer . We consider the representation of each s′ as a particle in the representation space and the reward function for each transition is given by r ( s , a , s′ ) = log c+ 1 k ∑ z ( j ) ∈Nk ( z=fθ ( s ) ) ‖fθ ( s ) − z ( j ) ‖nZ ( 5 ) In order to keep the rewards on a consistent scale , we normalize the intrinsic reward by dividing it by a running estimate of the mean of the intrinsic reward . See Figure 2 for illustration of the formulation .

Summary: The authors introduce Active Pretraining (APT), wherein an agent is encouraged to maximization the entropy of the state distribution represented in its replay buffer. The is accomplished by having rewards that are proportional to distance to nearest neighbours in some embedding space. This space is constructed via a separate contrastive learning objective.  This method is tested in a 2 stage process where the agent initially maximizes this state entropy entropy for a long time, and then finetunes on a down-stream time for a much shorter time, with the final performance on this task being the key metric. By this methodology, APT surpasses both data-efficient methods that train from scratch, as well as other methods that use the same 2 stage training scheme.