# INFOrmation Prioritization through
EmPOWERment in Visual Model-Based RL

Homanga Bharadhwaj
Carnegie Mellon University &Mohammad Babaeizadeh
Google Research, Brain Team &Dumitru Erhan
Google Research, Brain Team &Sergey Levine
Google Research, Brain Team
University of California Berkeley Work done during Homanga’s internship at
Google. Correspondence to<EMAIL_ADDRESS>
###### Abstract

Model-based reinforcement learning (RL) algorithms designed for handling
complex visual observations typically learn some sort of latent state
representation, either explicitly or implicitly. Standard methods of this sort
do not distinguish between functionally relevant aspects of the state and
irrelevant distractors, instead aiming to represent all available information
equally. We propose a modified objective for model-based RL that, in
combination with mutual information maximization, allows us to learn
representations and dynamics for visual model-based RL without reconstruction
in a way that explicitly prioritizes functionally relevant factors. The key
principle behind our design is to integrate a term inspired by variational
empowerment into a state-space model based on mutual information. This term
prioritizes information that is correlated with action, thus ensuring that
functionally relevant factors are captured first. Furthermore, the same
empowerment term also promotes faster exploration during the RL process,
especially for sparse-reward tasks where the reward signal is insufficient to
drive exploration in the early stages of learning. We evaluate the approach on
a suite of vision-based robot control tasks with natural video backgrounds,
and show that the proposed prioritized information objective outperforms
state-of-the-art model based RL approaches with higher sample efficiency and
episodic returns. https://sites.google.com/view/information-empowerment

##  1 Introduction

Model-based reinforcement learning (RL) provides a promising approach to
accelerating skill learning: by acquiring a predictive model that represents
how the world works, an agent can quickly derive effective strategies, either
by planning or by simulating synthetic experience under the model. However, in
complex environments with high-dimensional observations (e.g., images),
modeling the full observation space can present major challenges. While large
neural network models have made progress on this problem (Finn & Levine, 2017;
Ha & Schmidhuber, 2018; Hafner et al., 2019a; Watter et al., 2015; Babaeizadeh
et al., 2017), effective learning in visually complex environments
necessitates some mechanism for learning representations that _prioritize_
functionally relevant factors for the current task. This needs to be done
without wasting effort and capacity on _irrelevant_ distractors, and without
detailed reconstruction. Several recent works have proposed contrastive
objectives that maximize mutual information between observations and latent
states (Hjelm et al., 2018; Ma et al., 2020; Oord et al., 2018; Srinivas et
al., 2020). While such objectives avoid reconstruction, they still do not
distinguish between relevant and irrelevant factors of variation. We thus pose
the question: can we devise non-reconstructive representation learning methods
that explicitly prioritize information that is most likely to be functionally
relevant to the agent?

In this work, we derive a model-based RL algorithm from a combination of
representation learning via mutual information maximization (Poole et al.,
2019) and empowerment (Mohamed & Rezende, 2015). The latter serves to drive
both the representation and the policy toward exploring and representing
functionally relevant factors of variation. By integrating an empowerment-
based term into a mutual information framework for learning state
representations, we effectively _prioritize_ information that is most likely
to have functional relevance, which mitigates distractions due to irrelevant
factors of variation in the observations. By integrating this same term into
policy learning, we further improve exploration, particularly in the early
stages of learning in sparse-reward environments, where the reward signal
provides comparatively little guidance.

Our main contribution is InfoPower, a model-based RL algorithm for control
from image observations that integrates empowerment into a mutual information
based, non-reconstructive framework for learning state space models. Our
approach explicitly prioritizes information that is most likely to be
functionally relevant, which significantly improves performance in the
presence of time-correlated distractors (e.g., background videos), and also
accelerates exploration in environments when the reward signal is weak. We
evaluate the proposed objectives on a suite of simulated robotic control tasks
with explicit video distractors, and demonstrate up to 20% better performance
in terms of cumulative rewards at 1M environment interactions with 30% higher
sample efficiency at 100k interactions.

Figure 1: Overview of InfoPower. $I({\mathcal{O}}_{t};Z_{t})$ is the
contrastive learning objective for learning an encoder to map from image
${\mathcal{O}}$ to latent $Z$. $I(A_{t-1};Z_{t}|Z_{t-1})$ is the empowerment
objective that prioritizes encoding controllable representations in $Z$.
$-I(i_{t+1};Z_{t+1}|Z_{t},A_{t})$ helps learn a latent forward dynamics model
so that future $Z_{t+k}$ can be predicted from current $Z_{t}$.
$I(R_{t};Z_{t})$ helps learn a reward prediction model, such that the agent
can learn a plan $A_{t},..A_{t+k},..$ through latent rollouts. Together, this
combination of terms produces a latent state space model for MBRL that
captures all necessary information at convergence, while prioritizing the most
functionally relevant factors via the empowerment term.

##  2 Problem Statement and Notation

A partially observed Markov decision process (POMDP) is a tuple
$({\mathcal{S}},{\mathcal{A}},T,R,{\mathcal{O}})$ that consists of states
$s\in{\mathcal{S}}$, actions $a\in{\mathcal{A}}$, rewards $r\in R$,
observations $o\in{\mathcal{O}}$, and a state-transition distribution
$T(s^{\prime}|s,a)$. In most practical settings, the agent interacting with
the environment doesn’t have access to the actual states in ${\mathcal{S}}$,
but to some partial information in the form of observations ${\mathcal{O}}$.
The underlying state-transition distribution $T$ and reward distribution $R$
are also unknown to the agent.

In this paper, we consider the observations $o\in{\mathcal{O}}$ to be high-
dimensional images, and so the agent should learn a compact representation
space $Z$ for the latent state-space model. The problem statement is to learn
effective representations from observations ${\mathcal{O}}$ when there are
visual distractors present in the scene, and plan using the learned
representations to maximize the cumulative sum of discounted rewards,
$J=\mathbb{E}[\sum_{t}\gamma^{t-1}r_{t}]$. The value of a state $V(Z_{t})$ is
defined as the expected cumulative sum of discounted rewards starting at state
$Z_{t}$.

We use $q(\cdot)$ to denote parameterized variational approximations to
learned distributions. We denote random variables with capital letters and use
lowercase letters to denote particular realizations (e.g., $z_{t}$ denotes the
value of $Z_{t}$). Since the underlying distributions are unknown, we evaluate
all expectations through Monte-Carlo sampling with observed state-transition
tuples $(o_{t},a_{t-1},o_{t-1},z_{t},z_{t-1},r_{t})$.

##  3 Information Prioritization for the Latent State-Space Model

Our goal is to learn a latent state-space model with a representation $Z$ that
prioritizes capturing functionally relevant parts of observations
${\mathcal{O}}$, and devise a planning objective that explores with the
learned representation. To achieve this, our key insight is integration of
empowerment in the visual model-based RL pipeline. For representation learning
we maximize MI $\max_{Z}I({\mathcal{O}},Z)$ subject to a prioritization of the
empowerment objective $\max_{Z}I(A_{t-1};Z_{t}|Z_{t-1})$. For planning, we
maximize the empowerment objective along with reward-based value with respect
to the policy $\max_{A}I(A_{t-1};Z_{t}|Z_{t-1})+I(R_{t};Z_{t})$.

In the subsequent sections, we elaborate on our approach, InfoPower, and
describe lower bounds to MI that yield a tractable algorithm.

###  3.1 Learning Controllable Factors and Planning through Empowerment

Figure 2: PGM showing decomposition of state $S$ into controllable parts
$S^{+}$ (directly influenced by actions $A$), parts not influenced by actions
that still influence the reward, $\tilde{S}^{-}$, and distractors $DS^{-}$.
$\tilde{S}_{t+1}^{-}$ may be influenced by $\tilde{S}_{t}^{-}$ (arrow not
shown to reduce clutter) but not by $\tilde{S}_{t}^{+}$.

Controllable representations are features of the observation that correspond
to entities which the agent can influence through its actions. For example, in
quadrupedal locomotion, this could include the joint positions, velocities,
motor torques, and the configurations of any object in the environment that
the robot can interact with. For robotic manipulation, it could include the
joint actuators of the robot arm, and the configurations of objects in the
scene that it can interact with. Such representations are denoted by $S^{+}$
in Fig. 2, which we can formally define through conditional independence as
the smallest subspace of $S$, $S^{+}\leq S$, such that
$I(A_{t-1};S_{t}|S^{+}_{t})=0$. This conditional independence relation can be
seen in Fig. 2. We explicitly prioritize the learning of such representations
in the latent space by drawing inspiration from variational empowerment
(Mohamed & Rezende, 2015).

Initialize dataset $\mathcal{D}$ with random episodes. ;

Initialize model parameters $\phi,\chi,\psi,\eta$. ;

Initialize dual variable $\lambda$. ;

while  _not converged_ do

for  _update step $c=1..C$_ do

// Model learning

Sample data $\\{(a_{t},o_{t},r_{t})\\}_{t=k}^{k+L}\sim\mathcal{D}$. ;

Compute latents $z_{t}\sim p_{\phi}(z_{t}|z_{t-1},a_{t-1},o_{t})$. ;

Calculate ${\mathcal{L}}$ based on section 3.4. ;

$(\phi,\chi,\psi,\eta)\leftarrow(\phi,\chi,\psi,\eta)+\nabla_{\phi,\chi,\psi,\eta}{\mathcal{L}}$
;

$\lambda\leftarrow\lambda-\nabla_{\lambda}{\mathcal{L}}$;

// Behavior learning

Rollout latent plan,
$\mathcal{S}\leftarrow\mathcal{S}\cup\\{z_{t},a_{t},r_{t}\\}$;

$V(z_{t})\approx\mathbb{E}_{\pi}[\ln q_{\eta}(r_{t}|z_{t})+\ln
q_{\psi}(a_{t-1}|z_{t},z_{t-1})]$;

Update policy $\pi$ and value model;

end for

// Environment interaction

for  _time step $t=0..T-1$_ do

$z_{t}\sim p_{\phi}(z_{t}|z_{t-1},a_{t-1},o_{t})$;
$a_{t}\sim\pi(a_{t}|z_{t})$;

$r_{t},o_{t+1}\leftarrow\texttt{env.step(}a_{t}\texttt{)}.$ ;

end for

Add data
$\mathcal{D}\leftarrow\mathcal{D}\cup\\{(o_{t},a_{t},r_{t})_{t=1}^{T}\\}.$ ;

end while

Algorithm 1 Information Prioritization in Visual Model-based RL (InfoPower)

The empowerment objective can be cast as maximizing a conditional information
term
$I(A_{t-1};Z_{t}|Z_{t-1})={\mathcal{H}}(A_{t-1}|Z_{t-1})-{\mathcal{H}}(A_{t-1}|Z_{t},Z_{t-1})$.
The first term ${\mathcal{H}}(A_{t-1}|Z_{t-1})$ encourages the chosen actions
to be as diverse as possible, while the second term
$-{\mathcal{H}}(A_{t-1}|Z_{t},Z_{t-1})$ encourages the representations $Z_{t}$
and $Z_{t+1}$ to be such that the action $A_{t}$ for transition is
predictable. While prior approaches have used empowerment in the model-free
setting to learn policies by exploration through intrinsic motivation (Mohamed
& Rezende, 2015), we specifically use this objective in combination with MI
maximization for prioritizing the learning of controllable representations
from distracting images in the latent state-space model.

We include the same empowerment objective in _both_ representation learning
and policy learning. For this, we augment the maximization of the latent value
function that is standard for policy learning in visual model-based RL
(Sutton, 1991), with $\max_{A}I(A_{t-1};Z_{t}|Z_{t-1})$. This objectives
complements value based-learning and further improves exploration by seeking
controllable states. We empirically analyze the benefits of this in sections
4.3 and 4.5.

In Appendix 1 we next describe two theorems regarding learning controllable
representations. We observe that the $\max\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})$
objective for learning latent representations $Z$, when used along with the
planning objective, provably recovers controllable parts of the observation
${\mathcal{O}}$, namely $S^{+}$. This result in Theorem 1 is important because
in practice, we may not be able to represent every possible factor of
variation in a complex environment. In this situation, we would expect that
when $|Z|\ll|{\mathcal{O}}|$, learning $Z$ under the objective
$\max\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})$ would encode $S^{+}$.

We further show through Theorem 2 that the inverse information objective alone
can be used to train a latent-state space model and a policy through an
alternating optimization algorithm that converges to a local minimum of the
objective $\max\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})$ at a rate inversely
proportional to the number of iterations. In Section 4.3 we empirically show
how this objective helps achieve higher sample efficiency compared to pure
value-based policy learning.

###  3.2 Mutual Information maximization for Representation Learning

For visual model-based RL, we need to learn a representation space $Z$, such
that a forward dynamics model defining the probability of the next state in
terms of the current state and the current action can be learned. The
objective for this is $\sum_{t}-I(i_{t};Z_{t}|Z_{t-1},A_{t-1})$. Here, $i_{t}$
denotes the dataset indices that determine the observations
$p(o_{t}|i_{t})=\delta(o_{t}-o_{t^{\prime}})$. In addition to the forward
dynamics model, we need to learn a reward predictor by maximizing
$\sum_{t}I(R_{t};Z_{t})$, such that the agent can plan ahead in the future by
rolling forward latent states, without having to execute actions and observe
rewards in the real environment.

Finally, we need to learn an encoder for encoding observations ${\mathcal{O}}$
to latents $Z$. Most successful prior works have used reconstruction-loss as a
natural objective for learning this encoder (Babaeizadeh et al., 2017; Hafner
et al., 2019b; a). A reconstruction-loss can be motivated by considering the
objective $I({\mathcal{O}},Z)$ and computing its BA lower bound (Agakov,
2004). $I(o_{t};z_{t})\geq\mathbb{E}_{p(o_{t},z_{t})}[\log
q_{\phi^{\prime}}(o_{t}|z_{t})]+{\mathcal{H}}(p(o_{t}))$. The first term here
is the reconstruction objective, with $q_{\phi^{\prime}}(o_{t}|z_{t})$ being
the decoder, and the second term can be ignored as it doesn’t depend on $Z$.
However, this reconstruction objective explicitly encourages encoding the
information from every pixel in the latent space (such that reconstructing the
image is possible) and hence is prone to not ignoring distractors.

In contrast, if we consider other lower bounds to $I({\mathcal{O}},Z)$, we can
obtain tractable objectives that do not involve reconstructing high-
dimensional images. We can obtain an NCE-based lower bound (Hjelm et al.,
2018): $I(o_{t};z_{t})\geq\mathbb{E}_{q_{\phi}(z_{t}|o_{t})p(o_{t})}[\log
f_{\theta}(z_{t},o_{t})-\log\sum_{t^{\prime}\neq
t}f_{\theta}(z_{t},o_{t^{\prime}})],$ where $q_{\phi}(z_{t}|o_{t})$ is the
learned encoder, $o_{t}$ is the observation at timestep $t$ (positive sample),
and all observations in the replay buffer $o_{t^{\prime}}$ are negative
samples. $f_{\theta}(z_{t},o_{t})=\exp{(\hat{z}_{t}^{T}W_{\theta}z_{t})}$,
where $\hat{z}_{t}^{T}$ is an embedding for $o_{t}$. The lower-bound is a form
of contrastive learning as it maximizes compatibility of $z_{t}$ with the
corresponding observation $o_{t}$ while minimizing compatibility with all
other observations across time and batch.

Although prior work has explored NCE-based bounds for contrastive learning in
RL (Srinivas et al., 2020), to the best of our knowledge, prior work has not
used this in conjunction with empowerment for prioritizing information in
visual model-based RL. Similarly, the Nguyen-Wainwright-Jordan (NWJ) bound
(Nguyen et al., 2010), which to the best our knowledge has not been used by
prior works in visual model-based RL, can be obtained as,

$I(o_{t};z_{t})\geq\mathbb{E}_{q_{\phi}(z_{t}|o_{t})p(o_{t})}[f_{\theta}(z_{t},o_{t})]-e^{-1}\mathbb{E}_{q_{\phi}(z_{t}|o_{t})p(o_{t})}e^{f_{\theta}(z_{t},o_{t})},\vspace*{0.1cm}$

where $f_{\theta}$ is a _critic_. There exists an optimal critic function for
which the bound is tightest and equality holds.

We refer to the InfoNCE and NWJ lower bound based objectives as contrastive
learning, in order to distinguish them from a reconstruction-loss based
objective, though both are bounds on mutual information. We denote a lower
bound to MI by $I(o_{t},z_{t})$. We empirically find the NWJ-bound to perform
slightly better than the NCE-bound for our approach, explained in section 4.5.

###  3.3 Overall Objective

We now motivate the overall objective, which consists of maximizing mutual
information while prioritizing the learning of controllable representations
through empowerment in a latent state-space model. Based on the discussions in
Sections 3.1 and 3.2, we define the overall objective for representation
learning as

$\max_{Z_{0:H-1}}\sum_{t=0}^{H-1}I({\mathcal{O}}_{t};Z_{t})\;\;\text{s.t.}\;\;\sum_{t=0}^{H-1}\overbrace{\left(-I(i_{t};Z_{t}|Z_{t-1},A_{t-1})+I(A_{t-1};Z_{t}|Z_{t-1})+I(R_{t};Z_{t})\right)}^{{\mathcal{C}}_{t}}\geq
c_{0}.$

The objective is to maximize a MI term $I({\mathcal{O}}_{t};Z_{t})$ through
contrastive learning such that a constraint on ${\mathcal{C}}_{t}$ holds for
prioritizing the encoding of forward-predictive, reward-predictive and
controllable representations. We define the overall planning objective as

$\max_{A_{0:H-1}}\sum_{t=0}^{H-1}I(A_{t-1};Z_{t}|Z_{t-1})+V(Z_{t})\;\;\;;A_{t}=\pi(Z_{t})\;\;\;;V(Z_{t})\approx\sum_{t}R_{t}.$

The planning objective is to learn a policy as a function of the latent state
$Z$ such that the empowerment term and a reward-based value term are maximized
over the horizon $H$.

We can perform the constrained optimization for representation learning
through the method of Lagrange Multipliers, by the primal and dual updates
shown in Section A.2. In order to analyze this objective, let
$|{\mathcal{O}}|=n$ and $|Z|=d$. Since, ${\mathcal{O}}$ corresponds to images
and $Z$ is a bottlenecked latent representation, $d\ll n$.

$I({\mathcal{O}},Z)$ is maximized when $Z$ contains all the information
present in ${\mathcal{O}}$ such that $Z$ is a sufficient statistic of
${\mathcal{O}}$. However, in practice, this is not possible because
$|Z|\ll|{\mathcal{O}}|$. When $c_{0}$ is sufficiently large, and the
constraint $\sum_{t}{\mathcal{C}}_{t}\geq c_{0}$ is satisfied,
$Z=[S^{+},\tilde{S}^{-}]$. Hence, the objective
$\max\sum_{t}I({\mathcal{O}}_{t},Z_{t})\;\;\text{s.t.}\;\;\sum_{t}{\mathcal{C}}_{t}\geq
c_{0}$ cannot encode anything else in $Z$, in particular it cannot encode
distractors $DS^{-}$.

To understand the importance of prioritization through ${\mathcal{C}}_{t}\geq
c_{0}$, consider $\max\sum_{t}I({\mathcal{O}}_{t},Z_{t})$ without the
constraint. This objective would try to make $Z$ a sufficient statistic of
${\mathcal{O}}$, but since $|Z|\ll|{\mathcal{O}}|$, there are no guarantees
about which parts of ${\mathcal{O}}$ are getting encoded in $Z$. This is
because both distractors $S^{-}$ and non-distractors $S^{+},D\tilde{S}^{-}$
are equally important with respect to $I({\mathcal{O}},Z)$. Hence, the
constraint helps in prioritizing the type of information to be encoded in $Z$.

###  3.4 Practical Algorithm and Implementation Details

To arrive at a practical algorithm, we optimize the overall objective in
section 3.3 through lower bounds on each of the MI terms. For
$I({\mathcal{O}},Z)$ we consider two variants, corresponding to the NCE and
NWJ lower bounds described in Section 3.2. We can obtain a variational lower
bound on each of the terms in ${\mathcal{C}}_{t}$ as follows, with detailed
derivations in Appendix A.3:

$\displaystyle-I(i_{t};z_{t}|a_{t-1})\geq-\sum_{t}\mathbb{E}[D_{\text{KL}}(p(z_{t}|z_{t-1},a_{t-1},o_{t})||q_{\chi}(z_{t}|z_{t-1},a_{t-1}))]$
$\displaystyle I(r_{t};z_{t})\geq\mathbb{E}_{p(r_{t}|o_{t})}[\log
q_{\eta}(r_{t}|z_{t})]+{\mathcal{H}}(p(r_{t}))$ $\displaystyle
I(a_{t-1};z_{t}|z_{t-1})\geq\mathbb{E}_{p(o_{t}|z_{t-1},a_{t-1})q_{\phi}(z_{t}|o_{t})}[\log
q_{\psi}(a_{t-1}|z_{t},z_{t-1})]+\mathbb{E}[{\mathcal{H}}(\pi(a_{t-1}|z_{t-1}))]$

Here, $q_{\phi}(z_{t}|o_{t})$ is the observation encoder,
$q_{\psi}(a_{t-1}|z_{t},z_{t-1})$ is the inverse dynamics model,
$q_{\eta}(r_{t}|z_{t})$ is the reward decoder, and
$q_{\chi}(z_{t}|z_{t-1},a_{t-1}))$ is the forward dynamics model. The inverse
model helps in learning representations such that the chosen action is
predictable from successive latent states. The forward dynamics model helps in
predicting the next latent state given the current latent state and the
current action, without having the next observation. The reward decoder
predicts the reward (a scalar) at a time-step given the corresponding latent
state. We use dense networks for all the models, with details in Appendix A.7.
We denote by ${\mathcal{C}}_{t}$, the lower bound to ${\mathcal{C}}_{t}$ based
on the sum of the terms above. We construct a Lagrangian ${\mathcal{L}}=$
$\sum_{t}$ $I(o_{t};z_{t})$ $+$ $\lambda$ (${\mathcal{C}}_{t}$ \- $c_{0}$) and
optimize it by primal-dual gradient descent. An outline of this is shown in
Algorithm 1.

Planning Objective. For planning to choose actions at every time-step, we
learn a policy $\pi(a|z)$ through value estimates of task reward and the
empowerment objective. We learn value estimates with
$V(z_{t})\approx\mathbb{E}_{\pi}[\ln q_{\eta}(r_{t}|z_{t})+\ln
q_{\psi}(a_{t-1}|z_{t},z_{t-1})]$. We estimate $V(z_{t})$ similar to Equation
6 of (Hafner et al., 2019a). The empowerment term
$q_{\psi}(a_{t-1}|z_{t},z_{t-1})$ in policy learning incentivizes choosing
actions $a_{t-1}$ such that they can be predicted from consecutive latent
states $z_{t-1},z_{t}$. This biases the policy to explore controllable regions
of the state-space. The policy is trained to maximize the estimate of the
value, while the value model is trained to fit the estimate of the value that
changes as the policy is updated.

Finally, we note that the difference in value function of the underlying MDP
$Q^{\pi}(o,a)$ and the latent MDP $\hat{Q}^{\pi}(z,a)$, where $z\sim
q_{\phi}(z|o)$ is bounded, under some regularity assumptions. We provide this
result in Theorem 3 of the Appendix Section A.4. The overall procedure for
model learning, planning, and interacting with the environment is outlined in
Algorithm 1.

##  4 Experiments

Through our experiments, we aim to understand the following research
questions:

  1. 1.

How does InfoPower compare with the baselines in terms of episodic returns in
environments with explicit background video distractors?

  2. 2.

How sample efficient is InfoPower when the reward signal is weak ($<100k$ env
steps when the learned policy typically doesn’t achieve very high rewards)?

  3. 3.

How does InfoPower compare with baselines in terms of behavioral similarity of
latent states?

Please refer to the website for a summary and qualitative visualization
results https://sites.google.com/view/information-empowerment

###  4.1 Setup

We perform experiments with modified DeepMind Control Suite environments
(Tassa et al., 2018), with natural video distractors from ILSVRC dataset
(Russakovsky et al., 2015) in the background. The agent receives only image
observations at each time-step and does not receive the ground-truth simulator
state. This is a very challenging setting, because the agents must learn to
ignore the distractors and abstract out representations necessary for control.
While natural videos that are unrelated to the task might be easy to ignore,
realistic scenes might have other elements that resemble the controllable
elements, but are not actually controllable (e.g., other cars in a driving
scenario). To emulate this, we also add distractors that resemble other
potentially controllable robots, as shown for example in Fig. 3 (1st and 6th),
but are not actually influenced by actions.

Figure 3: Some of the natural video background distractors in our experiments.
The videos change every 50 time-steps. Some backgrounds (for e.g. the top left
and the bottom right) have more complex distractors in the form of agent-
behind-agent i.e. the background has pre-recorded motion of a similar agent
that is being controlled.

In addition to this setting, we also perform experiments with gray-scale
distractors based on videos from the Kinetics dataset (Kay et al., 2017). This
setting is adapted exactly based on prior works (Fu et al., 2021; Zhang et
al., 2020) and we compare directly to the published results.

We compare InfoPower with state-of-the art baselines that also learn world
models for control: Dreamer (Hafner et al., 2019a), C-Dreamer that is a
contrastive version of Dreamer similar to Ma et al. (2020), TIA (Fu et al.,
2021) that learns explicit reconstructions for both the distracting background
and agent separately, DBC (Zhang et al., 2020), and DeepMDP (Gelada et al.,
2019). We also compare variants of InfoPower with different MI bounds (NCE and
NWJ), and ablations of it that remove the empowerment objective from policy
learning.

###  4.2 A behavioral similarity metric based on graph kernels

In order to measure how good the learned latent representations are at
capturing the functionally relevant elements in the scene, we introduce a
behavioral similarity metric with details in A.5. Given the sets
${\mathcal{S}}_{z}=\\{z^{i}\\}_{i=1}^{n}$ and
${\mathcal{S}}_{z_{\text{gt}}}=\\{z^{i}_{\text{gt}}\\}_{i=1}^{n}$, we
construct complete graphs by connecting every vertex with every other vertex
through an edge. The weight of an edge is the Euclidean distance between the
respective pairs of vertices. The label of each vertex is an unique integer
and corresponding vertices in the two graphs are labelled similarly.

Let the resulting graphs be denoted as $G_{z}=(V_{z},E_{z})$ and
$G_{z_{\text{gt}}}=(V_{z_{\text{gt}}},E_{z_{\text{gt}}})$ respectively. Note
that both these graphs are shortest path graphs, by construction. Let,
$e_{i}=\\{u_{i},v_{i}\\}$ and $e_{j}=\\{u_{j},v_{j}\\}$. We define a kernel
$\hat{k}$ that measures similarity between the edges $e_{i},e_{j}$ and the
labels on respective vertices. $\hat{k}(e_{i},e_{j})=$

$k_{v}(l(u_{i}),l(u_{j}))k_{e}(l(e_{i}),l(e_{j}))k_{v}(l(v_{i}),l(v_{j}))+k_{v}(l(u_{i}),l(v_{j}))k_{e}(l(e_{i}),l(e_{j}))k_{v}(l(v_{i}),l(u_{j}))$

Here, the function $l(\cdot)$ denotes the labels of vertices and the weights
of edges. $k_{v}$ is a Dirac delta kernel i.e. $k_{e}(x,y)=1-\delta(x,y)$
(Note that $\delta(x,y)=1$ iff $x=y$, else $\delta(x,y)=0$) and $k_{e}$ is a
Brownian ridge kernel, $k_{e}(x,y)=\frac{1}{c}\max(0,c-|x-y|)$, where $c$ is a
large number. We define the shortest path kernel to measure similarity between
the two graphs, as,
$k(G_{z},G_{z_{\text{gt}}})=\frac{1}{|E_{z}|}\sum_{e_{i}\in
E_{z}}\sum_{e_{j}\in E_{z_{\text{gt}}}}\hat{k}(e_{i},e_{j})$. The value of
$k(G_{z},G_{z_{\text{gt}}})$ is low when a large number of pairs of
corresponding vertices in both the graphs have the same edge length. We expect
methods that recover latent state representations which better reflect the
true underlying simulator state (i.e., the positions of the joints) to have
higher values according to this metric.

###  4.3 Sample efficiency analysis

In Fig. 4, we compare InfoPower and the baselines in terms of episodic
returns. This version of InfoPower corresponds to an NWJ bound on MI which we
find works slightly better than the NCE bound variant analyzed in section 4.5.
It is evident that InfoPower achieves higher returns before 1M steps of
training quickly compared to the baselines, indicating higher sample
efficiency. This suggests the effectiveness of the empowerment model in
helping capture controllable representations early on during training, when
the agent doesn’t take on actions that yield very high rewards.

Figure 4: Evaluation of InfoPower and baselines in a suite of DeepMind Control
tasks with natural video distractors in the background. The x-axis denotes the
number of environment interactions and the y-axis shows the episodic returns.
The S.D is over 4 random seeds. Higher is better.

###  4.4 Behavioral similarity of latent states

In this section, we analyze how similar are the learned latent representations
with respect to the underlying simulator states. The intuition for this
comparison is that the proprioceptive features in the simulator state abstract
out distractors in the image, and so we want the latent states to be
behaviorally similar to the simulator states.

Quantitative results with the defined metric. In Table 1, we show results for
behavioral similarity of latent states (Sim), based on the metric in section
4.2. We see that the value of Sim for InfoPower is around $20\%$ higher than
the most competitive baseline, indicating high behavioral similarity of the
latent states with respect to the corresponding ground-truth simulator states.

Qualitative visualizations with t-sne. Fig. 5 shows a t-SNE plot of latent
states $z\sim q_{\phi}(z|o)$ for InfoPower and the baselines with
visualizations of 3 nearest neighbors for two randomly chosen latent states.
We see that the state of the agent is similar is each group for InfoPower,
although the background scenes are significantly different. However, for the
baselines, the nearest neighbor states are significantly different in terms of
the pose of the agent, indicating that the latent representations encode
significant background information.

###  4.5 Ablation Studies

In Fig. 6, we compare different ablations of InfoPower. Keeping everything
else the same, and changing only the MI lower bound to NCE, we see that the
performance is almost similar or slightly worse. However, when we remove the
empowerment objective from policy optimization (the versions with ‘-Policy’ in
the plot), we see that performance drops. The drop is significant in the
regions $<200k$ environment interactions, particularly in the sparse reward
environments - cartpole balance and ball in a cup, indicating the necessity of
the empowerment objective in exploration for learning controllable
representations, when the reward signal is weak.

Figure 5: t-SNE plot of latent states $z\sim q_{\phi}(z|o)$ with
visualizations of three nearest neighbors for two randomly sampled points (in
red frame). We see that the state of the agent is similar is each set for
InfoPower, whereas for Dreamer, and the most competitive baseline C-Dreamer,
the nearest neighbor frames have significantly different agent configurations.
Figure 6: Evaluation of InfoPower and ablated variants in a suite of DeepMind
Control tasks with natural video distractors in the background. The x-axis
denotes the number of environment interactions and the y-axis shows the
episodic returns. InfoPower-NWJ and InfoPower-NCE are full versions of our
method differing only in the MI lower bound. The versions with - Policy do not
include the empowerment objective in policy learning, but only use it from
representation learning. The S.D is over 4 random seeds. Higher is better.

##  5 Related Works

Visual model-based RL. Recent developments in video prediction and contrastive
learning have enabled learning of world-models from images (Watter et al.,
2015; Babaeizadeh et al., 2017; Finn & Levine, 2017; Hafner et al., 2019a; Ha
& Schmidhuber, 2018; Hafner et al., 2019b; Xie et al., 2020). All of these
approaches learn latent representations through reconstruction objectives that
are amenable for planning. Other approaches have used similar reconstruction
based objectives for control, but not for MBRL (Lee et al., 2019; Gregor et
al., 2019).

MI for representation learning. Mutual Information measures the dependence
between two random variables. The task of learning latent representations $Z$
from images ${\mathcal{O}}$ for downstream applications, has been very
successful with MI objectives of the form
$\max_{f_{1},f_{2}}I(f_{1}({\mathcal{O}}),f_{2}(Z))$ (Hjelm & Bachman, 2020;
Tian et al., 2020; Oord et al., 2018; Tschannen et al., 2019; Nachum & Yang,
2021). Since calculating MI exactly is intractable optimizing MI based
objectives, it is important to construct appropriate MI estimators that lower-
bound the true MI objective (Hjelm et al., 2018; Nguyen et al., 2010; Belghazi
et al., 2018; Agakov, 2004). The choice of the estimator is crucial, as shown
by recent works (Poole et al., 2019), and different estimators yield very
different behaviors of the algorithm. We incorporate MI maximization through
the NCE (Hjelm et al., 2018) and NWJ (Nguyen et al., 2010) lower bounds, such
that typical reconstruction objectives for representation learning which do
not perform well with visual distractors, can be avoided.

Inverse models and empowerment. Prior approaches have used inverse dynamics
models for regularization in representation learning (Agrawal et al., 2016;
Zhang et al., 2018) and as bonuses for improving policy gradient updates in RL
(Shelhamer et al., 2016; Pathak et al., 2017). The importance of information
theoretic approaches for learning representations that maximize predictive
power has been discussed in prior work (Still, 2009) and more recently in Lee
et al. (2020). Empowerment (Mohamed & Rezende, 2015) has been used as
exploration bonuses for policy learning (Leibfried et al., 2019; Klyubin et
al., 2008), and for learning skills in RL (Gregor et al., 2016; Sharma et al.,
2019; Eysenbach et al., 2018). In contrast to prior work, we incorporate
empowerment both for state space representation learning and policy learning,
in a visual model-based RL framework, with the aim of prioritizing the most
functionally relevant information in the scene.

RL with environment distractors. Some recent RL frameworks have studied the
problem of abstracting out only the task relevant information from the
environment when there are explicit distractors (Hansen & Wang, 2020; Zhang et
al., 2020; Fu et al., 2021; Ma et al., 2020). Zhang et al. (2020) constrain
the latent states by enforcing a bisimulation metric, without a reconstruction
objective. Our approach is different from this primarily because of the
empowerment objective that provides useful signal for dealiasing controllable
vs. uncontrollable representations independent of how strong is the observed
reward signal. Fu et al. (2021) model both the relevant and irrelevant aspects
of the environment separately, and differ from our approach that prioritizes
learning only the relevant aspects. Shu et al. (2020) used contrastive
representations for control. More recently, Nguyen et al. (2021) used temporal
predictive coding, and Ma et al. (2020) used InfoNCE based contrastive loss
for learning representations while ignoring distractors in visual MBRL.
Incorporating data augmentations for improving robustness with respect to
environment variations is an orthogonal line of work (Laskin et al., 2020;
Hansen & Wang, 2020; Raileanu et al., 2020; Srinivas et al., 2020; Kostrikov
et al., 2020), complementary to our approach.

##  6 Conclusion

In this paper we derived an approach for visual model-based RL such that an
agent can learn a latent state-space model and a policy by explicitly
prioritizing the encoding of functionally relevant factors. Our prioritized
information objective integrates a term inspired by variational empowerment
into a non-reconstructive objective for learning state space models. We
evaluate our approach on a suite of vision-based robot control tasks with two
sets of challenging video distractor backgrounds. In comparison to state-of-
the-art visual model-based RL methods, we observe higher sample efficiency and
episodic returns across different environments and distractor settings.

## References

  * Agakov (2004) David Barber Felix Agakov.  The im algorithm: a variational approach to information maximization.  _Advances in neural information processing systems_ , 16(320):201, 2004.
  * Agrawal et al. (2016) Pulkit Agrawal, Ashvin Nair, Pieter Abbeel, Jitendra Malik, and Sergey Levine.  Learning to poke by poking: Experiential learning of intuitive physics.  _arXiv preprint arXiv:1606.07419_ , 2016.
  * Babaeizadeh et al. (2017) Mohammad Babaeizadeh, Chelsea Finn, Dumitru Erhan, Roy H Campbell, and Sergey Levine.  Stochastic variational video prediction.  _arXiv preprint arXiv:1710.11252_ , 2017.
  * 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_ , pp. 531–540. PMLR, 2018.
  * Blahut (1972) Richard Blahut.  Computation of channel capacity and rate-distortion functions.  _IEEE transactions on Information Theory_ , 18(4):460–473, 1972.
  * Cover & Thomas (2006) Thomas M. Cover and Joy A. Thomas.  _Elements of Information Theory 2nd Edition (Wiley Series in Telecommunications and Signal Processing)_.  Wiley-Interscience, July 2006.  ISBN 0471241954.
  * Eysenbach et al. (2018) Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine.  Diversity is all you need: Learning skills without a reward function.  _arXiv preprint arXiv:1802.06070_ , 2018.
  * Finn & Levine (2017) Chelsea Finn and Sergey Levine.  Deep visual foresight for planning robot motion.  In _2017 IEEE International Conference on Robotics and Automation (ICRA)_ , pp. 2786–2793. IEEE, 2017.
  * Fu et al. (2021) Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola.  Learning task informed abstractions.  In _International Conference on Machine Learning_ , pp. 3480–3491. PMLR, 2021.
  * Gelada et al. (2019) Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare.  Deepmdp: Learning continuous latent space models for representation learning.  In _International Conference on Machine Learning_ , pp. 2170–2179. PMLR, 2019.
  * Gregor et al. (2016) Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra.  Variational intrinsic control.  _arXiv preprint arXiv:1611.07507_ , 2016.
  * Gregor et al. (2019) Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord.  Shaping belief states with generative environment models for rl.  _arXiv preprint arXiv:1906.09237_ , 2019.
  * Ha & Schmidhuber (2018) David Ha and Jürgen Schmidhuber.  World models.  _arXiv preprint arXiv:1803.10122_ , 2018.
  * Hafner et al. (2019a) Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  _arXiv preprint arXiv:1912.01603_ , 2019a.
  * Hafner et al. (2019b) Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson.  Learning latent dynamics for planning from pixels.  In _International Conference on Machine Learning_ , pp. 2555–2565. PMLR, 2019b.
  * Hansen & Wang (2020) Nicklas Hansen and Xiaolong Wang.  Generalization in reinforcement learning by soft data augmentation.  _arXiv preprint arXiv:2011.13389_ , 2020.
  * Hjelm & Bachman (2020) R Devon Hjelm and Philip Bachman.  Representation learning with video deep infomax.  _arXiv preprint arXiv:2007.13278_ , 2020.
  * Hjelm et al. (2018) 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.  _arXiv preprint arXiv:1808.06670_ , 2018.
  * 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.
  * Klyubin et al. (2008) Alexander S Klyubin, Daniel Polani, and Chrystopher L Nehaniv.  Keep your options open: an information-based driving principle for sensorimotor systems.  _PloS one_ , 3(12):e4018, 2008.
  * Kostrikov et al. (2020) Ilya Kostrikov, Denis Yarats, and Rob Fergus.  Image augmentation is all you need: Regularizing deep reinforcement learning from pixels.  _arXiv preprint arXiv:2004.13649_ , 2020.
  * Laskin et al. (2020) Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas.  Reinforcement learning with augmented data.  _arXiv preprint arXiv:2004.14990_ , 2020.
  * Lee et al. (2019) Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine.  Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model.  _arXiv preprint arXiv:1907.00953_ , 2019.
  * Lee et al. (2020) Kuang-Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama.  Predictive information accelerates learning in rl.  _arXiv preprint arXiv:2007.12401_ , 2020.
  * Leibfried et al. (2019) Felix Leibfried, Sergio Pascual-Diaz, and Jordi Grau-Moya.  A unified bellman optimality principle combining reward maximization and empowerment.  _arXiv preprint arXiv:1907.12392_ , 2019.
  * Ma et al. (2020) Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee.  Contrastive variational model-based reinforcement learning for complex observations.  _arXiv e-prints_ , pp. arXiv–2008, 2020.
  * Mohamed & Rezende (2015) Shakir Mohamed and Danilo Jimenez Rezende.  Variational information maximisation for intrinsically motivated reinforcement learning.  _arXiv preprint arXiv:1509.08731_ , 2015.
  * Nachum & Yang (2021) Ofir Nachum and Mengjiao Yang.  Provable representation learning for imitation with contrastive fourier features.  _arXiv preprint arXiv:2105.12272_ , 2021.
  * Nakagawa et al. (2021) Kenji Nakagawa, Yoshinori Takei, Shin-ichiro Hara, and Kohei Watabe.  Analysis of the convergence speed of the arimoto-blahut algorithm by the second-order recurrence formula.  _IEEE Transactions on Information Theory_ , 2021.
  * Nguyen et al. (2021) Tung Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon.  Temporal predictive coding for model-based planning in latent space.  _arXiv preprint arXiv:2106.07156_ , 2021.
  * Nguyen et al. (2010) XuanLong Nguyen, Martin J Wainwright, and Michael I Jordan.  Estimating divergence functionals and the likelihood ratio by convex risk minimization.  _IEEE Transactions on Information Theory_ , 56(11):5847–5861, 2010.
  * 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.
  * Pathak et al. (2017) Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  In _International conference on machine learning_ , pp. 2778–2787. PMLR, 2017.
  * 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_ , pp. 5171–5180. PMLR, 2019.
  * Raileanu et al. (2020) Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus.  Automatic data augmentation for generalization in reinforcement learning.  2020\.
  * Rakelly et al. (2021) Kate Rakelly, Abhishek Gupta, Carlos Florensa, and Sergey Levine.  Which mutual-information representation learning objectives are sufficient for control?  _arXiv preprint arXiv:2106.07278_ , 2021.
  * Russakovsky et al. (2015) Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  _International journal of computer vision_ , 115(3):211–252, 2015.
  * Sharma et al. (2019) Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman.  Dynamics-aware unsupervised discovery of skills.  _arXiv preprint arXiv:1907.01657_ , 2019.
  * Shelhamer et al. (2016) Evan Shelhamer, Parsa Mahmoudieh, Max Argus, and Trevor Darrell.  Loss is its own reward: Self-supervision for reinforcement learning.  _arXiv preprint arXiv:1612.07307_ , 2016.
  * Shu et al. (2020) Rui Shu, Tung Nguyen, Yinlam Chow, Tuan Pham, Khoat Than, Mohammad Ghavamzadeh, Stefano Ermon, and Hung Bui.  Predictive coding for locally-linear control.  In _International Conference on Machine Learning_ , pp. 8862–8871. PMLR, 2020.
  * Srinivas et al. (2020) Aravind Srinivas, Michael Laskin, and Pieter Abbeel.  Curl: Contrastive unsupervised representations for reinforcement learning.  _arXiv preprint arXiv:2004.04136_ , 2020.
  * Still (2009) Susanne Still.  Information-theoretic approach to interactive learning.  _EPL (Europhysics Letters)_ , 85(2):28005, 2009\.
  * Sutton (1991) Richard S Sutton.  Dyna, an integrated architecture for learning, planning, and reacting.  _ACM Sigart Bulletin_ , 2(4):160–163, 1991.
  * 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.
  * Tian et al. (2020) Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive multiview coding.  In _Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part XI 16_ , pp. 776–794. Springer, 2020.
  * Tschannen et al. (2019) Michael Tschannen, Josip Djolonga, Paul K Rubenstein, Sylvain Gelly, and Mario Lucic.  On mutual information maximization for representation learning.  _arXiv preprint arXiv:1907.13625_ , 2019.
  * 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.  _arXiv preprint arXiv:1506.07365_ , 2015.
  * Xie et al. (2020) Kevin Xie, Homanga Bharadhwaj, Danijar Hafner, Animesh Garg, and Florian Shkurti.  Latent skill planning for exploration and transfer.  In _International Conference on Learning Representations_ , 2020.
  * Zhang et al. (2018) Amy Zhang, Harsh Satija, and Joelle Pineau.  Decoupling dynamics and reward for transfer learning.  _arXiv preprint arXiv:1804.10689_ , 2018.
  * Zhang et al. (2020) Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine.  Learning invariant representations for reinforcement learning without reconstruction.  _arXiv preprint arXiv:2006.10742_ , 2020.

##  Appendix A Appendix

###  A.1 Theoretical results on controllable representations

Let $\Phi(\cdot)$ denote the encoder such that $Z=\Phi({\mathcal{O}})$.
${\mathcal{O}}$ is the observation seen by the agent and $S$ is the underlying
state. $S^{+}$ is defined as that part of underlying state $S$ which is
directly influenced by actions $A$ i.e. $S^{+}\leq S$ s.t.
$I(A_{t-1};S_{t}|S^{+}_{t})=0$.

We next describe two theorems regarding learning controllable representations,
with proofs in the Appendix. We observe that the
$\max\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})$ objective alone for learning latent
representations $Z$, along with the planning objective provably recovers
controllable parts of the observation ${\mathcal{O}}$, namely $S^{+}$.

######  Theorem 1.

The objective $\max\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})$ provably recovers
controllable parts $S^{+}$ of the observation ${\mathcal{O}}$. $S^{+}$ is
defined as that part of underlying state $S$ which is directly influenced by
actions $A$ i.e. $S^{+}\subset S$ s.t. $I(S_{t};A_{t-1}|S^{+}_{t})=0$.

###### Proof.

Based on Fig. 2 and Fig. 7, we have the following

$\displaystyle\max I(A_{t-1};Z_{t})$ $\displaystyle\leq\max
I(A_{t-1};{\mathcal{O}}_{t})\;\;\;\;\texttt{Data-Processing Inequality}$
$\displaystyle\leq I(A_{t-1};S_{t})\;\;\;\;\texttt{Data-Processing
Inequality}$ $\displaystyle\leq I(A_{t-1};S_{t},S^{+}_{t})$ $\displaystyle\leq
I(A_{t-1};S_{t}|S^{+}_{t})+I(A_{t-1};S^{+}_{t})\;\;\;\;\texttt{Chain Rule of
MI}$ $\displaystyle\leq I(A_{t-1};S^{+}_{t})\;\;\;\;\texttt{Conditional
independence}$

So, the maximum above can be obtained when the encoder $\Phi$ is an identity
function over $S^{+}$ and a zero function over $[\tilde{S}^{-},S^{-}]$. So,
$Z_{t}=S^{+}_{t}$. Hence, given $Z_{t-1}$, $I(A_{t-1};Z_{t}|Z_{t-1})$ is
maximized when $\Phi$ is an identity function over $S^{+}$. So,
$\max\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})$ learns the encoder $\Phi$ such that
controllable representations are recovered.

∎

This result is important because in practice, we may not be able to represent
every possible factor of variation in a complex environment. In this
situation, we would expect that when $|Z|\ll|{\mathcal{O}}|$, learning $Z$
under the objective $\max\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})$ would encode
$S^{+}$. We next show that the inverse information objective alone can be used
to train a latent-state space model and a policy through an alternating
optimization algorithm that converges to a local minimum of the objective
$\max\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})$ at a rate inversely proportional to the
number of iterations.

######  Theorem 2.

$\max_{\pi,\psi}\sum_{t}I(A_{t-1};Z_{t}|Z_{t-1})=\sum_{t}\mathbb{E}_{\pi(a_{t-1}|z_{t-1})p(z_{t}|z_{t-1},a_{t-1},o_{t})}\log\frac{q_{\psi}(a_{t-1}|z_{t},z_{t-1})}{\pi(a_{t-1}|z_{t-1})}$
can be optimized through an alternating optimization scheme that has a
convergence rate of ${\mathcal{O}}(1/N)$ to a local minima of the objective,
where $N$ is the number of iterations.

###### Proof.

Let $z\sim q_{\phi}(z|o)$ denote a latent state sampled from the encoder
distribution. Learning a world model such that reward prediction and inverse
information are maximized can be summarized as:

$\displaystyle\max_{\phi,\psi,\eta}\sum_{t}\mathbb{E}_{\pi(a_{t}|z_{t}),q_{\phi}(z_{t}|o_{t})}\left[\mathbb{E}_{p(r_{t}|z_{t},a_{t})}\log
q_{\eta}(r_{t}|z_{t},a_{t})+\mathbb{E}_{p(z_{t+1}|z_{t},a_{t})}\log\frac{q_{\psi}(a_{t}|z_{t+1},z_{t})}{\pi(a_{t}|z_{t})}\right]$

The policy $\pi(a_{t}|z_{t})$ is learned by

$\displaystyle\max_{\pi}\sum_{t}\mathbb{E}_{\pi(a_{t}|z_{t}),q_{\phi}(z_{t}|o_{t})}\left[\mathbb{E}_{p(r_{t}|z_{t},a_{t})}\log
q_{\eta}(r_{t}|z_{t},a_{t})+\mathbb{E}_{p(z_{t+1}|z_{t},a_{t})}\log\frac{q_{\psi}(a_{t}|z_{t+1},z_{t})}{\pi(a_{t}|z_{t})}\right]$

Model-based RL in this case involves alternating optimization between policy
learning and learning the world model. In the case where there is no reward
signal from the environment, i.e. $p(r_{t}|z_{t})$, we have the following
objective

$\displaystyle\max_{\pi}\max_{\phi,\psi}\sum_{t}\mathbb{E}_{\pi(a_{t}|z_{t}),q_{\phi}(z_{t}|s_{t})p(z_{t+1}|z_{t},a_{t})}\log\frac{q_{\psi}(a_{t}|z_{t+1},z_{t})}{\pi(a_{t}|z_{t})}$

The optimal $q^{*}_{\psi}$ give a fixed $\pi$ and $\phi$ can be derived as in
page 335 of Cover & Thomas (2006).

$\displaystyle
q^{*}_{\psi}(a_{t}|z_{t+1},z_{t})=\frac{q_{\phi}(z_{t}|s_{t})p(z_{t+1}|z_{t},a_{t})\pi(a_{t}|z_{t})}{\int_{a}q_{\phi}(z_{t}|s_{t})p(z_{t+1}|z_{t},a_{t})\pi(a_{t}|z_{t})}$

The optimal $\pi*$ given a fixed $\psi$ and $\phi$ can be derived as

$\displaystyle\pi^{*}(a_{t}|z_{t})=\frac{\exp{(\mathbb{E}_{q_{\phi}(z_{t}|s_{t})p(z_{t+1}|z_{t},a_{t})}[\log
q_{\psi}(a_{t}|z_{t+1},z_{t})]})}{\int_{a}\exp{(\mathbb{E}_{q_{\phi}(z_{t}|s_{t})p(z_{t+1}|z_{t},a_{t})q}[\log
q_{\psi}(a_{t}|z_{t+1},z_{t})]})}$

We note that these alternating iterations correspond to an instantiation of
the Blahut-Arimoto algorithm (Blahut, 1972) for determining the information
theoretic capacity of a channel. Based on (Nakagawa et al., 2021) we see that
this iterative optimization procedure between $q^{*}_{\psi}$ and $\pi^{*}$
converges, and the worst-case rate of convergence is ${\mathcal{O}}(1/N)$
where $N$ is the number of iterations.

∎

This result is useful because it shows that even in the absence of rewards
from the environment ($r_{t}=0\;\forall t)$, when planning to minimize regret
is not possible, the inverse information objective can be used to train a
policy that explores to seek out controllable parts of the state-space. When
rewards from the environment are present, we can train the policy with this
objective and the value estimates obtained from the cumulative rewards during
planning. In Section 4.3 we empirically show how this objective helps achieve
higher sample efficiency compared to pure value-based policy learning.

Figure 7: PGM of the MDP with distractor states. The state observed by the
agent ${\mathcal{O}}$ consists of three parts $S^{+}$, $\tilde{S}^{-}$ and
$DS^{-}$. $S^{+}$ is the controllable part of the state i.e. it is affected by
the actions of the agent, and in turn affects the reward $R$; $\tilde{S}^{-}$
is not controllable by the agent but affects the reward $R$ and future
$S^{+}$; $DS^{-}$ is not controllable by the agent and doesn’t affect the
reward $R$ and future $S^{+}$.

###  A.2 Primal Dual updates

$Z_{0:H-1}=Z_{0:H-1}+\frac{\partial\sum_{t=0}^{H-1}I({\mathcal{O}}_{t};Z_{t})+\lambda\left(-I(i_{t};Z_{t}|Z_{t-1},A_{t-1})+I(A_{t-1};Z_{t}|Z_{t-1})+I(R_{t};Z_{t})-c_{0}\right)}{\partial
Z_{0:H-1}}$

$\lambda=\lambda-\left(\underbrace{-I(i_{t};Z_{t}|Z_{t-1},A_{t-1})}_{\text{fwd
dynamics model}}+\underbrace{I(A_{t-1};Z_{t}|Z_{t-1})}_{\text{inv dynamics
model}}+\underbrace{I(R_{t};Z_{t})}_{\text{rew model}}-c_{0}\right)$

###  A.3 Forward, Inverse, and Reward Models

$\displaystyle-I(i_{t};Z_{t}|Z_{t-1},A_{t-1})$ $\displaystyle=-\int
p(z_{t},z_{t-1},a_{t-1},i_{t},o_{t})\log\frac{p(z_{t}|z_{t-1},a_{t-1},o_{t})}{p(z_{t}|z_{t-1},a_{t-1})}$
$\displaystyle=-\int
p(z_{t},z_{t-1},a_{t-1},o_{t})\left(\log\frac{p(z_{t}|z_{t-1},a_{t-1},o_{t})}{q_{\chi}(z_{t}|z_{t-1},a_{t-1})}+\log\frac{q_{\chi}(z_{t}|z_{t-1},a_{t-1})}{p(z_{t}|z_{t-1},a_{t-1})}\right)$
$\displaystyle\geq-\int
p(z_{t},z_{t-1},a_{t-1},o_{t})\log\frac{p(z_{t}|z_{t-1},a_{t-1},o_{t})}{q_{\chi}(z_{t}|z_{t-1},a_{t-1})}$
$\displaystyle=-D_{\text{KL}}(p(z_{t}|z_{t-1},a_{t-1},o_{t})||q_{\chi}(z_{t}|z_{t-1},a_{t-1}))$

We can obtain a variational lower bound on the empowerment

$\displaystyle I(a_{t};z_{t+k}|z_{t})$ $\displaystyle=\int
p(a_{t},z_{t+k},z_{t})\log\frac{p(a_{t}|z_{t+k},z_{t})}{p(a_{t}|z_{t})}$
$\displaystyle=\int
p(a_{t},z_{t+k},z_{t})\left(\log\frac{q(a_{t}|z_{t+k},z_{t})}{p(a_{t}|z_{t})}+\log\frac{p(a_{t}|z_{t+k},z_{t})}{q(a_{t}|z_{t+k},z_{t})}\right)$
$\displaystyle\geq\int
p(a_{t},z_{t+k},z_{t})\log\frac{q(a_{t}|z_{t+k},z_{t})}{p(a_{t}|z_{t})}$
$\displaystyle=\int p(a_{t},z_{t+k},z_{t})\log q(a_{t}|z_{t+k},z_{t})-\int
p(z_{t})p(z_{t+k}|a_{t},z_{t})p(a_{t}|z_{t})\log p(a_{t}|z_{t})$
$\displaystyle=\mathbb{E}_{p(a_{t},z_{t+k},z_{t})}[\log
q(a_{t}|z_{t+k},z_{t})]+\mathbb{E}_{p(z_{t})p(z_{t+k}|a_{t},z_{t})}[{\mathcal{H}}(p(a_{t}|z_{t}))]$

Here, $q(a_{t}|z_{t+k},z_{t})$ is the inverse dynamics model, and
${\mathcal{H}}(p(a_{t}|z_{t}))$ is the entropy of the policy.

We can obtain a variational lower bound on the reward model as

$\displaystyle I(r_{t};z_{t})$ $\displaystyle=\int
p(r_{t},z_{t})\log\frac{p(r_{t}|z_{t})}{p(r_{t})}$ $\displaystyle=\int
p(r_{t},z_{t})\left(\log\frac{q(r_{t}|z_{t})}{p(r_{t})}+\log\frac{p(r_{t}|z_{t})}{q(r_{t}|z_{t})}\right)$
$\displaystyle=\int
p(r_{t},z_{t})\log\frac{q(r_{t}|z_{t})}{p(r_{t})}+KL[p(r_{t},z_{t})||q(r_{t},z_{t})]$
$\displaystyle\geq\int p(r_{t},z_{t})\log\frac{q(r_{t}|z_{t})}{p(r_{t})}$
$\displaystyle=\int p(r_{t},z_{t})\log q(r_{t}|z_{t})-\int p(r_{t},z_{t})\log
p(r_{t})$ $\displaystyle=\mathbb{E}_{p(r_{t},z_{t})}[\log
q(r_{t}|z_{t})]+Constant$

Here, $q(r_{t}|z_{t})$ is the reward decoder. Since the reward is a scalar,
reconstructing it is computationally simpler compared to reconstructing high
dimensional observations as in the BA bound for the observation model.

###  A.4 Value Difference Result

We show that the difference in value function of the underlying MDP
$Q^{\pi}(o,a)$ and the latent MDP $\hat{Q}^{\pi}(z,a)$, where $z\sim
q_{\phi}(z|o)$ is bounded, under some regularity assumptions. Similar to the
assumption in (Gelada et al., 2019), let the policy $\pi$ have bounded semi-
norm value functions under the total variation distance $D_{\text{TV}}$ i.e.
$|\hat{V}^{\pi}(z)|_{D_{\text{TV}}}\leq K$ and $|\mathbb{E}_{z_{1}\sim
P}V(z_{1})-\mathbb{E}_{z_{2}\sim Q}V(z_{2})|\leq KD_{\text{TV}}(P,Q)$. Define
the forward dynamics learning objective as
$\min_{\chi}L_{T}=\min_{\chi}D_{\text{KL}}(p(o^{\prime}|o,a)||q_{\chi}(z^{\prime}|z,a))$
and the reward model objective as
$\min_{\eta}L_{R}=\min_{\eta}D_{\text{KL}}(p(r|o)||q_{\eta}(r|z))$

######  Theorem 3.

The difference between the Q-function in the latent MDP and that in the
original MDP is bounded by the loss in reward predictor and forward dynamics
model.

$\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim
q_{\phi}(z|o)}[Q^{\pi}(o,a)-\hat{Q}^{\pi}(z,a)]\leq\frac{\sqrt{L_{R}}+\gamma
K\sqrt{L_{T}}}{1-\gamma}$

###### Proof.

$\displaystyle\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim
q_{\phi}(z|o)}[Q^{\pi}(o,a)-\hat{Q}^{\pi}(z,a)]$
$\displaystyle\leq\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim
q_{\phi}(z|o)}(|r(o,a)-\hat{r}(z,a)|+\gamma|\mathbb{E}_{o^{\prime}\sim
p(o^{\prime}|o,a)}V(o^{\prime})-\mathbb{E}_{z^{\prime}\sim
q_{\chi}(z^{\prime}|z,a)}\hat{V}(z^{\prime})|)$
$\displaystyle\leq\sqrt{D_{\text{KL}}(p(r|o)||q_{\eta}(r|z))}+\mathbb{E}_{(o,a)\sim
d^{\pi}(o,a),z\sim q_{\phi}(z|o)}(\gamma|\mathbb{E}_{o^{\prime}\sim
p(o^{\prime}|o,a),z^{\prime}\sim
q_{\phi}(z^{\prime}|o^{\prime})}[V(o^{\prime})-\hat{V}(z^{\prime})]|$
$\displaystyle+\gamma|\mathbb{E}_{o^{\prime}\sim
p(o^{\prime}|o,a),z^{\prime}\sim
q_{\chi}(z^{\prime}|z,a)}[V(o^{\prime})-\hat{V}(z^{\prime})]|)$
$\displaystyle\leq\sqrt{L_{R}}+\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim
q_{\phi}(z|o)}(\gamma|\mathbb{E}_{o^{\prime}\sim
p(o^{\prime}|o,a),z^{\prime}\sim
q_{\phi}(z^{\prime}|o^{\prime})}[V(o^{\prime})-\hat{V}(z^{\prime})]|)$
$\displaystyle+\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim q_{\phi}(z|o)}(\gamma
KD_{\text{TV}}(p(o^{\prime}|o,a)||q_{\chi}(z^{\prime}|z,a)))$
$\displaystyle\leq\sqrt{L_{R}}+\gamma\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim
q_{\phi}(z|o)}\mathbb{E}_{o^{\prime}\sim p(o^{\prime}|o,a),z^{\prime}\sim
q_{\phi}(z^{\prime}|o^{\prime})}[V(o^{\prime})-\hat{V}(z^{\prime})])$
$\displaystyle+\gamma
K\sqrt{D_{\text{KL}}(p(o^{\prime}|o,a)||q_{\chi}(z^{\prime}|z,a))}$
$\displaystyle\leq\sqrt{L_{R}}+\gamma\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim
q_{\phi}(z|o)}[V(o)-\hat{V}(z)])+\gamma K\sqrt{L_{T}}$
$\displaystyle\leq\sqrt{L_{R}}+\gamma\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim
q_{\phi}(z|o)}[Q(o,a)-\hat{Q}(z,a)])+\gamma K\sqrt{L_{T}}$

So, we have shown that

$\mathbb{E}_{(o,a)\sim d^{\pi}(o,a),z\sim
q_{\phi}(z|o)}[Q^{\pi}(o,a)-\hat{Q}^{\pi}(z,a)]\leq\frac{\sqrt{L_{R}}+\gamma
K\sqrt{L_{T}}}{1-\gamma}$

∎

We can see that, as $L_{R},L_{T}\rightarrow 0$, the Q-function in the
representation space of the MDP, $\hat{Q}$, becomes increasingly closer to
that of the original MDP, $Q$. Since this result holds for all Q-functions, it
also holds for the optimal $Q^{*}$ and $\hat{Q}^{*}$. This result extends
sufficiency results in prior works Rakelly et al. (2021) to a setting with
stochastic encoders and KL-divergence losses on forward dynamics and reward,
as opposed to Wasserstein metrics (Gelada et al., 2019), and bisimulation
metrics (Zhang et al., 2020).

###  A.5 Behavioral Similarity metric details

Given the sets ${\mathcal{S}}_{z}=\\{z^{i}\\}_{i=1}^{n}$ and
${\mathcal{S}}_{z_{\text{gt}}}=\\{z^{i}_{\text{gt}}\\}_{i=1}^{n}$, we
construct complete graphs by connecting every vertex with every other vertex
through an edge. The weight of an edge is the Euclidean distance between the
respective pairs of vertices. The label of each vertex is an unique integer
and corresponding vertices in the two graphs are labelled similarly. Let the
resulting graphs be denoted as $G_{z}=(V_{z},E_{z})$ and
$G_{z_{\text{gt}}}=(V_{z_{\text{gt}}},E_{z_{\text{gt}}})$ respectively. Note
that both these graphs are shortest path graphs, by construction.

Since the Euclidean distances in both the graphs cannot be directly compared,
we scale the distances such that the shortest edge weight among all edge
weights are equal in both the graphs. Scaling every edge weight by the same
number doesn’t change the structure of the graph. Now, we define the shortest
path kernel to measure similarity between the two graphs, as follows

$k(G_{z},G_{z_{\text{gt}}})=\frac{1}{|E_{z}|}\sum_{e_{i}\in
E_{z}}\sum_{e_{j}\in E_{z_{\text{gt}}}}\hat{k}(e_{i},e_{j})$

Let, $e_{i}=\\{u_{i},v_{i}\\}$ and $e_{j}=\\{u_{j},v_{j}\\}$. Here, the kernel
$\hat{k}$ measures similarity between the edges $e_{i},e_{j}$ and the labels
on respective vertices. Let the function $l(\cdot)$ denote the labels of
vertices and the weights of edges.

$\hat{k}(e_{i},e_{j})=k_{v}(l(u_{i}),l(u_{j}))k_{e}(l(e_{i}),l(e_{j}))k_{v}(l(v_{i}),l(v_{j}))+k_{v}(l(u_{i}),l(v_{j}))k_{e}(l(e_{i}),l(e_{j}))k_{v}(l(v_{i}),l(u_{j}))$

Here, $k_{v}$ is a Dirac delta kernel i.e. $k_{e}(x,y)=1-\delta(x,y)$ (Note
that $\delta(x,y)=1$ iff $x=y$, else $\delta(x,y)=0$) and $k_{e}$ is a
Brownian ridge kernel, $k_{e}(x,y)=\frac{1}{c}\max(0,c-|x-y|)$, where $c$ is a
large number.

So, in summary, the value of $\hat{k}(e_{i},e_{j})$ is low when the edge
lengths between corresponding pairs of vertices in both the graphs are close.
Hence, the value of $k(G_{z},G_{z_{\text{gt}}})$ is low when a large number of
pairs of corresponding vertices in both the graphs have the same edge length.

We expect methods that recover latent state representations which better
reflect the true underlying simulator state (i.e., the positions of the
joints) to have higher values according to this metric.

###  A.6 Description of the distractors

ILSVRC dataset (Russakovsky et al., 2015) The distractors used for the main
experiments are natural video backgrounds with RGB images from the ILSVRC
dataset. We use 200 videos during training, and reserve 50 videos for testing.
An illustration of this for different environments are shown in Fig. 8. These
images are snapshots of what is received by the algorithm (64x64x3 images) and
no information about proprioceptive features is provided. While natural videos
that are unrelated to the task might be easy to ignore, realistic scenes might
have other elements that resemble the controllable elements, but are not
actually controllable (e.g., other cars in a driving scenario).

To emulate this, we also add distractors that resemble other potentially
controllable robots, as shown for example in Fig. 8 (top left and bottom right
frames), but are not actually influenced by the actions of the algorithm.
These are particularly challenging because the background video has a pre-
recorded motion of the same (or similar) agent that is being controlled. We
include such challenging agent-behind-agent background videos with a frequency
of 1 in very 3 videos.

For the experiments in Table 1, examples of different levels of distractors
are shown in Fig. 9. The three different levels have distractor windows of
size 32x32, 40x40, and 64x64 respectively.

Kinetics dataset (Kay et al., 2017) For comparing directly with results in
prior works DBC (Zhang et al., 2020), and TIA (Fu et al., 2021), we evaluate
on the setting where the random videos are grayscale images from the Kinetics
dataset driving car class. The settings are same as in the prior works and we
compare directly to the results in the respective papers.

###  A.7 Training and Network details

The agent observations have shape 64 x 64 x 3, action dimensions range from 1
to 12, rewards per time-step are scalars between 0 and 1. All episodes have
randomized initial states, and the initial distractor video background is
chosen randomly.

We implement our approach with TensorFlow 2 and use a single Nvidia V100 GPU
and 10 CPU cores for each training run. The training time for 1e6 environment
steps is 4 hours on average. In comparison, the average training time for 1e6
steps for Dreamer is 3 hours, for CDreamer is 4 hours, for TIA is 4 hours, for
DBC is 4.5 hours, and for DeepMDP is 5 hours.

The encoder consists of 4 convolutional layers with kernel size 4 and channel
numbers 32, 65, 128, 256. The forward model consists of deterministic and
stochastic parts, as in standard in visual model-based RL (Hafner et al.,
2019b; a). The stochastic states have size 30 and deterministic state has size
200. The compatibility function
$f_{\theta}(z_{t},o_{t})=\exp{(\hat{z}_{t}^{T}W_{\theta}z_{t})}$ for
contrastive learning is learned with a 200 x 200 $W_{\theta}$ matrix. Here,
$\hat{z}_{t}$ is the encoding of $o_{t}$. All other models are implemented as
three dense layers of size 300 with ELU activations.

We use ADAM optimizer, with learning rate of 6e-4 for the latent-state space
model, and 8e-5 for the value function and policy optimization. The hyper-
parameter $c_{0}$ for the prioritized information constraint is set to $1000$,
after doing a grid-search over the range $[100,10000]$ and observing similar
performance for $1000$ and $10000$. 100 training steps are followed by 1
episode of interacting with the environment with the currently trained policy,
for observing real transitions. The dataset is initialized with 7 episodes
collected with random actions in the environment. We kept the hyper-parameters
of InfoPower same across all environments and distractor types.

For fair comparisons, we kept all the common hyper-parameter values same as
Dreamer (Hafner et al., 2019a). This was the protocol followed in prior works
TIA (Fu et al., 2021) and Ma et al. (2020). For the baseline TIA, we tuned the
environment-specific parameters $\lambda_{R_{adv}}$ and $\lambda_{O_{s}}$ as
mentioned in Table 2 of (Fu et al., 2021). For $\lambda_{R_{adv}}$, we
performed a gridsearch over the range $10k-50k$ and for $\lambda_{O_{s}}$ we
performed a gridsearch over the range $1-3$ to obtain the parameters for best
performance, which we used for the plots in Fig. 4 and Fig. 11. Respectively
for walker-walk, cheetah-run, finger-spin, quadruped-walk, hopper-stand, ball-
in-a-cup-catch, cartpole-balance-sparse, quadruped-run, the values of
$\lambda_{R_{adv}}$ are 30k,30k,20k,40k,30k,40k,20k,30k. The values of
$\lambda_{O_{s}}$ are 2,2,1.5,2.5,2.5,2,2,2.

For baseline DBC (Zhang et al., 2020), we kept all the parameters same as in
Table 2 of the paper, because DBC does not have any environment-specific
parameters and the same values were used for all environments in the DBC
paper. The DeepMDP agent and its hyperparameters are adapted from the
implementation provided in the github repo of DBC (Zhang et al., 2020).

Figure 8: Illustration of the natural video background distractors used in our
experiments. The videos change after every 50 time-steps. Some video
backgrounds (for example the top left and the bottom right) have more complex
distractors in the form of agent-behind-agent i.e. an agent of similar
morphology moves in the background of the agent being controlled. Figure 9:
Illustration of the different levels of distractors. L1, L2, and L3
respectively have distractor windows of size 32x32, 40x40, and 64x64.

###  A.8 Contrastive Learning

Since the distracting backgrounds are changing videos of a certain duration,
for contrastive learning, we contrast both across time and across batches.
Contrasting across time helps model invariance against temporally contiguous
distractors (for example in the same video) while contrasting across batches
helps in learning invariance across different videos. Concretely, we sample
batches from the dataset ${}^{(i)}\\{(a_{t},o_{t},r_{t})\\}_{t=1}^{H}$, where
$i$ denotes a batch index, and for each observation $o_{t}$ in batch $i$, all
observations $o_{t^{\prime}}$ both in batch $i$ and in other batches $j\neq i$
are considered negative samples.

###  A.9 Result on varying distractor levels

We consider different levels of distractors by varying the size of the window
where distractors in the background are active. Fig. 9 in the Appendix
illustrates this visually. Table 1 shows results for the different approaches
on varying distractor levels. The size of the frame for distractors at levels
L1, L2, and L3 respectively are 32x32, 40x40, and 64x64. We observe that the
baselines perform worse as the distractor window size increases. While, for
InfoPower, the performance decrease with increasing level is minimal. This
indicates the effectiveness of InfoPower in filtering out background
distractors from the observations while learning latent representations. Note
that the distractors in Figs. 4 and 6 are of level 3 i.e. same size as the
observations 64x64.

Table 1: DM Control Walker Stand with natural video distractors at different levels. We tabulate the rewards (Rew) and behavioral similarity (Sim) at 500k and 1M steps of training. Higher is better for both Rew and Sim. Name | Levels | Rew@500k | Rew@1M | Sim@500k | Sim@1M
---|---|---|---|---|---
Dreamer | L1 | $197\pm 31$ | $240\pm 27$ | 0.73$\pm 0.02$ | 0.74$\pm 0.01$
DBC | L1 | $261\pm 25$ | $390\pm 34$ | 0.75$\pm 0.03$ | 0.74$\pm 0.02$
C-Dreamer | L1 | $291\pm 34$ | $590\pm 26$ | 0.73$\pm 0.02$ | 0.79$\pm 0.01$
DeepMDP | L1 | $263\pm 31$ | $340\pm 22$ | 0.74$\pm 0.02$ | 0.75$\pm 0.03$
InfoPower | L1 | $\mathbf{397\pm 22}$ | $\mathbf{650\pm 100}$ | $\mathbf{0.82}\pm 0.01$ | $\mathbf{0.84}\pm 0.03$
Dreamer | L2 | $180\pm 36$ | $197\pm 20$ | 0.56$\pm 0.03$ | 0.59$\pm 0.02$
DBC | L2 | $221\pm 28$ | $320\pm 33$ | 0.65$\pm 0.02$ | 0.66$\pm 0.02$
C-Dreamer | L2 | $282\pm 30$ | $550\pm 66$ | 0.63$\pm 0.01$ | 0.71$\pm 0.02$
DeepMDP | L2 | $213\pm 34$ | $300\pm 25$ | 0.59$\pm 0.02$ | 0.61$\pm 0.04$
InfoPower | L2 | $\mathbf{394\pm 30}$ | $\mathbf{644\pm 101}$ | $\mathbf{0.77}\pm 0.03$ | $\mathbf{0.77}\pm 0.03$
Dreamer | L3 | $140\pm 52$ | $157\pm 10$ | 0.32$\pm 0.02$ | 0.33$\pm 0.03$
DBC | L3 | $165\pm 20$ | $221\pm 24$ | 0.45$\pm 0.01$ | 0.48$\pm 0.01$
C-Dreamer | L3 | $231\pm 39$ | $485\pm 86$ | 0.58$\pm 0.02$ | 0.64$\pm 0.01$
DeepMDP | L3 | $153\pm 23$ | $212\pm 28$ | 0.44$\pm 0.02$ | 0.49$\pm 0.01$
InfoPower | L3 | $\mathbf{389\pm 20}$ | $\mathbf{624\pm 80}$ | $\mathbf{0.71}\pm 0.01$ | $\mathbf{0.74}\pm 0.03$

###  A.10 Results on distractors from the Kinetics dataset

In this section, we evaluate InfoPower on the same distractors used in prior
works, DBC and TIA. The distractor videos are from the Kinetics dataset (Kay
et al., 2017) Driving Car class, and converted to grayscale. From Table 2 we
see that InfoPower slightly outperforms the baselines. We believe the
performance gap is not as significant as Fig. 4 because the distractors being
grayscale images, and not RGB might be easier to ignore by the baselines as
well as InfoPower.

Table 2: Comparison of baselines TIA and DBC with InfoPower on environments with video distractors from the Kinetics dataset Driving Car class. The results for TIA and DBC are from the respective papers. | TIA | DBC | InfoPower
---|---|---|---
| 500k | 800k | 500k | 800k | 500k | 800k
Walker-Run | 389$\pm$45 | 602$\pm$40 | 165$\pm$92 | 210$\pm$32 | 417$\pm$22 | 630$\pm$33
Cheetah-Run | 384$\pm$149 | 579$\pm$172 | 261$\pm$64 | 277$\pm$81 | 425$\pm$97 | 572$\pm$158
Hopper-Stand | 338$\pm$221 | 581$\pm$214 | 110$\pm$92 | 207$\pm$120 | 395$\pm$140 | 615$\pm$176
Ball-in-a-Cup | 201$\pm$197 | 115$\pm$110 | – | – | 255$\pm$104 | 263$\pm$116
Walker-Walk | 878$\pm$37 | 948$\pm$28 | 545$\pm$57 | 630$\pm$68 | 925$\pm$18 | 972$\pm$26
Finger-Spin | 371$\pm$96 | 487$\pm$107 | 801$\pm$10 | 832$\pm$4 | 795$\pm$34 | 787$\pm$58
Hopper-Hop | 47$\pm$44 | 72$\pm$68 | 5$\pm$12 | 0$\pm$0 | 85$\pm$26 | 77$\pm$74

###  A.11 Setting with only shifted agent behind agent distractors

Figure 10: Ilustration of the distractor setting for results in Table 3.

In this section we evaluate InfoPower with the baselines only for the
challenging setting of agent-behind-agent distrators where the background
contains the same agent being controlled (a walker or a cheetah). An
illustration of this is shown in Fig. 10. From the results in Table 3, we see
that InfoPower significantly outperforms the competitive baselines both at
500k and 1M environment interactions. This suggests the utility of InfoPower
in separating out informative elements of the observation from uninformative
ones even in challenging settings.

Table 3: Evaluation of InfoPower and the most competitive baselines in a suite of challenging distractors where the background contains a shifted version of the agent being controlled. Results are averaged over 4 random seeds. Fig. 10 shows a visualization of the distractors for this setting. | Walker-Walk-shifted | Cheetah-Run-shifted
---|---|---
| 500k | 1M | 500k | 1M
DBC | 35$\pm$15 | 104$\pm$28 | 41$\pm$10 | 78$\pm$38
CDreamer | 108$\pm$36 | 164$\pm$45 | 73$\pm$44 | 110$\pm$52
TIA | 79$\pm$42 | 152$\pm$33 | 65$\pm$24 | 115$\pm$28
InfoPower | 168$\pm$48 | 278$\pm$42 | 130$\pm$46 | 207$\pm$35

###  A.12 Comparison of baselines with empowerment in policy learning

In this section, we compare against modified versions of the baselines, such
that we include empowerment in the policy learning of the baselines. This is
the only modification we make to the baselines, and keep everything else
unchanged. We described in section 3.3 that inclusion of empowerment in the
objective for visual model-based RL (by modifying both the representation
learning objective and the policy learning objective) is an important
contribution of our paper. The aim of this experiment is to disentangle the
benefits of empowerment in representation learning, for InfoPower with that in
policy learning.

We have plotted the results in Fig. 11. We can see that the performance of the
baselines is slightly better than Fig. 4 because of the inclusion of
empowerment in the policy learning objective which as we motivated previously
in section 3.4 helps with exploring controllable regions of the state-space.
Note that by adding empowerment, we have effectively modified the baselines.

Figure 11: Evaluation of InfoPower and baselines with empowerment in policy
learning. The setting corresponds to DeepMind Control tasks with RGB natural
video distractors in the background (same as that in Figure 5 of the main
paper). The x-axis denotes the number of environment interactions and the
y-axis shows the episodic returns.

###  A.13 Comparison on environments without distractors

We compare the performance of InfoPower with a state-of-the-art visual model-
based RL method, Dreamer in the default DM Control environments without
distractors. We see from Table 4 that InfoPower performs slightly better or
similar to Dreamer in all the environments consistently. This result
complements our main paper result on environments with distractors and
confirms that the benefit of InfoPower in challenging distracting environments
does not come at a cost of performance in simpler non-distracting
environments.

Table 4: Comparison of InfoPower with Dreamer in the default DM Control environments without distractors. We tabulate reward at 500k and 1M environment interactions. Results are averaged over 4 random seeds. Higher is better.  | Dreamer | InfoPower
---|---|---
| Rew @ 500k | Rew @ 1M | Rew @ 500k |  1M
Walker Walk | 890$\pm$52 | 985$\pm$ 10 | 910$\pm$41 | 980$\pm$ 17
Cheetah Run | 580$\pm$190 | 811$\pm$95 | 586$\pm$142 | 810$\pm$102
Finger Spin | 570$\pm$187 | 573$\pm$96 | 590$\pm$172 | 598$\pm$85
Quadruped Walk | 592$\pm$45 | 646$\pm$50 | 623$\pm$65 | 666$\pm$47
Hopper Stand | 701$\pm$68 | 906$\pm$33 | 715$\pm$60 | 930$\pm$39
Ball in a Cup | 871$\pm$102 | 910$\pm$48 | 913$\pm$67 | 939$\pm$52
Cart-Pole Sparse | 742$\pm$58 | 849$\pm$88 | 776$\pm$40 | 866$\pm$80
Quadruped Run | 450$\pm$47 | 515$\pm$65 | 468$\pm$42 | 510$\pm$61

###  A.14 Additional ablations

In this section, we consider additional ablations, namely a version of of
InfoPower that does not have the empowerment objective, a version of InfoPower
that does not have the MI maximization objective. We see from Table 5 that
removing either of these terms drops performance.

Table 5: DM Walker Stand with natural video distractors Name | Obs bound | Rew@500k | Rew@1M | Sim@500k | Sim@1M
---|---|---|---|---|---
Dreamer | BA | $140\pm 52$ | $157\pm 10$ | 0.32$\pm 0.02$ | 0.33$\pm 0.03$
Contrastive-Dreamer | NCE | $231\pm 39$ | $485\pm 86$ | 0.58$\pm 0.02$ | 0.64$\pm 0.01$
NonGenerative-Dreamer | NWJ | $280\pm 25$ | $480\pm 93$ | 0.59$\pm 0.01$ | 0.62$\pm 0.05$
InfoPower | NCE | $254\pm 27$ | $481\pm 30$ | 0.69$\pm 0.03$ | 0.72$\pm 0.02$
InfoPower | NWJ | $\mathbf{389\pm 20}$ | $\mathbf{624\pm 80}$ | 0.71$\pm 0.01$ | 0.74$\pm 0.03$
Inverse only | None | $282\pm 70$ | $367\pm 24$ | 0.63$\pm 0.04$ | 0.66$\pm 0.03$