Dynamic Bottleneck for Robust Self-Supervised Exploration

1 Introduction . The tradeoff between exploration and exploitation has long been a major challenge in reinforcement learning ( RL ) [ 35 , 50 , 58 ] . Generally , excessive exploitation of the experience suffers from the potential risk of being suboptimal , whereas excessive exploration of novel states hinders the improvement of the policy . A straightforward way to tackle the exploration-exploitation dilemma is to enhance exploration efficiency while keeping exploitation in pace . When the extrinsic rewards are dense , reward shaping is commonly adopted for efficient exploration . However , in many real-world applications such as autonomous driving [ 34 ] , the extrinsic rewards are sparse , making efficient exploration a challenging task in developing practical RL algorithms . The situations become even worse when the extrinsic rewards are entirely unavailable . In such a scenario , the task of collecting informative trajectories from exploration is known as the self-supervised exploration [ 11 ] . An effective approach to self-supervised exploration is to design a dense intrinsic reward that motivates the agent to explore novel transitions . Previous attempts include count-based [ 9 ] and curiosity-driven [ 39 ] explorations . The count-based exploration builds a density model to measure the pseudo-count of state visitation and assign high intrinsic rewards to less frequently visited states . In contrast , the 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . curiosity-driven methods maintain a predictive model of the transitions and encourage the agent to visit transitions with high prediction errors . However , all these methods becomes unstable when the states are noisy , e.g. , containing dynamics-irrelevant information . For example , in autonomous driving tasks , the states captured by the camera may contain irrelevant objects , such as clouds that behave similar to Brownian movement . Hence , if we measure the novelty of states or the curiosity of transitions through raw observed pixels , exploration are likely to be affected by the dynamics of these irrelevant objects . To encourage the agent to explore the most informative transitions of dynamics , we propose a Dynamic Bottleneck ( DB ) model , which generates a dynamics-relevant representation Zt of the current state-action pair ( St , At ) through the Information-Bottleneck ( IB ) principle [ 55 ] . The goal of training DB model is to acquire dynamics-relevant information and discard dynamics-irrelevant features simultaneously . To this end , we maximize the mutual-information I ( Zt ; St+1 ) between a latent representation Zt and the next state St+1 through maximizing its lower bound and using contrastive learning . Meanwhile , we minimize the mutual-information I ( [ St , At ] ; Zt ) between the state-action pair and the corresponding representation to compress dynamics-irrelevant information . Based on our proposed DB model , we further construct a DB-bonus for exploration . DB-bonus measures the novelty of state-action pairs by their information gain with respect to the representation computed from the DB model . We show that the DB-bonus are closely related to the provably efficient UCB-bonus in linear Markov Decision Processes ( MDPs ) [ 1 ] and the visiting count in tabular MDPs [ 3 , 22 ] . We further estimate the DB-bonus by the learned dynamics-relevant representation from the DB model . We highlight that exploration based on DB-bonus directly utilize the information gain of the transitions , which filters out dynamics-irrelevant noise . We conduct experiments on the Atari suit with dynamics-irrelevant noise injected . Results demonstrate that our proposed selfsupervised exploration with DB-bonus is robust to dynamics-irrelevant noise and outperforms several state-of-the-art exploration methods . 2 Related Work . Our work is closely related to previous exploration algorithms that construct intrinsic rewards to quantify the novelty of states and transitions . Several early approaches directly define the pseudocount by certain statistics to measure the novelty of states [ 41 , 30 ] ; more recent methods utilize density model [ 9 , 38 ] or hash map [ 52 , 42 ] for state statistics . Nevertheless , these approaches are easily affected by dynamics-irrelevant information such as white-noise . The contingency awareness method [ 16 ] addresses such an issue by using an attentive model to locate the agent and computes the pseudocount based on regions around the agent . However , such an approach could ignore features that are distant from the agent but relevant to the transition dynamics . Another line of research measures the novelty through learning a dynamics model and then use the prediction error to generate an intrinsic reward . These methods are known as the curiosity-driven exploration algorithms . Similar to the pseudo-count based methods , curiosity-driven methods become unstable in the presence of noises , because the prediction model is likely to yield high error for stochastic inputs or targets . Some recent attempts improve the curiosity-driven approach by learning the inverse dynamics [ 39 ] and variational dynamics [ 7 ] to define curiosity , or utilizes the prediction error of a random network to construct intrinsic rewards [ 12 ] . However , without explicitly removing dynamics-irrelevant information , these methods are still vulnerable to noises in practice [ 11 ] . The entropy-based exploration uses state entropy as the intrinsic reward . VISR [ 19 ] , APT [ 29 ] and APS [ 28 ] use unsupervised skill discovery for fast task adaptation . In the unsupervised stage , they use k-nearest-neighbor entropy estimator to measure the entropy of state , and then use it as the intrinsic reward . RE3 [ 47 ] and ProtoRL [ 59 ] use random encoder and prototypes to learn the representation and use state-entropy as bonuses in exploration . Nevertheless , the state entropy will increase significantly if we inject noises in the state space . The entropy-based exploration will be misled by the noises . Previous approaches also quantify the epistemic uncertainty of dynamics through Bayesian network [ 21 ] , bootstrapped Q-functions [ 37 , 8 ] , ensemble dynamics [ 40 ] , and Stein variational inference [ 43 ] to tackle noisy environments . However , they typically require either complicated optimization methods or large networks . In contrast , DB learns a dynamics-relevant representation and encourages exploration by directly accessing the information gain of new transitions via DB-bonus . Another closely related line of studies uses the mutual information to promote exploration in RL . Novelty Search ( NS ) [ 53 ] proposes to learn a representation through IB . Curiosity Bottleneck ( CB ) [ 26 ] also performs exploration based on IB by measuring the task-relevant novelty . However , both NS and CB require extrinsic rewards to learn a value function and are not applicable for selfsupervised exploration . Moreover , NS contains additional k-nearest-neighbor to generate intrinsic reward and representation loss to constrain the distance of consecutive states , which are costly for computation . In contrast , our DB model handles self-supervised exploration without accessing extrinsic rewards . EMI [ 25 ] learns a representation by maximizing the mutual information in the forward dynamics and the inverse dynamics , which is different from the IB principle used in our method . In addition , we aim to perform robust exploration to overcome the white-noise problem , while EMI does not have an explicit mechanism to address the noise . Our work is also related to representation learning in RL . DrQ [ 60 ] , RAD [ 27 ] and CURL [ 49 ] learn the state representation by data augmentation and contrastive learning [ 14 , 20 , 36 ] to improve the data-efficiency of DRL . Deep InfoMax [ 33 ] and Self-Predictive Representation ( SPR ) [ 46 ] learn the contrastive and predictive representations of dynamics , respectively , and utilize such representations as auxiliary losses for policy optimization . However , none of these existing approaches extracts information that benefits exploration . In contrast , we show that the dynamics-relevant representation learned by DB can be utilized for efficient exploration . 3 The Dynamic Bottleneck . In this section , we introduce the objective function and architecture of the DB model . We consider an MDP that can be described by a tuple ( O , A , P , r , γ ) , which consists of the observation space O , the action space A , the transition dynamics P , the reward function r , and the discount factor γ ∈ ( 0 , 1 ) . At each time step , an agent decides to perform an action at ∈ A after observing ot ∈ O , and then the observation transits to ot+1 with a reward rt received . In this paper , we use upper letters , such as Ot , to denote random variables and the corresponding lower case letter , such as ot , to represent their corresponding realizations . We first briefly introduce the IB principle [ 55 ] . In supervised setting that aims to learn a representation Z of a given input source X with the target source Y , IB maximizes the mutual information between Z and Y ( i.e . max I ( Z ; Y ) ) and restricts the complexity of Z by using the constrain as I ( Z ; X ) < Ic . Combining the two terms , the objective of IB is equal to max I ( Z ; Y ) − αI ( Z ; X ) with the introduction of a Lagrange multiplier . DB follows the IB principle [ 55 ] to learn dynamics-relevant representation . The input variable of the DB model is a tuple ( Ot , At ) that contains the current observation and action , and the target is the next observation Ot+1 . We denote by St and St+1 the encoding of observations Ot and Ot+1 . The goal of the DB model is to obtain a compressed latent representation Zt of ( St , At ) , that preserves the information that is relevant to St+1 only . Specifically , we use fSo and f S m as the encoders of two consecutive observations ot and ot+1 , respectively . We parameterize the dynamics-relevant representation zt by a Gaussian distribution with parameter φ , and it takes ( st , at ) as input . We summarize the DB model as follows , st = f S o ( ot ; θo ) , st+1 = f S m ( ot+1 ; θm ) , zt ∼ gZ ( st , at ; φ ) . ( 1 ) Following the IB principle , the objective of the DB model seeks to maximize the mutual information I ( Zt ; St+1 ) while minimizing the mutual information I ( [ St , At ] ; Zt ) . To this end , we propose the DB objective by following the IB Lagrangian [ 55 ] , which takes the form of min−I ( Zt ; St+1 ) + α1I ( [ St , At ] ; Zt ) . ( 2 ) Here α1 is a Lagrange multiplier that quantifies the amount of information about the next state preserved in Zt . Fig . 1 illustrates the DB objective . We minimize I ( [ St , At ] ; Zt ) and consider it as a regularizer in the representation learning . Then the representation learning is done by maximizing the mutual information I ( Zt , St+1 ) . Maximizing I ( Zt , St+1 ) ensures that we do not discard useful information from ( St , At ) . In DB , the mutual information is estimated by several variational bounds parameterized by neural networks to enable differentiable and tractable computations . In what follows , we propose a lower bound of ( 2 ) , which we optimize to train the DB model . 3.1 Maximizing the lower bound of I ( Zt ; St+1 ) As directly maximizing I ( Zt ; St+1 ) is intractable , we propose to optimize a predictive objective , which is a lower bound of I ( Zt ; St+1 ) [ 2 ] . It holds that I ( Zt ; St+1 ) = Ep ( zt , st+1 ) [ log p ( st+1|zt ) p ( st+1 ) ] = E [ log q ( st+1|zt ; ψ ) p ( st+1 ) ] +DKL [ p ( st+1|zt ) ‖q ( st+1|zt ; ψ ) ] , ( 3 ) where p ( st+1|zt ) is an intractable conditional distribution and q ( st+1|zt ; ψ ) is a tractable variational decoder with parameter ψ . By the non-negativity of the KL-divergence , we obtain the following lower bound , I ( Zt ; St+1 ) ≥ Ep ( zt , st+1 ) [ log q ( st+1|zt ; ψ ) ] +H ( St+1 ) , whereH ( · ) is the entropy . SinceH ( St+1 ) is irrelevant to the parameter ψ , maximizing I ( Zt ; St+1 ) is equivalent to maximizing the following lower bound , Ipred , Ep ( zt , st+1 ) [ log q ( st+1|zt ; ψ ) ] . ( 4 ) Ipred can be interpreted as the log-likelihood of the next-state encoding st+1 given the dynamicsrelevant representation zt . In practice , we parameterize the prediction head q ( st+1|zt ; ψ ) by a neural network that outputs a diagonal Gaussian random variable . Since st , st+1 and zt are all low-dimensional vectors instead of raw image pixels , optimizing Ipred is computationally efficient . Momentum Encoder To encode the consecutive observations ot and ot+1 , we adopt the Siamese architecture [ 10 ] that uses the same neural network structures for the two encoders . Nevertheless , we observe that if we train both the encoders by directly maximizing Ipred , the Siamese architecture tends to converge to a collapsed solution . That is , the generated encodings appears to be uninformative constants . A simple fact is that if both the encoders generate zero vectors , predicting zeros conditioning on zt ( or any variables ) is a trivial solution . To address such issue , we update the parameter θo of fSo in ( 1 ) by directly optimizing Ipred . Meanwhile , we update the parameter θm of fSm by a momentum moving average of θo , which takes the form of θm ← τθm + ( 1 − τ ) θo . In the sequel , we call fSo the online encoder and f S m the momentum encoder , respectively . Similar techniques is also adopted in previous study [ 18 , 46 ] to avoid the mode collapse . 3.2 Contrastive Objective for Maximizing I ( Zt ; St+1 ) In addition to the lower bound of I ( Zt ; St+1 ) in §3.1 , we also investigate the approach of maximizing the mutual information by contrastive learning ( CL ) [ 36 ] . CL classifies positive samples and negative samples in the learned representation space . An advantage of adopting CL is that training with negative samples plays the role of regularizer , which avoids collapsed solutions . Moreover , the contrastive objective yields a variational lower bound of the mutual information I ( Z ; St+1 ) . To see such a fact , note that by the Bayes rule , we have I ( Zt ; St+1 ) ≥ Ep ( zt , st+1 ) ES− [ log exp ( h ( zt , st+1 ) ) ∑ sj∈S−∪st+1 exp ( h ( zt , sj ) ) ] , Ince . ( 5 ) Here h is a score function which assigns high scores to positive pairs and low score to negative pairs . We refer to Appendix A for a detailed proof of ( 5 ) . The right-hand side of ( 5 ) is known as the InfoNCE objective [ 36 ] . The positive samples are obtained by directly sampling the transitions ( s , a , s′ ) . In contrast , the negative samples are obtained by first sampling a state-action pair ( s , a ) , and then sampling a state s̃ independently . Then a negative sample is obtained by concatenating them together to form a tuple ( s , a , s̃ ) . The negative samples do not follow the transition dynamics . In practice , we collect the negative sample by sampling observation encodings randomly from the batch . We remark that comparing with methods that require data augmentation to construct negative samples [ 20 , 49 ] , DB utilizes a simple scheme to obtain positive and negative samples from on-policy experiences . In ( 5 ) , we adopt the standard bilinear function as the score function h , which is defined as follows , h ( zt , st+1 ) = f P o ( q̄ ( zt ; ψ ) ) > WfPm ( st+1 ) , ( 6 ) where fPo ( · ; ϕo ) and fPm ( · ; ϕm ) project st+1 and the mean value of next-state prediction q ( st+1|zt ; ϕ ) , i.e. , q̄ ( · ; ψ ) , to a latent space to apply the contrastive loss Ince in ( 5 ) , andW is the parameter of the score function . Similar to the observation encoder and MoCo-based architectures [ 49 , 20 , 15 ] , we also adopt an online projector fPo and a momentum projector f P m for zt and st+1 , respectively . The momentum projector is updated by ϕm ← τϕm + ( 1− τ ) ϕo . 3.3 Minimizing the Upper Bound of I ( [ St , At ] ; Zt ) We minimize the mutual information I ( [ St , At ] ; Zt ) through minimizing a tractable upper bound of the mutual information . To this end , we introduce a variational approximation q ( zt ) to the intractable marginal p ( zt ) = ∫ p ( st , at ) p ( zt|st , at ) dstat . Specifically , the following upper-bound of I ( [ St , At ] ; Zt ) holds , I ( [ St , At ] ; Zt ) = Ep ( st , at ) [ p ( zt|st , at ) p ( zt ) ] = Ep ( st , at ) [ p ( zt|st , at ) q ( zt ) ] −DKL [ p ( zt ) ‖q ( zt ) ] ≤ Ep ( st , at ) [ DKL [ p ( zt|st , at ) ‖q ( zt ) ] ] , Iupper , ( 7 ) where the inequality follows from the non-negativity of the KL divergence , and q ( zt ) is an approximation of the marginal distribution of Zt . We follow Alemi et al . [ 2 ] and use a standard spherical Gaussian distribution q ( zt ) = N ( 0 , I ) as the approximation . The expectation of Iupper is estimated by sampling from on-policy experiences .

Summary: This paper proposes a new self-supervised exploration method for RL that encourages an agent to explore states which are more relevant in terms of information gain. To this aim, the Information-Bottleneck (IB) principle has been utilized to learn a model that can obtain dynamics-relevant information and discard dynamics-irrelevant features together. The paper evaluates the proposed technique in multiple Atari environments without accessing the extrinsic rewards.