DRIBO: ROBUST DEEP REINFORCEMENT LEARNING
Year: 2021

Abstract:
Deep reinforcement learning (DRL) agents are often sensitive to visual changes that were unseen in their training environments. To address this problem, we leverage the sequential nature of RL to learn robust representations that encode only task-relevant information from observations based on the unsupervised multiview setting. Specifically, we introduce a novel contrastive version of Multi-View Information Bottleneck (MIB) objective for temporal data. We train RL agents from pixels with this auxiliary objective to learn robust representations that can compress away task-irrelevant information and are predictive of task-relevant dynamics. This approach enables us to train high-performance policies that are robust to visual distractions and can generalize well to unseen environments. We demonstrate that our approach can achieve SOTA performance on diverse visual control tasks on the DeepMind Control Suite when the background is replaced with natural videos. In addition, we show that our approach outperforms well-established baselines for generalization to unseen environments on the Procgen benchmark.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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 cannot 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 cannot 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.

heading: 3 PRELIMINARIES
text: 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.

heading: 4 DRIBO
text: DRIBO learns robust representations that are predictive of future representations while discarding task-irrelevant information. To learn such representations, we construct a new MIB objective that (i) reduces the problem of maximizing the mutual information between sequences of observations and representations to maximizing the mutual information between them at each timestep and (ii) quantifies the amount of task-irrelevant information in the representations in the multi-view setting.

heading: 4.1 MUTUAL INFORMATION MAXIMIZATION
text: To capture the temporal evolution of observations and representations given any action sequence, we consider maximizing the conditional mutual information I(S1:T ;O1:T |A1:T ) which is a lower bound of I(O1:T ;S1:T ) (see Appendix A.1). The observations O1:T are generated sequentially in the environment by executing the actions A1:T . The conditional mutual information not only estimates the sufficiency of the representations but also maintains the temporal structure of RL problems.
To tackle the challenges of estimating the mutual information between sequences, we first factorize the mutual information between two sequential data to the mutual information at each timestep. Theorem 1. Let O1:T be the sequential observations obtained by executing action sequence A1:T . If S1:T is a sequence of sufficient representations for O1:T , we have:
I(S1:T ;O1:T |A1:T ) ≥ T∑ t=1 I(St;Ot|St−1, At−1) (3)
The proof is included in Appendix A.1. Theorem 1 shows that the sum of mutual information I(St;Ot|St−1, At−1) over multiple timesteps is a lower bound of I(S1:T ;O1:T |A1:T ). I(St;Ot|St−1, At−1) is defined with conditional probabilities, p(st,ot|st−1,at−1), p(st|st−1,at−1) and p(ot|st−1,at−1). This lower bound models the dynamics and temporal structure of RL since the first conditional probability is the composition of transition probability and observation model and the second conditional probability is the transition probability of the underlying MDP. Thus, the factorized mutual information explicitly retains the predictive property of representations. Even when the representations S1:T are not sufficient, maximizing the mutual information between St and Ot (I(St;Ot|St−1, At−1)) encourages encoding more detailed features from Ot into St and makes St sufficient.

heading: 4.2 MULTI-VIEW SETTING
text: To learn sufficient representations with minimal task-irrelevant information, we consider a two-view setting to identify the task-irrelevant information without supervision. Consider o(1)t and o (2) t to be two visual images of the control scenario from two different viewpoints. Under the multi-view assumption, any representation st containing all information accessible from both views and is predictive of future representations would contain sufficient task-relevant information. Furthermore, if st captures only the details that are visible from both observations, it would eliminate the viewspecific details and reduce the sensitivity of the representation to view-changes.
A sufficient representation in RL retains all the information that is shared by mutually redundant observations O(1)t and O (2) t . We refer to Appendix A for the sufficiency condition of representations and mutually redundancy condition (Federici et al., 2020) between O(1)t and O (2) t . Intuitively, with the mutual redundancy condition, any representation that contains all the information shared by both views is as task-relevant as the joint observation. By factorizing the mutual information between S(1)t and O(1)t as in Eq. (2), we can identify two components:
I(S(1)t ;O (1) t |S (1) t−1, At−1) = I(S (1) t ;O (1) t |S (1) t−1, At−1, O (2) t )+I(O (2) t ;S (1) t |S (1) t−1, At−1) (4)
Here, S(1)t−1 is a representation of visual observation O (1) t−1. Since we assume mutual redundancy between the two views, the information shared between O(1)t and S (1) t conditioned on O (2) t must be irrelevant to the task, which can be quantified as I(S(1)t ;O (1) t |S (1) t−1, At−1, O (2) t ) (first term in Eq. (4)). Then, I(O(2)t ;S (1) t |S (1) t−1, At−1) has to be maximal if the representation is sufficient. A formal description of the above statement can be found in Appendix A.
The less the two views have in common, the less task-irrelevant information can be encoded into the representations without violating sufficiency, and consequently, the less sensitive the resulting representation is to task-irrelevant nuisances. In the extreme, s(1)t is the underlying states of MDP if o (1) t and o (2) t share only task-relevant information. With Eq. (3) and 4, we have L (1) IB to maximize the mutual information between representations and observations and compress away task-irrelevant information I(O(2)t ;S (1) t |S (1) t−1, At−1) based on the information bottleneck principle. λ1 is a Lagrange multiplier. This loss also retains the temporally evolving information of the underlying dynamics.
L(1)IB = − ∑ t (I(S(1)t ;O (1) t |s (1) t−1, At−1, O (2) t )− λ1I(O (2) t ;S (1) t |S (1) t−1, At−1)) (5)
Symmetrically, we define a loss L(2)IB for representations and observations from view 2: L(2)IB = − ∑ t (I(S(2)t ;O (2) t |s (2) t−1, At−1, O (1) t )− λ2I(O (1) t ;S (2) t |S (2) t−1, At−1)) (6)
The above losses extend MIB to RL and minimizing it learns representations that are robust to task-irrelevant distractors and predictive of the future. The multi-view observations can be easily obtained with random data augmentation techniques so that each view is augmented differently.

heading: 4.3 DRIBO LOSS FUNCTION
text: By re-parameterizing the Lagrangian multipliers (details in Appendix B), the average of two loss functions L(1)IB and L (2) IB from two views at timestep t can be upper bounded as follows:
Lt(θ;β)=−Iθ(S(1)t ;S (2) t |St−1, At−1)+ (7)
βDSKL(pθ(s (1) t |o (1) t , s (1) t−1,at−1)||pθ(s (2) t |o (2) t , s (2) t−1,at−1))
where θ denotes the parameters of an encoder pθ(s (1) t |o (1) t , st−1,at−1) (details in Section 4.4), st−1 is a sufficient representation, DSKL represents the symmetrized KL divergence obtained by averaging the expected values of DKL(pθ(s (1) t |o (1) t , s (1) t−1,at−1)||pθ(s (2) t |o (2) t , s (2) t−1,at−1)), DKL(pθ(s (2) t |o (2) t , s (2) t−1,at−1)||pθ(s (1) t |o (1) t , s (1) t−1,at−1)), and the coefficient β represents the tradeoff between sufficiency and sensitivity to task-irrelevant information. β is a hyper-parameter.
Algorithm 1 DRIBO Loss 1: Input: Batch B storing N sequential observations and actions with length T from replay buffer. 2: Apply random augmentation transformations on B to obtain multi-view batches B(1) and B(2). 3: for i, (o(1)1:T ,o (2) 1:T ,a1:T ) in enumerate (B(1),B(2)) do
4: for t = 1 to T do 5: s(1)t ∼ pθ(s (1) t |o (1) t , s (1) t−1,at−1), s (2) t ∼ pθ(s (2) t |o (2) t , s (2) t−1,at−1) 6: (s(1),t+T (i−1), s(2),t+T (i−1))← (s(1)t , s (2) t ) 7: end for 8: LiSKL = 1T ∑T t=1DDKL(pθ(s (1) t |o (1) t , s (1) t−1,at−1)||pθ(s (2) t |o (2) t , s (2) t−1,at−1))
9: end for 10: return −Îψ({(s(1),i, s(2),i)}T∗Ni=1 ) + β N ∑N i=1 LiSKL
To generalize the above loss to sequential data in RL, we apply Theorem 1 to obtain the DRIBO loss: LDRIBO= 1T ∑T t=1 Lt(θ;β). We summarize the batch-based computation of the loss function in Algorithm 1. We sample s(1)t and s (2) t from pθ(s (1) t |o (1) t , s (1) t−1,at−1) and pθ(s (2) t |o (2) t , s (2) t−1,at−1) respectively. Though the first term in Eq. (7) is conditioned on st−1, we prove that the sampling process does not affect its effectiveness based on the multi-view assumption in Appendix B. The symmetrized KL divergence term can be computed from the probability density of S(1)t and S (2) t using the encoder. The mutual information between the two representations Iθ(S(1)t ;S (2) t |St−1, At−1) can be maximized by using any sample-based differentiable mutual information lower bound Îψ(s (1) t , s (2) t ), where ψ represents the learnable parameters. We use InfoNCE (Oord et al., 2018) to estimate mutual information since the multi-view setting provides a large number of negative examples. The positive pairs are the representations (s(1)t , s (2) t ) of the multi-view observations generated from the same observation. The remaining pairs of representations within the same batch are used as negative pairs. The full derivation of the DRIBO loss function can be found in Appendix B.

heading: 4.4 ENCODER ARCHITECTURE AND INCORPORATING DRIBO IN RL
text: The encoder pθ(st|ot, st−1,at−1) approximates the posterior representation given the current observation, and representation and action from the previous timestep. The posteriors can also be viewed as a reparameterization of pθ(s1:T |o1:T ,a1:T ) = ∏ t pθ(st|ot, st−1,at−1), which reflects the inherent temporal structure of RL. We implement the encoder as a recurrent space model (RSSM (Hafner et al., 2019)) which leverages recurrent neural networks to perform accurate long-term predictions. More details can be found in Appendix D.1. Training an RSSM encoder with DRIBO enables the representations to be predictive of future states.
We simultaneously train our representation learning models and the RL agent by adding LDRIBO (Algorithm 1) as an auxiliary objective during training. The multi-view observations can be easily obtained using the same experience replay of RL agents through data augmentation. We demonstrate the effectiveness of DRIBO by building the agents on top of SAC (Haarnoja et al., 2018), an off-policy RL algorithm, and PPO (Schulman et al., 2017), an on-policy RL algorithm, in Section 5.1 and Section 5.3 respectively. More details can be found in Appendix D.

heading: 5 EXPERIMENTS
text: We experimentally evaluate DRIBO on a variety of visual control tasks. We designed the experiments to compare DRIBO to the current best methods in the literature on: (i) the effectiveness of solving visual control tasks, (ii) their robustness against task-irrelevant distractors, and (iii) the ability to generalize to unseen environments. For effectiveness and robustness, we demonstrate DRIBO’s performance on DeepMind Control Suite (DMC (Tassa et al., 2018)) with task-irrelevant visual distractors in backgrounds. The backgrounds are replaced with natural videos from the Kinetics dataset (Kay et al., 2017) (middle column in Figure 2). In Appedix C.1, we also show comparison between DRIBO and SOTAs on DMC environments without the background distractors (left column in Figure 2).
For generalization, we present results on Procgen (Cobbe et al., 2020) which provides different levels of the same game to test how well agents generalize to unseen levels. We use single-step observations to train DRIBO without assuming that observation at each timestep provides full observability of the underlying dynamics. By contrast, current SOTA approaches require the use of consecutive observations 1 to capture predictive properties of the underlying states.
We also conduct careful ablations analysis to show that the benefit of DRIBO is due primarily to learning from the temporal structure of RL and the DRIBO loss.
For the DMC suite, all agents are built on top of SAC. For the Procgen suite, we augment PPO, a RL baseline for Procgen, with DRIBO. Implementation details are in Appendix D.

heading: 5.1 EFFECTIVENESS AND ROBUSTNESS
text: We compare DRIBO against several SOTA methods. The first is RAD (Laskin et al., 2020b), a recent method that uses augmented data to train pixel-based policies on DMC benchmarks. The second is SLAC (Lee et al., 2020a), a SOTA representation learning method for RL that learns a dynamic model using a reconstruction loss. The third is CURL (Laskin et al., 2020a), an approach that leverages contrastive learning to maximize the mutual information between representations of augmented versions of the same observation but does not distinguish between relevant and irrelevant features. The fourth is DBC (Zhang et al., 2021) which shares a similar goal with DRIBO. DBC learns an invariant representation based on bisimulation metrics without requiring reconstruction. Finally, we compare with PI-SAC (Lee et al., 2020b) which leverages the Predictive Information to compress away task-irrelevant information with CEB. We apply random crop+grayscale to obtain the augmented data for RAD which achieves the best performance in (Laskin et al., 2020b). For CURL and DRIBO, we apply random crop to obtain augmented data and multi-view observations.
Natural Video Setting. To investigate the effectiveness and robustness of RL agents in DMC environments, we introduce high-dimensional visual distractors by using natural videos from the Kinetics dataset (Kay et al., 2017) as backgrounds (Zhang et al., 2018a) (Figure 2: middle column). We use the class of “arranging flowers” videos to replace the background in training. During testing, we use the test set from the Kinetics dataset to replace the background, which contains videos from various different classes. Note that a single run of the DMC task may have multiple videos playing sequentially in the background. See Appendix C.4 for snapshots under the natural video setting.
1All other methods compared in this paper use stack frames of 3 consecutive observations.
In Figure 2, spatial attention maps (Zagoruyko & Komodakis, 2017) of the trained DRIBO encoder demonstrate that DRIBO trains agents to focus on the robot body while ignoring irrelevant scene details in the background. Figure 3 shows that DRIBO performs substantially better than RAD, SLAC and CURL which do not discard task-irrelevant information explicitly. Compared to PI-SAC and DBC which are recent state-of-art methods aimed at learning representations invariant to task-irrelevant information, DRIBO outperforms them consistently – DRIBO achieves on average 30% and at the maximum 77% higher returns at 88e4 steps compared to the second best performing method.
Visualization. We visualize the representations learned by DRIBO and CURL with t-SNE (Van der Maaten & Hinton, 2008). Figure 4 shows that even when the background looks drastically different, DRIBO learns to disregard irrelevant information and maps observations with similar robot configurations to the neighborhoods of one another. The color code represents value of the reward for each representation. The rewards can be viewed as task-relevant signals provided by the environments. DRIBO learns representations that are close in the latent space with similar reward values.

heading: 5.2 ABLATIONS
text: Temporal Structure of RL. To investigate whether DRIBO captures the temporal structure of RL, we conducted further experiments on DRIBO agents trained using sequences of different lengths under the natural video setting. Longer sequences carry more temporal information for DRIBO to learn. By default, we train DRIBO with sequences of length 32. In this ablation study, we present results of DRIBO trained using sequences of lengths 3, 6 and 16 respectively in Figure 5. Using sequences of length 3 is similar to stacking 3 consecutive frames which is a common choice for training in DMC.
Using sequences of length 6 is similar to the design choice made in PI-SAC (3 steps for the past and 3 steps for the future). We also include results on using 16-step sequential observations to investigate how DRIBO’s performance changes as the length of the sequences increases. Theoretically, the DRIBO loss provides a lower bound on the mutual information between sequences of observations and sequences of (latent) representations with the same length. It can be observed that DRIBO performs significantly better when training with longer sequences.
Learning Objective. In the DRIBO loss, we use the SKL term (second term in Eq. (7)) to compress away task-irrelevant information. Figure 5 studies the effect of removing the SKL term in the DRIBO loss (annotated as DRIBO-no-skl). The objective without the SKL term is equivalent to an InfoMax objective. We use identical training configurations (e.g. sequence length of 32) for DRIBO-no-skl and DRIBO. The only difference is whether the SKL term is included in the learning objective. Figure 5 shows that DRIBO outperforms DRIBO-no-skl substantially in the natural video setting.

heading: 5.3 GENERALIZATION
text: The natural video setting of DMC is suitable for benchmarking robustness to high-dimensional visual distractors. However, the task-relevant information and the task difficulties are unchanged. Thus, we use the ProcGen suite to investigate the generalization capabilities of DRIBO. For each game, agents are trained on the first 200 levels, and evaluated w.r.t. their zero-shot performance averaged over unseen levels during testing. Unseen levels typically have different backgrounds or layouts, which are relatively easy for humans to adapt to but challenging for RL agents.
We compare DRIBO with recent methods that incorporate data augmentation. In addition to comparing with RAD, we compare DRIBO with DrAC (Raileanu et al., 2020) which applies two regularization terms for policy and value function using augmented data. UCB-DrAC is built on top of DrAC, which automatically selects the best type of data augmentation for DrAC. For RAD and DrAC, we use the best reported augmentation types for different environments. DRIBO selects the same augmentation types except for a few games. The details can be found in Appendix D. We also compare the Procgen results with DAAC (Raileanu & Fergus, 2021) and IDAAC (Raileanu & Fergus, 2021), two state-of-art methods on the Procgen suite that do not apply data augmentation. DAAC decouples the learning of the policy and value function in RL to improve the generalization of RL. IDAAC is built on top of DAAC by adding an auxiliary loss based on an adversarial framework.
Table 1 shows that DRIBO achieves higher averaged testing returns compared to the PPO baseline and other augmentation-based methods. The few environments, in which our approach does not outperform the other augmentation-based methods, share the commonality that task-relevant layouts remain static throughout the same run of the game. Since the current version of DRIBO only considers the mutual information between the complete input and the encoder output (global MI (Hjelm et al., 2019)), it may fail to capture local features. The representations for a sequence of observations within the same run of the game are treated as
globally negative pairs in DRIBO but they may be locally positive pairs. Thus, the performance of DRIBO can be further improved by considering local features (e.g. positions of the layouts) shared between representations as positive pairs in the mutual information estimation. We leave this investigation to future work. DRIBO also outperforms DAAC and IDAAC in 9 of the 16 games.

heading: 6 CONCLUSION
text: We introduce a novel robust representation learning approach based on the multi-view information bottleneck principle for RL problems. Our experimental results show that (1) DRIBO learns representations that are robust against task-irrelevant distractions and boosts the RL agent’s performance even when complex visual distractors are introduced, and (2) DRIBO improves generalization performance compared to well-established baselines on the large-scale Procgen benchmarks.
Reproducibility Statement. The implementation code can be found in Supplementary Material.zip, and we will also release it on GitHub once the paper is published. All datasets we use are public. In addition, we also provide detailed experiment parameters in the Appendix D.

heading: APPENDIX

heading: A THEOREMS AND PROOFS
text: In this section, we first list properties of the mutual information we use in our proof. For any random variables X , Y and Z.
(P.1) Positivity: I(X;Y ) ≥ 0, I(X;Y |Z) ≥ 0
(P.2) Chain rule: I(XY ;Z) = I(Y ;Z) + I(X;Z|Y )
(P.3) Chain rule (Multivariate Mutual Information):
I(X;Y ;Z) = I(Y ;Z)− I(Y ;Z|X)
(P.4) Entropy and Mutual Information:
I(X;Y ) = H(X)−H(X|Y )
(P.5) Chain rule for Entropy:
H(X1, X2, . . . , Xn)= n∑ i=1 H(Xi|Xi−1, . . ., X1)
(P.6) The mutual information among three variables is bounded by:
−min{I(X;Y |Z), I(X;Z|Y ), I(Y ;Z|X)} ≤ I(X;Y ;Z)≤min{I(X;Y ), I(X;Z), I(Y ;Z)}

heading: A.1 THEOREM 1
text: We first show that I(S1:T ;O1:T |A1:T ) is a lower bound of I(S1:T ;O1:T ).
I(S1:T ;O1:T |A1:T ) (P.3) = I(O1:T ;S1:T )− I(S1:T ;O1:T ;A1:T )
Since S1:T are representations of O1:T , we have I(S1:T ;A1:T |O1:T ) = 0. With (P.1) and (P.6), we have that:
I(S1:T ;O1:T ;A1:T ) ≥ 0
Thus, we have:
I(O1:T ;S1:T ) ≥ I(S1:T ;O1:T |A1:T )
Furthermore, since I(S1:T ;O1:T ;A1:T ) is lower bounded, maximizing I(S1:T ;O1:T |A1:T ) also maximizes I(O1:T ;S1:T ).
Theorem A.1. Let O1:T be the observation sequence obtained by executing action sequence A1:T . If S1:T is a sequence of sufficient representations for O1:T , then we have:
I(S1:T ;O1:T |A1:T ) ≥ T∑ t=1 I(St;Ot|St−1, At−1) (8)
Proof. We indicate the property we use for each step of the derivation below.
I(S1:T ;O1:T |A1:T ) (P.4) = H(S1:T |A1:T )−H(S1:T |O1:T , A1:T )
(P.5) = ∑ t (H(St|A1:T , S1:t−1)−H(St|A1:T , O1:T , S1:t−1))
(P.4) = ∑ t I(St;O1:T |A1:T , S1:t−1)
(P.4) = ∑ t (H(O1:T |A1:T , S1:t−1)−H(O1:T |St, A1:T , S1:t−1))
(P.5) = ∑ t ∑ τ (H(Oτ |A1:T , S1:t−1, O1:τ−1)
−H(Oτ |St, A1:T , S1:t−1, O1:τ−1)) (P.4) = ∑ t ∑ τ I(St;Oτ |A1:T , S1:t−1, O1:τ−1)
(P.1) ≥ ∑ t I(St;Ot|A1:T , S1:t−1, O1:t−1)
= ∑ t I(St;Ot|St−1, At−1)
Here, we provide a formal proof for the last step of the derivation above. Let τ represent a1:T , s1:t−1,o1:t−1. We have
I(St;Ot|A1:T , S1:t−1, O1:t−1)
= ∫ τ ∫ st ∫ ot p(τ)p(st,ot|τ) log [ p(st,ot|a1:T , s1:t−1,o1:t−1) p(st|a1:T , s1:t−1,o1:t−1)p(ot|a1:T , s1:t−1,o1:t−1) ] dstdotdτ
With the density of multi-view trajectories Eq. (1), we can observe that ot and st are generated by p(ot|st)p(st|st−1,at−1) = p(st,ot|st−1,at−1). Thus, we have:
p(st,ot|a1:T , s1:t−1,o1:t−1) = p(st,ot|st−1,at−1) p(st|a1:T , s1:t−1,o1:t−1) = p(st|st−1,at−1) p(ot|a1:T , s1:t−1,o1:t−1) = p(ot|st−1,at−1)
This in turn implies that:
I(St;Ot|A1:T , S1:t−1, O1:t−1) = I(St;Ot|St−1, At−1)
As a result, we have a lower bound of I(S(1)1:T ;O (1) 1:T |A1:T ):
I(S (1) 1:T ;O (1) 1:T |A1:T ) ≥ ∑ t (I(S (1) t ;O (1) t |S (1) t−1, At−1, O (2) t )+I(O (2) t ;S (1) t |S (1) t−1, At−1))
With the information bottleneck principle, we have losses in Eq. (5) and Eq. (6) to compress away task-irrelevant information in representations while maximizing the mutual information between representations and observations.

heading: A.2 SUFFICIENT REPRESENTATIONS IN RL
text: In this section, we first present the sufficiency condition for sequential data. Then, we prove that if the sufficiency condition on the sequential data holds, then the sufficiency condition on each corresponding individual representation and observation holds as well.
Theorem A.2. Let O1:T and A∗1:T be random variables with joint distribution p(o1:T ,a∗1:T ). Let S1:T be the representation of O1:T , then S1:T is sufficient for A∗1:T if and only if I(O1:T ;A ∗ 1:T ) = I(S1:T ;A ∗ 1:T ). Also, St is a sufficient representation of Ot since I(Ot;A ∗ t |St, St−1, At−1) = 0.

heading: Hypothesis:
text: (H.1) S1:T is a sequence of sufficient representations for O1:T :
I(O1:T ;A ∗ 1:T |S1:T ) = 0
Proof.
I(O1:T ;A ∗ 1:T |S1:T )
(P.3) = I(O1:T ;A ∗ 1:T )− I(O1:T ;A∗1:T ;S1:T )
(P.3) = I(O1:T ;A ∗ 1:T )− I(A∗1:T ;S1:T )− I(A∗1:T ;S1:T |O1:T )
With S1:T as a representation of O1:T , we have I(S1:T ;A∗1:T |O1:T )=0. The reason is that O1:T shares the same level of information as A∗1:T and S1:T . Then,
I(O1:T ;A ∗ 1:T |S1:T ) =I(O1:T ;A∗1:T )− I(A∗1:T ;S1:T ) (9)
So the sufficiency condition I(O1:T ;A∗1:T |S1:T ) = 0 holds if and only if I(O1:T ;A∗1:T ) = I(A∗1:T ;S1:T ).
We factorize the mutual information between sequential observations and optimal actions
I(O1:t;A ∗ 1:t)
(P.2) = I(Ot;A ∗ 1:t|O1:t−1) + I(O1:t−1;A∗1:t)
(H.1) = I(Ot;A ∗ 1:t|O1:t−1) + I(S1:t−1;A∗1:t)
I(S1:t;A ∗ 1:t)
(P.2) = I(St;A ∗ 1:t|S1:t−1) + I(S1:t−1;A∗1:t)
(H.1) = I(St;A ∗ 1:t|S1:t−1) + I(O1:t−1;A∗1:t)
Then we obtain the following relation:
I(Ot;A ∗ 1:t|O1:t−1) = I(St;A∗1:t|S1:t−1) (10)
We also have
I(Ot;A ∗ 1:t|O1:t−1)
(P.2) = I(O1:t;A ∗ 1:t)− I(O1:t−1;A∗1:t)
(P.2) = I(O1:t−1;A ∗ 1:t|Ot) + I(Ot;A∗1:t)
− I(O1:t−1;A∗1:t) (H.1) = I(S1:t−1;A ∗ 1:t|Ot) + I(Ot;A∗1:t)
− I(S1:t−1;A∗1:t) (P.2) = I(OtS1:t−1;A ∗ 1:t)− I(S1:t−1;A∗1:t)
(P.2) = I(Ot;A ∗ 1:t|S1:t−1)
Eq. (10) = I(St;A ∗ 1:t|S1:t−1)
(P.2)⇐⇒ I(Ot;A ∗ t |A∗1:t−1, S1:t−1) + I(Ot;A∗1:t−1|S1:t−1)
=I(St;A ∗ t |A∗1:t−1, S1:t−1) + I(St;A∗1:t−1|S1:t−1)
Eq. (10)⇐⇒ I(Ot;A ∗ t |A∗1:t−1, S1:t−1)=I(St;A∗t |A∗1:t−1, S1:t−1)
Eq. (9)⇐⇒ I(Ot, A ∗ t |St, S1:t−1, A∗1:t−1) = 0
With the above derivation and Markov property, we have I(Ot;A∗t |St, St−1, At−1) = 0. We can generalize A∗t−1 to any At−1 by assuming A ∗ t as the optimal action for state St whose last timestep state-action pair is (St−1, At−1). Thus, we have St is a sufficient representation for Ot if and only if S1:T is a sufficient representation of O1:T .

heading: A.3 MULTI-VIEW REDUNDANCY AND SUFFICIENCY
text: Proposition A.1. O(1)1:T is a redundant view with respect to O (2) 1:T to obtain A ∗ 1:T if only if I(O (1) 1:T ;A ∗ 1:T |O (2) 1:T ) = 0. Any representation S (1) 1:T ofO (1) 1:T that is sufficient forO (2) 1:T is also sufficient for A∗1:T .
Proof. See proof of Proposition B.3 in the MIB paper (Federici et al., 2020).
Corollary A.1. Let O(1)1:T and O (2) 1:T be two mutually redundant views for A ∗ 1:T . Let S (1) 1:T be a representation of O(1)1:T . If S (1) 1:T is sufficient for O (2) 1:T , S (1) t can derive A ∗ t as the joint observation of the two views (I(O(1)t O (2) t ;A ∗ t |St−1, At−1)=I(S (1) t ;A ∗ t |St−1, At−1)), where St−1 is any sufficient representation at timestep t− 1.
Proof. For the sequential data, see proof of Corollary B.2.1 in the MIB paper (Federici et al., 2020) to prove
I(O (1) 1:TO (2) 1:T ;A ∗ 1:T )=I(S (1) 1:T ;A ∗ 1:T )
According to Theorem A.2, if S(1)1:T is a sufficient representation of O (2) 1:T , S (1) t is a sufficient representation of O(2)t . Similar to proof on sequential data, we can use Corollary B.2.1 in the MIB paper (Federici et al., 2020) to show that
I(O (1) t O (2) t ;A ∗ t |St−1, At−1) = I(S (1) t ;A ∗ t |St−1, At−1)
Theorem A.3. Let the two views o(1)1:T and o (2) 1:T of observation o1:T are obtained by data augmentation transformation sequences t(1)1:T and t (2) 1:T respectively (o (1) 1:T=t (1) 1:T (o1:T ) and o (2) 1:T=t (2) 1:T (o1:T )). We abuse the notation t for simplicity to represent t(1)1:T (O1:T ) and t (2) 1:T (O1:T ) as random variables for augmented observations. Whenever I(t(1)1:T (O1:T );A ∗ 1:T )=I(t (2) 1:T (O1:T );A ∗ 1:T )=I(O1:T ;A ∗ 1:T ), the two views O(1)1:T and O (2) 1:T must be mutually redundant for A ∗ 1:T . Besides, the two views O (1) t and O (2) t must be mutually redundant for A ∗ t .
Proof. Let s1:T be a sufficient representation for both original and multi-view observations. We first factorize the mutual information and refer A.2 as Theorem A.2.
I(t (1) 1:t (O1:t);A ∗ 1:t) = I(O (1) 1:t ;A ∗ 1:t)
(P.2) = I(O (1) t ;A ∗ 1:t|O (1) 1:t−1)+I(O (1) 1:t−1;A ∗ 1:t)
A.2 = I(O (1) t ;A ∗ 1:t|S1:t−1)+I(S1:t−1;A∗1:t)
I(t (2) 1:t (O1:t);A ∗ 1:t) = I(O (2) 1:t ;A ∗ 1:t)
(P.2) = I(O (2) t ;A ∗ 1:t|O (2) 1:t−1)+I(O (2) 1:t−1;A ∗ 1:t)
A.2 = I(O (2) t ;A ∗ 1:t|S1:t−1)+I(S1:t−1;A∗1:t)
I(O1:t;A ∗ 1:t) = I(O1:t;A ∗ 1:t)
(P.2) = I(Ot;A ∗ 1:t|O1:t−1)+I(O1:t−1;A∗1:t)
A.2 = I(Ot;A ∗ 1:t|S1:t−1)+I(S1:t−1;A∗1:t)
Then, we have the following equality
I(O (1) t ;A ∗ 1:t|S1:t−1)=I(O (2) t ;A ∗ 1:t|S1:t−1)=I(Ot;A∗1:t|S1:t−1)
Similar as derivation in Theorem A.2
I(O (1) t ;A ∗ 1:t|S1:t−1)
(P.2) = I(O (1) t ;A ∗ t |A∗1:t−1, S1:t−1) + I(O (1) t ;A ∗ 1:t−1|S1:t−1)
Eq. (10) = I(O (1) t ;A ∗ t |A∗1:t−1, S1:t−1) + I(St;A∗1:t−1|S1:t−1)
We apply the same derivation for o(2) and o, we have the following with Markov property
I(t (1) t (Ot);A ∗ t |St−1, At−1)
=I(t (2) t (Ot);A ∗ t |St−1, At−1) =I(Ot;A ∗ t |St−1, At−1)
We show that the condition on sequential data can be expressed at each timestep with the similar form. See proof of Proposition B.4 in the MIB paper (Federici et al., 2020) for mutual redundancy between sequential views and individual pairs of views.
Theorem A.4. Suppose the mutually redundant condition holds, i.e. I(t
(1) 1:T (O1:T );A ∗ 1:T )=I(t (2) 1:T (O1:T );A ∗ 1:T )=I(O1:T ;A ∗ 1:T ). If S (1) 1:T is a sufficient representa-
tion for t(2)1:T (O1:T ) then I(Ot;A ∗ t |St−1, At−1) = I(S (1) t ;A ∗ t |St−1, At−1).
Proof. Since t(1)1:T (O1:T ) is redundant for t (2) 1:T (O1:T ) (Theorem A.3), any representation S (1) t of t (1) 1:T (O1:T ) that is sufficient for t (2) 1:T (O1:T ) must also be sufficient for A ∗ t (Theorem A.2 and Proposition A.1). Using Theorem A.2 we have I(S(1)t ;A ∗ t |St−1, At−1)=I(t (1) t (Ot);A ∗ t |St−1, At−1). With I(t(1)t (Ot);A ∗ t |St−1, At−1) = I(Ot;A∗t |St−1, At−1), we conclude I(Ot;A∗t |St−1, At−1) = I(S (1) t ;A ∗ t |St−1, At−1).
We finally show the proposition for the Multi-Information Bottleneck principle in RL with the generalization of sufficiency and mutually redundancy condition from sequential data to each individual pairs of data.
Proposition A.2. Let O(1)t and O (2) t be mutually redundant views for A ∗ t that share only optimal action information. Then a sufficient representation of S(1)t of O 1 t for O (2) t that is minimal for O (2) t is also a minimal representation for A∗t .
Proof. See proof of Proposition E.1 in the MIB paper (Federici et al., 2020).

heading: B DRIBO LOSS COMPUTATION
text: We consider the average of the information bottleneck losses from the two views.
L 1+2 2
(11)
= I(S
(1) t ;O (1) t |S (1) t−1,At−1,O (2) t )+I(S (2) t ;O (2) t |S (2) t−1,At−1,O (1) t )
2
− λ1I(S
(1) t ;O (2) t |S (1) t−1, At−1)+λ2I(S (2) t ;O (1) t |S (2) t−1, At−1)
2 (12)
Consider s(1)t and s (2) t on the same domain S, I(S (1) t ;O (1) t |S (1) t−1, At−1, O (2) t ) can be expressed as:
I(S (1) t ;O (1) t |S (1) t−1, At−1, O (2) t )
=E [ log pθ(s (1) t |o (1) t , s (1) t−1,at−1)
pθ(s (1) t |o (2) t , s (1) t−1,at−1)
]
=E [ log pθ(s (1) t |o (1) t , s (1) t−1,at−1)
pθ(s (2) t |o (2) t , s (2) t−1,at−1)
pθ(s (2) t |o (2) t , s (2) t−1,at−1) pθ(s (1) t |o (2) t , s (1) t−1,at−1) ] =DKL(pθ(s (1) t |o (1) t , s (1) t−1,at−1)||pθ(s (2) t |o (2) t , s (2) t−1,at−1))
−DKL(pθ(s(1)t |o (2) t , s (1) t−1,at−1)||pθ(s (2) t |o (2) t , s (2) t−1,at−1))
≤DKL(pθ(s(1)t |o (1) t , s (1) t−1,at−1)||pθ(s (2) t |o (2) t , s (2) t−1,at−1)) (13)
Note that equality holds if the two distributions coincide. Analogously I(S(2)t ;O (2) t |S (2) t−1, At−1, O (1) t ) is upper bounded by DKL(pθ(s (2) t |o (2) t , s (2) t−1,at−1)||pθ(s (1) t |o (1) t , s (1) t−1,at−1)).
Assume St−1 is a sufficient representation of Ot−1. Then, S (1) t−1 provides task-relevant information no more than the sufficient representation St−1. I(S (1) t ;O (2) t |S (1) t−1, At−1) can be thus re-expressed as:
I(S (1) t ;O (2) t |S (1) t−1, At−1)
≥I(S(1)t ;O (2) t |St−1, At−1)
(P.2) = I(S (1) t ;S (2) t O (2) t |St−1, At−1)−I(S (1) t ;S (2) t |O (2) t , St−1, At−1)
∗ =I(S (1) t ;S (2) t O (2) t |St−1, At−1)
=I(S (1) t ;S (2) t |St−1, At−1) + I(S (1) t ;O (2) t |S (2) t , St−1, At−1)
≥I(S(1)t ;S (2) t |St−1, At−1) (14)
Where ∗ follows from S(2)t being the representation of O (2) t . The bound is tight whenever S (2) t is sufficient from S(1)t (I(S (1) t ;O (2) t |S (2) t , St−1, At−1)=0). This happens whenever S (2) t contains all the information regarding S(1)t . Once again, we can have I(S (2) t ;O (1) t |S (2) t−1, At−1) ≥ I(S (1) t ;S (2) t |St−1, At−1). Therefore, the averaged loss functions can be upper-bounded by
L 1+2 2 ≤−λ1 + λ2 2 I(S (1) t ;S (2) t |St−1, At−1) (15)
+DSKL(pθ(s (1) t |o (1) t ,s (1) t−1,at−1)||pθ(s (2) t |o (2) t , s (2) t−1,at−1))
Lastly, by re-parametrizing the objective, we obtain:
L(θ;β)=− Iθ(S(1)t ;S (2) t |St−1, At−1) (16)
+βDSKL(pθ(s (1) t |o (1) t ,s (1) t−1,at−1)||pθ(s (2) t |o (2) t ,s (2) t−1,at−1))
In Algorithm 1, we use s(1)t ∼ pθ(s (1) t |o (1) t , s (1) t−1,at−1) and s (2) t ∼ pθ(s (2) t |o (2) t , s (2) t−1,at−1) to obtain representations for multi-view observations. We argue that the substitution does not affect the effectiveness of the averaged objective. With the multi-view assumption, we have that representations s (1) t−1 and s (2) t−1 do not share any task-irrelevant information. So, the representations at timestep t conditioned on them do not share any task-irrelevant information. Maximizing the mutual information between s(1)t and s (2) t (first term in Eq. (16)) will encourage the representations to share maximal task-relevant information. Similar argument also works for the second term in Eq. (16). Since s(1)t−1 and s(2)t−1 do not share any task-irrelevant information, any task-irrelevant information introduced from the conditional probability will be also identified as task-irrelevant information by KL divergence, which will be reduced through minimizing the DRIBO loss.

heading: C ADDITIONAL RESULTS

heading: C.1 ADDITIONAL DMC RESULTS
text: For clean setting, the pixel observations have simple backgrounds as shown in Figure 2 (left column). Figure 6 shows that RAD, SLAC, CURL and PI-SAC generally perform the best, whereas DRIBO consistently outperforms DBC and matches SOTA.
For natural video setting, Figure 7 shows that DRIBO performs substantially better than RAD, SLAC and CURL which do not discard task-irrelevant information explicitly. Compared to PI-SAC and DBC, recent state-of-art methods that aim at learning representations that are invariant to task-irrelevant information, DRIBO outperforms them consistently.

heading: C.2 DRIBO LOSS VS. INFOMAX

heading: C.2.1 NATURAL VIDEO SETTING
text: To quantify the amount of task-irrelevant information retained in the representations, we compare the SKL term values between DRIBO and DRIBO-no-skl during training in Figure 8. As described in Section 5, DRIBO-no-skl is trained with an InfoMax-based objective without a SKL term. The gap between the SKL values explains the performance gap between the two approaches (as shown in Figure 5). The models trained with DRIBO take advantage of the information bottleneck to map observations from different views close to each other in the latent space. Figure 8 shows that minimizing the DRIBO loss consistently reduces the KL divergence between representations from different views. On the other hand, the models trained with DRIBO-no-skl fail to discard the task-irrelevant information contained in observations from different views even though the RSSM model helps to learn predictive representations. For DRIBO-no-skl, the KL divergence between representations from different views is consistently larger than the one learned by DRIBO.

heading: C.2.2 CLEAN SETTING
text: We provide additional results on comparing DRIBO with DRIBO-no-skl under the clean setting Figure 9 and Figure 10. It can be observed that DRIBO and DRIBO-no-skl perform similarly in terms of average returns. Recall our earlier plot on comparing DRIBO with DRIBO-no-skl under the natural video setting Figure 5 which shows DRIBO substantially outperforms DRIBO-no-skl (in other words, the performance of DRIBO-no-skl drops when the setting is changed from clean to natural video). Similar results can be seen in Figure 6 and Figure 7, where the performance of approaches like RAD, SLAC, CURL and PI-SAC significantly degrades when the backgrounds of the environments are changed to different natural videos during training and testing.

heading: C.3 TRAINING AND TESTING PERFORMANCE IN DMC UNDER THE NATURAL VIDEO SETTING
text: We further compare DRIBO with RAD, CURL and PI-SAC on the DMC environments under the natural video setting. We observe that RAD, CURL and PI-SAC could achieve high scores during training but failed to achieve the same high scores (with a substantial gap) during testing.

heading: C.4 ADDITIONAL VISUALIZATION
text: Figure 14 are the snapshots of training environments of DMC under the natural video setting. The background videos are randomly sampled from the class of "arranging flower" which are drastically different from backgrounds used during testing (Figure 15).
In Figure 15, we show sequential observations in test environments of DMC under the natural video setting. The backgrounds videos are randomly sampled from the test data of Kinetic dataset which contain various classes of videos. Note that a single run of the DMC task may have multiple videos playing sequentially in the background.
Figure 15 also visualizes the corresponding spatial attention maps of DRIBO trained encoders. For different tasks, DRIBO encoders’ activations concentrate on entire edges of the body, providing a more complete and robust representation of the visual observations.
In addition to Figure 2, we also show spatial attention maps for each convolutional layers in Figure 16. We can observe that DRIBO trained encoders filter the complex distractors in the backgrounds gradually layer by layer.
To further compare DRIBO with frame-stacked CURL, we visualize the learned representations using t-SNE plots in Figure 17. We see that even when the backgrounds are drastically different, our encoder learns to map observations with similar robot configurations near each other, whereas CURL’s encoder maps similar robot configurations far away from each other. This shows that CURL does not discard as many irrelevant features in the background as DRIBO does despite leveraging data augmentations and backpropagating RL objectives to the encoder.
D IMPLEMENTATION DETAILS

heading: D.1 RSSM
text: We split the representation st into a stochastic part zt and a deterministic part ht, where st = (zt,ht). The generative and inference models of RSSM are defined as:
Deterministic state transition: ht = f(ht−1, zt−1,at−1) Stochastic state transition: zt = p(zt|ht)
Observation model: ot = p(ot|ht, zt) where f(ht−1, zt−1,at−1) is implemented as a recurrent neural network (RNN) that carries the dependency on the stochastic and deterministic parts at the previous timestep. Then, we obtain the representation with the encoder pθ((s1:T |o1:T ,a1:T )= ∏ t pθ(st|ot,ht), where ht retains information from st−1 = (zt−1,ht−1) and at−1.

heading: D.2 DRIBO + SAC
text: We first show how we train SAC agent given the representations of DRIBO. Let φ(o) = s ∼ pθ(s|o, s′,a′) denote the encoder, where s′ and a′ as the representation and action at last timestep.
Algorithm 2 SAC + DRIBO Encoder
1: Input RL batch BRL = {(φ(oi), ai, ri, φ(o′i))} (T−1)∗N i=1 with (T −1)∗N pairs of representation,
action, reward and next representation. 2: Get value: V = mini=1,2 Q̂i(φ̂(o′))− α log π(a|φ̂(o′)) 3: Train critics: J(Qi, φ) = (Qi(φ(o))− r − γV )2 4: Train actor: J(π)=α log π(a|φ(o))−mini=1,2Qi(φ(o)) 5: Train alpha: J(α) = −α log π(a|φ(o))−αH(a|φ(o)) 6: Update target critics: Q̂i = τQQi + (1− τQ)Q̂i 7: Update target encoder: φ̂← τφφ+ (1− τφ)φ
Then we incorporate the above SAC algorithm into minimizing DRIBO loss in Algorithm 3
Algorithm 3 DRIBO + SAC 1: Input: Replay buffer D storing sequential observations and actions with length T . The batch
size is N . The number of total training step is K. The number of total episodes is E. 2: for e = 1, . . . , E do 3: Sample sequential observations and actions from the environment and append new samples
to D. 4: for each step k = 1, . . . ,K do 5: Sample a sequential batch B ∼ D. 6: Compute the representations batch BRL which has the shape (T,N) using the encoder pθ(s1:T |o1:T ,a1:T ) 7: Train SAC agent: EBRL [J(π,Q, φ)] . Algorithm 2 8: Update θ and ψ to minimize LDRIBO using B . Algorithm 1. 9: end for
10: end for

heading: D.3 DRIBO + PPO
text: The main difference between SAC and PPO is that PPO is an on-policy RL algorithm while SAC is an off-policy RL algorithm. With the update of the encoder, representations may not be consistent within each training step which breaks the on-policy sampling assumption. To address this issue, instead of obtaining st propagating from the initial observation of the observation sequence, we store the representations as soldt while sampling from the on-policy batch. Then, we use ϕ(o) = s ∼
pθ(s|o, sold,a′) to denote the representation from the encoder. Here, sold and a′ are the representation and action at the previous timestep. By treating the encoding process as a part of the policy and value function, the on-policy requirement is satisfied since the new action/value at timestep t depends only on (ot, soldt−1,at−1).
Algorithm 4 DRIBO + PPO 1: Input: Replay buffer D and on-policy replay buffer DPPO storing sequential observations and
actions with length T . The batch size is N . The minibatch size for PPO is M. The number of total episodes is E. 2: for e = 1, . . . , E do 3: Sample sequential observations and actions from the environment {(o1:T ,a1:T , r1:T , sold1:T }Ni=1. 4: Append new samples to D and update the on-policy replay buffer DPPO. 5: for j = 1, . . . ,M do 6: {(ϕ(oi), ai, ri)} bT∗NM c i=1 ∼ DPPO 7: Optimize PPO policy, value function and encoder using each sample (ϕ(oi), ai, ri) in the batch. 8: Sample a sequential batch B ∼ D. 9: Update θ and ψ to minimize LDRIBO using B . Algorithm 1.
10: end for 11: end for

heading: D.4 DMC
text: We use an almost identical encoder architecture as the encoder in the RSSM paper (Hafner et al., 2019), with two minor differences. Firstly, we deploy the encoder architecture in Tassa et al. (2018) as the observation embedder, with two more convolutional layers to the CNN trunk. Secondly, we add layer normalization to process the output of CNN, deterministic output and stochastic output. Deterministic part of the representation is a 200-dimensional vector. Stochastic part of the representation is a 30-dimensional diagonal Gaussian with predicted mean and standard deviation. Thus, the representation is a 230-dimensional vector. We implement Q-network and policy in SAC as MLPs with two fully connected layers of size 1024 with ReLU activations. We estimate the mutual information using a bi-linear inner product as the similarity measure (Oord et al., 2018).
Pixel Preprocessing. We construct an observational input as a single frame, where each frame is a RGB rendering of size 100× 100 from the 0th camera. We then divide each pixel by 255 to scale it down to [0, 1) range. For methods compared in our experiments, an observation inputs contains an 3-stack of consecutive frames.
Augmentations of Visual Observations. For RAD, we use random crop+ grayscale to generate augmented data. For our approach DRIBO and CURL, we use random crop to generate multi-view observations. We apply the implementation of RAD to do the augmentation. For random crop, it extracts a random patch from the original observation. In DMC, we render 100 × 100 pixel observations and crop randomly to 84× 84 pixels. For random grayscale, it converts RGB images to grayscale with a probability p = 0.3.
Hyperparameters. To facilitate the optimization, the hyperparameter β in the DRIBO loss Algorithm 1 is slowly increased during training. β value starts from a small value 1e− 4 and increases to 1e− 3 with an exponential scheduler. The same procedure is also used in the MIB paper (Federici et al., 2020). We show other hyperparameters for DMC experiments in Table 2.
KL Balancing. During training, we also incorporate KL balancing from a variant method described in DreamerV2 (Hafner et al., 2021) to train RSSM. KL balancing encourages learning an accurate prior over increasing posterior entropy, so that the prior better approximates the aggregate posterior. This help us to avoid regularizing the representations toward a poorly trained prior. KL balancing is orthogonal to our MIB objective (DRIBO Loss). Note that DreamerV2 deploy KL balancing based on a reconstruction loss. In our case, our results show that KL balancing can improve training of RSSM with a contrastive-learning-based or mutual-information-maximization objective. We implement
this technique as shown in Algorithm 5. β shares the same value as the coefficient for SKL term in DRIBO Loss.
Algorithm 5 DRIBO Loss + KL Balancing 1: Compute kl balancing term:
kl_balancing = 0.8 · compute_kl(stop_grad(approx_posterior) + prior) + 0.2 · compute_kl(approx_posterior + stop_grad(prior))
return DRIBO Loss + β· kl_balancing

heading: D.5 PROCGEN
text: For Procgen suite, the implementation of DRIBO is almost the same as DMC experiments. Better design choice could be found by validation. We use the same as the encoder architecture used in DMC experiments, except for the observation embedder, which we use the network from IMPALA paper to take the visual observations. In addition, since the actions in Procgen suite are discrete, we use an action embedder to embed discrete actions into continuous space. The action embedder is implemented as a simple one hidden layer resblock with 64 neurons. It maps a one-hot action vector to a 4-dimensional vector. The policy and value function share one hidden layer with 1024 neurons. The policy uses another fully connected layer to generate a categorical distribution to select the discrete action. The value function uses another fully connected layer to generate the value for an input representation. All activation functions are ReLU activations.
Augmentation of Visual Observations. We select augmentation types based on the best reported augmentation types for each environment. DrAC (Raileanu et al., 2020) reported best augmentation types for RAD and DrAC in Table 4 and 5 of the DrAC paper. We list the augmentation types used in DRIBO in Table 3 and 4. We use the same settings for each augmentation type as DrAC. Note that we only performed limited experiments to select the augmentations reported in the tables due to time constraints. So, the tables do not show the best augmentation types in each environment for DRIBO.
Hyperparameters. We use the same β scheduler as the DMC experiments. The starting β value is 1e− 8 and the final β value is 1e− 3. We show other hyperparameters for Procgen environments in Table 5.
Discussion. Here, we extend the discussion on why our method underperforms on some environments, whose screenshots are shown in Figure 18.
Our approach, DRIBO, only consider global MI (Hjelm et al., 2019) in the current implementation. As a result, local structures can be easily ignored in the representation. More specifically, representations containing the same static local features within a single execution but at different timesteps are treated as negative examples in the mutual information maximization. Then, the information of these local features shared between representations is not maximized. Negative pairs of representations sharing this type of local features are globally negative pairs but locally positive pairs.
For Plunder, the goal is to destroy moving enemy pirate ships by firing cannonballs. The enemy ships can be identified with the color of the target in the bottom left corner. The background of the game maintains the same within a single execution. Wooden obstacles capable of blocking the player’s cannonballs. In a single execution, critical features like the target label and wooden obstacles remain unchanged. Failing to capture these local features results in poor performance of an agent.
For Chaser, the goal is to collect all green orbs as well as stars. A collision with an enemy that is not vulnerable (red) results in the death of the player. The background remains the same across different executions. Walls in the environment are dense and remain static during a single execution. Failing to capture walls’ positions in representations hinders the agent from navigating to avoid enemies and collect orbs/stars. In Maze and Leaper, walls also remain static, but the backgrounds are different across different executions. This difference reduces the influence introduced by globally negative pairs but locally positive pairs. By contrast, walls in DodgeBall remain static but more critical since hitting at a wall ends the game.

heading: D.6 COMPUTE RESOURCES AND LICENSE
text: Compute Resources. We used a desktop with a 12-core CPU and a single GTX 1080 Ti GPU for benchmarking. Each seed for DMC benchmarks takes 3 days to finish. For Procgen suite, it takes 2 days to finish experiments on each seed.
License. DMC benchmarks are simulated in MuJoCo 2.0. We perform the experiments in this paper under an Academic Lab License. We perform experiments in Procgen suite with its open source code under MIT License.

Review summaries:
This paper tackles the problem of generalization amidst visual distractors for control tasks. In particular, the distractors have no dependence on the optimal policy and thus clearly form a task-irrelevant component. The proposal is to use mutual information between two views as a proxy for how much task-relevant information is present in the constructed representation. This objective is adapted for RL to consider the long term sequential nature. Finally, the authors test this approach on ProcGen and DMC Suite with distractors, while also performing certain ablations. 

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). 

This paper studies the problem of pixel-based control with reinforcement learning. The authors categorize observation changes into task-relevant and task-irrelevant changes. They define task-irrelevant changes as those do not have casual relations with actions, and introduce a conditional prior to compress task-irrelevant information, such that agent can focus on task-relevant information. The authors evaluated the proposed DRIBO method on pixel-based control tasks (DM Control) with natural video background and the procgen suite.