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

Jiameng Fan
Department of Electrical and Computer Engineering
Boston University
<EMAIL_ADDRESS>
&Wenchao Li
Department of Electrical and Computer Engineering
Boston University
<EMAIL_ADDRESS>

###### 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 multi-view setting. Specifically, we introduce an auxiliary
objective based on the multi-view information bottleneck (MIB) principle which
quantifies the amount of task-irrelevant information and encourages learning
representations that are both predictive of the future and less sensitive to
task-irrelevant distractions. This enables us to train high-performance
policies that are robust to visual distractions and can generalize to unseen
environments. We demonstrate that our approach can achieve SOTA performance on
diverse visual control tasks on the DeepMind Control Suite, even 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. Our code is open-sourced and available
at https://github.com/BU-DEPEND-Lab/DRIBO.

##  1 Introduction

Deep reinforcement learning (DRL) methods have been shown to be successful in
learning high-quality controllers directly from raw images in an end-to-end
fashion [28, 24, 3]. Recent studies, however, show that DRL agents can suffer
from the problem of overfitting to training environments [42, 40, 38]. In
particular, 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 [7, 5]. By contrast, humans routinely adapt to new, unseen
environments that contain similar underlying dynamics. 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 underlying task-relevant dynamics and thus can be
distracted by task-irrelevant visual details [15, 41, 1]

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 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 [10, 23] 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.

Figure 1:  Robust Deep Reinforcement Learning via Multi-View Infomration
BOttleneck (DRIBO) incorporates the inherent sequential structure of RL and
unsupervised multi-view settings into robust representation learning in RL. We
consider sequential multi-view observations, ${\bm{o}}_{1:T}^{(1)}$ and
${\bm{o}}_{1:T}^{(2)}$, of original observations ${\bm{o}}_{1:T}$ sharing the
same task-relevant information while any information not shared by them are
task-irrelevant. DRIBO uses a multi-view information bottleneck loss to ensure
that $s_{1:T}^{(1)}$ and $s_{1:T}^{(2)}$, 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.

We tackle this problem by considering state representations for RL that are
robust under the multi-view setting [26, 8], where each view is assumed to
provide the same amount of task-relevant information while all the information
not shared by them is deemed task-irrelevant. Data augmentation can be easily
leveraged to generate such multi-view observations without requiring
additional new data. Incorporating data augmentation into RL has delivered
promising results for visual control tasks [21, 19, 20]. 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 [20], we
propose to learn a predictive model that captures the temporal evolution of
representations from a sequence of observations and actions. Concretely, we
introduce a new multi-view information bottleneck (MIB) objective that
maximizes the mutual information between sequences of observations and
representations while reducing the task-irrelevant information identified from
the multi-view observations. We incorporate this MIB objective into RL by
optimizing RL objectives on top of the learned encoder. 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 the multi-view information bottleneck principle.

  * •

We leverage the sequential nature of RL to learn representations better suited
for RL with a new non-reconstruction-based, MIB objective 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 trained autoencoders
to learn representations to reconstruct raw observations first. Then, the RL
agent was trained from the learned representations [18, 19]. 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 [36, 37].
More recently, Hafner et al. [10, 11, 12] and Lee et al. [23] 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 [40]. Experimentally, we show that our non-reconstructive
approach, DRIBO, is substantially more robust against this type of visual
changes than prior works. We also compare DRIBO with another non-
reconstructive method, DBC [41], which uses bisimulation metrics to learn
representations in RL that contain only task-relevant information without
requiring reconstruction.

Contrastive Representations Learning. Contrastive representation learning
methods train an encoder that obeys similarity constraints in a dataset
typically organized by similar and dissimilar pairs. The similar examples are
typically obtained from nearby image patches [29, 13] or through data
augmentation [4]. A scoring function that lower-bounds mutual information is
one of the typical objects to be maximized [2, 29, 14, 30]. A number of works
have applied the above ideas to RL settings to extract predictive signals. EMI
[17] applies a Jensen-Shannon divergence-based lower bound on mutual
information across subsequent frames as an exploration bonus. DRIML [27] uses
an auxiliary contrastive objective to maximize concordance between
representations to increase predictive properties of the representations
conditioned on actions. CURL [20] 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 [8].

Multi-View Information Bottleneck (MIB). In classification, Federici et al.
[8] 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. However, MIB cannot be
directly used in RL settings due to the sequential nature of decision making
problems. Task-relevant information in RL is relevant because they influence
not only current control and reward but also states and rewards in the future,
which requires representations to be predictive of the future representations.
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. In addition to works in representation learning, MVRL [25] uses
the multi-view setting to form a generalization of partially observable Markov
decision process. which can substantially reduce sample complexity when
training RL agents on it.

##  3 Preliminaries

We denote a Markov Decision Process (MDP) as ${\mathcal{M}}$, with state
${\bm{s}}$, action ${\bm{a}}$, and reward $r$. We denote a policy on
${\mathcal{M}}$ as $\pi$. The agent’s goal is to learn a policy $\pi$ that
maximizes the cumulative rewards. We define
${\mathcal{S}}{\subseteq}{\mathbb{R}}^{d}$ as the state-representation space.
The visual observations are $o{\in}{\mathcal{O}}$, where we denote multi-view
observations from the viewpoint $i$ as $o^{(i)}$. We introduce a multi-view
trajectory
$\tau^{{\mathcal{M}}}{=}[{\bm{s}}_{1},{\bm{o}}_{1}^{(i)},{\bm{a}}_{1},\dots,{\bm{s}}_{T},{\bm{o}}_{T}^{(i)},{\bm{a}}_{T}]$
where $T$ is the length. Knowing that the trajectory density is defined over
joint observations, states, and actions, we write:

$p_{\pi}(\tau^{{\mathcal{M}}}){=}\pi({\bm{a}}_{T}|{\bm{s}}_{T}){\mathcal{P}}_{\text{obs}}^{(i)}({\bm{o}}_{T}^{(i)}|{\bm{s}}_{T}){\mathcal{P}}({\bm{s}}_{T}|{\bm{s}}_{T-1},{\bm{a}}_{T-1})\cdots\pi({\bm{a}}_{1}|{\bm{s}}_{1}){\mathcal{P}}_{\text{obs}}^{(i)}({\bm{o}}_{1}^{(i)}|{\bm{s}}_{1}){\mathcal{P}}_{0}({\bm{s}}_{1})$
(1)

with ${\mathcal{P}}_{0}$ being the initial state distribution, ${\mathcal{P}}$
being the transition model and ${\mathcal{P}}_{\text{obs}}^{(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 Equation 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 ${\bm{a}}^{*}_{1:T}$ be the optimal action sequence for some
${\bm{o}}_{1:T}$ which is obtained by executing the action sequence
${\bm{a}}_{1:T}$. We assume that ${\bm{o}}_{1:T}$ contains enough information
to obtain ${\bm{a}}^{*}_{1:T}$ which maximizes the cumulative rewards. With
the above assumption, we define task-relevant information if it is necessary
to derive ${\bm{a}}^{*}_{1:T}$. By contrast, task-irrelevant information does
not contribute to the choice of ${\bm{a}}^{*}_{1:T}$. We first consider
sufficient representations that are discriminative enough to obtain
${\bm{a}}^{*}$ at each timestep. This property can be quantified by the amount
of mutual information between ${\bm{o}}_{1:T}$ and ${\bm{a}}^{*}_{1:T}$ and
mutual information between ${\bm{s}}_{1:T}$ and ${\bm{a}}^{*}_{1:T}$.

######  Definition 1.

Representations ${\bm{s}}_{1:T}$ of ${\bm{o}}_{1:T}$ are sufficient for RL iff
$I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T}){=}I({\bm{s}}_{1:T};{\bm{a}}^{*}_{1:T})$.

RL agents that have access to a sufficient representation ${\bm{s}}_{t}$ at
timestep $t$ must be able to generate ${\bm{a}}^{*}_{t}$ as if it has access
to the original observations. This can be better understood by subdividing
$I({\bm{o}}_{1:T};{\bm{s}}_{1:T})$ into two components using the chain rule of
mutual information:

$\displaystyle
I({\bm{o}}_{1:T};{\bm{s}}_{1:T}){=}I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}^{*}_{1:T})+I({\bm{s}}_{1:T};{\bm{a}}^{*}_{1:T})$
(2)

Conditional mutual information
$I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}^{*}_{1:T})$ quantifies the
information in ${\bm{s}}_{1:T}$ that is task-irrelevant.
$I({\bm{s}}_{1:T};{\bm{a}}^{*}_{1:T})$ quantifies task-relevant information
that is accessible from ${\bm{s}}_{1:T}$. Note that the last term is
independent of the representation as long as ${\bm{s}}_{t}$ is sufficient for
${\bm{a}}^{*}_{t}$ (see Definition 1). Thus, a representation contains minimal
task-irrelevant information whenever
$I({\bm{o}}_{1:T};{\bm{s}}_{1:T}|{\bm{a}}^{*}_{1:T})$ is minimized. To obtain
a sufficient representation, we can maximize the mutual information
$I({\bm{o}}_{1:T};{\bm{s}}_{1:T})$. With the information bottleneck principle
[34], we can construct an objective to maximize
$I({\bm{o}}_{1:T};{\bm{s}}_{1:T})$ while minimizing
$I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}^{*}_{1:T})$ to reduce task-
irrelevant information.

However, estimating the mutual information between sequences is difficult due
to the high dimensionality of the problem. In addition, the minimization of
$I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}^{*}_{1:T})$ can only be done
directly in supervised settings where ${\bm{a}}^{*}_{1:T}$ are observed. One
option is to use MIB which can reduce task-irrelevant information in the
representations in unsupervised settings [8]. 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 sequential 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 information
$I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})$ which is a lower bound of
$I({\bm{o}}_{1:T};{\bm{s}}_{1:T})$ (see Appendix A.1). The observations
${\bm{o}}_{1:T}$ are generated sequentially in the environment by executing
the actions ${\bm{a}}_{1:T}$. The conditional mutual information not only
estimates the sufficiency of the representations but also maintains the
sequential 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 ${\bm{o}}_{1:T}$ be the sequential observations obtained by executing
action sequence ${\bm{a}}_{1:T}$. If ${\bm{s}}_{1:T}$ is a sequence of
sufficient representations for ${\bm{o}}_{1:T}$, we have:

$I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})\geq\sum_{t=1}^{T}I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
(3)

Proof in Appendix A.1. Theorem 3 shows that the sum of mutual information
$I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$ over timesteps is
a lower bound of the mutual information
$I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})$. Even when the
representations ${\bm{s}}_{1:T}$ are not sufficient, maximizing
$I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$ encodes more
detailed features into ${\bm{s}}_{t}$ which makes it sufficient. The
factorized mutual information is conditioned on the representation and the
action at $t{-}1$, which explicitly retains the predictive property of
representations.

###  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 ${\bm{o}}_{t}^{(1)}$ and ${\bm{o}}_{t}^{(2)}$ to
be two visual images of the control scenario from two different viewpoints.
Under the multi-view assumption, any representation ${\bm{s}}_{t}$ containing
all information accessible from both views and being predictive of future
representations would contain the sufficient task-relevant information.
Furthermore, if ${\bm{s}}_{t}$ captures only the details that are visible from
both observations, it would eliminate the view-specific details and reduce the
sensitivity of the representation to view-changes.

A sufficient representation in RL maintains all the information that is shared
by mutually redundant observations ${\bm{o}}_{t}^{(1)}$ and
${\bm{o}}_{t}^{(2)}$. We refer to Appendix A for the sufficiency condition of
representations and mutually redundancy condition between ${\bm{o}}_{t}^{(1)}$
and ${\bm{o}}_{t}^{(2)}$. Intuitively, with the mutual redundancy condition,
any representation which contains all information shared by both views is as
task-relevant as the joint observation. By factorizing the mutual information
between $s_{t}^{(1)}$ and ${\bm{o}}_{t}^{(1)}$ as in Equation 2, we can
identify two components:

$\displaystyle
I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})=I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(2)}){+}I({\bm{o}}_{t}^{(2)};{\bm{s}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})$
(4)

Here, ${\bm{s}}_{t-1}^{(1)}$ is a representation of visual observation
${\bm{o}}_{t-1}^{(1)}$. Since we assume mutual redundancy between the two
views, the information shared between ${\bm{o}}_{t}^{(1)}$ and
${\bm{s}}_{t}^{(1)}$ conditioned on ${\bm{o}}_{t}^{(2)}$ must be irrelevant to
the task, which can be quantified as
$I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(2)})$
(first term in Equation 4). Then,
$I({\bm{o}}_{t}^{(2)};{\bm{s}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})$
has to be maximal if the representation is sufficient. A formal description
for 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 to task-
irrelevant nuisances. In the extreme, ${\bm{s}}_{t}^{(1)}$ is the underlying
states of MDP if ${\bm{o}}_{t}^{(1)}$ and ${\bm{o}}_{t}^{(2)}$ share only
task-relevant information. With Equation 3 and 4, we have
${\mathcal{L}}_{\text{InfoMax}}$ to maximizing the mutual information between
representations and observations in the multi-view setting, which retains the
temporally evolving information of the underlying dynamics.

$\displaystyle{\mathcal{L}}_{\text{InfoMax}}=-\sum_{t}(I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(2)})+I({\bm{o}}_{t}^{(2)};{\bm{s}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1}))$
(5)

The above loss extends mutual information maximization to RL and minimizing it
learns representations that are sufficient 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

In Section 4.2, we show how to factorize the mutual information between
sequences of observations and representations into mutual information at each
timestep. This enables us to obtain sufficient representations by maximizing
the mutual information between ${\bm{s}}_{t}^{(1)}$ and ${\bm{o}}_{t}^{(1)}$
while discarding task-irrelevant information by reducing
$I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(2)})$.
The predictive properties are ensured by conditioning the mutual information
on the previous timestep representation and action. With the information
bottleneck principle, we can construct a relaxed Lagrangian loss to obtain a
sufficient representation ${\bm{s}}_{t}^{(1)}$ for ${\bm{o}}_{t}^{(1)}$ with
minimal task-irrelevant information:

$\displaystyle{\mathcal{L}}_{1}(\theta;\lambda_{1})=I_{\theta}({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(2)})-\lambda_{1}I_{\theta}({\bm{o}}_{t}^{(2)};{\bm{s}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})$
(6)

where $\theta$ denotes the parameters of an encoder
$p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
(details in Section 4.4), and $\lambda_{1}$ is the Lagrangian multiplier.
Symmetrically, we define a loss ${\mathcal{L}}_{2}$ to learn a sufficient
representation ${\bm{s}}_{t}^{(2)}$ for ${\bm{o}}_{t}^{(2)}$ with minimal
task-irrelevant information:

$\displaystyle{\mathcal{L}}_{2}(\theta;\lambda_{2})=$ $\displaystyle
I_{\theta}({\bm{s}}_{t}^{(2)};{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(1)})-\lambda_{2}I_{\theta}({\bm{o}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})$
(7)

By re-parameterizing the Lagrangian multipliers, the average of two loss
functions ${\mathcal{L}}_{1}$ and ${\mathcal{L}}_{2}$ from two views at
timestep $t$ can be upper bounded as follows:

$\displaystyle{\mathcal{L}}_{t}(\theta;\beta){=}$
$\displaystyle{-}I_{\theta}({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})+$
(8) $\displaystyle\beta
D_{\text{SKL}}(p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})||p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1}))$

where ${\bm{s}}_{t-1}$ is a sufficient representation, $D_{\text{SKL}}$
represents the symmetrized KL divergence obtained by averaging the expected
values of
$D_{\text{KL}}(p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})||p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1}))$,
$D_{\text{KL}}(p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})||p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1}))$,
and the coefficient $\beta$ represents the trade-off between sufficiency and
sensitivity to task-irrelevant information. $\beta$ is a hyper-parameter.

To generalize the above loss to sequential data in RL, we apply Theorem 3 to
obtain the DRIBO loss:
${\mathcal{L}}_{\text{DRIBO}}{=}\frac{1}{T}\sum_{t=1}^{T}{\mathcal{L}}_{t}(\theta;\beta)$.
We summarize the batch-based computation of the loss function in Algorithm 1.
We sample ${\bm{s}}_{t}^{(1)}$ and ${\bm{s}}_{t}^{(2)}$ from
$p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})$
and
$p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})$
respectively. Though first term in Equation 8 is conditioned on
${\bm{s}}_{t-1}$, we prove that the sampling process does not affect its
effectiveness based on the multi-view assumption in Appendix B. The
symmetrized KL divergence term can be computed from the probability density of
${\bm{s}}_{t}^{(1)}$ and ${\bm{s}}_{t}^{(2)}$ using the encoder. The mutual
information between the two representations
$I_{\theta}({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
can be maximized by using any sample-based differentiable mutual information
lower bound $\hat{I}_{\psi}({\bm{s}}_{t}^{(1)},{\bm{s}}_{t}^{(2)})$, where
$\psi$ represents the learnable parameters. We use InfoNCE [29] to estimate
mutual information since the multi-view setting provides a large number of
negative examples. The positive pairs are the representations
$({\bm{s}}_{t}^{(1)},{\bm{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.

Algorithm 1 DRIBO Loss

1:Input: Batch ${\mathcal{B}}$ storing $N$ sequential observations and actions
with length $T$ from replay buffer.

2:Apply random augmentation transformations on ${\mathcal{B}}$ to obtain
multi-view batches ${\mathcal{B}}^{(1)}$ and ${\mathcal{B}}^{(2)}$.

3:for $i,({\bm{o}}_{1:T}^{(1)},{\bm{o}}_{1:T}^{(2)},{\bm{a}}_{1:T})$ in
enumerate $({\mathcal{B}}^{(1)},{\mathcal{B}}^{(2)})$ do

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

5: ${\bm{s}}_{t}^{(1)}\sim
p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})$

6: ${\bm{s}}_{t}^{(2)}\sim
p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})$

7:
$({\bm{s}}^{(1),t+T(i-1)},{\bm{s}}^{(2),t+T(i-1)})\leftarrow({\bm{s}}_{t}^{(1)},{\bm{s}}_{t}^{(2)})$

8: end for

9:
${\mathcal{L}}_{\text{SKL}}^{i}=\frac{1}{T}\sum_{t=1}^{T}D_{\text{DKL}}(p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})||p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1}))$

10:end for

11:return
$-\hat{I}_{\psi}(\\{({\bm{s}}^{(1),i},{\bm{s}}^{(2),i})\\}_{i=1}^{T*N})+\frac{\beta}{N}\sum_{i=1}^{N}{\mathcal{L}}_{\text{SKL}}^{i}$

###  4.4 Encoder Architecture

The encoder
$p_{\theta}({\bm{s}}_{t}|{\bm{o}}_{t},{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
approximates representation posteriors from the current observation, the
previous timestep representation and action. The posteriors can also be seen
as a reparameterization of
$p_{\theta}({\bm{s}}_{1:T}|{\bm{o}}_{1:T},{\bm{a}}_{1:T})=\prod_{t}p_{\theta}({\bm{s}}_{t}|{\bm{o}}_{t},{\bm{s}}_{t-1},{\bm{a}}_{t-1})$,
which reflects the inherent sequential structure of RL.

We implement the encoder as a recurrent space model (RSSM [10]) with a
convolutional neural network (CNN [22]) applied to the visual observations.
RSSM is a latent dynamics model with an expressive recurrent neural network to
perform accurate long-term prediction. We split the representation
${\bm{s}}_{t}$ into a stochastic part ${\bm{z}}_{t}$ and a deterministic part
${\bm{h}}_{t}$, where ${\bm{s}}_{t}=({\bm{z}}_{t},{\bm{h}}_{t})$. The
generative and inference models of RSSM are defined as:

Deterministic state transition:
$\displaystyle{\bm{h}}_{t}=f({\bm{h}}_{t-1},{\bm{z}}_{t-1},{\bm{a}}_{t-1})$
Stochastic state transition:
$\displaystyle{\bm{z}}_{t}=p({\bm{z}}_{t}|{\bm{h}}_{t})$ Observation model:
$\displaystyle{\bm{o}}_{t}=p({\bm{o}}_{t}|{\bm{h}}_{t},{\bm{z}}_{t})$

where $f({\bm{h}}_{t-1},{\bm{z}}_{t-1},{\bm{a}}_{t-1})$ is implemented as a
recurrent neural network (RNN) that carries the dependency on the stochastic
and deterministic parts at the previous timestep. Then, we obtain the
representation with the encoder
$p_{\theta}(({\bm{s}}_{1:T}|{\bm{o}}_{1:T},{\bm{a}}_{1:T}){=}\prod_{t}p_{\theta}({\bm{s}}_{t}|{\bm{o}}_{t},{\bm{h}}_{t})$,
where ${\bm{h}}_{t}$ retains information from
${\bm{s}}_{t-1}=({\bm{z}}_{t-1},{\bm{h}}_{t-1})$ and ${\bm{a}}_{t-1}$. The
encoder architecture based on the RSSM model encourages the representations to
be predictive of future states, one of the key properties of DRIBO.

###  4.5 Incorporating DRIBO in RL

We simultaneously train our representation learning models with the RL agent
by adding ${\mathcal{L}}_{\text{DRIBO}}$ (Algorithm 1) as an auxiliary
objective during training. The multi-view observations required by DRIBO can
be easily obtained using the same experience replay of RL agents with data
augmentation. The policy and (or) value function in RL will directly take the
representation of the visual observation to allow the models in RL to backprop
to our encoder. This further improves the performance by accommodating the
learned representations with policy/value function. As a result, DRIBO
encourages the agent to learn the underlying task-relevant dynamics of the
environments while being robust against visual changes that are task-
irrelevant.

We demonstrate the effectiveness of DRIBO by building the agents on top of SAC
[9] and PPO [32] in Section 5.1 and Section 5.2 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.

Figure 2: The left images are DMC observations in the clean setting. The
center images are observations with natural videos as backgrounds during
testing. The right images are the spatial attention maps of encoders for the
center images.

For effectiveness, we demonstrate DRIBO’s performance on the DeepMind Control
Suite with no distractors (DMC [33]). The DMC suite provides qualitatively
different visual control challenges. For robustness, we investigate whether
our DRIBO agent can ignore high-dimensional visual distractors that are task-
irrelevant in the DMC environments when the backgrounds are replaced with
natural videos from the Kinetics dataset [16]. For generalization, we present
results on Procgen [6] which provides different levels of the same game to
test how well agents generalize to unseen levels. Since DRIBO does not assume
that observation at each timestep provides full observability of the
underlying dynamics, we use single-step observations to train the DRIBO
representations. By contrast, current SOTA approaches require the use of
consecutive observations111All other methods stack 3 consecutive observations
while DRIBO does not use frame stack. to capture predictive properties of the
underlying representations implicitly.

For the DMC suite, all agents are built on top of SAC, an off-policy RL
algorithm. For the Procgen suite, we augment PPO, an on-policy RL baseline for
Procgen, with DRIBO. Implementation details are given in Appendix D.

###  5.1 Effectiveness and Robustness

First, we study the effectiveness and robustness of DRIBO agents. We evaluate
our approach on DMC under "clean" settings (without distractors), as well as
much more difficult settings with distractors.

We compare DRIBO against several state-of-the-art methods. The first is RAD
[21], a recent method that uses augmented data to train pixel-based policies
on DMC benchmarks. The second is SLAC [23], a SOTA representation learning
method for RL that learns a dynamic model using a reconstruction loss. The
third is CURL [20], 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. Finally, we compare with DBC [41] which is the most
similar work to ours. DBC learns an invariant representation based on
bisimulation metrics without requiring reconstruction. We apply random
crop+grayscale to obtain the augmented data for RAD which achieves the best
performance in [21]. For CURL and DRIBO, we apply random crop to obtain
augmented data and multi-view observations across different settings.

Clean Setting. For clean setting, the pixel observations have simple
backgrounds as shown in Figure 2 (left column). Figure 3 shows that RAD, SLAC
and CURL generally perform the best, whereas DRIBO consistently outperforms
DBC and matches SOTA. However, since the testing and training environments are
identical in this setting, the RL agents may overfit to the training
environments.

Natural Video Setting. To investigate the issue of overfitting, we introduce
high-dimensional visual distractors by using natural videos from the Kinetics
dataset [16] as new backgrounds [40] (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 classes. See Appendix C.3 for
snapshots under the natural video setting.

(a)  Figure 3: Results for DMC tasks over 5 seeds with one standard error
shaded in the clean setting. (a)  Figure 4: Results for DMC tasks over 5 seeds
with one standard error shaded in the natural video setting.

In Figure 2, spatial attention maps [39] of the trained DRIBO encoder
demonstrate that DRIBO trains agents to focus on the robot body while ignoring
irrelevant scene details in the background. Figure 4 shows that DRIBO performs
substantially better than RAD, SLAC and CURL which do not discard task-
irrelevant information explicitly. Compared to DBC, a recent state-of-art
method for learning representations that are invariant to task-irrelevant
information, DRIBO outperforms it consistently.

Visualizing Learned Representations. We visualize the representations learned
by DRIBO and CURL with t-SNE [35]. Figure 5 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.

Figure 5: 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.

###  5.2 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. For this reason, we use the ProcGen
suite [6] 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 different 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 [31] 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.

Table 1: Procgen returns on test levels after training on 25M environment
steps. The mean and standard deviation are computed over 10 seeds.

Env PPO RAD DrAC UCB-DrAC DRIBO BigFish 4.0 ${\pm}$ 1.2 9.9 ${\pm}$ 1.7 8.7
${\pm}$ 1.4 9.7 ${\pm}$ 1.0 10.9 ${\pm}$ 1.6 StarPilot 24.7 ${\pm}$ 3.4 33.4
${\pm}$ 5.1 29.5 ${\pm}$ 5.4 30.2 ${\pm}$ 2.8 36.5 ${\pm}$ 3.0 FruitBot 26.7
${\pm}$ 0.8 27.3 $\pm$ 1.8 28.2 ${\pm}$ 0.8 28.3 ${\pm}$ 0.9 30.8 ${\pm}$ 0.8
BossFight 7.7 $\pm$ 1.0 7.9 ${\pm}$ 0.6 7.5 $\pm$ 0.8 8.3 ${\pm}$ 0.8 12.0
${\pm}$ 0.5 Ninja 5.9 ${\pm}$ 0.7 6.9 $\pm$ 0.8 7.0 ${\pm}$ 0.4 6.9 $\pm$ 0.6
9.7 ${\pm}$ 0.7 Plunder 5.0 $\pm$ 0.5 8.5 ${\pm}$ 1.2 9.5 $\pm$ 1.0 8.9
${\pm}$ 1.0 5.8 ${\pm}$ 1.0 CaveFlyer 5.1 ${\pm}$ 0.9 5.1 $\pm$ 0.6 6.3
${\pm}$ 0.8 5.3 $\pm$ 0.9 7.5 ${\pm}$ 1.0 CoinRun 8.5 $\pm$ 0.5 9.0 ${\pm}$
0.8 8.8 $\pm$ 0.2 8.5 ${\pm}$ 0.6 9.2 ${\pm}$ 0.7 Jumper 5.8 ${\pm}$ 0.5 6.5
$\pm$ 0.6 6.6 ${\pm}$ 0.4 6.4 $\pm$ 0.6 8.4 ${\pm}$ 1.6 Chaser 5.0 $\pm$ 0.8
5.9 ${\pm}$ 1.0 5.7 ${\pm}$ 0.6 6.7 ${\pm}$ 0.6 4.8 ${\pm}$ 0.8 Climber 5.7
${\pm}$ 0.8 6.9 $\pm$ 0.8 7.1 $\pm$ 0.7 6.5 $\pm$ 0.8 8.1 ${\pm}$ 1.6
DodgeBall 11.7 $\pm$ 0.3 2.8 ${\pm}$ 0.7 4.3 ${\pm}$ 0.8 4.7 ${\pm}$ 0.7 3.8
${\pm}$ 0.9 Heist 2.4 ${\pm}$ 0.5 4.1 $\pm$ 1.0 4.0 $\pm$ 0.8 4.0 ${\pm}$ 0.7
7.7 ${\pm}$ 1.6 Leaper 4.9 $\pm$ 0.7 4.3 ${\pm}$ 1.0 5.3 ${\pm}$ 1.1 5.0
${\pm}$ 0.3 5.3 ${\pm}$ 1.5 Maze 5.7 ${\pm}$ 0.6 6.1 ${\pm}$ 1.0 6.6 ${\pm}$
0.8 6.3 ${\pm}$ 0.6 8.5 ${\pm}$ 1.6 Miner 8.5 ${\pm}$ 0.5 9.4 ${\pm}$ 1.2 9.8
${\pm}$ 0.6 9.7 ${\pm}$ 0.7 9.8 ${\pm}$ 0.9

Results in Table 1 show that DRIBO attains higher averaged testing returns
compared to the PPO baseline and other augmentation-based methods. The few
environments, in which our approach does not outperform the others, share the
commonality that task-relevant layouts remain static throughout the same run
of the game. Since the current version of DRIBO only considers the mutual
information between the complete input and the encoder output (global MI
[14]), it may fail to capture local features. The representations for a
sequence of observations within the same run of the game are treated as
globally negative pairs in DRIBO but they may be locally positive pairs. Thus,
the performance of DRIBO can be further improved by considering local features
(e.g. positions of the layouts) shared between representations as positive
pairs in the mutual information estimation.

##  6 Conclusion

In this paper, we introduce a novel robust representation learning approach in
RL based on the multi-view information bottleneck principle. Visual
observations are encoded into representations that are robust against
different task-irrelevant details and predictive of the future, a property
central to the sequential aspect of RL. Our experimental results show that 1)
DRIBO learns representations that are robust against task-irrelevant
distractions and boosts training performance when complex visual distractors
are introduced and 2) DRIBO improves generalization performance compared to
well-established baselines on the large-scale Procgen benchmarks. We plan to
explore the direction of incorporating knowledge about locality in the
observations into DRIBO in the future.

## References

  * Agarwal et al. [2021] 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.
  * Belghazi et al. [2018] 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_ , pages 531–540. PMLR, 2018.
  * Bojarski et al. [2016] 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.
  * Chen et al. [2020] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pages 1597–1607. PMLR, 2020.
  * Cobbe et al. [2019] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman.  Quantifying generalization in reinforcement learning.  In _International Conference on Machine Learning_ , pages 1282–1289. PMLR, 2019.
  * Cobbe et al. [2020] Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman.  Leveraging procedural generation to benchmark reinforcement learning.  In _International Conference on Machine Learning_. PMLR, 2020.
  * Farebrother et al. [2018] Jesse Farebrother, Marlos C Machado, and Michael Bowling.  Generalization and regularization in dqn.  _arXiv preprint arXiv:1810.00123_ , 2018.
  * Federici et al. [2020] 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.
  * Haarnoja et al. [2018] 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_ , pages 1861–1870. PMLR, 2018.
  * Hafner et al. [2019] 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_ , pages 2555–2565. PMLR, 2019.
  * Hafner et al. [2020] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  _International Conference on Learning Representation_ , 2020.
  * Hafner et al. [2021] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba.  Mastering atari with discrete world models.  _International Conference on Learning Representation_ , 2021.
  * Hénaff et al. [2020] 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_ , pages 4182–4192. PMLR, 2020.
  * Hjelm et al. [2019] 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.
  * Jonschkowski and Brock [2015] Rico Jonschkowski and Oliver Brock.  Learning state representations with robotic priors.  _Autonomous Robots_ , 39(3):407–428, 2015.
  * Kay et al. [2017] 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.
  * Kim et al. [2019] Hyoungseok Kim, Jaekyeom Kim, Yeonwoo Jeong, Sergey Levine, and Hyun Oh Song.  Emi: Exploration with mutual information.  In _International Conference on Machine Learning_ , pages 3360–3369, 2019.
  * Lange and Riedmiller [2010] Sascha Lange and Martin Riedmiller.  Deep auto-encoder neural networks in reinforcement learning.  In _The 2010 International Joint Conference on Neural Networks (IJCNN)_ , pages 1–8. IEEE, 2010.
  * Lange et al. [2012] 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)_ , pages 1–8. IEEE, 2012.
  * Laskin et al. [2020a] Michael Laskin, Aravind Srinivas, and Pieter Abbeel.  Curl: Contrastive unsupervised representations for reinforcement learning.  In _International Conference on Machine Learning_ , pages 5639–5650. PMLR, 2020a.
  * Laskin et al. [2020b] 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.
  * LeCun et al. [1989] Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne Hubbard, and Lawrence D Jackel.  Backpropagation applied to handwritten zip code recognition.  _Neural computation_ , 1(4):541–551, 1989.
  * Lee et al. [2020] 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, 2020.
  * Levine et al. [2016] 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.
  * Li et al. [2019] Minne Li, Lisheng Wu, WANG Jun, and Haitham Bou Ammar.  Multi-view reinforcement learning.  In _Advances in neural information processing systems_ , pages 1420–1431, 2019.
  * Li et al. [2018] 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.
  * Mazoure et al. [2020] 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.
  * Mnih et al. [2015] 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.
  * Oord et al. [2018] Aaron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding.  _arXiv preprint arXiv:1807.03748_ , 2018.
  * Poole et al. [2019] 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_ , pages 5171–5180. PMLR, 2019.
  * Raileanu et al. [2020] 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.
  * Schulman et al. [2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ , 2017.
  * Tassa et al. [2018] 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.
  * Tishby et al. [2000] Naftali Tishby, Fernando C Pereira, and William Bialek.  The information bottleneck method.  _arXiv preprint physics/0004057_ , 2000.
  * Van der Maaten and Hinton [2008] Laurens Van der Maaten and Geoffrey Hinton.  Visualizing data using t-sne.  _Journal of machine learning research_ , 9(11), 2008.
  * Wahlström et al. [2015] 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.
  * Watter et al. [2015] 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_ , pages 2746–2754, 2015.
  * Yu et al. [2019] Wenhao Yu, C Karen Liu, and Greg Turk.  Policy transfer with strategy optimization.  In _International Conference on Learning Representations_ , 2019.
  * Zagoruyko and Komodakis [2017] 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.
  * Zhang et al. [2018a] Amy Zhang, Yuxin Wu, and Joelle Pineau.  Natural environment benchmarks for reinforcement learning.  _arXiv preprint arXiv:1811.06032_ , 2018a.
  * Zhang et al. [2021] 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.
  * Zhang et al. [2018b] Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio.  A study on overfitting in deep reinforcement learning.  _arXiv preprint arXiv:1804.06893_ , 2018b.

## Appendix

##  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$.

  1. (P.1)

Positivity:

$I(X;Y)\geq 0,I(X;Y|Z)\geq 0$

  2. (P.2)

Chain rule:

$I(XY;Z)=I(Y;Z)+I(X;Z|Y)$

  3. (P.3)

Chain rule (Multivariate Mutual Information):

$I(X;Y;Z)=I(Y;Z)-I(Y;Z|X)$

  4. (P.4)

Entropy and Mutual Information:

$I(X;Y)=H(X)-H(X|Y)$

  5. (P.5)

Chain rule for Entropy:

$H(X_{1},X_{2},\dots,X_{n}){=}\sum_{i=1}^{n}H(X_{i}|X_{i-1},{\dots},X_{1})$

  6. (P.6)

The mutual information among three variables is bounded by:

${-}\min\\{I(X;Y|Z),I(X;Z|Y),I(Y;Z|X)\\}\leq
I(X;Y;Z){\leq}\min\\{I(X;Y),I(X;Z),I(Y;Z)\\}$

###  A.1 Theorem 3

We first show that $I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})$ is a
lower bound of $I({\bm{s}}_{1:T};{\bm{o}}_{1:T})$.

$\displaystyle
I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})\overset{\ref{p:multi_chain}}{=}I({\bm{o}}_{1:T};{\bm{s}}_{1:T})-I({\bm{s}}_{1:T};{\bm{o}}_{1:T};{\bm{a}}_{1:T})$

Since ${\bm{s}}_{1:T}$ are representations of ${\bm{o}}_{1:T}$, we have
$I({\bm{s}}_{1:T};{\bm{a}}_{1:T}|{\bm{o}}_{1:T})=0$. With (P.1) and (P.6), we
have that:

$\displaystyle I({\bm{s}}_{1:T};{\bm{o}}_{1:T};{\bm{a}}_{1:T})\geq 0$

Thus, we have:

$\displaystyle I({\bm{o}}_{1:T};{\bm{s}}_{1:T})\geq
I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})$

Furthermore, since $I({\bm{s}}_{1:T};{\bm{o}}_{1:T};{\bm{a}}_{1:T})$ is lower
bounded, maximizing $I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})$ also
maximizes $I({\bm{o}}_{1:T};{\bm{s}}_{1:T})$.

######  Theorem A.1.

Let ${\bm{o}}_{1:T}$ be the observation sequence obtained by executing action
sequence ${\bm{a}}_{1:T}$. If ${\bm{s}}_{1:T}$ is a sequence of sufficient
representations for ${\bm{o}}_{1:T}$, then we have:

$I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})\geq\sum_{t=1}^{T}I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
(9)

###### Proof.

We indicate the property we use for each step of the derivation below.

$\displaystyle I({\bm{s}}_{1:T};{\bm{o}}_{1:T}|{\bm{a}}_{1:T})$
$\displaystyle\overset{\ref{p:m and h}}{=}$ $\displaystyle
H({\bm{s}}_{1:T}|{\bm{a}}_{1:T})-H({\bm{s}}_{1:T}|{\bm{o}}_{1:T},{\bm{a}}_{1:T})$
$\displaystyle\overset{\ref{p:h_chain}}{=}$
$\displaystyle\sum_{t}\left(H({\bm{s}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1})-H({\bm{s}}_{t}|{\bm{a}}_{1:T},{\bm{o}}_{1:T},{\bm{s}}_{1:t-1})\right)$
$\displaystyle\overset{\ref{p:m and h}}{=}$
$\displaystyle\sum_{t}I({\bm{s}}_{t};{\bm{o}}_{1:T}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1})$
$\displaystyle\overset{\ref{p:m and h}}{=}$
$\displaystyle\sum_{t}\left(H({\bm{o}}_{1:T}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1})-H({\bm{o}}_{1:T}|{\bm{s}}_{t},{\bm{a}}_{1:T},{\bm{s}}_{1:t-1})\right)$
$\displaystyle\overset{\ref{p:h_chain}}{=}$
$\displaystyle\sum_{t}\sum_{\tau}(H({\bm{o}}_{\tau}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:\tau-1})$
$\displaystyle-H({\bm{o}}_{\tau}|{\bm{s}}_{t},{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:\tau-1}))$
$\displaystyle\overset{\ref{p:m and h}}{=}$
$\displaystyle\sum_{t}\sum_{\tau}I({\bm{s}}_{t};{\bm{o}}_{\tau}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:\tau-1})$
$\displaystyle\overset{\ref{p:positive}}{\geq}$
$\displaystyle\sum_{t}I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})$
$\displaystyle=$
$\displaystyle\sum_{t}I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$

Here, we provide a formal proof for the last step of the derivation above. Let
$\tau$ represent ${\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1}$. We have

$\displaystyle
I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})$
$\displaystyle=\int_{\tau}\int_{{\bm{s}}_{t}}\int_{{\bm{o}}_{t}}p(\tau)p({\bm{s}}_{t},{\bm{o}}_{t}|\tau)\log\left[\frac{p({\bm{s}}_{t},{\bm{o}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})}{p({\bm{s}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})p({\bm{o}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})}\right]d{\bm{s}}_{t}d{\bm{o}}_{t}d\tau$

With the density of multi-view trajectories Equation 1, we can observe that
${\bm{o}}_{t}$ and ${\bm{s}}_{t}$ are generated by
$p({\bm{o}}_{t}|{\bm{s}}_{t})p({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})=p({\bm{s}}_{t},{\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$.
Thus, we have:

$\displaystyle
p({\bm{s}}_{t},{\bm{o}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})$
$\displaystyle=p({\bm{s}}_{t},{\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
$\displaystyle
p({\bm{s}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})$
$\displaystyle=p({\bm{s}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$ $\displaystyle
p({\bm{o}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})$
$\displaystyle=p({\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$

This in turn implies that:

$\displaystyle
I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{a}}_{1:T},{\bm{s}}_{1:t-1},{\bm{o}}_{1:t-1})=I({\bm{s}}_{t};{\bm{o}}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$

∎

As a result, minimizing ${\mathcal{L}}_{\text{InfoMax}}$ is maximizing a lower
bound of $I({\bm{s}}_{1:T}^{(1)};{\bm{o}}_{1:T}^{(1)}|{\bm{a}}_{1:T})$ since

$\displaystyle
I({\bm{s}}_{1:T}^{(1)};{\bm{o}}_{1:T}^{(1)}|{\bm{a}}_{1:T})\geq\sum_{t}(I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(2)}){+}I({\bm{o}}_{t}^{(2)};{\bm{s}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1}))$

###  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 ${\bm{o}}_{1:T}$ and ${\bm{a}}^{*}_{1:T}$ be random variables with joint
distribution $p({\bm{o}}_{1:T},{\bm{a}}^{*}_{1:T})$. Let ${\bm{s}}_{1:T}$ be
the representation of ${\bm{o}}_{1:T}$, then ${\bm{s}}_{1:T}$ is sufficient
for ${\bm{a}}^{*}_{1:T}$ if and only if
$I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T})=I({\bm{s}}_{1:T};{\bm{a}}^{*}_{1:T})$.
Also, ${\bm{s}}_{t}$ is a sufficient representation of ${\bm{o}}_{t}$ since
$I({\bm{o}}_{t};{\bm{a}}^{*}_{t}|{\bm{s}}_{t},{\bm{s}}_{t-1},{\bm{a}}_{t-1})=0$.

Hypothesis:

  1. (H.1)

${\bm{s}}_{1:T}$ is a sequence of sufficient representations for
${\bm{o}}_{1:T}$:

$I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T}|{\bm{s}}_{1:T})=0$

###### Proof.

$\displaystyle I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T}|{\bm{s}}_{1:T})$
$\displaystyle\overset{\ref{p:multi_chain}}{=}$ $\displaystyle
I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T})-I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T};{\bm{s}}_{1:T})$
$\displaystyle\overset{\ref{p:multi_chain}}{=}$ $\displaystyle
I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T})-I({\bm{a}}^{*}_{1:T};{\bm{s}}_{1:T})-I({\bm{a}}^{*}_{1:T};{\bm{s}}_{1:T}|{\bm{o}}_{1:T})$

With ${\bm{s}}_{1:T}$ as a representation of ${\bm{o}}_{1:T}$, we have
$I({\bm{s}}_{1:T};{\bm{a}}^{*}_{1:T}|{\bm{o}}_{1:T}){=}0$. The reason is that
${\bm{o}}_{1:T}$ shares the same level of information as ${\bm{a}}^{*}_{1:T}$
and ${\bm{s}}_{1:T}$. Then,

$\displaystyle I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T}|{\bm{s}}_{1:T})=$
$\displaystyle
I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T})-I({\bm{a}}^{*}_{1:T};{\bm{s}}_{1:T})$
(10)

So the sufficiency condition
$I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T}|{\bm{s}}_{1:T})=0$ holds if and only if
$I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T})=I({\bm{a}}^{*}_{1:T};{\bm{s}}_{1:T})$.

We factorize the mutual information between sequential observations and
optimal actions

$\displaystyle I({\bm{o}}_{1:t};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{o}}_{1:t-1})+I({\bm{o}}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{h:sufficiency}}{=}$ $\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{o}}_{1:t-1})+I({\bm{s}}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle I({\bm{s}}_{1:t};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{s}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1})+I({\bm{s}}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{h:sufficiency}}{=}$ $\displaystyle
I({\bm{s}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1})+I({\bm{o}}_{1:t-1};{\bm{a}}^{*}_{1:t})$

Then we obtain the following relation:

$\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{o}}_{1:t-1})=I({\bm{s}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1})$
(11)

We also have

$\displaystyle I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{o}}_{1:t-1})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}_{1:t};{\bm{a}}^{*}_{1:t})-I({\bm{o}}_{1:t-1};{\bm{a}}_{1:t}^{*})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}_{1:t-1};{\bm{a}}^{*}_{1:t}|{\bm{o}}_{t})+I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t})$
$\displaystyle-I({\bm{o}}_{1:t-1};{\bm{a}}_{1:t}^{*})$
$\displaystyle\overset{\ref{h:sufficiency}}{=}$ $\displaystyle
I({\bm{s}}_{1:t-1};{\bm{a}}^{*}_{1:t}|{\bm{o}}_{t})+I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t})$
$\displaystyle-I({\bm{s}}_{1:t-1};{\bm{a}}_{1:t}^{*})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}_{t}{\bm{s}}_{1:t-1};{\bm{a}}^{*}_{1:t})-I({\bm{s}}_{1:t-1};{\bm{a}}_{1:t}^{*})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1})$
$\displaystyle\overset{\text{Equation~{}\ref{eq:equality}}}{=}$ $\displaystyle
I({\bm{s}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1})$
$\displaystyle\overset{\ref{p:chain}}{\Longleftrightarrow}$ $\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{t}|{\bm{a}}^{*}_{1:t-1},{\bm{s}}_{1:t-1})+I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t-1}|{\bm{s}}_{1:t-1})$
$\displaystyle=$ $\displaystyle
I({\bm{s}}_{t};{\bm{a}}^{*}_{t}|{\bm{a}}^{*}_{1:t-1},{\bm{s}}_{1:t-1})+I({\bm{s}}_{t};{\bm{a}}^{*}_{1:t-1}|{\bm{s}}_{1:t-1})$
$\displaystyle\overset{\text{Equation~{}\ref{eq:equality}}}{\Longleftrightarrow}$
$\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{t}|{\bm{a}}^{*}_{1:t-1},{\bm{s}}_{1:t-1}){=}I({\bm{s}}_{t};{\bm{a}}^{*}_{t}|{\bm{a}}^{*}_{1:t-1},{\bm{s}}_{1:t-1})$
$\displaystyle\overset{\text{Equation~{}\ref{eq:sufficiency}}}{\Longleftrightarrow}$
$\displaystyle
I({\bm{o}}_{t},{\bm{a}}^{*}_{t}|{\bm{s}}_{t},{\bm{s}}_{1:t-1},{\bm{a}}^{*}_{1:t-1})=0$

With the above derivation and Markov property, we have
$I({\bm{o}}_{t};{\bm{a}}^{*}_{t}|{\bm{s}}_{t},{\bm{s}}_{t-1},{\bm{a}}_{t-1})=0$.
We can generalize ${\bm{a}}^{*}_{t-1}$ to any ${\bm{a}}_{t-1}$ by assuming
$a^{*}_{t}$ as the optimal action for state ${\bm{s}}_{t}$ whose last timestep
state-action pair is $({\bm{s}}_{t-1},{\bm{a}}_{t-1})$. Thus, we have
${\bm{s}}_{t}$ is a sufficient representation for ${\bm{o}}_{t}$ if and only
if ${\bm{s}}_{1:T}$ is a sufficient representation of ${\bm{o}}_{1:T}$. ∎

###  A.3 Multi-View Redundancy and Sufficiency

######  Proposition A.1.

${\bm{o}}_{1:T}^{(1)}$ is a redundant view with respect to
${\bm{o}}_{1:T}^{(2)}$ to obtain ${\bm{a}}^{*}_{1:T}$ if only if
$I({\bm{o}}_{1:T}^{(1)};{\bm{a}}_{1:T}^{*}|{\bm{o}}_{1:T}^{(2)})=0$. Any
representation ${\bm{s}}_{1:T}^{(1)}$ of ${\bm{o}}_{1:T}^{(1)}$ that is
sufficient for ${\bm{o}}_{1:T}^{(2)}$ is also sufficient for
${\bm{a}}^{*}_{1:T}$.

###### Proof.

See proof of Proposition B.3 in the MIB paper [8]. ∎

######  Corollary A.1.

Let ${\bm{o}}_{1:T}^{(1)}$ and ${\bm{o}}_{1:T}^{(2)}$ be two mutually
redundant views for ${\bm{a}}^{*}_{1:T}$. Let ${\bm{s}}_{1:T}^{(1)}$ be a
representation of ${\bm{o}}_{1:T}^{(1)}$. If ${\bm{s}}_{1:T}^{(1)}$ is
sufficient for ${\bm{o}}_{1:T}^{(2)}$, ${\bm{s}}_{t}^{(1)}$ can derive
${\bm{a}}^{*}_{t}$ as the joint observation of the two views
($I({\bm{o}}_{t}^{(1)}{\bm{o}}_{t}^{(2)};{\bm{a}}_{t}^{*}|{\bm{s}}_{t-1},{\bm{a}}_{t-1}){=}I({\bm{s}}_{t}^{(1)};{\bm{a}}_{t}^{*}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$),
where ${\bm{s}}_{t-1}$ is any sufficient representation at timestep $t-1$.

###### Proof.

For the sequential data, see proof of Corollary B.2.1 in the MIB paper [8] to
prove

$\displaystyle
I({\bm{o}}_{1:T}^{(1)}{\bm{o}}_{1:T}^{(2)};{\bm{a}}_{1:T}^{*}){=}I({\bm{s}}_{1:T}^{(1)};{\bm{a}}_{1:T}^{*})$

According to Theorem A.2, if ${\bm{s}}_{1:T}^{(1)}$ is a sufficient
representation of ${\bm{o}}_{1:T}^{(2)}$, ${\bm{s}}_{t}^{(1)}$ is a sufficient
representation of ${\bm{o}}_{t}^{(2)}$. Similar to proof on sequential data,
we can use Corollary B.2.1 in the MIB paper [8] to show that

$\displaystyle
I({\bm{o}}_{t}^{(1)}{\bm{o}}_{t}^{(2)};{\bm{a}}_{t}^{*}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})=I({\bm{s}}_{t}^{(1)};{\bm{a}}_{t}^{*}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$

∎

######  Theorem A.3.

Let the two views ${\bm{o}}_{1:T}^{(1)}$ and ${\bm{o}}_{1:T}^{(2)}$ of
observation ${\bm{o}}_{1:T}$ are obtained by data augmentation transformation
sequences $t^{(1)}_{1:T}$ and $t^{(2)}_{1:T}$ respectively
(${\bm{o}}^{(1)}_{1:T}{=}t^{(1)}_{1:T}({\bm{o}}_{1:T})$ and
${\bm{o}}^{(2)}_{1:T}{=}t^{(2)}_{1:T}({\bm{o}}_{1:T})$). Whenever
$I(t^{(1)}_{1:T}({\bm{o}}_{1:T});{\bm{a}}^{*}_{1:T}){=}I(t^{(2)}_{1:T}({\bm{o}}_{1:T});{\bm{a}}^{*}_{1:T}){=}I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T})$,
the two views ${\bm{o}}_{1:T}^{(1)}$ and ${\bm{o}}_{1:T}^{(2)}$ must be
mutually redundant for ${\bm{a}}^{*}_{1:T}$. Besides, the two views
${\bm{o}}_{t}^{(1)}$ and ${\bm{o}}_{t}^{(2)}$ must be mutually redundant for
${\bm{a}}^{*}_{t}$.

###### Proof.

Let ${\bm{s}}_{1: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.

$\displaystyle
I(t^{(1)}_{1:t}({\bm{o}}_{1:t});{\bm{a}}^{*}_{1:t})=I({\bm{o}}^{(1)}_{1:t};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}^{(1)}_{t};{\bm{a}}^{*}_{1:t}|{\bm{o}}^{(1)}_{1:t-1}){+}I({\bm{o}}^{(1)}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{prop:sufficiency}}{=}$ $\displaystyle
I({\bm{o}}_{t}^{(1)};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1}){+}I({\bm{s}}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle
I(t^{(2)}_{1:t}({\bm{o}}_{1:t});{\bm{a}}^{*}_{1:t})=I({\bm{o}}^{(2)}_{1:t};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}^{(2)}_{t};{\bm{a}}^{*}_{1:t}|{\bm{o}}^{(2)}_{1:t-1}){+}I({\bm{o}}^{(2)}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{prop:sufficiency}}{=}$ $\displaystyle
I({\bm{o}}_{t}^{(2)};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1}){+}I({\bm{s}}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle
I(t^{(2)}_{1:t}({\bm{o}}_{1:t});{\bm{a}}^{*}_{1:t})=I({\bm{o}}^{(2)}_{1:t};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}^{(2)}_{t};{\bm{a}}^{*}_{1:t}|{\bm{o}}^{(2)}_{1:t-1}){+}I({\bm{o}}^{(2)}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle
I({\bm{o}}_{1:t};{\bm{a}}^{*}_{1:t})=I({\bm{o}}_{1:t};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{o}}_{1:t-1}){+}I({\bm{o}}_{1:t-1};{\bm{a}}^{*}_{1:t})$
$\displaystyle\overset{\ref{prop:sufficiency}}{=}$ $\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1}){+}I({\bm{s}}_{1:t-1};{\bm{a}}^{*}_{1:t})$

Then, we have the following equality

$\displaystyle
I({\bm{o}}_{t}^{(1)};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1}){=}I({\bm{o}}_{t}^{(2)};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1}){=}I({\bm{o}}_{t};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1})$

Similar as derivation in Theorem A.2

$\displaystyle I({\bm{o}}_{t}^{(1)};{\bm{a}}^{*}_{1:t}|{\bm{s}}_{1:t-1})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{o}}_{t}^{(1)};{\bm{a}}^{*}_{t}|{\bm{a}}_{1:t-1}^{*},{\bm{s}}_{1:t-1})+I({\bm{o}}_{t}^{(1)};{\bm{a}}_{1:t-1}^{*}|{\bm{s}}_{1:t-1})$
$\displaystyle\overset{\text{Equation~{}\ref{eq:equality}}}{=}$ $\displaystyle
I({\bm{o}}_{t}^{(1)};{\bm{a}}^{*}_{t}|{\bm{a}}_{1:t-1}^{*},{\bm{s}}_{1:t-1})+I({\bm{s}}_{t};{\bm{a}}_{1:t-1}^{*}|{\bm{s}}_{1:t-1})$

We apply the same derivation for ${\bm{o}}^{(2)}$ and ${\bm{o}}$, we have the
following with Markov property

$\displaystyle
I(t_{t}^{(1)}({\bm{o}}_{t});{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
$\displaystyle=$ $\displaystyle
I(t_{t}^{(2)}({\bm{o}}_{t});{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
$\displaystyle=$ $\displaystyle
I({\bm{o}}_{t};{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{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
[8] 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}({\bm{o}}_{1:T});{\bm{a}}^{*}_{1:T}){=}I(t^{(2)}_{1:T}({\bm{o}}_{1:T});{\bm{a}}^{*}_{1:T}){=}I({\bm{o}}_{1:T};{\bm{a}}^{*}_{1:T})$.
If ${\bm{s}}^{(1)}_{1:T}$ is a sufficient representation for
$t_{1:T}^{(2)}({\bm{o}}_{1:T})$ then
$I({\bm{o}}_{t};{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})=I({\bm{s}}_{t}^{(1)};{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$.

###### Proof.

Since $t_{1:T}^{(1)}({\bm{o}}_{1:T})$ is redundant for
$t_{1:T}^{(2)}({\bm{o}}_{1:T})$ (Theorem A.3), any representation
${\bm{s}}_{t}^{(1)}$ of $t_{1:T}^{(1)}({\bm{o}}_{1:T})$ that is sufficient for
$t_{1:T}^{(2)}({\bm{o}}_{1:T})$ must also be sufficient for ${\bm{a}}^{*}_{t}$
(Theorem A.2 and Proposition A.1). Using Theorem A.2 we have
$I({\bm{s}}_{t}^{(1)};{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1}){=}I(t_{t}^{(1)}({\bm{o}}_{t});{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$.
With
$I(t_{t}^{(1)}({\bm{o}}_{t});{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})=I({\bm{o}}_{t};{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$,
we conclude
$I({\bm{o}}_{t};{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})=I({\bm{s}}_{t}^{(1)};{\bm{a}}^{*}_{t}|{\bm{s}}_{t-1},{\bm{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 ${\bm{o}}_{t}^{(1)}$ and ${\bm{o}}_{t}^{(2)}$ be mutually redundant views
for ${\bm{a}}^{*}_{t}$ that share only optimal action information. Then a
sufficient representation of ${\bm{s}}_{t}^{(1)}$ of ${\bm{o}}_{t}^{{1}}$ for
${\bm{o}}_{t}^{(2)}$ that is minimal for ${\bm{o}}_{t}^{(2)}$ is also a
minimal representation for ${\bm{a}}^{*}_{t}$.

###### Proof.

See proof of Proposition E.1 in the MIB paper [8]. ∎

##  Appendix B DRIBO Loss Computation

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

$\displaystyle{\mathcal{L}}_{\frac{1{+}2}{2}}$ (12) $\displaystyle{=}$
$\displaystyle\frac{I({\bm{s}}_{t}^{(1)}{;}{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t{-}1}^{(1)}{,}{\bm{a}}_{t{-}1}{,}{\bm{o}}_{t}^{(2)}){+}I({\bm{s}}_{t}^{(2)}{;}{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t{-}1}^{(2)}{,}{\bm{a}}_{t{-}1}{,}{\bm{o}}_{t}^{(1)})}{2}$
$\displaystyle{-}\frac{\lambda_{1}I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1}){+}\lambda_{2}I({\bm{s}}_{t}^{(2)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})}{2}$
(13)

Consider ${\bm{s}}_{t}^{(1)}$ and ${\bm{s}}_{t}^{(2)}$ on the same domain
${\mathbb{S}}$,
$I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(2)})$
can be expressed as:

$\displaystyle
I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(2)})$
$\displaystyle{=}\mathbb{E}\left[\log\frac{p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})}{p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})}\right]$
$\displaystyle{=}\mathbb{E}\left[\log\frac{p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})}{p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})}\frac{p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})}{p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})}\right]$
$\displaystyle{=}D_{\text{KL}}(p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})||p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1}))$
$\displaystyle{-}D_{\text{KL}}(p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})||p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1}))$
$\displaystyle{\leq}D_{\text{KL}}(p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})||p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1}))$
(14)

Note that equality holds if the two distributions coincide. Analogously
$I({\bm{s}}_{t}^{(2)};{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1},{\bm{o}}_{t}^{(1)})$
is upper bounded by
$D_{\text{KL}}(p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})||p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1}))$.

Assume ${\bm{s}}_{t-1}$ is a sufficient representation of ${\bm{o}}_{t-1}$.
Then, ${\bm{s}}_{t-1}^{(1)}$ provides task-relevant information no more than
the sufficient representation ${\bm{s}}_{t-1}$.
$I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})$
can be thus re-expressed as:

$\displaystyle
I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})$
$\displaystyle\geq$ $\displaystyle
I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
$\displaystyle\overset{\ref{p:chain}}{=}$ $\displaystyle
I({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1}){-}I({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
$\displaystyle\overset{*}{=}$ $\displaystyle
I({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
$\displaystyle{=}$ $\displaystyle
I({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})+I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t}^{(2)},{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
$\displaystyle{\geq}$ $\displaystyle
I({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$ (15)

Where $*$ follows from ${\bm{s}}_{t}^{(2)}$ being the representation of
${\bm{o}}_{t}^{(2)}$. The bound is tight whenever ${\bm{s}}_{t}^{(2)}$ is
sufficient from ${\bm{s}}_{t}^{(1)}$
$(I({\bm{s}}_{t}^{(1)};{\bm{o}}_{t}^{(2)}|{\bm{s}}_{t}^{(2)},{\bm{s}}_{t-1},{\bm{a}}_{t-1}){=}0)$.
This happens whenever ${\bm{s}}_{t}^{(2)}$ contains all the information
regarding ${\bm{s}}_{t}^{(1)}$. Once again, we can have
$I({\bm{s}}_{t}^{(2)};{\bm{o}}_{t}^{(1)}|{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})\geq
I({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$.
Therefore, the averaged loss functions can be upper-bounded by

$\displaystyle{\mathcal{L}}_{\frac{1+2}{2}}{\leq}{-}\frac{\lambda_{1}+\lambda_{2}}{2}I({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
(16)
$\displaystyle{+}D_{\text{SKL}}(p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)}{,}{\bm{s}}_{t{-}1}^{(1)}{,}{\bm{a}}_{t{-}1}){||}p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t{-}1}^{(2)}{,}{\bm{a}}_{t{-}1}))$

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

$\displaystyle{\mathcal{L}}(\theta;\beta){=}-I_{\theta}({\bm{s}}_{t}^{(1)};{\bm{s}}_{t}^{(2)}|{\bm{s}}_{t-1},{\bm{a}}_{t-1})$
(17) $\displaystyle{+}\beta
D_{\text{SKL}}(p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)}{,}{\bm{s}}_{t{-}1}^{(1)}{,}{\bm{a}}_{t{-}1})||p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)}{,}{\bm{s}}_{t{-}1}^{(2)}{,}{\bm{a}}_{t{-}1}))$

In Algorithm 1, we use ${\bm{s}}_{t}^{(1)}\sim
p_{\theta}({\bm{s}}_{t}^{(1)}|{\bm{o}}_{t}^{(1)},{\bm{s}}_{t-1}^{(1)},{\bm{a}}_{t-1})$
and ${\bm{s}}_{t}^{(2)}\sim
p_{\theta}({\bm{s}}_{t}^{(2)}|{\bm{o}}_{t}^{(2)},{\bm{s}}_{t-1}^{(2)},{\bm{a}}_{t-1})$
to obtain representations for multi-view observations. We argue that the
substitution does not affect the effectiveness of the averaged objective. With
the multi-view assumption, we have that representations ${\bm{s}}_{t-1}^{(1)}$
and ${\bm{s}}_{t-1}^{(2)}$ 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
${\bm{s}}_{t}^{(1)}$ and ${\bm{s}}_{t}^{(2)}$ (first term in Equation 17) will
encourage the representations to share maximal task-relevant information.
Similar argument also works for the second term in Equation 17. Since
${\bm{s}}_{t-1}^{(1)}$ and ${\bm{s}}_{t-1}^{(2)}$ 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.

##  Appendix C Additional Results

###  C.1 InfoMax vs. Information Bottleneck

Figure 6: Average returns of cartpole swingup task in DMC environments with
natural videos as background.

We further investigate the effectiveness of using the information bottleneck
principle to reduce task-irrelevant information. We train an agent to only
maximize the mutual information between representations from different views,
which is equivalent to InfoMax [14, 27] and contrastive learning [20]. The
objective is similar to DRIBO loss but without the SKL term. The encoder uses
the same RSSM architecture as the one used in DRIBO. The hyperparameters are
also set to be the same. We show the performance gap on the cartpole swingup
task from DMC in Figure 6.

Figure 7: Average SKL values during training of cartpole swingup task 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 InfoMax
during training in Figure 7. The gap between the SKL values explains the
performance gap between the two approaches. 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 7 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 Additional DMC Results

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

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

###  C.3 Additional Visualization

Figure 10 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 11).

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

In Figure 11, 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.
Figure 11 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.

Figure 11: 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 addition to Figure 2, we also show spatial attention maps for each
convolutional layers in Figure 12. We can observe that DRIBO trained encoders
filter the complex distractors in the backgrounds gradually layer by layer.

Figure 12: 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 13. 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.

Figure 13: t-SNE of latent spaces learned with DRIBO (left) and CURL (right).
They are the same as t-SNE in Figure 5. 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 5). 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.

##  Appendix D Implementation Details

###  D.1 DRIBO + SAC

We first show how we train SAC agent given the representations of DRIBO. Let
$\phi({\bm{o}})={\bm{s}}\sim
p_{\theta}({\bm{s}}|{\bm{o}},{\bm{s}}^{\prime},{\bm{a}}^{\prime})$ denote the
encoder, where ${\bm{s}}^{\prime}$ and ${\bm{a}}^{\prime}$ as the
representation and action at last timestep.

Algorithm 2 SAC + DRIBO Encoder

1:Input RL batch
${\mathcal{B}}_{\text{RL}}=\\{(\phi({\bm{o}}_{i}),a_{i},r_{i},\phi({\bm{o}}_{i}^{\prime}))\\}_{i=1}^{(T-1)*N}$
with $(T-1)*N$ pairs of representation, action, reward and next
representation.

2:Get value:
$V=\min_{i=1,2}\hat{Q}_{i}(\hat{\phi}({\bm{o}}^{\prime}))-\alpha\log\pi(a|\hat{\phi}(o^{\prime}))$

3:Train critics: $J(Q_{i},\phi)=(Q_{i}(\phi({\bm{o}}))-r-\gamma V)^{2}$

4:Train actor:
$J(\pi){=}\alpha\log\pi(a|\phi({\bm{o}})){-}\min_{i=1,2}Q_{i}(\phi({\bm{o}}))$

5:Train alpha:
$J(\alpha)=-\alpha\log\pi(a|\phi({\bm{o}})){-}\alpha{\mathcal{H}}(a|\phi({\bm{o}}))$

6:Update target critics: $\hat{Q}_{i}=\tau_{Q}Q_{i}+(1-\tau_{Q})\hat{Q}_{i}$

7:Update target encoder:
$\hat{\phi}\leftarrow\tau_{\phi}\phi+(1-\tau_{\phi})\phi$

Then we incorporate the above SAC algorithm into minimizing DRIBO loss in
Algorithm 3

Algorithm 3 DRIBO + SAC

1:Input: Replay buffer ${\mathcal{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,\dots,E$ do

3: Sample sequential observations and actions from the environment and append
new samples to ${\mathcal{D}}$.

4: for each step $k=1,\dots,K$ do

5: Sample a sequential batch ${\mathcal{B}}\sim{\mathcal{D}}$.

6: Compute the representations batch ${\mathcal{B}}_{\text{RL}}$ which has the
shape $(T,N)$ using the encoder
$p_{\theta}({\bm{s}}_{1:T}|{\bm{o}}_{1:T},{\bm{a}}_{1:T})$

7: Train SAC agent: $\mathbb{E}_{{\mathcal{B}}_{\text{RL}}}[J(\pi,Q,\phi)]$
$\triangleright$ Algorithm 2

8: Update $\theta$ and $\psi$ to minimize ${\mathcal{L}}_{DRIBO}$ using
${\mathcal{B}}$ $\triangleright$ Algorithm 1.

9: end for

10:end for

###  D.2 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 ${\bm{s}}_{t}$ propagating from the initial observation of the
observation sequence, we store the representations as
${\bm{s}}_{t}^{\text{old}}$ while sampling from the on-policy batch. Then, we
use $\varphi({\bm{o}})={\bm{s}}\sim
p_{\theta}({\bm{s}}|{\bm{o}},{\bm{s}}^{\text{old}},{\bm{a}}^{\prime})$ to
denote the representation from the encoder. Here, ${\bm{s}}^{\text{old}}$ and
${\bm{a}}^{\prime}$ 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
$({\bm{o}}_{t},{\bm{s}}_{t-1}^{old},{\bm{a}}_{t-1})$.

Algorithm 4 DRIBO + PPO

1:Input: Replay buffer ${\mathcal{D}}$ and on-policy replay buffer
${\mathcal{D}}_{\text{PPO}}$ 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,\dots,E$ do

3: Sample sequential observations and actions from the environment
$\\{({\bm{o}}_{1:T},{\bm{a}}_{1:T},r_{1:T},{\bm{s}}_{1:T}^{\text{old}}\\}_{i=1}^{N}$.

4: Append new samples to ${\mathcal{D}}$ and update the on-policy replay
buffer ${\mathcal{D}}_{\text{PPO}}$.

5: for $j=1,\dots,M$ do

6:
$\\{(\varphi({\bm{o}}_{i}),a_{i},r_{i})\\}_{i=1}^{\lfloor\frac{T*N}{M}\rfloor}\sim{\mathcal{D}}_{\text{PPO}}$

7: Optimize PPO policy, value function and encoder using each sample
$(\varphi({\bm{o}}_{i}),a_{i},r_{i})$ in the batch.

8: Sample a sequential batch ${\mathcal{B}}\sim{\mathcal{D}}$.

9: Update $\theta$ and $\psi$ to minimize ${\mathcal{L}}_{DRIBO}$ using
${\mathcal{B}}$ $\triangleright$ Algorithm 1.

10: end for

11:end for

###  D.3 DMC

We use an almost identical encoder architecture as the encoder in the RSSM
paper [10], with two minor differences. Firstly, we deploy the encoder
architecture in Tassa et al. [33] 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
[29].

Pixel Preprocessing. We construct an observational input as a single frame,
where each frame is a RGB rendering of size 100 $\times$ 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\times 100$ pixel observations and
crop randomly to $84\times 84$ pixels. For random grayscale, it converts RGB
images to grayscale with a probability $p=0.3$.

Hyperparameters. To facilitate the optimization, the hyperparameter $\beta$ in
the DRIBO loss Algorithm 1 is slowly increased during training. $\beta$ 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 [8]. 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 [12] 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. $\beta$ 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
$\displaystyle=0.8\cdot\text{compute\\_kl}(\text{stop\\_grad}(\text{approx\\_posterior})+\text{prior})$
$\displaystyle+0.2\cdot\text{compute\\_kl}(\text{approx\\_posterior}+\text{stop\\_grad}(\text{prior}))$
return DRIBO Loss + $\beta\cdot$ kl_balancing

Table 2: Hyperparameters used for DMC experiments.

Hyperparameters Value Observation size $(100\times 100)$ Cropped size
$(84\times 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; $\alpha$ learning rate: 1e-4.
Batch size $8\times 32$, where $T=32$ Target update $\tau$ Q network: 0.01
encoder: 0.05 Target update freq 2 Discount $\gamma$ .99 Initial temperature
0.1 Total timesteps 88e4 $\beta$ scheduler start episode 10 $\beta$ scheduler
end episode 60

###  D.4 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 [31] 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 $\beta$ scheduler as the DMC experiments. The
starting $\beta$ value is $1e-8$ and the final $\beta$ value is $1e-3$. We
show other hyperparameters for Procgen environments in Table 5.

Table 5: Hyperparameters used for Procgen experiments.

Hyperparameters Value Observation size $(64\times 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;
$\alpha$ learning rate: 1e-4. Batch size $8\times 256$, where $T=256$ Entropy
bonus 0.01 PPO clip range 0.2 Discount $\gamma$ .99 GAE parameter$\lambda$
0.95 Reward normalization yes Num. of workers 1 Num. of environments per
worker 64 Total timesteps 25M $\beta$ scheduler start episode 10 $\beta$
scheduler end episode 110

Discussion. Here, we extend the discussion on why our method underperforms on
some environments, whose screenshots are shown in Figure 14.

Our approach, DRIBO, only consider global MI [14] 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.

###  D.5 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 14: 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.