DRIBO: Robust Deep Reinforcement Learning via Multi-View Information Bottleneck

1 INTRODUCTION . Deep reinforcement learning ( DRL ) methods have been shown to be successful in learning highquality controllers directly from raw images in an end-to-end fashion ( Mnih et al. , 2015 ; Levine et al. , 2016 ; Bojarski et al. , 2016 ) . However , it has been observed that DRL agents perform poorly in environments different from those where the agents were trained on , even when these environments contain semantically equivalent information relevant to the control task ( Farebrother et al. , 2018 ; Cobbe et al. , 2019 ; Zhang et al. , 2018b ; a ; Yu et al. , 2019 ) . By contrast , humans routinely adapt to new , unseen environments . For example , while visual scenes can be drastically different when driving in different cities , human drivers can quickly adjust to driving in a new city which they have never visited . We argue that this ability to adapt stems from the fact that driving skills are invariant to many visual details that are actually not relevant to driving . Conversely , DRL agents without this ability are hindered from understanding the temporal structure of task-relevant dynamics without being distracted by task-irrelevant visual details ( Jonschkowski & Brock , 2015 ; Zhang et al. , 2021 ; Agarwal et al. , 2021 ; Lee et al. , 2020b ) . Viewing from a representation learning perspective , a desired representation for RL should facilitate the prediction of future states ( beyond expected rewards ) on potential actions and discard excessive , task-irrelevant information from visual observations . An RL agent that learns from such representations has the advantage of learning an optimal policy more easily upon the prediction and being more robust to visual changes . In addition , the resulting policy is more likely to generalize to unseen environments if the task-relevant information in the new environment remains similar to that in the training environments . Prior works ( Hafner et al. , 2019 ; Lee et al. , 2020a ) that encode images into a low-dimensional latent space for RL typically rely on a reconstruction loss to learn representations that are sufficient to reconstruct the input images and predict ahead in the latent space While these approaches can learn representations that retain information in the visual observations , they do nothing to discard the irrelevant information . We tackle this problem by considering state representations for RL that are robust under the multi-view setting ( Li et al. , 2018 ; Federici et al. , 2020 ; Fischer , 2020 ) , where each view is assumed to provide the same amount of task-relevant information while all the information not shared by them is deemed task-irrelevant . Data augmentation can be easily leveraged to generate such multi-view observations without requiring additional new data . Data augmentation in RL has delivered promising results for visual control tasks ( Laskin et al. , 2020b ; Lange et al. , 2012 ; Laskin et al. , 2020a ) . However , these methods rarely exploit the sequential aspect of RL which requires an ideal representation to be predictive of future states given actions . In fact , the sequential nature of RL provides an additional temporal dimension for identifying task-irrelevant information when it is independent of actions . Instead of learning representations from each visual observation ( Laskin et al. , 2020a ) , we propose to learn a predictive model that captures the temporal evolution of representations from a sequence of observations and actions . Concretely , we introduce a new multi-view information bottleneck ( MIB ) objective that maximizes the mutual information between sequences of observations and representations while reducing the task-irrelevant information identified from the multi-view observations . We incorporate this MIB objective into RL by using it as an auxiliary learning objective . We illustrate our approach in Figure 1 . Our contributions are summarized below . • We propose DRIBO , a novel technique that learns robust representations in RL by identifying and discarding task-irrelevant information in the representations based on MIB . • We leverage the sequential nature of RL to learn representations better suited for RL with a non-reconstruction-based , DRIBO loss that maximizes the mutual information between sequences of observations and representations while disregarding task-irrelevant information . • Empirically , we show that our approach can ( i ) lead to better robustness against task-irrelevant distractors on the DeepMind Control Suite and ( ii ) significantly improve generalization on the Procgen benchmarks compared to current state-of-the-arts . 2 RELATED WORK . Reconstruction-based Representation Learning . Early works first trained autoencoders to learn representations to reconstruct raw observations . Then , the RL agent was trained from the learned representations ( Lange & Riedmiller , 2010 ; Lange et al. , 2012 ) . However , there is no guarantee that the agent will capture useful information for control . To address this problem , learning encoder and dynamics jointly has been proved effective in learning task-oriented and predictive representations ( Wahlström et al. , 2015 ; Watter et al. , 2015 ) . More recently , Hafner et al . ( 2019 ; 2020 ; 2021 ) and Lee et al . ( 2020a ) learn a latent dynamics model and train RL agents with predictive latent representations . However , these approaches suffer from the problem of embedding all details into representations even when they are task-irrelevant . The reason is that improving reconstruction quality from representations to visual observations forces the representations to retain more details . Despite success on many benchmarks , task-irrelevant visual changes can affect performance significantly ( Zhang et al. , 2018a ) . Experimentally , we show that our non-reconstructive approach , DRIBO , is substantially more robust against visual changes than prior works . We also compare DRIBO with another non-reconstructive method , DBC ( Zhang et al. , 2021 ) , which uses bisimulation metrics to learn representations in RL that contain only task-relevant information . Contrastive Representations Learning . Contrastive representation learning methods train an encoder that obeys similarity constraints in a dataset typically organized by similar and dissimilar pairs . The similar examples are typically obtained from nearby image patches ( Oord et al. , 2018 ; Hénaff et al. , 2020 ) or through data augmentation ( Chen et al. , 2020 ) . A scoring function that lower-bounds mutual information is one of the typical objects to be maximized ( Belghazi et al. , 2018 ; Oord et al. , 2018 ; Hjelm et al. , 2019 ; Poole et al. , 2019 ) . A number of works have applied the above ideas to RL settings to extract predictive signals . EMI ( Kim et al. , 2019 ) applies a Jensen-Shannon divergence-based lower bound on mutual information across subsequent frames as an exploration bonus . DRIML ( Mazoure et al. , 2020 ) uses an auxiliary contrastive objective to maximize concordance between representations to increase predictive properties of the representations conditioned on actions . CURL ( Laskin et al. , 2020a ) incorporates contrastive learning into RL algorithms to maximize similarity between augmented versions of the same observation . However , solely maximizing the lower-bound of mutual information retains all the information including those that are task-irrelevant ( Federici et al. , 2020 ; Fischer , 2020 ) . Multi-View Information Bottleneck ( MIB ) . MVRL ( Li et al. , 2019 ) uses the multi-view setting to tackle partially observable Markov decision processes with more than one observation model . For classification tasks , Federici et al . ( 2020 ) uses MIB by maximizing the mutual information between the representations of the two views while at the same time eliminating the label-irrelevant information identified by multi-view observations . Fischer ( 2020 ) describes a variant of the Conditional Entropy Bottleneck ( CEB ) which is mathematically equivalent to MIB . However , MIB/CEB can not be directly used in RL settings due to the sequential nature of decision making problems . PI-SAC ( Lee et al. , 2020b ) uses a contrastive version of CEB to model Predictive Information ( Bialek & Tishby , 1999 ) which is the mutual information between the past and the future to solve RL problems . However , this approach does not scale to long sequential data in RL and in practice only models short-term Predictive Information . Task-relevant information in RL is relevant because they influence not only current control decision and reward but also states and rewards well into the future . Our work , DRIBO , learns robust representations with a predictive model to maximize the mutual information between sequences of representations and observations , while eliminating task-irrelevant information based on the information bottleneck principle . Learning a predictive model also adopts richer learning signals than those provided by individual observation and reward alone . Philosophically and technically , our approach is different from PI-SAC which does not quantify task-irrelevant information from multi-view observations and can not capture long-term dependencies . Another line of work , IDAAC ( Raileanu & Fergus , 2021 ) , leverages an adversarial framework so that the learned representations yield features that are instance-independent and invariant to task-irrelevant changes . 3 PRELIMINARIES . We denote a Markov decision process ( MDP ) asM , with state s , action a , and reward r. S and A stand for the corresponding random variables . We denote a policy onM as π . The agent ’ s goal is to learn a policy π that maximizes the cumulative rewards . We define S⊆Rd as the state-representation space . The visual observations are o∈O , where we denote multi-view observations from the viewpoint i as o ( i ) . O stands for the random variable of the observation . We introduce a multi-view trajectory τM= [ s1 , o ( i ) 1 , a1 , . . . , sT , o ( i ) T , aT ] where T is the length . Knowing that the trajectory density is defined over joint observations , states , and actions , we write : pπ ( τ M ) =π ( aT |sT ) P ( i ) obs ( o ( i ) T |sT ) P ( sT |sT−1 , aT−1 ) · · ·π ( a1|s1 ) P ( i ) obs ( o ( i ) 1 |s1 ) P0 ( s1 ) ( 1 ) with P0 being the initial state distribution , P being the transition model and P ( i ) obs being the unknown observation model for view i. DRL agents learn from visual observations by treating consecutive observations as states to implicitly capture the predictive property . However , rich details in observations can easily distract the agent . An ideal representation should contain no task-irrelevant information and satisfy some underlying MDP which determines the distribution of the multi-view trajectory in Eq . ( 1 ) . Thus , instead of mapping a single-step observation to a representation , we consider learning a predictive model that correlates sequential observations and representations . Let a∗1 : T be the optimal action sequence for some o1 : T which is obtained by executing the action sequence a1 : T . We assume that o1 : T contains enough information to obtain a∗1 : T which maximizes the cumulative rewards . With this assumption , we define task-relevant information if it is necessary for deriving a∗1 : T . By contrast , task-irrelevant information does not contribute to the choice of a∗1 : T . We first consider sufficient representations that are discriminative enough to obtain A ∗ at each timestep . This property can be quantified by the amount of mutual information between O1 : T and A∗1 : T and mutual information between S1 : T and A ∗ 1 : T . Definition 1 . Representations S1 : T of O1 : T are sufficient for RL iff I ( O1 : T ; A∗1 : T ) =I ( S1 : T ; A∗1 : T ) . RL agents that have access to a sufficient representation St at timestep t must be able to generate A∗t as if it has access to the original observations . This can be better understood by subdividing I ( O1 : T ; S1 : T ) into two components using the chain rule of mutual information : I ( O1 : T ; S1 : T ) = I ( S1 : T ; O1 : T |A∗1 : T ) + I ( S1 : T ; A∗1 : T ) ( 2 ) Conditional mutual information I ( S1 : T ; O1 : T |A∗1 : T ) quantifies the information in S1 : T that is taskirrelevant . I ( S1 : T ; A∗1 : T ) quantifies task-relevant information that is accessible from S1 : T . The last term is independent of the representation as long as St is sufficient for A∗t ( see Definition 1 ) . Thus , a representation contains minimal task-irrelevant information whenever I ( O1 : T ; S1 : T |A∗1 : T ) is minimized . Maximizing I ( O1 : T ; S1 : T ) learns a sufficient representation . With the information bottleneck principle ( Tishby et al. , 2000 ) , we can construct an objective to maximize I ( O1 : T ; S1 : T ) while minimizing I ( S1 : T ; O1 : T |A∗1 : T ) to compress away task-irrelevant information . However , estimating the mutual information between long sequences is difficult due to the high dimensionality of the problem . In addition , the minimization of I ( S1 : T ; O1 : T |A∗1 : T ) can only be done directly in supervised settings whereA∗1 : T are observed . One option is to use MIB which can compress away task-irrelevant information in the representations in unsupervised settings ( Federici et al. , 2020 ) . The problem , however , is that MIB in its original form only considers a single observation and its representation and thus does not guarantee that the learned representations retain the important temporal structure of RL . In the next section , we describe how we extend MIB to RL settings .

Summary: This paper aims to learn robust representation from the replay buffer for reinforcement learning. The key idea is to leverage the concept of mutual information and the InfoNCE tool to compute the mutual information as a regularizer (a.k.a. DRIBO loss). The authors conducted experiments on some standard benchmarks (DeepMind Visual Control Suite and ProcGen).