Information Prioritization through Empowerment in Visual Model-based RL
Authors: Homanga Bharadhwaj
Mohammad Babaeizadeh
Dumitru Erhan
Sergey Levine
Venue: ICLR 2022 Poster (2022)
Keywords: model-based reinforcement learning
visual distractors
empowerment

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 learning 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 by an average of 20\% in terms of episodic returns at 1M environment interactions with 30\% higher sample efficiency at 100k interactions.

Published as a conference paper at ICLR 2022
INFORMATION PRIORITIZATION THROUGH EMPOWERMENT IN VISUAL MODEL-BASED RL
Homanga Bharadhwaj
Mohammad Babaeizadeh Google Research, Brain Team
Dumitru Erhan Google Research, Brain Team
Carnegie Mellon University
Sergey Levine
Google Research, Brain Team University of California Berkeley
ABSTRACT
Model-based reinforcement learning (RL) algorithms designed for handling com- plex visual observations typically learn some sort of latent state representation, either explicitly implicitly. Standard methods of this sort do not distinguish be- tween functionally relevant aspects of the state and irrelevant distractors, instead aiming to represent all available information equally. We propose modified ob- jective for model-based RL that, in combination with mutual information maxi- mization, allows 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 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 empower- ment 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-basse robot control tasks with natural video backgrounds, and show that the proposed prioritized information objective outperforms state-of-the-ee model based RL ap- proaches with higher sample efficiency and episodic returns.
INTRODUCTION
Model-based reinforcement learning (RL) provides promising approach to accelerating skill learn- ing: by acquiring 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, 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; Babacizadch et al., 2017), effecttve learning in visually complex environments necessitates some mechanism for learning representations that prioritize functionally relevant fac- tors for the current task. This needs 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; 8; Ma et al., 2020; Oord et al., 2018; Srinivas et al., 2020). While such objectives avoid recon- struction, they still do not distinguish between relevant and irrelevant factors of variation. We thus pose the question: can we devise reconnsttruee representation learning methods that explicitly prioritize information that is most likely to be functionally relevant the agent?
In this work, we derive a model-based RL algorithm from a combination of representation learning via mutual information maximization (Poole al., 2019) and empowerment (Mohamed Rezende, 2015). The latter serves to drive both the representation and the policy toward exploring and repre- senting functionally relevant factors of variation. By integrating an empowerment-based term into
Work done during Homanga's 's research internship at Google. hoha radh 8cs cmu.edu

---

Published as a conference paper at ICLR 2022
Ox
R
O++1
O2
Negetive Samples
A
Reward Objective
Empowerment
I(R4; Zz)
Empowarment Objective
[[A Z++1 Zx) ((z))
Enzodor
Encoder
Encodar
Encodor
I I(A Zet1 (Z4)
Au xt T(ee)
Arst ~ =(z+1))
Contrastive Oojective
Z
Ze
Zz
Ze+11
2012 ..
[1,
V1
V4+1
Representation learning
Planning
Figure I: Overview of InfoFo. I(O,; Z.) is the contrastive learning objective for learning an encoder
to map from image O to latent Z. I(At-1; Ze (rr-) is the empowerment objective that prioritizes encoding controllable representations in Z. [( Z,+1 Zz. Ar) helps learn a latent forward dynamics model so that future Z,+a can be predicted from current Ze. I(R;; Ze) helps learn reward prediction model, such that the agent learn a plan At, ...A.: through latent rollouts. Together, this combination of terms produces latent state space model for MBRL that captures all necessary information at convergence, while prioritizing the most functionally relevant factors via the empowerment term.
mutual information framework for learning state representations, we effectively prioritize informa- tion 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 fur- ther 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 ob- servations that integrates empowerment into mutual information based, non-reconstructive frame- work for learning state space models. Our approach explicitly prioritizes information that is most likely be functionally relevant, which significantly improves performance 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 suite of simulated robotic control tasks with explicit video distractors, and demonstrate up to 20% better performance in terms of cumulative rewards IM environment interactions with 30% higher sample efficiency at 100k interactions.
PROBLEM STATEMENT AND NOTATION
A partially observed Markov decision process (POMDP) is a tuple (S, R, that consists of states E S, actions a E A, rewards r E R. observations E O, and a state-transition distribu- tion [(s'|s, a). In most practical settings, the agent interacting with the environment doesn't 't have access to the actual states in S, but to some partial information in the form of observations O. The underlying state-transition distribution T and reward distribution R are also unknown the agent.
In this paper, we consider the observations E O to be high-dimensional images, and so the agent should learn a compact representation space Z for the latent sate-space model. The problem state- ment is to learn effective representations from observations 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 = E " -r))) The value of state V(Z)) is defined as the expected cumulative sum of discounted rewards starting at state Z,
We use q(-) to denote parameterized variational approximations to learned distributions. We denote random variables with capital letters and use lowercase letters to denote particular re- alizations (e.g.,
denotes the value of Z,). Since the underlying distributions unknown, we evaluate all expectations through Monte-Card sampling with observed stte-transition tuples (ot, 1 1 Or- 1 St, ct- 1, T.).
2

---

Published as a conference paper at ICLR 2022
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 O and devise 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 maxz I(O, Z) sub- ject to prioritization of the empowerment objective maxz [(A,-i; 1; Ze Zz) For planning, we maximize the empowerment objective along with reward-based value with respect to the policy max 1: (11)) + [[;;; Z,).
In the subsequent sections, we elaborate on our approach, InfoPower, and describe lower bounds to MI that yield tractable algorithm.
3.1 LEARNING CONTROLLABLE FACTORS AND PLANNING THROUGH EMPOWERMENT
Controllable representations are features of the observation that correspond to enti- ties 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 en- vironment 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 denoted by S in Fig. 2, which we can formally define through conditional inde- pendence as the smallest subspace of S. S+ < S, such that T(A,-1; Se St)
4-1
At
S
S
DS,
Se+
DS
R
O₂
Re+1
O++1
Figure 2: PGM showing decomposition of state S into con- trollable parts S (directly influenced by actions A), parts not influenced by actions that still influence the reward, S and distractors DS- ST may be influenced by 8, (arrow can be seen in Fig. 2. We explicitly priori- not shown to reduce clutter) but not by $
This conditional independence relation
tize the learning of such representations in the latent space by drawing inspiration from variational empowerment (Mohamed & Rezende, 2015).
The empowerment objective can be cast as maximizing a conditional information term H(A,-1 Z1-1). The first term (((A,1/Z,-1) encour-
[Ar-1: 1z)Z--1) H(A,-1) ages the chosen actions to be as diverse as possible, while the second term -H(A,-1)4 1l Z,-1) en- courages the representations Ze and Ze+1 to be such that the action A, for transition is predictable. While prior approaches have used empowerment in the model-free setting to learn policies by explo- ration 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 learn- ing in visual model-based RL (Sutton, 1991), with max, / (A,,1; Z, 1.) This objectives com- plements 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 A.1 we next describe two theorems regarding learning controllable representations. We objective for learning latent representations 2, when
observe that the max E. 7(A,1;; Zz/Z-1) ( used along with the planning objective, provably recovers controllable parts of the observation 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| encode 5*
O), learning Z under the objective max Σ, I(A, Z, Zz-1) would
We further show through Theorem 2 that the inverse information objective alone can be used train a latent-state space model and a policy through an alternating optimization algorithm that converges to local minimum of the objective max E. T(A,1;; 7, 74-1) at rate inversely proportional to
3

---

Published as a conference paper at ICLR 2022
the number of iterations. In Section 4.3 we empirically show how this objective helps achieve higher sample efficiency compared to pure value-base based policy learning.
3.2 MUTUAL INFORMATION MAXIMIZATION FOR REPRESENTATION LEARNING
Information Prioritization in Visual Model-based RL (InfoPower)
Algorithm I:
For visual model-based RL, need to learn a representation space Z. such that forward dynamies 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 -T;;;; Z--,, At-1). Here, it denotes the dataset indices that de- termine the observations p(o,lid) (ie) 8(or - 0r). In addition to the for- ward dynamics model, we need to learn a reward predictor by maximiz- ing E, T(Re; Ze), such that the agent can plan ahead in the future by rolling forward latent states, without having to execute actions and observe re- wards in the real environment.
Initialize dataset D with random episodes. Initialize model parameters x, 27. Initialize dual variable X. while not converged do for update step c 1..C do
// Model learning
Sample data {(a,, oz, re)) ~ D.
Compute latents 2, ~ Pos 2-,, 0z).
Calculate L based on section 3.4.
(ø. P. n) + (0. (6.x. + x.
À A
// Behavior learning
Rollout latent plan, S 4 SU (2,, a,, rz}
V(a)) = E, [In gn (re (r) + In qo (at-1 22, 2-1)
Update policy ι and value model
Finally, we need to learn an en- coder for encoding observations O to latents Z. Most successful prior works have used reconstruction- loss as natural objective for learning this encoder (Babaeizadch et al., 2017; Hafner et al., 2019b;a). A reconstruction-los. can be moti- vated by considering the objective end I(O, Z) and computing its BA lower
end
// Envi ronment interaction for time step = 0.. - do
21 ~ p ( 1:-1, or): ; or ~ (arl2) Tt,0+11 t env.step (ax).
end Add data D + DU {(on a. rt
bound (Agakov, 2004). [(o;; zx) 2 Ep(or...) (log g (or]==)) + H(p(o:)). The first term here the reconstruction objective, with qar (or12)) being the decoder, and the second term can be ignored as it doesn't '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 consider other lower bounds I(O, Z), can obtain tractable objectives that do not involve reconstructing high-dimensional images. We can obtain an NCE-based lower bound (Hjelm al., 2018): [(o 21) 2 (=) [log fo(ze, o;) log Ee ftz fe or)], where (2 lo:) is the learned encoder, o2 the observation timestep (positive sample), and all observa- tions in the replay buffer oy are negative samples. fo(2,, oy) = exp (oTWozr.) The lower-bound is a form of contrastive learning as it maximizes compatibility of with the corresponding observation or 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 conjunction with empowerment for prioritizing information in visual model-based RL. Similarly, the Nguyen- wainee--aiiiihhh--ooree (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,
[(o 21) 2 Bns(e(lo)ppoo) [f 0z)] - "Enoo (p(o)))) )e (o(2
where fo 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 reconstruction-los. based objective, though both are bounds on mutual
4

---

Published as a conference paper at ICLR 2022
information. We denote lower bound to MI by I(o..2)). 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 pri- oritizing 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 repre- sentation learning as
H-1
Σ 1(On s.t. Σ 70..
max
T(i 1, + I(A I Z.)) co.
The objective to maximize a MI term T(O,; Z+) through contrastive learning such that constraint on C. holds for prioritizing the encoding of forward- pwwerrd-reeiittve, reward-predichte and controllable representations. We define the overall planning objective as
H-1
(AA--; + V(2,) - (()) V((2) R Σ Rr.
T118X
The planning objective is to learn a policy as a function of the latent state such that the empower- ment term and reward-based value term 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 |O] T2 and |2| = d. Since, O corresponds to images and Z is bottlenecked latent representation, n.
I(0, Z) is maximized when Z contains all the information present in O such that Z is suffi- cient statistic of O. However, in practice, this is not possible because [Z] < |O]. When co is sufficiently large, and the constraint E, C. 2 co is satisfied, max >. (O,, Z.) s.t. C. 2 co cannot encode anything else in Z, in particular it cannot encode distractors DS
[S+, S-]. Hence, the objective
To understand the importance of prioritization through C, consider max E. I(Or, Z,) without the constraint. This objective would try to make Z a sufficient statistic O, but since |Z| < [O], there are guarantees about which parts of O are getting encoded in Z. This is because both distractors S and non-distractors S++ DD are equally important with respect to I). 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 each of the MI terms. For T(O, Z) consider 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 C, as follows, with detailed derivations in Appendix A.3:
-I(ig 2(1
Σ
1(rei z)) 2 Ep(raje)) [lo (Te|ze) +((
((1-11;2/2(-11 2 E flog ds (
Here, (2(2))0)) is the observation encoder, qu (a 1 (2. 1-11 is the inverse dynamics model, (n(r1]2) is the reward decoder, and 4x (2/ at-1)) 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 C,, the lower bound to C₁ based on the sum of the terms above. We construct a Lagrangian L E. [(o.:,)) + > (C. cp) and optimize by primal-dual gradient descent. An outline of this is shown in Algorithm
11) is the forward dynamics model. The inverse
5

---

Published as a conference paper at ICLR 2022
Planning Objective. For planning to choose actions at every time-step, we learn a policy π (al2) through value estimates of task reward and the empowerment objective. We learn value estimates with (2+) E [I q, (ra(2e) + In qu (4t-1 12,,, 2t-1]]. We estimate ((+) similar to Equation 6 of (Hafner et al., 2019a), The empowerment term Zv (1- [ 1 in policy learning incentivizes choosing actions - such that they be predicted from consecutive latent states 2 2t. 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 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" (o. a) and the latent MDP (z, a). where 2 ~ (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 outlined in Algorithm
EXPERIMENTS
Through our experiments, we aim to understand the following research questions:
1. How does InfoPower compare with the baselines in terms of episodic returns in environments with explicit background video distractors?
2.
How sample efficient is InfoPowez when the reward signal is weak (< 100k env steps when the learned policy typically doesn't I achieve very high rewards)?
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://sits..oogge
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 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 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.
In addition to this setting, we also perform exper- iments with gray-scale distractors based on videos from the Kinetics dataset (Kay et al., 2017). This setting is adapted exactly based prior works (Fu et al., 2021; Zhang et al., 2020) and we compare di- rectly to the published results.
We compare InfoPower with state-of-the art base- lines that also leam world models for control: Dreamer (Hafner et al., 2019a), C-Dreamer that is a contrastive version of Dreamer similar 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 al., 2019). We also compare variants of In fo Power with different MI bounds (NCE and NWJ), and ablations of it that remove the empowerment objective from policy learning.
Figure Some natural video background distractors in our experiments. The videos change every 50 time-steps. Some backgrounds (for e.g. com-
the top left and the bottom right) have plex distractors the form of agent-behind behind-agenn i.c. background has pre-recorded motion of a similar agent that is being controlled.
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 behavioral similarity metric with details in A.5. Given

---

Published as a conference paper at ICLR 2022
walker walk
cheetsh run
finger spin
quadruped waik
Episodic Raturns
0010
600-
600
400
400
200-
200
Environiment Step
1ve
Environment Step
0.0
0.3
Environment Step
1.0
1.3
2.0
1/8 Environment Step
10e
hopper stand
ball in a cup catch
cartpole balance sparse
quadruped run
Episodic Returns
5000
300
1000
200-
750
200
200
500
100
0.3
0.5
Environment Step
1.0
2.0
Environment Step
led
0.0
Environment Step
0.5
1.5
le6
2,0
0.0
0.5
Environment Step
1.5
2.5
InfoPower
dreamer
cdreamer
ta
dbc
deopmdp
Figure 4: Evaluation of InfoPower and baselines in a suite of DeepMind Control tasks with natural vidco 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.
the sets S₂- = [= and S.y { 11 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 unique integer and corresponding vertices in the two graphs are labelled similarly.
Let the resulting graphs be denoted G- (Vz, F.)) and Gzp. (Vagr F=g)) respectively. Note that both these graphs are shortest path graphs, by construction. Let, e {u, U/} and e, {uz, 2y): We define a kernel à that measures similarity between the edges c,, eg and the labels on respective vertices. k(ei,es)
ke ((u,) (1(e)), (((o), + ky (l(u,) (((u)),(w))))) ((yy)) (I(e,), ((e)))
Here, the function 2(-) denotes the labels of vertices and the weights of edges. ko a Dirac delta 8(x, v) (Note that 8(2,y)
1 iff Brownian ridge kernel, (x, v) = max(0, c - j - yl). where e is a large number. We define the shortest path kernel to measure similarity between the two graphs as, k(G,, Gzy)
kernel i.e. he ke y)
y. else 8(x, v) = 0) and ke
F.. EeEE, EejEB,,, e(),, The value of k(G,, G...) 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
1
positions of the joints) have higher values according to this metric.
4.3 SAMPLE EFFICIENCY ANALYSIS
In Fig. 4, we compare In InfoPower ower and the baselines in terms of episodic returns. This version of foPowrr corresponds to 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 IM steps of training quickly compared to the baselines, indicating higher sample efficiency. This suggests the effectiveness of the empowerment model in helping capture controllable represen- tations early on during training, when the agent doesn't take on actions that yield very high rewards.
4.4 BEHAAIIRAA 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 simi- larity of latent states (Sim), based on the metric in section 4.2. We see that the value of Sim for InfoPower r is around 20% higher than the most competitive baseline, indicating high behavioral similarity of the latent states with respect to the corresponding ground-trum -tunt simulator states.
Qualitative visualizations with t-sne, Fig. 5 shows t-SNE plot of latent states {s(2)o) for EnfoPower 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 Infooowee r, although the
7

---

Published as a conference paper at ICLR 2022
InfoPower
Dreamer
C-Dreamer
Figure 5: (-SNE plot of latent states (z(o)) o) with visualizations of three nearest neighbors for two ran- domly sampled points (in red frame). We that the state of the agent similar is each for InfoPowe whereas for Dreamer, and the most competitive baseline C-Dreamer, the nearest neighbor frames have signifi- cantly different agent configurations.
Wab welliter walk
finger spin
quadruped
1000
600
800
250
400
500
400
250
Step
1eb
Environient Step
0.0
0.5
Environment Stop
1.0
5.5
2.0 Lo6
0.0
0.5
Environmmnt Step
1.0
1.3
2.0
hopper stand
cartpole balance sparse
ball cup catch
400
300
300
Egtsadic
200
100
100
300
0.
Enutronment Stap
0.5
1.0
1.5 2.0
0.0
Environment Step
0.5
100
1.5
lee
2.0
Environmen: Step
1e6
Envtronment Stop
0.5
1.3
2.0
108
InfoPower
InfoPower NCE
InfoPower NWI Policy
InfoPower NCE Pooiic
Figure 6: Evaluation of In fooore ablated variants in suite 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 In Infooower-NCC are full versions of our method differing only in the MI lower bound. The versions with Policy do include the empowerment objective in policy learning. but only use it from representation learning. The
over random seeds. Higher better.
background scenes are significantly different. However, for the baselines, the nearest neighbor states 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 Inooooeer. Keeping everything else the same, and changing only the MI lower bound to NCE, we see that the performance almost similar or slightly worse. However, when remove the empowerment objective from policy optimization (the ver- sions with 'PPliccy' the plot), we see that performance drops. The drop is significant 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.
8

---

Published as a conference paper at ICLR 2022
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 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 al., 2019; Gregor et al., 2019).
MI for representation learning. Mutual Information measures the dependence between two ran- dom variables. The task of learning latent representations Z from images for downstream appli- cations, has been very successful with MI objectives of the form fa T(fi (O), f2(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-boume 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 behav- iors of the algorithm. We incorporate MI maximization through the NCE (Hjelm al., 2018) and NWJ (Nguyen et al., 2010) lower bounds, such that typical reconstruction objectives for representa- tion learning which do not perform well with visual distractors, can be avoided.
Inverse models and empowerment. Prior approaches have used inverse dynamics models for reg- ularization in representation learning (Agrawal et al., 2016; Zhang et al., 2018) and 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 predic- tive 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 learn- ing (Leibfried 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 empower- ment both for state space representation learning and policy learning, in visual model-basee 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 ab- stracting out only the task relevant information from the environment when there are explicit dis- tractors (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 ob- jective. Our approach 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 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) was and earlier approach that 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 environ- ment 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 approach.
6 CONCLUSION
In this paper we derived approach for visual model-based RL such that an agent can learn latent state-space model and a policy by explicitly prioritizing the encoding of functionally relevant factors. Our prioritized information objective integrates term inspired by variational empowerment into 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-att visual model-based RL methods, we observe higher sample efficiency and episodic returns across different environments and distractor settings.

---

Published as a conference paper at ICLR 2022
REFERENCES
David Barber Felix Agakov. The im algorithm: variational approach to information maximization. Advances in neural information processing systems, 16(320):201, 2004.
Pulkit Agrawal, Ashvin Nair, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Learning to poke by poking: Experiential learning of intuitive physies. arXiv preprint (rXxv:/606.074199 2016.
Mohammad Babacizadch, Chelsea Finn, Dumitru Erhan, Roy H Campbell, and Sergey Levine.
Stochastic variational video prediction. arXiv preprint arXiv:1710.1122 2017.
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.
Richard Blahut. Computation of channel capacity and rate-distortion functions. IEEE transactions on Information Theory, 18(4):460-473, 1972.
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.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint arXiv: /802000002 2018.
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.
Xiang Fu, Ge Yang. Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, pp. 3480-3491. PMLR, 2021.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pp. 2170-2179, PMLR, 2019.
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. arXiv preprint 161 11/6611..770 1.07507, 2016.
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: arXiv:1906.000 096.799377 2019.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint
2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 2019a.
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.
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmen- tation. arXiv preprint arXiv:201 133389 2020.
Devon Hjelm and Philip Bachman. Representation learning with video deep infomax. arXiv preprint arXiv:2007.1322 2020.
Devon Hjelm, Alex Fedorov, Samuel Moroii-Marchhldon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. arXiv preprint IrXiv:I800s 0607000670, 2018.
Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijaya- narasimhan, Fabio Viola. Tim Green, Trevor Back, Paul Natsev, et al. The kinetics human action video dataset. arXiv preprint arXiv: 7777000.0950 2017.
10

---

Published as a conference paper at ICLR 2022
Alexander S Klyubin, Daniel Polani, and Chrystopher L Nehaniv. Keep your options open: information- based driving principle for sensorimotor systems. PloS one, 3(12):e4018, 2008
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint urXiv:2004.13643 2020.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. arXiv preprint arXiv:2004. /4990, 2020.
Alex X Lee, Anusha Nagabandi, Picter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. arXiv preprint arXiv: 19077009970 2019.
Kuang-Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in rl. arXiv preprint urxiv:2007.12401, 2020.
Felix Leibfried, Sergio Pascual-Dias Diaz, and Jordi Grau-Moya. A unified bellman optimality principle combining reward maximization and empowerment. arXiv preprint arXiv:1907.123999 2019.
Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee. Contrastive variational model-based reinforce- ment learning for complex observations. arXiv e-prints, pp. arXiv-2008, 2020.
Shakir Mohamed and Danilo Jimenez Rezende. Variational information maximisation for intrinsi- cally motivated reinforcement learning. arXiv preprint (rXXvv:509.087/) 2015.
Ofir Nachum and Mengjiao Yang. Provable representation learning for imitation with contrastive fourier features. arXiv preprint arXiv:2105. /2272, 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.
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.
XuanLong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11):5877-5861, 2010.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding, arXiv preprint arXiv: 11XX:::8077.0374: 2018.
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-2777 2787. PMLR, 2017.
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.
Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in reinforcement learning. 2020.
Kate Rakelly, Abhishek Gupta, Carlos Florensa, and Sergey Levine. Which mutualiinformation representation learning objectives are sufficient for control? arXiv preprint arXiv:2106.07278, 2021.
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.
Archit Sharma, Shixiang Gu, Sergey Levine. Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. arXiv preprint arXiv: /07.006707006 2019.

---

Published as a conference paper at ICLR 2022
Evan Shelhamer, Parsa Mahmoudich, Max Argus, and Trevor Darrell. Loss is its reward: Self- supervision for reinforcement learning. arXiv preprint 1rXiv:1612.073077 2016.
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.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. arXiv preprint urXiv:2004.041:0 2020.
Susanne Still. Information -hhooreticc approach to interactive learning. EPL (Europhysics Letters), 85(2):2005, 2009.
Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160-163, 1991.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite, arXiv 2018.
preprint 10Xiv:8801.000
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. In Computer Vision-ECCH 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XI 16. pp. 776-794. Springer, 2020.
Michael Tschannen, Josip Djolonga, Paul Rubenstein, Sylvain Gelly, and Mario Lucie. On mutual information maximization for representation learning. arXiv preprint arXiv: 2019.
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.
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.
Amy Zhang, Harsh Satija, and Joelle Pineau. Decoupling dynamics and reward for transfer learning. arXiv preprint arXiv:1804. 1804.1069,, 2018.
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.
12

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1 THEORETICAL RESULTS ON CONTROLLABLEE REPRESENTATIONS
Let (-)) denote the encoder such that Z ((O). 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 directly influenced by actions A i.e. S+ I 5 s.t. 1( S,15,) 0.
We next describe two theorems regarding learning controllable representations, with proofs in the Appendix. We observe that the max E, (A Ze(Z,-1) objective alone for learning latent repre- sentations Z, along with the planning objective provably recovers controllable parts of the observa- tion O, namely S
Theorem 1. The objective max Ee 1; Z(-1) provably recovers controllable parts S* of the observation O. S defined as that part of underlying state S which is directly influenced by actions A i.e. S+ C S 8.1. [Sr; S O.
Proof. Based on Fig. 2 and Fig. 7, we have the following
max max Z,) $ max (; Procesiog Inequality
((A--1;))) Proc-Procsssing Inequality
$ (1,
$ (AAc-1; + 1(A Chain Rule of MI
- 1(A,11;5;) Conditions tonal i independence
So, the maximum above be obtained when the encoder 4 is identity function over S and zero function over ST]. So, Z, S Hence, given -,-1, I(A,-;; Z, Z.-1) is maximized when an identity function over 5+ So, max Σ. 1(; Zi (1)) learns the encoder such
that controllable representations recovered.
This result is important because in practice, we may not be able represent every possible factor of variation in a complex environment. In this situation, we would expect that when |2| < 10), learning Z under the objective max Σ, [(A,,1; Ze| ((--) would encode St. We next show that the inverse information objective alone can be used to train a latent-state space model and policy through an alternating optimization algorithm that converges to local minimum of the objective max >, [(A,i; Z, 2z-1) at a rate inversely proportional to the number of iterations.
Theorem max ab E. 1(A--;; = E.
log golar 1121,21-11 ( 118e-1)
can be optimized through an alternating optimization scheme that has convergence rate of ((/NN) local minima of the objective, where N the number of iterations.
Proof. Let (2(2))) denote a latent state sampled from the encoder distribution. Learning world model such that reward prediction and inverse information are maximized can be summarized
log Zo (a:a (4,1
max 0.1.11)
Σ
Br(aa (x1lo))))))))
Eptrikana (rela a.) +
T(ae(2)))
The policy x (atlze) (2x) is learned by
gr (a 11 22)
max Σ E
E Ep(rizead log q (re at) + Eptz (l log
π T (2))
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(ra)=)), we have the following objective
ge (ag 2++1, π (a,
max
E (aa (l..)) log
1.V
13

---

Published as a conference paper at ICLR 2022
A-11
A++K-1
Sitk
X
S
Sith
St
Ri+k
Si+k
R
Ot
Ou+
Ze
Zetk
Figure 7: PGM of the MDP with distractor states. The state observed by the agent O consists of three parts S S and DS S* is the controllable part of the state i.e. is affected by the actions of the agent, and in turn affects the reward R: Š- not controllable by the agent doesn't affect reward R future S
controllable by the agent but affects the reward and future S+ DS
The optimal ghe give a fixed and & can be derived as in page 335 of Cover & Thomas (2006).
(a 2.) Ja ds (z (: [ az ( (a)
The optimal т* given a fixed y and & can derived as
exp
[log qo (ar| (x)))
x >' 24) Ja exp (aslss)) (a [log qu (ar (,)11,
We note that these alternating iterations correspond to an instantiation of the Blahut-Arimods al- gorithm (Blahut, 1972) for determining the information theoretic capacity of a channel. Based on (Nakagawa et al., 2021) see that this iterative optimization procedure between g.. and x converges, and the worst-case rate of convergence is O(1/N) where N is the number of iterations.
This result is useful because it shows that in the absence of rewards from the environment (rz = 0 Vi), when planning to minimize regret is not possible, the inverse information objective can be used to train policy that explores seek 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.
A.2 PRIMAL DUAL UPDATES
8 ~ 1(0,; Z,) + + Z,,, I(Ar-1: Z (R; Ze) - -o)
Zo.H-1 Zo.H-1t
0Zo.H-1
14

---

Published as a conference paper at ICLR 2022
(-liis 1,,
-
IAA-1 2.-1)
[RR;; 2.) -co rew model
fwd dynamics model
inv dynamics model
A.3 FORWARD, INVERSE, AND REWARD MODELS
pl 2 1,-1,,0-1,0
-I(ie Z1 Ac-1)=
pl 2 2-
/ 01,,,211,001111 log
p 2t 1,-1, 0z)
4x (21
log
9x (26 21-1, 0
4-11
2 f p( 20..2-1.-111111 log
1
9x (2) 2 2-11 a-11)
DeL (p(2) 21-1,
(z (11, 4-11)
We can obtain variational lower bound on the empowerment
p(ae|2e+k:21)
[ay 21)
log
p(ae(z))
f pla,, 2) log g(ael204k:20) 21)
+ log p(ae 21)
p(arl2))
g(ar| 21)
f p(a,,244x 21) log ((a ((412))
p(ar12))
log / 2og 21) - ((2e)p(=e+x)) la 21 2)p(as))) logp(a))<))
Ep a. [log g(0/) (1+x 2)] + (e ((p(((=()))))
Here, 2 2)) is the inverse dynamies model, and H(p(a, aa) the entropy of the policy.
We can obtain variational lower bound on the reward model as p(relza)
/
I(r 21)
p(re, 21) log
p(r))
/ 2ire log g(ralze) p(re)
log p(re12t) q(re)21)
/ 2()))2)) log
(1e((e)
p(re)
/
p(r., log ((re(ze)
p(r.)
/ p(re, log g(rel2)) - ff plre 2)) log p(r))
Ep(,. [log ((ra|2))] + Constant
Here, q(r: (2)) the reward decoder. Since the reward is scalar, reconstructing is computation- ally 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" a) and the latent MDP (= a). where 2 ~ 9s (z|o) is bounded, under some regularity assumptions. Similar to the assump- tion in (Gelada et al., 2019), let the policy æ have bounded semi-norm value functions under the total variation distance Drv i.e. V (z) K and VV
(Pv( Define the forward dynamics learning objective as miny Lr = min, Dxa. (p(o (o, a) q (2112, and the reward model objective as min, I.n min, DRL. (p(rlo) (r(()))
15

---

Published as a conference paper at ICLR 2022
Theorem 3. The difference between the Q-function in the latent MDP and that in the original MDP bounded by the loss in reward predictor and forward dynamics model.
+ vbn+gKvIn Lrr
(lo.a) (0.0) (0.0))
(e0 (o, (z)
1 1-7
Proof.
(z) - I -1r.::/:--rrrooooo--------------
$ +a
+ Y /ote +
V(z')l)
+ (4a.
Ln + y (oa)) or
V(('))
+ K v DKi (p(o 10x
/ Ln + (o. (,,)), (V(o) - (()))) LKVLT
S V + + +7 (o0)))))) (o)o)
a(), Q(: +7KLr
So, we have shown that
(z,,a) S vK
E (ooa))/dttoo))
We can see that, as LR, LT + 0, the Q-function in the representation space of the MDP, Q, becomes increasingly closer to that of the original MDP, Q. Since this result holds for all O-functions, it also holds for the optimal Q' and Q*. This result extends sufficiency results in prior works Rakelly al. (2021) to a setting with stochastic encoders and KL--divergence losses on forward dynamics and reward, as opposed to Wasserstein metries (Gelada al., 2019), and bisimulation metrics (Zhang et al., 2020).
A.5 BEHAVIORAL SIMILARITY METRIC DETAILS
Given the sets S. = {2 and S.y = {2 (ig 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 G2 = (Vz, Ez) and Gz, = (Vay. Exg.) 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,, Gz,)
1 CEEF, ef
1
Σ
Σ
Σ
k(e, ez)
EE
Let, e, {u, e and e, = {u,, fu,, Here, the kernel & measures similarity between the edges e, and the labels on respective vertices. Let the function 1(-) denote the labels of vertices and the weights of edges.
k((e,e))) ko (aleee/ahhaaaeeaeeeeeealealclaad
16

---

Published as a conference paper at ICLR 2022
Here, K, is Dirac delta kernel i.e. (z. y) = 1 8(x, y) (Note that 8(x,)) iff x y, else ((, y) = 0) and be is a Brownian ridge kernel, lig (z. y) max(0, c- -2 -y)), where c is large number.
So, in summary, the value of k(e;, e;) is low when the edge lengths between corresponding pairs of vertices in both the graphs are close. Hence, the value of k( Gzg)) ) is low when 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 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, 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 background videos with a frequency of I in very 3 videos.
ggent-behind-ung
For the experiments in Table I, 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 al., 2020), and TIA (Fu 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 64 x 3, action dimensions range from 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 and use a single Nvidia V100 GPU and 10 CPU cores for each training run. The training time for 1еб environment steps is hours on average. In comparison, the average training time for leó 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 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 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 fo 0r) exp (27 Woze) ) for contrastive learning is learned with 200 x 200 Wa matrix. Here, is the encoding of ov. 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 co 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 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-paramette parameters of InfoPowez same across all environments and distractor types.
For fair comparisons, we kept all the common syper-parameter values same as Dreamer (Hafner et al., 2019a). This the protocol followed in prior works TIA (Fu et al., 2021) and Ma et al.
17

---

Published as a conference paper at ICLR 2022
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- seattt-ehind-ggeee i.e. agent of similar morphology moves in the back- ground of the agent heing controlled.
Figure 9: Illustration of the different levels of distractors. L1, L2, and L3 respectively have distractor windows size 32x32, 40x40, and 64x64.
(2020). For the baseline TIA, we tuned the environment-speecfee parameters AR.a- and λo as mentioned in Table 2 of (Fu et al., 2021). For AR... we performed a gridsearch over the range 104 - 50k for λo, we performed a gridsearch over the range - to obtain the parame- ters for best performance, which we used for the plots in Fig. 4 and Fig. 11. Respectively for walker-walk, cheetah-run, finger-spin, quadruped-wald hopper-starn stand, aull-i--a-cup-atchh. cartpole- balance-sparse, quadruped-run, the values of λRaáv are 200,K,200.04033K4440000 ues of λo, are 2,2,1 1,2.5,2,,,222,,,,,2222
The val-
For baseline DBC (Zhang et al., 2020), kept all the parameters because DBC does not have any environment-specde specific parameters and the same values were used for all environments in the DBC paper. The DeepMDP agent and its hyperparameters the implementation provided the github repo of DBC (Zhang et al., 2020).
as in Table of the paper,
adapted from
A.8 CONTRASTIVE LEARNING
Since the distracting backgrounds changing videos of 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 learning invariance across different videos. Concretely, we sample batches from the dataset (4) { (at, o,, r where denotes a batch index, and for each observation o, in batch i, all considered negative samples.
1 observations ou both in batch i and in other batches
A.9 RESULT ON VARYING AISTRACTOR LEVELS
We consider different levels distractors by varying the size of the window where distractors in the background are active. Fig. 9 in the Appendix illustrates this visually. Table I shows results for
18

---

Published as a conference paper at ICLR 2022
Table I: 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 training, Higher better for both Rew and Sim.
Name
Levels
Rew @ 500k
Rew @ 1M
Sim@500k
Sim @ 1M
Dreamer
197 + 31
240 + 27
0.73=0.02
0.74=0.01
DBC
L1
261 + 25
390 = 34
0.75=0.03
0.74=0.02
C-Dreamer
291 34
590 + 26
0.73#0.02
0.79=0.01
DeepMDDP
L1
263 31
340 + 22
0.74+0.02
0.75+0.03
InfoPower
22
650 100
0.82 0.01
0.84 + 0.03
Dreamer
L22
180 + 36
197 + 20
0.56+#003
0.59=0.02
DBC
L2
221 28
320 + 33
0.65=0.02
0.66=0.02
C-Dreamer
L2
282
30
550 + 66
0.63
0.71=0.02
DeepMDP
L2
213 34
300 + 25
0.59
0.61=0.04
InfoPower
L2
394 + 30
644 + 101
0.77 0.03
0.77 0.03
Dreamer
L33
140 + 52
157 + 10
0.32+0.02
0.33
DBC
165 20
221 + 24
0.45+0.01
0.48+0.01
C-Dreamer
L3
231 39
485
+ 86
0.58=0.02
0.64=0.01
DeepMDD
L3
153 23
212
28
0.4410.02
0.49=0.01
InfoPower
L3
389 - 20
624
80
0.71 0.011
0.74 + 0.03
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
3891:45
602-40
165=92
2101:32
4171:22
630.133
Cheetah-Run
384:1149
57911172
261=644
277181
425:197
552
Hopper-Stand
338::221
581::214
110=92
207+:120
395+:140
615:176
Ball-in-a-Cup
201::197
115+110
255+:104
263+:116
Walker-Wall
878:1:37
948-28
545
630-168
925+118
972-226
Finger-Spin 371-1966
487::107
801-10
832.14
795+-34
787+-58
Hopper-Hop 47144
72:68
51122
0.:0
851:26
77174
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 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 i.. same size the observations 64x64.
A. RESULTS ON DISTRACTORS FROM THE KINETICS DATASET
In this section, we evaluate InfoPowed over 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 Fig. 4 because the distractors being grayscale images, and not RGB might be easier to ignore by the basclines as well as InfoP foor....
A.11 SETTING WITH ONLY SHIFTED AGENT BEHIND AGENT DISTRACTORS
In this section we evaluate InfoPower with the base- lines only for the challenging setting of agent-behind- agent distrators where the background contains the same agent being controlled (a walker cheetah). An il- lustration of this shown Fig. 10. From the results in Table 3, that Info significantly out-
19
Figure 10: Ilustration of the distractor set-
ting for results in Table 3.

---

Published as a conference paper at ICLR 2022
performs the competitive baselines both 500k and 1M environment interactions. This suggests the utility of InfoPower in separating out informative elements of the observation from uninformative ones even in chal- lenging settings.
Table 3: Evaluation of InfoPower and the most competitive baselines in suite of challenging distractors where the background contains a shifted version of the agent being controlled. Results are averaged over random seeds. Fig. 10 shows visualization the distractors for this setting.
Walker- walker-iakkshhifee
500k
1M
500k
1M
DBC
35:15
104428
41:10
78:138
CDreamer
108++36
1644-45
73+444
110+-52
TIA
79++42
152+:33
65++24
115++28
InfoPower
168+488
278:1:42
130-:46
207+:35
20

---

Published as a conference paper at ICLR 2022
cheetah run
spin
quadiuueed walk
Eptsotcc Retuime
701
509
400
400
250
200
200
200
0.0
0.5
Environment Step
1.0
1.5
2.0
Environment
1-5
2.0
Eriironment Step
LV
le0
0.0
0.5
Envirooment Step
1.0
1.5
2.0
14b
hopper stand
cartpole sparse
quadruped run
200
300
1000
200
200
200
Eptcodic
200
500
300
0.0
0.5
Environment Step
1.0
1.5
2.0 le6
Environientt Step
1.0
le0
Environment
2.0
Environment Step
InfoPower
dreamer
cdreamer
tla
dbc
Figure 11: Evaluation Inf ower and baselines with empowerment in policy learning. The setting cor- responds to DeepMind Control tasks with RGB natural video distractors in the background (same as that in Figure of the main paper). The x-axis denotes the number of environment interactions and the y-axis shows the episodic returns.
A.12 COMPARISON OF BASELINES WITH EMPOWERMENT IN POLICY LEARNING
In this section, compare against modified versions of the baselines, such that include empow- erment in the policy learning of the baselines. This 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 RI. (by modifying both the representation learning objective and the policy learning objective) is important contribution of paper. The aim of this exper- iment is to disentangle the benefits of empowerment in representation learning, for InfoP ower 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. 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.
A.13 COMPARISON ON ENVIRONMENTS WITHOUT DISTRACTORS
We compare the performance of InfoPower with state-of-the-arr visual model-based RL method, Dreamer in default DM Control environments without distractors. We see from Table that InfoPower performs slightly better 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 In fowor in challenging distracting environments does come cost of per- formance in simpler non-distracting environments.
Table 4: Comparison of InfoPower with state-of-the-art visual model-hased RI. method, Dreamer in the default DM Control environments without distractors. tabulate reward 500k and IM environment inter- actions. Results are averaged over random seeds. Higher is better.
Dreamer
InfoPower
Rew @ 500k
Rew @ IM
Rew @ 500k
1M
Walker Walk
890-152
985 10
910:441
980 17
Cheetah Run
180
811195
586-1142
810:1102
Finger Spin
570+188
573196
590-1172
598.-85
Quadruped Walk
5924-45
646+50
623:165
666-147
Hopper Stand
701+68
9061:33
715-160
930-139
Ball in Cup
871+102
910++48
913-167
939--22
Part-Pole Sparse
742+-58
849+88
776+440
866 80
Quadruped Run
450++47
515+66
468--42
510+61
21

---

Published as a conference paper at ICLR 2022
A.14 ADDITIONAL ABLATIONS
In this section, we consider additional ablations, namely a version of of InfoPower that does not have the empowerment objective, version of InfoPower that does not have the MI maximization objective. We see from Table that removing either of these terms drops performance.
Table 5: DM Walker Stand with natural video distractors
Name
Obs bound
Rew Rew@500k
Rew @ 1M
Sim@@ 500k
Sim @ 1M
Dreamer
BA
140 = 52
157 + 10
0.32+0002
0.33:5.008
Contrastive- Dreamer
NCE
231 = 39
485 + 86
0.58
0.64#0.01
NonGenerative- Dreamer
NWJ
280 =
480 + 93
0.59
0.62::0.05
InfoPower
NCE
254 + 27
481 + 30
0.69
0.72++0.02
InfoPower
NWJ
389 + 20
624 + 80
0.71+0.01
0.74+H0.03
Inverse only
None
282 + 70
367 24
0.6310.04
0.66100.03
22


Reviews:
[{"review_id": "oCNjjwiX9e", "reviewer": "Reviewer_sJgS", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n- I like the overall idea. It's not too complicated and the results seem very promising.\n- With one exception, the illustrations and figures are helpful.\n\nWeaknesses:\n\n- W1. The paper is structured badly. It feels like the whole paper was written, it came out as 11 pages and hours before the deadline, hastily important sections were moved to the appendix. For example there are 2 theorems in the main paper but no proof. Why format them as \"theorems\" at all in the main body of the paper if there's no proof right there? I'd just trim this down, only include the description of these in prose and move the \"theorem\" and proof both into the appendix. I also felt like the theory part in the main body of the paper was a bit shallow because it's long and verbose but doesn't feel complete and as easy to follow as I would have expected based on how much space it takes up. Another example of this is section 4.2 where you \"introduce\" a new similarity metric but never explain it in the main body of the paper. Figure 4 doesn't make any sense without consulting the appendix. This needs to go back into the main body of the paper!\n- W2. The writing could have definitely used a proofreading. There are tons of errors in the main doc and a lot more in the appendix. Also some expressions are a bit handwavy and could have used a bit more rigorous writing. E.g. in the intro, you wrote \"... modeling the observation space so as to most quickly capture the functionality relevant factors... \" - (a) You usually don't model the observation space but \"states\" or a \"latent representation\" that's inferred from the observations; (b) \"quickly\"? Why quickly? That's not really explained; (c) \"functionality relevant factors\" - you probably mean \"functionally-relevant factors\", right? (d), functionally relevant to what? To a task, to control, to an understanding of the robot's odometry? You have a good example of this later in Figure 2 but here, without any context, this is too imprecise language.\n- W3. Fundamentally, I think to show your assumption properly, you should've included a task where the observations consists quantifiably of those 3 different factors: S+: images of the robot itself that can be affected by actions, S-: elements that can't be controlled by the agent but influences the reward, DS-: distractors that don't affect the reward. In your current setup, there's only S+ and DS- because there is a random video playing in the background. How about this: Include obstacles on the ground (e.g. in front of halfcheetah and walker) that obviously cannot be affected by actions but would make the agent stumble if ignored. Or alternatively or additionally, instead of a random video, include a video of the agent itself but shifted up or down. These would be a harder but more informative test if your method can actually learn to separate out informative state elements from uninformative ones.\n- W4. Why did you include a model-based version of DBC and not DBC itself? If you compare the results from Fig.3 in the DBC paper with your results in Fig. 5, I can see that DBC is supposed to be significantly better than this. The whole purpose of DBC is to remove distractions from the observations when learning a policy and it is already very sample-efficient (it can reach the performance of your model in 1e5 steps instead of 1e6 or 2e6). And speaking of which, in the appendix you mention training for 10e6 steps. Why didn't you include these results? What made you pick 1e6/2e6 for comparison and not 1e5 or 1e7? Feels a bit cherry-picky.\n\nNitpicks & Questions:\n\n- Q1. Section 4.4 \"the window size for distractors\" is not explained.\n- Q2. Spacing of Fig.3's caption needs to be fixed.\n- Q3. The algorithm 1 block needs to be improved a lot. E.g. you initialize dataset D but not the policy and not lambda? Also what is lambda, why is it important there? Why is the environments not stepped until T but only until T-1? What is the line below \"Compute latents\" refer to? Update your model parameters? If that's the case then you need to specify that these are your model parameters. Where is the loss calculated? Computing the latents is not inherently also calculating the loss.", "summary_of_the_paper": "The paper introduces a variation on the DREAMER model-based RL architecture that aims to remove elements from the latent state representation that cannot be affected by the actor. To realize this, there is an inverse model objective included in the value function.", "summary_of_the_review": "UPDATE (2021-11-21): The authors have incorporated all of the suggested changes and addressed my concerns reasonably, which is why I'm now recommending acceptance.\n\nOLD: The paper is overall okay. I don't think it's remotely in its best possible shape (in terms of structure, writing, Figure 4, etc.) but this is nothing that can't be solved with a couple hours editing and rearranging things. The main idea is good, nothing in the paper appears to me a major flaw and so I'm mildly recommending acceptance and strongly recommend polishing!\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "dxv8ziaS8cp", "reviewer": "Reviewer_Kbjg", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nI like the idea of the paper in that it appears to be logical and\nalmost obvious, after reading about it here. The solution elegantly\naddresses the problems of representation, and policy learning, using\nthe same formalism. I also like the information-theoretic view as a general\nformulation independent of a specific model.\n\nResults appear to be very good on the chosen benchmarks. Results of\nthe ablation (appendix) and comparison without distractors is useful\nto see, and the description of the distractors at least give a good\nidea (but see below).\n\n\nWeaknesses:\n\nFrom the results table (in the main paper) it appears InfoPower comes\nwith much higher variance than other approaches. For results without\ndistractors (in the appendix), this is almost reversed. The variance\n(and the differences) is not much discussed in the paper, but it may\nbe worthwhile to.\n\nI cannot see the results being easily reproduced for a lack of detail\n/ availability on the set of distractors, and was wondering if future\nwork and comparisons could be helped by data and/or code allowing\nsimilar distractors as those used in the paper.\n\n\nOther feedback:\n\nWhile the approach and solution to the problem is new as far as I\nknow, relevant ideas have been discussed much earlier in a paper by S\nStill, Information theoretic approach to interactive learning, EPL 85,\n2009 https://arxiv.org/abs/0709.1948 that would be worth referencing.\nI'd like to suggest including a brief discussion of the differences\nand objectives of both ideas.\n\nIt might be useful to mention the relative difference in computation\nof the approaches (on similar hardware). Compute times are given in\nthe paper and useful.\n\nTypos etc\n- p2: problem statement section: \"captial\" -> capital\n- p4 bottom: \"summmarise\"\n- In table 1, should the C-Dreamer L2 Rew@1M be bolded (because of the overlap / variances) (as in the level 1\\\n case); the semantics of \"bold\" isn't explained so it's difficult to say.", "summary_of_the_paper": "The paper introduces a method to learn representations for model-based\nRL. The key idea is to find a representation that maximises\ninformation between the action at the previous step and the current\nlatent representation, given the representation at the previous step,\nthereby maximising the information that the representation encodes\nabout effects of an action. The work contributes a method and\nevaluation on tasks where high-dimensional visual inputs are used to\ncontrol an agent in an environment with complex backgrounds as\ndistractors.\n\n", "summary_of_the_review": "The paper is well written and motivated. I find the results compelling\nand useful. There is not much to complain, except it would be good to\nexplore differences in variance of the different methods. Minor issues\nare related to potential reproducibility / experimental setup, and a\nreference to prior work.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "RDcMcuV0Kgq", "reviewer": "Reviewer_8ezz", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I appreciate the idea of using the empowerment objective for both representation learning and policy learning. The agent-behind-agent distractors are also novel and more challenging than previous work. However, I have some concerns regarding the experiments and the clarity of the paper.\n- While most of the baselines were tested in previous work using the same set of videos from the Kinetics dataset as background distractions, this paper uses a different set of videos without providing enough details. For example, how many videos are there? Are the videos obtained from some specific dataset? Do training and evaluation use the same set of videos? This makes it hard to compare with published results. In particular, some baselines require quite extensive hyperparameter tuning depending on the dataset, and TIA even requires separate hyperparameters for each domain in DMC. The paper did not include any detail about how the baselines were tuned. From the experiments, I do not see the need to change the background videos from previous work, so I think it would be better to compare with the baselines on the same dataset used in DBC and TIA. As for the agent-behind-agent distractors, I think it would be better to have a separate dataset focusing exclusively on this challenging setting, rather than mixing the challenging cases with standard ones.\n- The comparison to baselines seems unfair, as the baselines do not use empowerment for policy learning. In fact, the ablation study in Figure 7 shows that removing the empowerment for policy learning significantly decreases performance of the proposed model, even on dense reward tasks like Walker Walk, where the ablated version underperforms C-Dreamer and TIA (compare to Figure 5). The paper can be strengthened by adding empowerment to the baselines.\n- The paper claims to \"outperform state-of-the-art model-based RL approaches by an average of 20% in terms of episodic returns at 1M environment interactions with 30% higher sample efficiency at 100k interactions\". I do not see how these numbers are obtained.\n- In Appendix A.10 Table 2, the proposed model achieves $1026 \\pm 21$ on Walker Walk without distractions. This seems impossible, and undermines the credibility of the reported results.\n- I find the paper a bit hard to follow. At the beginning of Section 3.1, I did not understand the relation between controllable representations and the conditional independence $I(A_{t-1}; S_t \\mid S_t^+)=0$. Is this a necessary and/or sufficient condition and why? Also the PGM in Figure 2 may need more justification/explanation. Given that both $S_t^+$ and $S_t^-$ affect the reward and transition, is it reasonable to only learn $S_t^+$ and discard $S_t^-$?\n\n**MINOR COMMENTS**\n- The paper can be strengthened by comparing to TPC [1], a contrastive approach that outperforms DBC and CVRL.\n- The total number of steps are different in Figure 5 and 7. It would be better to use the same number of steps for all tasks.\n- What distractor levels are used in Figure 5 and 7?\n\n[1] Temporal Predictive Coding For Model-Based Planning In Latent Space. Nguyen et al., ICML 2021.", "summary_of_the_paper": "This paper tackles the problem of prioritizing functionally relevant information from complex observations for model-based RL. To this end, previous work has proposed to replace the reconstruction loss with contrastive loss. Building on that, this paper introduces an additional empowerment objective that is used for both representation learning and policy learning. Experiments show that the proposed model outperforms baselines on a set of DeepMind Control tasks with custom background distractions, which include other visually similar but uncontrollable agents. It is also shown that the similarity between learned states can match well to the similarity between groundtruth simulator states, according to the proposed metric.", "summary_of_the_review": "I am leaning toward reject, due to potentially unfair comparisons and incorrect results.\n\n---\n\nMy concerns have been adequately addressed during rebuttal, so I now recommend accept.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "5gilzf8M3I7", "reviewer": "Reviewer_EU46", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "#### **Strengths**\n- The paper tackles a very important problem in the topic of learning to control from visual observations. In real-life scenarios, the observations are likely to contain distractors or task-irrelevant information, and being able to focus on relevant aspects of the environment in those cases is crucial.\n- The proposed use of information empowerment for learning representations and exploration is novel to the best of my knowledge.\n- The experimental results look promising.\n\n#### **Weaknesses**\n- The first major issue of the paper is the lack of details in the presentation of the practical algorithm. Specifically, it is not clear from the text what components/networks constitute the InfoPower model. In Section 3.4, the authors derived lower bounds of MI terms that will be used to optimize the model. However, it is very confusing to read because the authors have not specified what $p$, $q_{\\mathcal{X}}$, $q_{\\eta}$, and $q_{\\psi}$ are and how they are parameterized. Moreover, these lower bounds also need to be derived in more detail as they are not straightforward from the text. In Appendix A.8, the authors presented the training and network details but again it does not provide enough information to fully understand the components of InfoPower.\n- The second major issue lies in the overall objective function, as there are multiple terms being optimized together. First of all, why do we need to maximize $I(O_t; Z_t)$ if we are already minimizing $I(i_t; Z_t \\mid Z_{t-1}, A_{t-1})$, which encourages encoding the predictable parts of the observations? In previous works [1,2,3,4] people have found that maximizing $I(Z_t; Z_{t-1}, A_{t-1})$, which is similar to minimizing $I(i_t; Z_t \\mid Z_{t-1}, A_{t-1})$, allows to learn a much better representation compared to maximizing $I(O_t; Z_t)$. While being a contrastive loss, $I(O_t; Z_t)$ does not behave differently from a reconstructions loss, as it requires the encoder to retain as much information in the observation as possible, which is contradictory to the motivation of the paper.\n- The use of both $I(i_t; Z_t \\mid Z_{t-1}, A_{t-1})$ and $I(A_{t-1}; Z_{t} \\mid Z_{t-1})$ also seems redundant. In fact, $I(A_{t-1}; Z_{t} \\mid Z_{t-1}) = I(A_{t-1}, Z_{t-1}; Z_t) - I(Z_t; Z_{t-1})$ and therefore maximizing $I(A_{t-1}; Z_{t} \\mid Z_{t-1})$ will also maximize $I(A_{t-1}, Z_{t-1}; Z_t)$, which has similar effects to minimizing $I(i_t; Z_t \\mid Z_{t-1}, A_{t-1})$.\n\n#### **Other comments and suggestions**\n- While the assumption that the features of the observations influenceable by the agent through its actions are important for control is true in the experiments that the paper considers, it is not true in general. Consider a robot trying to navigate from A to B by using visual information, its actions such as turning left or right can actually influence the distractors (trees, irrelevant subjects, etc.) in the scene that it observes. However, this can be left for future work.\n- The paper is missing some important references. PC3 [2] is one of the first methods that proposed to use mutual information/contrastive learning to learn representations for control. TPC [1] is a very relevant paper that was also proposed to address the problem of learning representations for visual model-based RL with distractors in the observations. TPC outperforms Dreamer, CDreamer and DBC and should serve as a stronger baseline. PI-SAC [3] is also a method related to the line of work that uses mutual information to learn representations for RL. In [4] the authors theoretically investigated different mutual information objectives used to learn representations for RL.\n- I want to see the performance of InfoPower vs the baselines in the standard setting (no natural backgrounds), as a good method should work well in both standard and natural background settings.\n\n[1] Nguyen, T., Shu, R., Pham, T., Bui, H. and Ermon, S., 2021. Temporal Predictive Coding For Model-Based Planning In Latent Space. arXiv preprint arXiv:2106.07156.\n\n[2] Shu, R., Nguyen, T., Chow, Y., Pham, T., Than, K., Ghavamzadeh, M., Ermon, S. and Bui, H., 2020, November. Predictive coding for locally-linear control. In International Conference on Machine Learning (pp. 8862-8871). PMLR.\n\n[3] Lee, K. H., Fischer, I., Liu, A., Guo, Y., Lee, H., Canny, J., & Guadarrama, S. (2020). Predictive information accelerates learning in rl. arXiv preprint arXiv:2007.12401.\n\n[4] Rakelly, K., Gupta, A., Florensa, C. and Levine, S., 2021. Which Mutual-Information Representation Learning Objectives are Sufficient for Control?. arXiv preprint arXiv:2106.07278.\n", "summary_of_the_paper": "The paper proposes a new non-reconstruction method for model-based RL from high-dimensional observations. The main idea is to make use of an information empowerment objective that prioritizes encoding parts of the environment that are influenceable by the actions. This allows the model to focus on functionally relevant information and filter out distractors. The same empowerment term can also be used to promote faster exploration when the reward signal is sparse. The proposed method outperformed the existing baselines in difficult Deepmind control tasks with natural video backgrounds.", "summary_of_the_review": "In general, the proposed method is promising, but the current presentation of the paper has major issues that need to be addressed and clarified.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agree that the proposed method to create a more robust representation of a task for model-based-RL is interesting and has significant merits. After some revision, more critical reviewers improved their ratings of the paper, such that there is unanimous agreement that the paper can be accepted to ICLR.
Decision: Accept (Poster)