# DRIBO: ROBUST DEEP REINFORCEMENT LEARNING
## VIA MULTI-VIEW INFORMATION BOTTLENECK

**Anonymous authors**
Paper under double-blind review

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.

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_


-----

Sequential Observations


#!:#(!) Learning ObjectiveReinforcement

#!:#(') Multi-View Information Bottleneck Loss


%!:#(!)

%!:#(')


!!:#


View 1

View 2


Encoder

!$(#!:#|%!:#, '!:#)


Figure 1: Robust Deep Reinforcement Learning via Multi-View Infomration BOttleneck (DRIBO)
incorporates the inherent temporal structure of RL and unsupervised multi-view settings into robust
representation learning in RL. We consider sequential multi-view observations, o[(1)]1:T [and][ o]1:[(2)]T [, of]
original observations o1:T sharing the same task-relevant information while any information not
shared by them are task-irrelevant. For example, natural video backgrounds of sequential observations
are task-irrelevant and can be drastically different between training and testing environments. DRIBO
uses a multi-view information bottleneck loss to ensure that s[(1)]1:T [and][ s]1:[(2)]T [, the representations of]
multi-view observations, shares maximal task-relevant information while reducing the task-irrelevant
information. DRIBO trains the RL policy and (or) value function on top of the encoder.

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 ob-_**
**_servations. 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 en-**
coder 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.

3 PRELIMINARIES

We denote a Markov decision process (MDP) as M, with state s, action a, and reward r. S and A
stand for the corresponding random variables. We denote a policy on M as π. The agent’s goal is to
learn a policy π that maximizes the cumulative rewards. We define S⊆R[d] 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[(]1[i][)][,][ a][1][, . . .,][ s][T][,][ o]T[(][i][)][,][ a][T][ ]][ where][ T][ is the length. Knowing that the trajectory density is]_
defined over joint observations, states, and actions, we write:

_pπ(τ_ _[M])=π(aT |sT )Pobs[(][i][)][(][o]T[(][i][)][|][s][T][ )][P][(][s][T][ |][s][T][ −][1][,][ a][T][ −][1][)][ · · ·][ π][(][a][1][|][s][1][)][P]obs[(][i][)][(][o]1[(][i][)][|][s][1][)][P][0][(][s][1][)]_ (1)

with P0 being the initial state distribution, P being the transition model and Pobs[(][i][)] [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][ o][1:][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][ S][1:][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][(][S][1:][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][(][S][1:][T] [;][ A][∗]1:T [)] (2)

Conditional mutual informationirrelevant. I(S1:T ; A[∗]1:T [)][ quantifies] I(S1:[ task-relevant]T ; O1:T |A[∗]1:T[ information that is accessible from][)][ quantifies the information in][ S][1:][T][ that is][ S][1:][T][ . The][ task-]
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 wheneveris minimized. Maximizing I(O1:T ; S1:T ) learns a sufficient representation. With the information I(O1:T ; S1:T |A[∗]1:T [)]
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 ofdirectly in supervised settings where A[∗]1:T [are observed. One option is to use MIB which can compress] I(S1:T ; O1:T |A[∗]1:T [)][ can only be done]
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.

4 DRIBO

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.

4.1 MUTUAL INFORMATION MAXIMIZATION

To capture the temporal evolution of observations and representations given any action sequence,
we consider maximizing the conditional mutual informationbound of I(O1:T ; S1:T ) (see Appendix A.1). The observations I O(1:ST1: are generated sequentially in theT ; O1:T |A1:T ) which is a lower
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(St; Ot|St−1, At−1)_ (3)
_t=1_

X


_I(S1:T ; O1:T |A1:T ) ≥_


The proof is included in Appendix A.1. Theorem 1 shows that the sum of mutual information (St; Ot _St_ 1, At 1) over multiple timesteps is a lower bound of (S1:T ; O1:T _A1:T )._
_I(S It; Ot|St−1|, A−t−1)_ _−is_ defined with conditional probabilities, _Ip(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 ( (St; Ot _St_ 1, At 1)) encourages encoding more detailed features
_I_ _|_ _−_ _−_
from Ot into St and makes St sufficient.


-----

4.2 MULTI-VIEW SETTING

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 Ot[(1)] and Ot[(2)]. We refer to Appendix A for the sufficiency condition of representations
and mutually redundancy condition (Federici et al., 2020) between Ot[(1)] and Ot[(2)][. 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 St[(1)]
and Ot[(1)] as in Eq. (2), we can identify two components:

(St[(1)]; Ot[(1)] _t_ 1[, A][t][−][1][) =][ I][(][S]t[(1)]; Ot[(1)] _t_ 1[, A][t][−][1][, O]t[(2)][)+][I][(][O]t[(2)][;][ S]t[(1)] _St[(1)]1[, A][t][−][1][)]_ (4)
_I_ _[|][S][(1)]−_ _[|][S][(1)]−_ _|_ _−_

Here, St[(1)]1 [is a representation of visual observation][ O]t[(1)]1[. Since we assume mutual redundancy]
_−_ _−_
between the two views, the information shared between Ot[(1)] and St[(1)] conditioned on Ot[(2)] must
be irrelevant to the task, which can be quantified as (St[(1)]; Ot[(1)] _t_ 1[, A][t][−][1][, O]t[(2)][)][ (first term in]
_I_ _[|][S][(1)]−_
Eq. (4)). Then, (Ot[(2)][;][ S]t[(1)] _St[(1)]1[, A][t][−][1][)][ has to be maximal if the representation is sufficient. A]_
_I_ _|_ _−_
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 LIB[(1)] [to maximize the]
mutual information between representations and observations and compress away task-irrelevant information (Ot[(2)][;][ S]t[(1)] _St[(1)]1[, A][t][−][1][)][ based on the information bottleneck principle.][ λ][1][ is a Lagrange]_
_I_ _|_ _−_
multiplier. This loss also retains the temporally evolving information of the underlying dynamics.


_LIB[(1)]_ [=][ −]


( (St[(1)]; Ot[(1)] _t_ 1[, A][t][−][1][, O]t[(2)][)][ −] _[λ][1][I][(][O]t[(2)][;][ S]t[(1)]_ _St[(1)]1[, A][t][−][1][))]_ (5)
_I_ _[|][s][(1)]−_ _|_ _−_


Symmetrically, we define a loss LIB[(2)] [for representations and observations from view][ 2][:]

IB [=][ −] ( (St[(2)]; Ot[(2)] _t_ 1[, A][t][−][1][, O]t[(1)][)][ −] _[λ][2][I][(][O]t[(1)][;][ S]t[(2)]_ _St[(2)]1[, A][t][−][1][))]_ (6)
_L[(2)]_ _t_ _I_ _[|][s][(2)]−_ _|_ _−_
X

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.

4.3 DRIBO LOSS FUNCTION

By re-parameterizing the Lagrangian multipliers (details in Appendix B), the average of two loss
functions LIB[(1)] [and][ L]IB[(2)] [from two views at timestep][ t][ can be upper bounded as follows:]

_Lt(θ; β)=−Iθ(St[(1)]; St[(2)]|St−1, At−1)+_ (7)

_βDSKL(pθ(s[(1)]t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)][||][p][θ][(][s]t[(2)]_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][))]_

_[|][o][(1)]_ _−_ _[|][o][(2)]_ _−_

wherest 1 is a sufficient representation, θ denotes the parameters of an encoder DSKL represents the symmetrized KL divergence obtained pθ(s[(1)]t _[|][o]t[(1)][,][ s][t][−][1][,][ a][t][−][1][)][ (details in Section][ 4.4][),]_
_−_
by averaging the expected values of DKL(pθ(s[(1)]t _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)][||][p][θ][(][s]t[(2)]_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][))][,]_

_[|][o][(1)]_ _−_ _[|][o][(2)]_ _−_
_DKL(pθ(s[(2)]t_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][)][||][p][θ][(][s]t[(1)]_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][))][, and the coefficient][ β][ represents the trade-]_
off between sufficiency and sensitivity to task-irrelevant information.[|][o][(2)] _−_ _[|][o][(1)]_ _−_ _β 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]1:[(2)]T_ _[,][ a][1:][T][ )][ in enumerate][ (][B][(1)][,][ B][(2)][)][ do]_

4: **for t = 1 to T do**

5: **_s[(1)]t_** _pθ(s[(1)]t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)][,][ s]t[(2)]_ _pθ(s[(2)]t_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][)]_
_∼_ _[|][o][(1)]_ _−_ _∼_ _[|][o][(2)]_ _−_

6: (s[(1)][,t][+][T][ (][i][−][1)], s[(2)][,t][+][T][ (][i][−][1)]) ← (s[(1)]t _[,][ s]t[(2)][)]_

7: **end for**

8: SKL [=][ 1]T _Tt=1_ _[D][DKL][(][p][θ][(][s]t[(1)]_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)][||][p][θ][(][s]t[(2)]_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][))]_

9: end forL[i] _[|][o][(1)]_ _−_ _[|][o][(2)]_ _−_

_N_

10: return _Iψ(_ (Ps[(1)][,i], s[(2)][,i]) _i=1_ [) +][ β]N _i=1_ SKL
_−_ [ˆ] _{_ _}[T][ ∗][N]_ _[L][i]_

P

To generalize the above loss to sequential data in RL, we apply Theorem 1 to obtain the DRIBO
loss: DRIBO= _T[1]_ _Tt=1_
_L_ _[L][t][(][θ][;][ β][)][. We summarize the batch-based computation of the loss function in]_

Algorithmrespectively. Though the first term in Eq. ( 1. We sampleP **_s[(1)]t_** and s[(2)]t from p7) is conditioned onθ(s[(1)]t _[|][o]t[(1)][,][ s]t[(1)]−1[,][ a] s[t][−]t−[1][)]1[ and], we prove that the sampling[ p][θ][(][s]t[(2)][|][o]t[(2)][,][ s]t[(2)]−1[,][ a][t][−][1][)]_
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 St[(1)] and St[(2)] using the
encoder. The mutual information between the two representations Iθ(St[(1)]; St[(2)]|St−1, At−1) can be
maximized by using any sample-based differentiable mutual information lower bound _I[ˆ]ψ(s[(1)]t_ _[,][ s]t[(2)][)][,]_
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]t[(2)][)][ 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.


4.4 ENCODER ARCHITECTURE AND INCORPORATING DRIBO IN RL

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 oftemporal structure of RL. We implement the encoder as a recurrent space model (RSSM ( pθ(s1:T |o1:T, a1:T ) = _t_ _[p][θ][(][s][t][|][o][t][,][ s][t][−][1][,][ a][t][−][1][)][, which reflects the inherent]Hafner_

et al., 2019)) which leverages recurrent neural networks to perform accurate long-term predictions.

[Q]

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.

5 EXPERIMENTS

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.

5.1 EFFECTIVENESS AND ROBUSTNESS


Figure 2: Left: DMC observations
without visual distractors. Middle: observations with natural videos as backgrounds. Right: spatial attention maps
of encoders for the middle images.


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.


cartpole/swingup

1 2 3 4 5 6 7

400

300

200

100

hopper/hop


cheetah/run

2 3 4 5 6

|Col1|100 80 60 40 20|
|---|---|


walker/walk


finger/spin
1000

800

600

400

200


800

600

400

200

0

150

100

50


400

300

200

100

0

|Col1|Col2|
|---|---|


3 4 5 6

walker/run


|Col1|40 30 20 10|
|---|---|

|Col1|Col2|
|---|---|


800

600

400

200


Environment Steps1e5 Environment Steps1e5 Environment Steps1e5

DRIBO (Ours) RAD SLAC CURL DBC PI-SAC

Figure 3: Results for DMC over 5 seeds with one standard error shaded in the natural video setting.

**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
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
higher returns at 88e4 steps compared to the second best performing method.

Figure 4: t-SNE of latent spaces learned with DRIBO (left t-SNE) and CURL (right t-SNE). We
color-code the embedded points with reward values (higher value yellow, lower value green). Each
pair of solid lines indicates the corresponding embedded points for observations with an identical
foreground but different backgrounds. DRIBO learns representations that are neighboring in the
embedding space with similar reward values. This property holds even if the backgrounds are
drastically different (see middle images). By contrast, CURL maps the same image pairs to points far
away from each other in the embedding space.

We visualize the representations learned by DRIBO and CURL with t-SNE (
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.

5.2 ABLATIONS

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





cartpole/swingup cheetah/run finger/spin

1000

800

400

750

600 300

500

400 200

250

200 100

0

Average Returns 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

hopper/hop walker/walk 1e5 walker/run

400

150 800

300

600

100

400 200

50

200 100

Average Returns 0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

Environment Steps1e5 Environment Steps1e5 Environment Steps1e5

seq-len-3 seq-len-6 seq-len-16 seq-len-32 DRIBO-no-skl


Figure 5: DRIBO achieves better performance by capturing the temporal structure of RL from longer
training sequences. Compressing away task-irrelevant information using the SKL term in DRIBO
loss improves performance when the architecture choice and training configurations are the same. We
perform 5 runs for each method under the natural video setting. More results are in Appendix C.2.


-----

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.

5.3 GENERALIZATION

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 aver- Table 1: Procgen returns on test levels after trainaged testing returns compared to the PPO baseline ing on 25M environment steps. The mean and
and other augmentation-based methods. The few standard deviation are computed over 10 seeds.
environments, in which our approach does not outperform the other augmentation-based methods,
The representations for a sequence of observationsshare the commonality that task-relevant layoutsremain static throughout the same run of the game.Since the current version of DRIBO only consid-ers the mutual information between the completeinput and the encoder output (global MI (et al., 2019)), it may fail to capture local features.Hjelm DodgeBallCaveFlyerBossFightCoinRunStarPilotFruitBotClimberBigFishPlunderJumperChaserLeaperMinerNinjaMazeHeist 24.726.711.74.07.75.95.05.18.55.85.05.72.44.95.78.5 ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± 1.2 1.0 0.7 0.5 0.9 0.5 0.5 0.8 0.8 0.5 0.7 0.6 0.5 3.4 0.8 0.3 33.427.39.97.96.98.55.19.06.55.96.94.14.36.19.42.8 ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± 0.7 1.7 0.6 0.8 1.2 0.6 0.8 0.6 1.0 0.8 1.0 1.0 1.0 1.2 5.1 1.8 29.528.24.38.77.57.09.56.38.86.65.77.14.05.36.69.8 ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± 1.4 0.8 0.4 1.0 0.8 0.2 0.4 0.6 0.7 0.8 0.8 1.1 0.8 0.6 5.4 0.8 30.228.39.79.78.36.98.95.38.56.46.76.54.74.05.06.3 ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± 1.0 0.8 0.6 1.0 0.9 0.6 0.6 0.6 0.8 0.7 0.7 0.3 0.6 0.7 2.8 0.9 17.836.428.620.79.66.84.69.26.56.67.83.33.37.35.58.6 ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± 0.5 0.4 0.2 0.2 0.4 1.2 0.2 0.5 0.2 1.1 0.2 0.9 1.4 2.8 0.6 3.3 27.918.537.023.39.46.88.37.79.86.85.06.33.33.55.69.5 ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± 0.6 0.4 0.6 0.1 0.2 1.0 0.4 0.3 0.2 1.0 0.3 0.4 1.2 2.3 0.5 1.4 10.936.530.812.09.75.87.59.28.44.88.13.87.75.38.59.8 ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± ± 0.7 1.0 1.0 0.7 1.6 0.8 1.6 0.9 1.6 1.5 1.6 0.9 1.6 3.0 0.8 0.5
within the same run of the game are treated as

|Env PPO RAD DrAC UCB-DrAC|Col2|DAAC IDAAC|DRIBO|
|---|---|---|---|
|BigFish StarPilot FruitBot BossFight Ninja Plunder CaveFlyer CoinRun Jumper Chaser Climber DodgeBall Heist Leaper Maze Miner|4.0 ± 1.2 9.9 ± 1.7 8.7 ± 1.4 9.7 ± 1.0 24.7 ± 3.4 33.4 ± 5.1 29.5 ± 5.4 30.2 ± 2.8 26.7 ± 0.8 27.3 ± 1.8 28.2 ± 0.8 28.3 ± 0.9 7.7 ± 1.0 7.9 ± 0.6 7.5 ± 0.8 8.3 ± 0.8 5.9 ± 0.7 6.9 ± 0.8 7.0 ± 0.4 6.9 ± 0.6 5.0 ± 0.5 8.5 ± 1.2 9.5 ± 1.0 8.9 ± 1.0 5.1 ± 0.9 5.1 ± 0.6 6.3 ± 0.8 5.3 ± 0.9 8.5 ± 0.5 9.0 ± 0.8 8.8 ± 0.2 8.5 ± 0.6 5.8 ± 0.5 6.5 ± 0.6 6.6 ± 0.4 6.4 ± 0.6 5.0 ± 0.8 5.9 ± 1.0 5.7 ± 0.6 6.7 ± 0.6 5.7 ± 0.8 6.9 ± 0.8 7.1 ± 0.7 6.5 ± 0.8 11.7 ± 0.3 2.8 ± 0.7 4.3 ± 0.8 4.7 ± 0.7 2.4 ± 0.5 4.1 ± 1.0 4.0 ± 0.8 4.0 ± 0.7 4.9 ± 0.7 4.3 ± 1.0 5.3 ± 1.1 5.0 ± 0.3 5.7 ± 0.6 6.1 ± 1.0 6.6 ± 0.8 6.3 ± 0.6 8.5 ± 0.5 9.4 ± 1.2 9.8 ± 0.6 9.7 ± 0.7|17.8 ± 1.4 18.5 ± 1.2 36.4 ± 2.8 37.0 ± 2.3 28.6 ± 0.6 27.9 ± 0.5 9.6 ± 0.5 9.8 ± 0.6 6.8 ± 0.4 6.8 ± 0.4 20.7 ± 3.3 23.3 ± 1.4 4.6 ± 0.2 5.0 ± 0.6 9.2 ± 0.2 9.4 ± 0.1 6.5 ± 0.4 6.3 ± 0.2 6.6 ± 1.2 6.8 ± 1.0 7.8 ± 0.2 8.3 ± 0.4 3.3 ± 0.5 3.3 ± 0.3 3.3 ± 0.2 3.5 ± 0.2 7.3 ± 1.1 7.7 ± 1.0 5.5 ± 0.2 5.6 ± 0.3 8.6 ± 0.9 9.5 ± 0.4|10.9 ± 1.6 36.5 ± 3.0 30.8 ± 0.8 12.0 ± 0.5 9.7 ± 0.7 5.8 ± 1.0 7.5 ± 1.0 9.2 ± 0.7 8.4 ± 1.6 4.8 ± 0.8 8.1 ± 1.6 3.8 ± 0.9 7.7 ± 1.6 5.3 ± 1.5 8.5 ± 1.6 9.8 ± 0.9|

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.

6 CONCLUSION

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.

REFERENCES

Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Marc G Bellemare. Contrastive
behavioral similarity embeddings for generalization in reinforcement learning. International
_Conference on Learning Representation, 2021. 1_

Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron
Courville, and Devon Hjelm. Mutual information neural estimation. In International Conference
_on Machine Learning, pp. 531–540. PMLR, 2018. 3_

William Bialek and Naftali Tishby. Predictive information. arXiv preprint cond-mat/9902341, 1999.

3

Mariusz Bojarski, Davide Del Testa, Daniel Dworakowski, Bernhard Firner, Beat Flepp, Prasoon
Goyal, Lawrence D Jackel, Mathew Monfort, Urs Muller, Jiakai Zhang, et al. End to end learning
for self-driving cars. arXiv preprint arXiv:1604.07316, 2016. 1

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for
contrastive learning of visual representations. In International conference on machine learning, pp.
1597–1607. PMLR, 2020. 3

Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization
in reinforcement learning. In International Conference on Machine Learning, pp. 1282–1289.
PMLR, 2019. 1

Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation
to benchmark reinforcement learning. In International Conference on Machine Learning. PMLR,
2020. 7

Jesse Farebrother, Marlos C Machado, and Michael Bowling. Generalization and regularization in
dqn. arXiv preprint arXiv:1810.00123, 2018. 1

Marco Federici, Anjan Dutta, Patrick Forré, Nate Kushman, and Zeynep Akata. Learning robust
representations via multi-view information bottleneck. International Conference on Learning
_Representation, 2020. 1, 3, 4, 5, 16, 17, 18, 27_

Ian Fischer. The conditional entropy bottleneck. Entropy, 22(9):999, 2020. 1, 3

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy
maximum entropy deep reinforcement learning with a stochastic actor. In International Conference
_on Machine Learning, pp. 1861–1870. PMLR, 2018. 6_

Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James
Davidson. Learning latent dynamics for planning from pixels. In International Conference on
_Machine Learning, pp. 2555–2565. PMLR, 2019. 1, 2, 6, 27_

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning
behaviors by latent imagination. International Conference on Learning Representation, 2020. 2

Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete
world models. International Conference on Learning Representation, 2021. 2, 27

Olivier J Hénaff, Aravind Srinivas, Jeffrey De Fauw, Ali Razavi, Carl Doersch, SM Eslami, and
Aaron van den Oord. Data-efficient image recognition with contrastive predictive coding. In
_International Conference on Machine Learning, pp. 4182–4192. PMLR, 2020. 3_

R Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Phil Bachman, Adam
Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation
and maximization. International Conference on Learning Representation, 2019. 3, 9, 29


-----

Rico Jonschkowski and Oliver Brock. Learning state representations with robotic priors. Autonomous
_Robots, 39(3):407–428, 2015. 1_

Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan,
Fabio Viola, Tim Green, Trevor Back, Paul Natsev, et al. The kinetics human action video dataset.
_arXiv preprint arXiv:1705.06950, 2017. 6, 7_

Hyoungseok Kim, Jaekyeom Kim, Yeonwoo Jeong, Sergey Levine, and Hyun Oh Song. Emi:
Exploration with mutual information. In International Conference on Machine Learning, pp.
3360–3369, 2019. 3

Sascha Lange and Martin Riedmiller. Deep auto-encoder neural networks in reinforcement learning.
In The 2010 International Joint Conference on Neural Networks (IJCNN), pp. 1–8. IEEE, 2010. 2

Sascha Lange, Martin Riedmiller, and Arne Voigtländer. Autonomous reinforcement learning on raw
visual input data in a real world application. In The 2012 international joint conference on neural
_networks (IJCNN), pp. 1–8. IEEE, 2012. 2_

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations
for reinforcement learning. In International Conference on Machine Learning, pp. 5639–5650.
PMLR, 2020a. 2, 3, 7

Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in Neural Information Processing Systems, 33,
2020b. 2, 7

Alex Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep
reinforcement learning with a latent variable model. Advances in Neural Information Processing
_Systems, 33, 2020a. 1, 2, 7_

Kuang-Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio
Guadarrama. Predictive information accelerates learning in rl. Advances in Neural Information
_Processing Systems, 33:11890–11901, 2020b. 1, 3, 7_

Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep
visuomotor policies. The Journal of Machine Learning Research, 17(1):1334–1373, 2016. 1

Minne Li, Lisheng Wu, WANG Jun, and Haitham Bou Ammar. Multi-view reinforcement learning.
In Advances in neural information processing systems, pp. 1420–1431, 2019. 3

Yingming Li, Ming Yang, and Zhongfei Zhang. A survey of multi-view representation learning.
_IEEE transactions on knowledge and data engineering, 31(10):1863–1883, 2018. 1_

Bogdan Mazoure, Remi Tachet des Combes, Thang Long DOAN, Philip Bachman, and R Devon
Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing
_Systems, 33, 2020. 3_

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare,
Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control
through deep reinforcement learning. nature, 518(7540):529–533, 2015. 1

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive
coding. arXiv preprint arXiv:1807.03748, 2018. 3, 6, 27

Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational
bounds of mutual information. In International Conference on Machine Learning, pp. 5171–5180.
PMLR, 2019. 3

Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization in reinforcement
learning. In Proceedings of the 38th International Conference on Machine Learning, pp. 1334–
1373, 2021. 3, 9

Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data
augmentation for generalization in deep reinforcement learning. arXiv preprint arXiv:2006.12862,
2020. 9, 29


-----

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy
optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. 6

Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden,
Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint
_arXiv:1801.00690, 2018. 6, 27_

Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv
_preprint physics/0004057, 2000. 4_

Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine
_learning research, 9(11), 2008. 8_

Niklas Wahlström, Thomas B Schön, and Marc P Desienroth. From pixels to torques: Policy learning
with deep dynamical models. In Deep Learning Workshop at the 32nd International Conference
_on Machine Learning (ICML 2015), July 10-11, Lille, France, 2015. 2_

Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to
control: a locally linear latent dynamics model for control from raw images. In Proceedings
_of the 28th International Conference on Neural Information Processing Systems-Volume 2, pp._
2746–2754, 2015. 2

Wenhao Yu, C Karen Liu, and Greg Turk. Policy transfer with strategy optimization. In International
_Conference on Learning Representations, 2019. 1_

Sergey Zagoruyko and Nikos Komodakis. Paying more attention to attention: Improving the
performance of convolutional neural networks via attention transfer. International Conference on
_Learning Representation, 2017. 8_

Amy Zhang, Yuxin Wu, and Joelle Pineau. Natural environment benchmarks for reinforcement
learning. arXiv preprint arXiv:1811.06032, 2018a. 1, 2, 7

Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant
representations for reinforcement learning without reconstruction. International Conference on
_Learning Representation, 2021. 1, 2, 7_

Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. A study on overfitting in deep
reinforcement learning. arXiv preprint arXiv:1804.06893, 2018b. 1


-----

APPENDIX

A THEOREMS AND PROOFS

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)=_


_H(Xi|Xi−1, . . ., X1)_
_i=1_

X


**(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)}

A.1 THEOREM 1

We first show that I(S1:T ; O1:T |A1:T ) is a lower bound of I(S1:T ; O1:T ).


**(P.3)**
= I(O1:T ; S1:T ) − _I(S1:T ; O1:T ; A1:T )_


_I(S1:T ; O1:T |A1:T )_


Sincehave that: S1:T are representations of O1:T, we have I(S1:T ; A1:T |O1:T ) = 0. With (P.1) and (P.6), we

_I(S1:T ; O1:T ; A1:T ) ≥_ 0

Thus, we have:

_I(O1:T ; S1:T ) ≥_ _I(S1:T ; O1:T |A1:T )_

Furthermore, sincemaximizes I(O1:T ; S I1:(ST )1:.T ; O1:T ; A1:T ) is lower bounded, maximizing I(S1:T ; O1:T |A1:T ) also

**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(St; Ot|St−1, At−1)_ (8)
_t=1_

X


_I(S1:T ; O1:T |A1:T ) ≥_


-----

_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)**
= (H(St _A1:T, S1:t_ 1) _H(St_ _A1:T, O1:T, S1:t_ 1))

_|_ _−_ _−_ _|_ _−_
_t_

X

**(P.4)**
= _I(St; O1:T_ _A1:T, S1:t_ 1)

_t_ _|_ _−_

X

**(P.4)**
= (H(O1:T _A1:T, S1:t_ 1) _H(O1:T_ _St, A1:T, S1:t_ 1))

_t_ _|_ _−_ _−_ _|_ _−_

X

**(P.5)**
= (H(Oτ _A1:T, S1:t_ 1, O1:τ 1)

_|_ _−_ _−_

_t_ _τ_

X X

_H(Oτ_ _St, A1:T, S1:t_ 1, O1:τ 1))
_−_ _|_ _−_ _−_

**(P.4)**
= _I(St; Oτ_ _A1:T, S1:t_ 1, O1:τ 1)

_|_ _−_ _−_

_t_ _τ_

X X

**(P.1)**

_I(St; Ot_ _A1:T, S1:t_ 1, O1:t 1)

_≥_ _|_ _−_ _−_

_t_

X

= _I(St; Ot|St−1, At−1)_
X


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)
_|_ _−_ _−_

_p(st, ot_ **_a1:T, s1:t_** 1, o1:t 1)

= _p(τ_ )p(st, ot _τ_ ) log _|_ _−_ _−_

_τ_ **_st_** **_ot_** _|_ _p(st_ **_a1:T, s1:t_** 1, o1:t 1)p(ot **_a1:T, s1:t_** 1, o1:t 1)

Z Z Z  _|_ _−_ _−_ _|_ _−_ _−_


_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(S1:[(1)]T [;][ O]1:[(1)]T _[|][A][1:][T][ )][:]_

_I(S1:[(1)]T_ [;][ O]1:[(1)]T (I(St[(1)]; Ot[(1)] _t_ 1[, A][t][−][1][, O]t[(2)][)+][I][(][O]t[(2)][;][ S]t[(1)] _St[(1)]1[, A][t][−][1][))]_

_[|][A][1:][T][ )][ ≥]_ _t_ _[|][S][(1)]−_ _|_ _−_
X

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.

A.2 SUFFICIENT REPRESENTATIONS IN RL

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][(][o][1:][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][(][O][1:][T][ ;][ A]1:[∗]_ _T_ [) =]
_I(S1:T ; A[∗]1:T_ [)][. Also,][ S][t][ is a sufficient representation of][ O][t][ since][ I][(][O][t][;][ A]t[∗][|][S][t][, S][t][−][1][, A][t][−][1][) = 0][.]

_Hypothesis:_

**_(H.1) S1:T is a sequence of sufficient representations for O1:T :_**

_I(O1:T ; A[∗]1:T_ [) = 0]

_[|][S][1:][T]_

_Proof._

_I(O1:T ; A[∗]1:T_ [)]

_[|][S][1:][T]_

**(P.3)**
= I(O1:T ; A[∗]1:T [)][ −] _[I][(][O][1:][T]_ [;][ A][∗]1:T [;][ S][1:][T] [)]

**(P.3)**
= I(O1:T ; A[∗]1:T [)][ −] _[I][(][A]1:[∗]_ _T_ [;][ S][1:][T] [)][ −] _[I][(][A][∗]1:T_ [;][ S][1:][T] _[|][O][1:][T]_ [)]


With S1:T as a representation of O1:T, we have I(S1:T ; A[∗]1:T
shares the same level of information as A[∗]1:T [and][ S][1:][T][ . Then,] _[|][O][1:][T][ )=0][. The reason is that][ O][1:][T]_

_I(O1:T ; A[∗]1:T_ _[|][S][1:][T]_ [) =][I][(][O][1:][T] [;][ A][∗]1:T [)][ −] _[I][(][A]1:[∗]_ _T_ [;][ S][1:][T] [)] (9)

So the sufficiency conditionI(A[∗]1:T [;][ S][1:][T][ )][.] _I(O1:T ; A[∗]1:T_ _[|][S][1:][T][ ) = 0][ holds if and only if][ I][(][O][1:][T][ ;][ A]1:[∗]_ _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[|][O][1:][t][−][1][) +][ I][(][O][1:][t][−][1][;][ A][∗]1:t[)]

**(H.1)**
= I(Ot; A[∗]1:t[|][O][1:][t][−][1][) +][ I][(][S][1:][t][−][1][;][ A][∗]1:t[)]

_I(S1:t; A[∗]1:t[)]_


**(P.2)**
= I(St; A[∗]1:t[|][S][1:][t][−][1][) +][ I][(][S][1:][t][−][1][;][ A][∗]1:t[)]

**(H.1)**
= I(St; A[∗]1:t[|][S][1:][t][−][1][) +][ I][(][O][1:][t][−][1][;][ A][∗]1:t[)]

Then we obtain the following relation:

_I(Ot; A[∗]1:t[|][O][1:][t][−][1][) =][ I][(][S][t][;][ A][∗]1:t[|][S][1:][t][−][1][)]_ (10)


-----

We also have

_I(Ot; A[∗]1:t[|][O][1:][t][−][1][)]_

**(P.2)**
= I(O1:t; A[∗]1:t[)][ −] _[I][(][O][1:][t][−][1][;][ A][∗]1:t[)]_

**(P.2)**
= I(O1:t 1; A[∗]1:t[|][O][t][) +][ I][(][O][t][;][ A][∗]1:t[)]
_−_
_I(O1:t_ 1; A[∗]1:t[)]
_−_ _−_

**(H.1)**
= I(S1:t 1; A[∗]1:t[|][O][t][) +][ I][(][O][t][;][ A][∗]1:t[)]
_−_
_I(S1:t_ 1; A[∗]1:t[)]
_−_ _−_

**(P.2)**
= I(OtS1:t 1; A[∗]1:t[)][ −] _[I][(][S][1:][t][−][1][;][ A][∗]1:t[)]_
_−_

**(P.2)**
= I(Ot; A[∗]1:t[|][S][1:][t][−][1][)]

Eq. (10)
= I(St; A[∗]1:t[|][S][1:][t][−][1][)]

**(P.2)**
_⇐⇒_
_I(Ot; A[∗]t_ _[|][A]1:[∗]_ _t−1[, S][1:][t][−][1][) +][ I][(][O][t][;][ A]1:[∗]_ _t−1[|][S][1:][t][−][1][)]_

Eq. (10=)I(St; A[∗]t _[|][A]1:[∗]_ _t−1[, S][1:][t][−][1][) +][ I][(][S][t][;][ A]1:[∗]_ _t−1[|][S][1:][t][−][1][)]_
_⇐⇒_

Eq. (9)I(Ot; A[∗]t _[|][A]1:[∗]_ _t−1[, S][1:][t][−][1][)=][I][(][S][t][;][ A][∗]t_ _[|][A]1:[∗]_ _t−1[, S][1:][t][−][1][)]_
_⇐⇒_
_I(Ot, A[∗]t_ _[|][S][t][, S][1:][t][−][1][, A]1:[∗]_ _t−1[) = 0]_

With the above derivation and Markov property, we have I(Ot; A[∗]t
generalize A[∗]t−1 [to any][ A][t][−][1] [by assuming][ A]t[∗] [as the optimal action for state][|][S][t][, S][t][−][ S][1][, A][t] [whose last timestep][t][−][1][) = 0][. We can]
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 ._


A.3 MULTI-VIEW REDUNDANCY AND SUFFICIENCY

**Proposition A.1. O1:[(1)]T** [is a redundant view with respect to][ O]1:[(2)]T [to obtain][ A]1:[∗] _T_ [if only if]
_Ifor(O A1:[(1)][∗]1:TT[;][ A][.]_ 1:[∗] _T_ _[|][O]1:[(2)]T_ [) = 0][. Any representation][ S]1:[(1)]T [of][ O]1:[(1)]T [that is sufficient for][ O]1:[(2)]T [is also sufficient]

_Proof. See proof of Proposition B.3 in the MIB paper (Federici et al., 2020)._

**Corollary A.1. Let O1:[(1)]T** [and][ O]1:[(2)]T [be two mutually redundant views for][ A]1:[∗] _T_ [. Let][ S]1:[(1)]T [be a]
representation of O1:[(1)]T [. If][ S]1:[(1)]T [is sufficient for][ O]1:[(2)]T [,][ S]t[(1)] can derive A[∗]t [as the joint observation of]
the two views (representation at timestepI(Ot[(1)][O]t[(2)] t[;] −[ A]t[∗]1[|][S]. _[t][−][1][, A][t][−][1][)=][I][(][S]t[(1)]; A[∗]t_ _[|][S][t][−][1][, A][t][−][1][)][), where][ S][t][−][1]_ [is any sufficient]

_Proof. For the sequential data, see proof of Corollary B.2.1 in the MIB paper (Federici et al., 2020)_
to prove

_I(O1:[(1)]T_ _[O]1:[(2)]T_ [;][ A]1:[∗] _T_ [)=][I][(][S]1:[(1)]T [;][ A]1:[∗] _T_ [)]

According to Theorem A.2, if S1:[(1)]T [is a sufficient representation of][ O]1:[(2)]T [,][ S]t[(1)] is a sufficient representation of Ot[(2)][. 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(Ot[(1)][O]t[(2)][;][ A]t[∗][|][S][t][−][1][, A][t][−][1][) =][ I][(][S]t[(1)]; A[∗]t_ _[|][S][t][−][1][, A][t][−][1][)]_


-----

**Theorem A.3. Let the two views o[(1)]1:T** _[and][ o]1:[(2)]T_ _[of observation][ o][1:][T][ are obtained by data augmenta-]_
_tion transformation sequences t[(1)]1:T_ _[and][ t]1:[(2)]T_ _[respectively (][o]1:[(1)]T_ [=][t]1:[(1)]T [(][o][1:][T][ )][ and][ o]1:[(2)]T [=][t]1:[(2)]T [(][o][1:][T][ )][).]
_We abuse the notation t for simplicity to represent t[(1)]1:T_ [(][O][1:][T][ )][ and][ t]1:[(2)]T [(][O][1:][T][ )][ as random variables]
_for augmented observations. Whenever I(t[(1)]1:T_ [(][O][1:][T][ );][ A]1:[∗] _T_ [)=][I][(][t]1:[(2)]T [(][O][1:][T][ );][ A]1:[∗] _T_ [)=][I][(][O][1:][T][ ;][ A]1:[∗] _T_ [)][,]
_the two views O1:[(1)]T_ _[and][ O]1:[(2)]T_ _[must be mutually redundant for][ A]1:[∗]_ _T_ _[. Besides, the two views][ O]t[(1)]_ _and_
_Ot[(2)]_ _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_ [(][O][1:][t][);][ A]1:[∗] _t[) =][ I][(][O]1:[(1)]t_ [;][ A]1:[∗] _t[)]_

**(P.2)**
= I(Ot[(1)][;][ A]1:[∗] _t[|][O]1:[(1)]t_ 1[)+][I][(][O]1:[(1)]t 1[;][ A]1:[∗] _t[)]_
_−_ _−_

_A.2_
= I(Ot[(1)][;][ A]1:[∗] _t[|][S][1:][t][−][1][)+][I][(][S][1:][t][−][1][;][ A][∗]1:t[)]_

_I(t[(2)]1:t_ [(][O][1:][t][);][ A]1:[∗] _t[) =][ I][(][O]1:[(2)]t_ [;][ A]1:[∗] _t[)]_

**(P.2)**
= I(Ot[(2)][;][ A]1:[∗] _t[|][O]1:[(2)]t_ 1[)+][I][(][O]1:[(2)]t 1[;][ A]1:[∗] _t[)]_
_−_ _−_

_A.2_
= I(Ot[(2)][;][ A]1:[∗] _t[|][S][1:][t][−][1][)+][I][(][S][1:][t][−][1][;][ A]1:[∗]_ _t[)]_

_I(O1:t; A[∗]1:t[) =][ I][(][O][1:][t][;][ A]1:[∗]_ _t[)]_


**(P.2)**
= I(Ot; A[∗]1:t[|][O][1:][t][−][1][)+][I][(][O][1:][t][−][1][;][ A]1:[∗] _t[)]_

_A.2_
= I(Ot; A[∗]1:t[|][S][1:][t][−][1][)+][I][(][S][1:][t][−][1][;][ A]1:[∗] _t[)]_

Then, we have the following equality

_I(Ot[(1)][;][ A]1:[∗]_ _t[|][S][1:][t][−][1][)=][I][(][O]t[(2)][;][ A]1:[∗]_ _t[|][S][1:][t][−][1][)=][I][(][O][t][;][ A]1:[∗]_ _t[|][S][1:][t][−][1][)]_

Similar as derivation in Theorem A.2


_I(Ot[(1)][;][ A]1:[∗]_ _t[|][S][1:][t][−][1][)]_

**(P.2)**
= I(Ot[(1)][;][ A]t[∗] 1:t 1[, S][1:][t][−][1][) +][ I][(][O]t[(1)][;][ A]1:[∗] _t_ 1[|][S][1:][t][−][1][)]

Eq. (10) _[|][A][∗]_ _−_ _−_
= I(Ot[(1)][;][ A]t[∗] 1:t 1[, S][1:][t][−][1][) +][ I][(][S][t][;][ A]1:[∗] _t_ 1[|][S][1:][t][−][1][)]

_[|][A][∗]_ _−_ _−_

We apply the same derivation for o[(2)] and o, we have the following with Markov property

_I(t[(1)]t_ [(][O][t][);][ A]t[∗][|][S][t][−][1][, A][t][−][1][)]

=I(t[(2)]t [(][O][t][);][ A]t[∗][|][S][t][−][1][, A][t][−][1][)]
=I(Ot; A[∗]t

_[|][S][t][−][1][, A][t][−][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_ [(][O][1:][T][ );][ A]1:[∗] _T_ [)=][I][(][t]1:[(2)]T [(][O][1:][T][ );][ A]1:[∗] _T_ [)=][I][(][O][1:][T][ ;][ A]1:[∗] _T_ [)][.] _If S1:[(1)]T_ _[is a sufficient representa-]_
_tion for t[(2)]1:T_ [(][O][1:][T][ )][ then][ I][(][O][t][;][ A]t[∗][|][S][t][−][1][, A][t][−][1][) =][ I][(][S]t[(1)]; A[∗]t _[|][S][t][−][1][, A][t][−][1][)][.]_

_Proof. Since t[(1)]1:T_ [(][O][1:][T][ )][ is redundant for][ t]1:[(2)]T [(][O][1:][T][ )][ (Theorem][ A.3][), any representation][ S]t[(1)] of
_t[(1)]1:T_ [(][O][1:][T][ )][ that is sufficient for][ t]1:[(2)]T [(][O][1:][T][ )][ must also be sufficient for][ A]t[∗] [(Theorem][ A.2][ and Propo-]
sition A.1). Using Theorem A.2 we have I(St[(1)]; A[∗]t _[|][S][t][−][1][, A][t][−][1][)=][I][(][t]t[(1)][(][O][t][);][ A]t[∗][|][S][t][−][1][, A][t][−][1][)][.]_
With I(t[(1)]t [(][O][t][);][ A]t[∗][|][S][t][−][1][, A][t][−][1][) =][ I][(][O][t][;][ A]t[∗][|][S][t][−][1][, A][t][−][1][)][, we conclude][ I][(][O][t][;][ A]t[∗][|][S][t][−][1][, A][t][−][1][) =]
_I(St[(1)]; A[∗]t_

_[|][S][t][−][1][, A][t][−][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 Ot[(1)]** and Ot[(2)] be mutually redundant views for A[∗]t [that share only optimal]
action information. Then a sufficient representation of St[(1)] of Ot[1] [for][ O]t[(2)] that is minimal for Ot[(2)] is
also a minimal representation for A[∗]t [.]

_Proof. See proof of Proposition E.1 in the MIB paper (Federici et al., 2020)._

B DRIBO LOSS COMPUTATION

We consider the average of the information bottleneck losses from the two views.


(11)

(12)


1+2
_L_ 2

_t_ ;Ot[(1)] _t_ 1[,A][t][−][1][,O]t[(2)][)+][I][(][S]t[(2)];Ot[(2)] _t_ 1[,A][t][−][1][,O]t[(1)][)]
= _[I][(][S][(1)]_ _[|][S][(1)]−_ _[|][S][(2)]−_

2

_t_ ; Ot[(2)] _t_ 1[, A][t][−][1][)+][λ][2][I][(][S]t[(2)]; Ot[(1)] _t_ 1[, A][t][−][1][)]

_[|][S][(1)]−_ _[|][S][(2)]−_
_−_ _[λ][1][I][(][S][(1)]_ 2


Consider s[(1)]t and s[(2)]t on the same domain S, I(St[(1)]; Ot[(1)] _t_ 1[, A][t][−][1][, O]t[(2)][)][ can be expressed as:]

_[|][S][(1)]−_

_I(St[(1)]; Ot[(1)]_ _t_ 1[, A][t][−][1][, O]t[(2)][)]

_[|][S][(1)]−_

_t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)]_

=E log _[p][θ][(][s][(1)][|][o][(1)]_ _−_

" _pθ(s[(1)]t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)]_ #

_[|][o][(2)]_ _−_

_t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)]_ _pθ(s[(2)]t_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][)]_

=E log _[p][θ][(][s][(1)][|][o][(1)]_ _−_ _[|][o][(2)]_ _−_

" _pθ(s[(2)]t_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][)]_ _pθ(s[(1)]t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)]_ #

_[|][o][(2)]_ _−_ _[|][o][(2)]_ _−_

=DKL(pθ(s[(1)]t _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)][||][p][θ][(][s]t[(2)]_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][))]_

_[|][o][(1)]_ _−_ _[|][o][(2)]_ _−_

_DKL(pθ(s[(1)]t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)][||][p][θ][(][s]t[(2)]_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][))]_
_−_ _[|][o][(2)]_ _−_ _[|][o][(2)]_ _−_

_DKL(pθ(s[(1)]t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)][||][p][θ][(][s]t[(2)]_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][))]_ (13)
_≤_ _[|][o][(1)]_ _−_ _[|][o][(2)]_ _−_


Note that equality holds if the two distributions coincide. Analogously I(St[(2)]; Ot[(2)] _t_ 1[, A][t][−][1][, O]t[(1)][)]

_[|][S][(2)]−_
is upper bounded by DKL(pθ(s[(2)]t _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][)][||][p][θ][(][s]t[(1)]_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][))][.]_

_[|][o][(2)]_ _−_ _[|][o][(1)]_ _−_

Assume St−1 is a sufficient representation of Ot−1. Then, St[(1)]−1 [provides task-relevant information]
no more than the sufficient representation St 1. I(St[(1)]; Ot[(2)] _t_ 1[, A][t][−][1][)][ can be thus re-expressed]
as: _−_ _[|][S][(1)]−_

_I(St[(1)]; Ot[(2)]_ _t_ 1[, A][t][−][1][)]

_[|][S][(1)]−_

_I(St[(1)]; Ot[(2)]_
_≥_ _[|][S][t][−][1][, A][t][−][1][)]_

**(P.2)**
= I(St[(1)]; St[(2)]Ot[(2)][|][S][t][−][1][, A][t][−][1][)][−][I][(][S]t[(1)]; St[(2)]|Ot[(2)][, S][t][−][1][, A][t][−][1][)]

=∗ _I(St[(1)]; St[(2)]Ot[(2)]_

_[|][S][t][−][1][, A][t][−][1][)]_

=I(St[(1)]; St[(2)]|St−1, At−1) + I(St[(1)]; Ot[(2)][|][S]t[(2)], St−1, At−1)

_≥I(St[(1)]; St[(2)]|St−1, At−1)_ (14)


-----

Where ∗ follows from St[(2)] being the representation of Ot[(2)][. The bound is tight whenever][ S]t[(2)]
is sufficient from St[(1)] (I(St[(1)]; Ot[(2)] _t_ _, St−1, At−1)=0)._ This happens whenever St[(2)] con
_[|][S][(2)]_
tains all the information regarding St[(1)]. Once again, we can have I(St[(2)]; Ot[(1)] _t_ 1[, A][t][−][1][)][ ≥]

_[|][S][(2)]−_
_I(St[(1)]; St[(2)]|St−1, At−1). Therefore, the averaged loss functions can be upper-bounded by_

_L 1+22_ _[≤−]_ _[λ][1][ +]2_ _[ λ][2]_ _I(St[(1)]; St[(2)]|St−1, At−1)_ (15)

+DSKL(pθ(s[(1)]t _t_ _[,][s]t[(1)]1[,][a][t][−][1][)][||][p][θ][(][s]t[(2)]_ _t_ _[,][ s]t[(2)]1[,][a][t][−][1][))]_

_[|][o][(1)]_ _−_ _[|][o][(2)]_ _−_


Lastly, by re-parametrizing the objective, we obtain:

_L(θ; β)= −_ _Iθ(St[(1)]; St[(2)]|St−1, At−1)_ (16)

+βDSKL(pθ(s[(1)]t _t_ _[,][s]t[(1)]1[,][a][t][−][1][)][||][p][θ][(][s]t[(2)]_ _t_ _[,][s]t[(2)]1[,][a][t][−][1][))]_

_[|][o][(1)]_ _−_ _[|][o][(2)]_ _−_

In Algorithm 1, we use s[(1)]t _pθ(s[(1)]t_ _t_ _[,][ s]t[(1)]1[,][ a][t][−][1][)][ and][ s]t[(2)]_ _pθ(s[(2)]t_ _t_ _[,][ s]t[(2)]1[,][ a][t][−][1][)][ to]_
obtain representations for multi-view observations. We argue that the substitution does not affect the∼ _[|][o][(1)]_ _−_ _∼_ _[|][o][(2)]_ _−_
effectiveness of the averaged objective. With the multi-view assumption, we have that representations
**_s[(1)]t_** 1 [and][ s]t[(2)]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.

C ADDITIONAL RESULTS

C.1 ADDITIONAL DMC RESULTS

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.


-----

cartpole/swingup cheetah/run finger/spin

800 1000

800

600 800

600

600

400

400 400

200

200 200

Average Returns

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

ball_in_cup/catch reacher/easy hopper/hop

1000 1000 250

800 800 200

600 600 150

400 400 100

200 200 50

Average Returns 0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

walker/stand 1000 walker/walk 600 walker/run

1000 800 500

800 400

600

600 300

400

400 200

Average Returns 200 200 100

0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

Environment Steps 1e5 Environment Steps 1e5 Environment Steps 1e5

DRIBO (Ours) RAD SLAC CURL@1e6 DBC PI-SAC@1e6


Figure 6: Average returns on DMC tasks over 5 seeds with mean and one standard error shaded in
the clean setting.





cartpole/swingup cheetah/run finger/spin

1000

800 400 800

600 300 600

400 200 400

200 100 200

Average Returns

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

ball_in_cup/catch reacher/easy hopper/hop

1000

800

800 150

600

600

400 100

400

200 200 50

Average Returns 0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

walker/stand walker/walk walker/run

1000 400

800

800

300

600

600

200

400

400

200 100

Average Returns 200

0 1 2 3 4 5 6 7 8 00 1 2 3 4 5 6 7 8 00 1 2 3 4 5 6 7 8

Environment Steps 1e5 Environment Steps 1e5 Environment Steps 1e5

DRIBO (Ours) RAD SLAC CURL DBC PI-SAC


Figure 7: Average returns on DMC tasks over 5 seeds with mean and one standard error shaded in
the natural video setting.


-----

C.2 DRIBO LOSS VS. INFOMAX

C.2.1 NATURAL VIDEO SETTING




cartpole/swingup cheetah/run finger/spin

200

300

150 200

200

100

100

SKL Value100 50

0

0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

hopper/hop walker/walk hopper/hop

250

500

150 200

400

150

100 300

100 200

SKL Value 50

50 100

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

1e5 1e5 1e5

Environment Steps Environment Steps Environment Steps

DRIBO (Ours) DRIBO-no-skl


Figure 8: Average SKL values during training in DMC environments with natural videos as background.

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.


C.2.2 CLEAN SETTING

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.


-----

cartpole/swingup cheetah/run finger/spin

1000

800 600

800

600

400 600

400 400

200

200 200

Average Returns 0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

hopper/hop walker/walk walker/run

200 1000

500

800

150 400

600

100 300

400

200

50

200 100

Average Returns 0

0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

1e5 1e5 1e5

Environment Steps Environment Steps Environment Steps

DRIBO (Ours) DRIBO-no-skl


Figure 9: Average returns of DRIBO and DRIBO-no-skl on DMC tasks over 5 seeds with mean and
one standard error shaded in the clean setting.





cartpole/swingup cheetah/run finger/spin

500

300

800

400

600 300 200

400

200

100

SKL Value200 100

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

hopper/hop walker/walk hopper/hop

500

400 400

400

300 300

300

200 200

200

SKL Value

100 100 100

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

1e5 1e5 1e5

Environment Steps Environment Steps Environment Steps

DRIBO (Ours) DRIBO-no-skl


Figure 10: Average SKL values during training in DMC environments in the clean setting.


-----

C.3 TRAINING AND TESTING PERFORMANCE IN DMC UNDER THE NATURAL VIDEO SETTING





cartpole/swingup 500 cheetah/run 1000 finger/spin

800 400 800

600 300 600

400 200 400

200 100 200

Average Returns

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

hopper/hop 1000 walker/walk walker/run 1e5

200 400

800

150 300

600

100

200

400

50

200 100

Average Returns 0

0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

Environment Steps 1e5 Environment Steps 1e5 Environment Steps 1e5

DRIBO (Ours) Train DRIBO (Ours) Test RAD Trian RAD Test

Figure 11: Training and testing performance of RAD and DRIBO.





1000 cartpole/swingup 500 cheetah/run 1000 finger/spin

800 400 800

600 300 600

400 200 400

200 100 200

Average Returns

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

hopper/hop 1000 walker/walk walker/run 1e5

200 400

800

150 300

600

100 200

400

50 200 100

Average Returns 0

0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

Environment Steps 1e5 Environment Steps 1e5 Environment Steps 1e5

DRIBO (Ours) Train DRIBO (Ours) Test RAD Trian RAD Test

Figure 12: Training and testing performance of CURL and DRIBO.





cartpole/swingup 500 cheetah/run 1000 finger/spin

800 400 800

600 300 600

400 200 400

200 100 200

Average Returns

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

hopper/hop 1000 walker/walk walker/run 1e5

200 400

800

150 300

600

100 200

400

50 200 100

Average Returns

0 0 0

0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 8

Environment Steps 1e5 Environment Steps 1e5 Environment Steps 1e5

DRIBO (Ours) Train DRIBO (Ours) Test PI-SAC Test PI-SAC Train

Figure 13: Training and testing performance of PI-SAC and DRIBO.

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.


-----

C.4 ADDITIONAL VISUALIZATION

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

Figure 14: Training environments of DMC under the natural video setting: The background
videos are sampled from arranging flower class in Kinetics dataset.

Figure 15: Test environments of DMC under the natural video setting: We show the sequential
observations with natural videos as background in DMC and their corresponding spatial attention
maps of the DRIBO trained encoder.

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.

Figure 16: Spatial attention maps for each convolutional layers of the DRIBO trained encoders. The
observations are the same as ones in Figure 2 from the snapshots during testing.

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

Background 1 Background 2

Figure 17: t-SNE of latent spaces learned with DRIBO (left) and CURL (right). They are the same as
t-SNE in Figure 4. But we color-code the embedded points corresponding to their backgrounds. The
observations are from the same trajectory but with different background natural videos (the same as
in Figure Figure 4). Points from different backgrounds are close to each other in the embedding space
learned by DRIBO, whereas no such structure is seen in the embedding space learned by CURL.


-----

D IMPLEMENTATION DETAILS

D.1 RSSM

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 encoderfrom st 1 = (zt 1, ht 1) and pθ a((ts1:1T. |o1:T, a1:T )= _t_ _[p][θ][(][s][t][|][o][t][,][ h][t][)][, where][ h][t][ retains information]_

_−_ _−_ _−_ _−_

[Q]

D.2 DRIBO + SAC

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 RL = (φ(oi), ai, ri, φ(o[′]i[))][}][(]i=1[T][ −][1)][∗][N] with (T 1) _N pairs of representation,_
_B_ _{_ _−_ _∗_
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,2 Qi(φ(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

7: _pθ(s1:T |Train SAC agent:o1:T, a1:T )_ E RL [J(π, Q, φ)] _▷_ Algorithm 2
_B_

8: Update θ and ψ to minimize LDRIBO using B _▷_ Algorithm 1.

9: **end for**

10: end for

D.3 DRIBO + PPO

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 s[old]t while sampling from the on-policy batch. Then, we use ϕ(o) = s
_∼_


-----

_pθ(s_ **_o, s[old], a[′]) to denote the representation from the encoder. Here, s[old]_** 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, s[old]t 1[,][ a][t][−][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, s[old]1:T _i=1[.]_

4: _{_ Append new samples to[}][N] _D and update the on-policy replay buffer DPPO._

5: **for j = 1, . . ., M do**

6: (ϕ(oi), ai, ri) _⌊i=1[T]M[ ∗][N]_ _[⌋]_ PPO
_{_ _}_ _∼D_

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

D.4 DMC

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 Algo-**
rithm 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**

Table 2: Hyperparameters used for DMC experiments.

**Hyperparameters** **Value**

Observation size (100 × 100)
Cropped size (84 × 84)
Replay buffer size 1000000
Initial steps 1000
Stacked frames No
Action repeat 2 finger, spin; walker, walk;
8 cartpole swingup
4 otherwise
Evaluation episodes 8
Optimizer Adam
Learning rate encoder learning rate: 1e-5;
policy/Q network learning rate: 1e-5;
_α learning rate: 1e-4._
Batch size 8 × 32, where T = 32
Target update τ Q network: 0.01
encoder: 0.05
Target update freq 2
Discount γ .99
Initial temperature 0.1
Total timesteps 88e4
_β scheduler start episode_ 10
_β scheduler end episode_ 60


-----

D.5 PROCGEN

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.

Table 3: Augmentation type used for each game.

|Env|BigFish|StarPilot|FruitBot|BossFight|Ninja|Plunder|CaveFlyer|CoinRun|
|---|---|---|---|---|---|---|---|---|
|Augmentation|crop|cutout|cutout|cutout|random-conv|crop|random-conv|random-conv|



Table 4: Augmentation type used for each game.

|Env|Jumper|Chaser|Climber|DodgeBall|Heist|Leaper|Maze|Miner|
|---|---|---|---|---|---|---|---|---|
|Augmentation|random-conv|crop|random-conv|cutout|crop|crop|crop|flip|



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


-----

Table 5: Hyperparameters used for Procgen experiments.

**Hyperparameters** **Value**

Observation size (64 × 64)
Replay buffer size 1000000
Num. of steps per rollout 256
Num. of epochs per rollout 3
Num. of minibatches per epoch 8
Stacked frames No
Evaluation episodes 10
Optimizer Adam
Learning rate encoder learning rate: 1e-4;
policy learning rate: 5e-4;
_α learning rate: 1e-4._
Batch size 8 × 256, where T = 256
Entropy bonus 0.01
PPO clip range 0.2
Discount γ .99
GAE parameterλ 0.95
Reward normalization yes
Num. of workers 1
Num. of environments per worker 64
Total timesteps 25M
_β scheduler start episode_ 10
_β scheduler end episode_ 110

D.6 COMPUTE RESOURCES AND LICENSE

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

Figure 18: Screenshots of the three Procgen games where our approach DRIBO does not improve
generalization performance compared to the other methods. From left to right, they are Plunder,
Chaser and DodgeBall.


-----

