Mutual Information Maximization for Robust Plannable Representations
Authors: Yiming Ding
Ignasi Clavera
Pieter Abbeel
Venue:  (2020)
Keywords: reinforcement learning
robust learning
model based
planning
representation learning

Abstract: Extending the capabilities of robotics to real-world complex, unstructured environments requires the capability of developing better perception systems while maintaining low sample complexity. When dealing with high-dimensional state spaces, current methods are either model-free, or model-based with reconstruction based objectives. The sample inefficiency of the former constitutes a major barrier for applying them to the real-world. While the latter present low sample complexity, they learn latent spaces that need to reconstruct every single detail of the scene. Real-world environments are unstructured and cluttered with objects. Capturing all the variability on the latent representation harms its applicability to downstream tasks. In this work, we present mutual information maximization for robust plannable representations (MIRO), an information theoretic representational learning objective for model-based reinforcement learning. Our objective optimizes for a latent space that maximizes the mutual information with future observations and emphasizes the relevant aspects of the dynamics, which allows to capture all the information needed for planning.
We show that our approach learns a latent representation that in cluttered scenes focuses on the task relevant features, ignoring the irrelevant aspects. At the same time, state-of-the-art methods with reconstruction objectives are unable to learn in such environments.

Under review as a conference paper at ICLR 2020
MUTUAL INFORMATION MAXIMIZATION FOR ROBUST PLANNABLE REPRESENTATIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Extending the capabilities of robotics to real-world complex, unstructured envi- ronments requires the capability of developing better perception systems while maintaining low sample complexity. When dealing with high-dimensions state spaces, current methods are either model-free, or model-based with reconstruc- tion based objectives. The sample inefficiency of the former constitutes major barrier for applying them to the real-world. While the latter present low sample complexity, they learn latent spaces that need to reconstruct every single detail of the scene. Real-world environments are unstructured and cluttered with objects. Capturing all the variability the latent representation harms its applicability to downstream tasks. In this work, we present mutual information maximization for robust plannable representations (MIRO), information theoretic representa- tional learning objective for model-based reinforcement learning. Our objective optimizes for latent space that maximizes the mutual information with future observations and emphasizes the relevant aspects of the dynamics, which allows to capture all the information needed for planning. We show that our approach learns a latent representation that in cluttered scenes focuses on the task relevant features, ignoring the irrelevant aspects. At the same time, state-of-the-art methods with reconstruction objectives are unable to learn in such environments.
INTRODUCTION
A fundamental challenge in applying reinforcement learning (RL) to real robotics is the need to define a suitable state space representation. Designing perception systems manually makes it difficult to apply the same reinforcement learning algorithm to a wide range of tasks. Real-world environments are unstructured, cultured, and changing. A suitable representation has to capture features relevant for the task, while ignoring the other elements in the scene. How do we impose inductive biases in the learned representation for real-world robotics?
One could, in principle, apply model-free free reinforcement algorithms from raw, low-level observations; however, these tend to be slow, sensitive to hyperparameters, and sample inefficient (Lee et al.) 2019). Current work has focused learning latent spaces for model-based RI. constraining them to be feature points (Finn et al.) 2016), or learning them with reconstruction objectives (Hafner et al. 2018 Watter et al. [0155] Ha & Schmidhuber] 2018). While feature points objectives provide good inductive bias, they are too rigid for unstructured tasks and cannot incorporate semantic knowledge of the environment. Learned representations with reconstruction objectives (Kingma Welling 2013) force the model to capture all the diversity of the environment in the learned representation in order to accurately reconstruct the observations. As result, in environments with variability, the relevant features for the task either represent small component of the latent space or are ignored, both of which prevents deriving a controller from them.
This work tackles the problem of representation learning from an information theoretic point of view: learning a latent space that maximizes the mutual information between the latent and future observations (van den Oord et al. 2018 Hénaff et al.) 2019). The mutual information objective induces an inductive bias towards the shared information with future observations, which allows the latent space to discard lowelevel information and local noise. Hence, by foregoing the reliance on reconstruction objectives, we obtain a latent space that is robust to the presence of variations on the scene if they unrelated to the task.

---

Under review as a conference paper at ICLR 2020
The main contribution of our work is a representational learning approach, MIRO (Mutual Information Maximization for RObust Plannable Representations), that jointly optimizes for the latent representa- tion and model, and results in latent space that is robust against disturbances. noisy observations, and achieves performance comparable to state-of-the-art algorithms. We show that previous approaches brittle even to minor changes in the environment. Finally, experimentally investigate these effects on the latent variables in our approach and previous work, pinpointing that our learned space invariant to such changes in the environment. The experimental evaluation carried out in standard DeepMind Control Suite (Tassa et al. 2018) environments.
2 RELATED WORK
Model-based Reinforcement Learning from Image Observations. Recent work has shown that model-based reinforcement learning RL is able to achieve the same performance model-free methods while being substantially more sample efficient (Wang et al.) 2019). However, such achievements have been limited to envitonments where compact state space representation is available, and extending them to raw sensory high-dimensional input spaces poses a major challenge. In such cases, learning compact and accurate latent space is crucial to relieve the dynamics model from directly modeling the raw sensor input space. Recent work on high- high-dimensional observations on model-based RL can be categorized in two main classes: 1) video prediction models (Jayaraman et al. 2018 Ebert et al. 2018 Kaiser et al. 2019). and 2) latent space learning with reconstruction objectives (Hafner al. 2018 Lee et al. 2019 Watter et al. 2015 Wahlström et al. 2015 Zhang
et al. 2018 Ha & Schmidhuber 2018).
Both video prediction models and latent space with reconstruction objectives share the commonality that the latent space is learned using loss on the raw pixel observations. As a result, the latent space model needs to incorporate all information to reconstruct every detail on the observations, which is redundant the task is usually represented by small fraction of the scene in real world environments.
Representation Learning. Variational autoencoders (VAE) (Kingma & Welling 2013 Higgins et al.; van den Oord et al., 2017) learn an embedding by optimizing the variational lower bound of the likelihood of data. The learned latent follows a prior distribution, usually Gaussian, which makes sampling and distance meaningful in the latent space. Our method uses a similar architecture to the extension of VAE to sequential data (Chung al. 2015 Lee et al. 2019 Hafner al. 2018 Fraccaro et al., 2016). However, instead of a reconstruction objective, we propose to learn latent space by maximizing the mutual information (MI) between the latent and future observations. The omission of reconstruction objective relieves the dynamics model from encoding unpredictable variation in the scene, which proves to be more robust when the observation visually noisy. Our method inspired by the InfoMax principle (Linsker 1988), learns a latent space by maximizing the mutual information (MI) between observation and latent representation. Despite the difficulty of estimating MI directly, recent work has shown that representations learned by maximizing noise contrastive estimation lower bound of MI are able to achieve statc-of-the performance in downstream benchmark tasks in various domains, including image classification, video captioning and natural language processing (van den Oord et al. 2018 Hénaff et al.] 2019 Hjelm et al. 2018 Sun et al.) 2019: Tian al.) 2019 Bachman al. 2019). The maximizing the MI has been applied to the context of reinforcement learning, but not in the context of learning robust representations, but to either speed up learning (van den Oord et al. 2018), 3), to achieve hierarchical represenations (Nachum et al. 2018), or model the believe of state space (Gregor et al., 2019 Guo et al. 2018)), which use a similar architecture to our method. (Anand et al. 2019) use an MI objective to learn representations for Atari games.
3 BACKGROUND
This work tackles the problem of learning a latent space that is suitable for planning from high- dimensional observations in POMDPs. In this work, the observations constitute single images, which are not enough to completely capture the dynamics of the system. For instance, a single image does not contain any information of the velocity object might be occluded.
2

---

Under review as a conference paper at ICLR 2020
Partially Observable Markov Decision Process. A discrete-time finite partially observable Markov decision process (POMDP) M is defined by the tuple (S, R, O, O, 7 Po. H). Here, is the set of states, A the action space, ([$+11/s, a.) the transition distribution (p($t+1|st, a,)), R(s the probability of obtaining the reward 7, at the state 8, (p(rz)s:)), O is the observation space, (00e(s)) (p(oe(s))) the discount factor, Po S R represents the initial state distribution, and H is the horizon of the process. We define the return as the sum of rewards Tr along a trajectory T = (so: 002 ..y 5H-1, 1:-1::--1-000 sH) The goal of reinforcement learning to find a controller x A R* that maximizes the expected return, i.c.: max, ( E scrup
E 7"T]].
Mutual Information. The mutual information between random variables X and Y, denoted by T(X; Y) is reparametrizatioo-nnaaaia measure of dependency. Specifically, it characterizes the Kullback- divergence between the joint distribution (X, Y) and the product of the marginals log p(z.g) p(z)p(u)
X and Y: I(X;Y) Bp(x,))
log p(alv)
E Ep(z.))
DRL ((X, l) @ Y).
p(z)
Estimating and optimizing the mutual information objective poses a challenging problem. In this work, we multi-sample unnormalized lower bound based on noise contrastive estimation, INCE van den Oord et al. (2018).
4 ROBUST PLANNABLE REPRESENTATIONS
Enabling complex real robotics tasks requires extending current model-based methods to low-level high-dimensional observations. However, in order to do so, we need to specify which space they should operate on. Real-world environments are unstructured, cluttered, and present distractors. Our approach, MIRO, is able learn latent representations that capture the relevant aspects of the dynamics and the task by framing the representational learning problem in information theoretic terms: maximizing the mutual information between the latent space and the future observations. This objective advocates for representing just the relevant aspects of the dynamics, removing the burden of reconstructing the entire pixel observation.
We first motivate the use of the mutual information objective representational learning objec- tive for control, then we derive our objective that entangles dynamics and reward learning with the representational objective, and finally we instantiate this objective in a concrete model-based algorithm.
4.1 ROBUSTNESS OF LEARNED LATENT SPACES
Representations that ignore the presence of elements unrelated to the task at hand would allow us to apply successfully off-the-shelf reinforcement learning algorithms in real-world environments. In contrast, current representational learning approaches for control are based on reconstruction objectives. Here we show that these methods are not suitable for real-world tasks, since even in the presence of simple distractors, they completely fail to capture the important aspects of the task.
Concretely, several prior work learns representations by reconstructing the observation using variational auto-encoder (VAE) (Kingma & Welling 2013):
max E [=( po [)] DRi (go((o)()))) (p(2))
0,0
here p(z) (0, 1), and 8, o are the parameters of the encoder and decoder, respectively. This objective captures the variation of the input in the latent space while the commonalities are stored in the weights of the auto-encoder. As result, an unstructured and changing environment will be encoded in the latent space even when this variation is irrelevant to the control problem.
The mutual information objective presents itself as an alternative to learn representations. We argue that this objective is more suitable for learning robust representations. In the POMDP setup, we propose to maximize the mutual information between past observations and future observations, ie., max; I [01,1; 0tth ((::++)) where 8 are the parameters the encoder (no decoder is needed). Then.given limited capacity of the encoder, the optimal representation should not capture the details on the scene that do not depend on the actions. Including more information in latent space always increases the mutual information, hence limiting the capacity adding an information bottleneck is
3

---

Under review as a conference paper at ICLR 2020
0.8
VAE
Regression Validation
0.6
CPC
0.4
0.2
0.0
Number of Distractors
Figure I: Left. Point-mass environment with distractors. Right. As function of the number of the distractors, we plot the accuracy of the regression the position of the point-mass when regressing from the latent representation learned variational autoencoder (VAE) and when regressing from the latent representation learned by contrastive predictive coding (CPC). Even in simple domains the reconstruction based objective for learning representations (VAE) discards the information needed for modelling the dynamics of the system, while the mutual information objective preserves it.
crucial for obtaining a robust representation. We estimate the mutual information using InfoNCE lower bound (van den Oord et al. 2018):
fo(o1:1- octh) En fo Or4N))
max E
Here, the expectation is on sets of observtions ( Oh 1 with N negative examples and one positive example. o₁, and fe represent score function. The mutual information objective allows us to encode an inductive bias towards the important aspects of the task. The inductive bias arises from the fact that VAE. Figure shows the effect of distractors on the learned representation in a simple point-mass environment. Both representations are learned using the same data, raw pixel images of the point mass with the distractors. Then, we regress on top of the learned representations to the position of the point-mass. The mutual information objective allows to maintain the accuracy even when more distractors are present. However, the autoencoder accuracy is increasingly harmed by the appearance of unrelated objects to the task.
optimizing diseriminative model, instead of generative such the
4.2 LEARNING LATENT SPACES FOR CONTROL
In this section, we develop the MIRO objective that entangles the representation, dynamics model, and reward predictor. Consequently, the objective not only learns a representation that is robust, but also emphasizes the important aspects of the dynamics and the task.
ao
a1
a2
So
S1
S2
SS
Figure shows the POMDP set-up, where the only observed variables are the observations o, and re- wards Tt, and we are trying to model the under- lying latent space and the underlying process that generates the observed variables. We model the latent space and reward to be Gaussian with diagonal covariance given the previous state and action, i.e., St41 1441/st.o ~ N(u++1 E++1) and re| a, ~ (V Ar). order to learn the underly- ing process of the POMDP and successfully use this representation for control, we learn four functions: an encoder, filter, transition dynamics, and reward predictor. In order to plan, we need to process the current observation (encoder). update the latent space with the current observation (filter), predict the
ro Oo
r1.01
T2
03.09
Figure 2: Probabilistic graphical model of the POMDP. The only variables observed (shaded nodes)
the actions, observations (high- dimensional images), and rewards. Our model has to infer latent space (dashed lines) as well as to model the conditional probabilities (solid lines).
4

---

Under review as a conference paper at ICLR 2020
D++1
Pi-2
11+3
6
C6
68
66
2t
St
2++
51+1
2++2
2++3
9
96
at
St
1++1
st+
f
8++1
SE+2
Figure 3: Proposed architecture that optimizes the MIRO ohjective. The shaded variables are observed, while the unshaded are latent, the yellow nodes represent parametric functions. The ohservation or is fed into the encoder resulting in the intermediate latent variable Together with the prior on the latent state & constitute the input of the filter, modelled feed-forward neural network. The filtered state s; and action a, is passed through the dynamics model, which is recurrent neural network, to obtain the predicted (prior) latent next state 8t-1. The dashed lines denote the variables between which we maximize the mutual information. (Circle: stochastic; square: deterministic). Reward predictor network is omitted.
next states (transition dynamics) and predict the future rewards (reward predictor). In the following, we describe detail the parametrization, and functionality each these components:
The encoder parametrizes map from high-dimensional observation to a lower dimensional manifold, Eg O C I!" We represent the encoded observation (0 with the notation & Z. Contrary to prior work in learning latent spaces for planning, do not make any assumption on the underlying distribution of 2 or og (such as they are Gaussian (Kingma & Welling 2013)).
Encoder.
Filter. This function, 9c Z S + S. filters the belief of the prior & state variable with the current encoded observation 2 resulting in the filtered variable 8t. Essentially, it updates the latent to be coherent with the most recent observation.
Dynamics model. The dynamics model approximates the transition function the learned latent space, f& S x A - S. The dynamics parametrizes the mean br and diagonal covariance Σ, of the underlying transition distribution of the latent process. We denote by 8 the prior latent variable resulting from sampling from the dynamics model distribution.
Reward predictor. The reward predictor regresses onto the observed rewards from the latent variable. Hence, it is mapping from latent states to rewards, Tv S + R. As said, we assume that the rewards follow a Gaussian distribution with unit variance.
We denote by θ = {0, C, } the parameters of our function approximators. These are learned altogether with the following constrained optimization objective:
max +
I(s o0tal ((:c++))
s.t.: DKL (2+41, Se+ (1+1)] Se Dxu(rellre(ss, ac)) Se
As seen in the previous section, the mutual information maximization term allows us to obtain latent spaces that just contain information about the relevant components of the scene. The enforcement of the constraints has two purposes: 1) ensures small prediction errors when planning, and 2) guides the state space capture the components for predicting rewards and future states.
In practice, the previous objective intractable to optimize. The exact maximization, or even evaluation, of the mutual information term is intractable, as well the of non-linear inequalities.
5

---

Under review as a conference paper at ICLR 2020
Algorithm Model-Based Reinforcement Learning with MIRO
1: Initialize encoder cø, filter 9c, dynamics model fø. reward predictor ry and dataset D 0.
2: repeat
3: T
4: for = K do
5: Sample initial state 50 Po
6:
for f = O...
.... H do
7:
Filter the state & with o, using 9ς to get 8t Get action a, from planner using fo and T Take action a, in the environment and obtain 1t T T TU {8x, dz, 0t)
8:
9:
10:
11:
Estimate next state $141 ~ fo(s;, as)
12:
end for
13:
DUDUT end for
14:
15:
Optimize Equation with mini-batches from dataset D
until Until desired performance is achieved
17:
return Optimal model parameters +
Instead, we optimize tractable lower bound of this objective. First, we replace the mutual information term with the noise constrative estimator lower bound INCE. Second, we formulate its Lagrangian and treat dual parameters as hyperparameters. Altogether, it gives an optimization objective that can be easily optimized with stochastic gradient descent
max TNCE (8;; OL+h (aa+h)) - x DKL (gc ( (+41)(18441 (rr ( at)
(1)
We optimize this objective by using Monte-Carlo estimates and backpropagating through the stochas- tie nodes with the path-wis, derivative, also known reparametrization trick (Kingma & Welling 2013). As yan den Oord et al. (2018), we use log-bilinear model for the score function, which leads to the following formulation of the InfoNCE objective. Be xt [se; aes+a):
exp(xr WAZ++h)
Isc 064h(0ca4R) = log
exp(x; Wh + Σ exp(x] Wwez))
Even though more complicated models can be used, as well as models that encode the future sequences of actions; This simple score model tends to work well in practice.
ROBUST LATENT SPACE MODEL-BASED REINFORCEMENT LEARNING
In this section, present architecture, implementation, and instantiation into model-based method. The presented algorithm uses planning, but our representational learning method for control agnostic the derived controller. Methods such as Lee al. (2019) could be used with MIRO. The pseudocode for the model-based approach is shown in Alg I and is typical in model- RL algorithms, it iterates between three steps: 1) data collection, model learning, and 3) policy improvement or planning.
each iteration, we first collect data by planning with the current dynamics model. On-policy data collection relieves the model from learning the entire space, and instead it just focuses on the regions that the agents visits. also allows to overcome the insufficient coverage of the initial data distribution. The data collected is stored in replay buffer that is used for training the model.
Data Collection.
Model Learning. The architecture of the model is shown in Figure 3 It is composed by an convolutional encoder, feed forward filter, and dynamics model that has the form of a recurrent feed forward neural
neural network. The reward predictor, not shown in Figure is modelled network that has as input the latent state variable. The model is trained using all the data collected so far the MIRO objective (Equation

---

Under review as a conference paper at ICLR 2020
Figure 4: An example trajectory of Cheetah environment with distractors. The distractor objects (red sphere and green cube) placed at random position each time step. Planning. In this work, we use moddll-predicte for action selection (Botey et al., 2013). The CEM algorithm selects the action sequence that maximizes the expected return under the learned dynamics. Specifically, it is population based procedure that iteratively refits Gaussian distribution, starting from unit Gaussian, on the best sequences of actions. The MPC component prevents overfitting to the model deficiencies by selecting the best CEM action and replanning at each step. Note that our represenational learning objective is orthogonal the planning method used.
control (MPC) with cross-entropy method (CEM)
RESULTS
In this section, empirically corroborate the claims in the previous sections. Specifically, the experiments are designed to address the following questions:
1. Is our approach able to maintain its performance front of distractors the scene?
2. How does our method compare with state-of-the-ar reconstruction objectives?
3. Does latent variable remain unchanged in the presence of visual noise?
To answer the posed questions, evaluate framework. in four continuous control benchmark tasks MuJoCo simulator: curtpole-bllanee reacher, finger-spin and half-cheetah (Todorov et al. Tassa al. 2018). We choose PlaNet (Hafner et al. 2018) the state-of-the-art reconstruction
2012
objective baseline for plannable representations.
6.1 ROBUST CONTINUOUS CONTROL COMPARISON
To test the robustness of MIRO and PlaNet visually noisy environments, we add distractors each of the four environments, as shown in Fig. 4
1001
1000
800
800
600
600
400
400
200
edc
800
100
200
300
400
(a) Cartpole Balance
(b) Reacher
600
400
400
200
300
400
zōo
600
800
(c) Finger
(d) Half Cheetah
PlaNet
PlaNet with distractor
MIRO
MIRO with distractor
Figure 5: Learning curves MIRO and PlaNet on environments and without distractors. y-axis plots cumulative reward while x-axis plots number rollouts collected. All curves represent mean and the shaded area represents one standard deviation among 3 seeds. MIRO's performance remains unchange in the presence of disctractor objects in the environments while PlaNet is drastically affected such objects. When distractors are present in the scene both approaches have comparable performance.

---

Under review as a conference paper at ICLR 2020
The results, depicted in Figure show that in all four environments the performance of MIRO is not undermined by the presence of objects unrelated to the task. Interestingly, in the half-cheetah environment, the performance presence of visual noise forces the embedding to focus even more on information relevant to the task, and thus makes the embedding more suitable for planning. In comparison, the performance of PlaNet struggles in face of distractors: the agent is unable to learn in the reacher environment, and struggle learn in the other environment obtaining sub-optimal performance. When distractors are present in the scene, Figure [S] shows that our approach achieves comparable or superior performance than its reconstruction- bassoddd----ddd---- counterpart.
improves with visual noise in the background. We argue that the
6.2 REALISTIC DISTRACTORS
In the previous section, the distractors added to the scene is placed at a random position at each time step. While this resembles particular types of random noise, the movement of majority of objects in real life do not conform to such behavior. To simulate more realistic visual noise, we add the same geometric objects in to the scene but move the i-th distractor according to the following linear rule:
(x) + 1. vt + 1) = (xt, yt) + Aü'
where / size of each step, û' is a unit vector denoting the direction of the movement and it is resampled with probability p 0.2 at each timestep. We term the resulting distractors "linear distractors" and term the distractors Figure "random distractors".
We compared two ways of obtaining dynamics model on tasks with linear distractors: 1. Transfer: Train with random distractors and test on linear distractors; 2. From scratch: Train with linear distractors and test linear distractors. The performance of both MIRO and PlaNet on two schemes are presented in table. The results showcase that MIRO, when presented with realistic distractors, is still able to perform the task. Since MIRO is maximizing an MI objective, the latent space is more robust to the distractors when it is exposed to higher entropy distractors at training time.
PlaNet (Transfer)
MIRO (Transfer)
PlaNet (From scratch)
MIRO (From scratch)
Cartpole Balance
615 + 31
984 + 6
583
57
780 + 71
Reacher
139 126
673 148
274 + 238
551 + 43
Finger
512 + 21
784 + 17
582 + 31
648 + 50
Half Cheetah
308 128
615 + 62
211 73
576 + 148
6.3
UNDERSTANDING THE LEARNED LATENT SPACE
Here, we provide insight into robustness of la- tent learned by MIRO. An ideal latent space should only encode information essential to the dynamies of the task. Thus, given observa- tions Oclean and Odis. the later being the same
Sequential VAE
MIRO
Normalized difference of latent
1.5
1.2-
as the former but with added distractors in the background, the desired behaviour of our en- coder would be to output the same latent varible;
0.9
0.6
ie., Selean Adis. being Zelean the encoded Oclcan (eo (Oeteam) )) and zdis the encoded Oas (Co (Q.is).
0.3
In order to quantify the discrepancy on the latent, we measure its normalized difference
0.0
100
200
300
400
500
600
Babao Fالl Figure shows the average nor-
Training Episodes
malized difference on the Cartpole Balance en- vironment when the latent space learned using the MIRO objective and a sequential VAE objec- tive (Hafner et al. 20188 Lee et al. 2019 Ha & [Schmidhuber] 2018). The results manifest that the embedding learned by MIRO excludes irrel-
Figure 6: Normalized different in the latent variable when encoding the same observation with and without distractors. MIRO learns representations that are invari- ant to distractors, resulting in lower difference, while the embedding learned by the sequential VAE fails to cap- ture the relevant aspects of the dynamics despite being trained jointly with the dynamics.
evant information over the course of training. In
8

---

Under review as a conference paper at ICLR 2020
contrast, with sequential VAE, the normalized enror increases and lands higher final value, which indicates that the model fails to understand the essential elements in the task despite being trained jointly with the dynamics.
CONCLUSION
In this paper, present MIRO, information thcoretic representational learning objective that tackles the problem of learning robust and generalizable representations for control. Our approach maximizes the mutual information between the learned latent space and future observations, empha- sizing the features that are relevant for the dynamics and the reward predictions. We compare our approach in standard benchmark control environments with and without visual disturbances. Our method remains unaffected by those disturbances while state-of-the ff-the-an art methods fail to complete the tasks. When disturbances are removed it attains comparable or superior performance to prior work. Furthermore, we experimentally investigate the effect of such disturbances in the latent variables, showing that MIRO's latents are invariant to visual noises. The development of this method has been motivated for the applicability of general reinforcement learning to real-world and unstructured environments. Consequently, an enticing direction for future work would be to analyze MIRO's behaviour in real-robotic agents outside structured environments.

---

Under review as a conference paper at ICLR 2020
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning atari, 2019.
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning Representatiins by Maximizing 00910
Mutual Information Across Views. jun 2019.
Zdravko I Botev, Dirk P Kroese, Reuven Y Rubinstein, and Pierre L'Ecuyer. The cross-entropy method for optimization. In Handbook of statistics, volume 31, pp. 35-59, Elsevier, 2013.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron Courville, and Yoshua Bengio. A recurrent latent variable model for sequential data, 2015.
Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control. dec 2018. URL http//aaxvvv...////./ org/ abs 1812. C0568
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Learning Visual Feature Spaces for Robotic Manipulation with Deep Spa- tial Autoencoders. Proceedings IEEE International Conference on Robotics and Automation, 2016-June:512-511 519, 2016. ISSN 10504729. doi: 10.1109/CRA. 2016.7487173. URL http::/rrll
Feature
OAhttp: http://wwwwwwwwww////ooo//sso
Marco Fraccaro, Søren Kaae Sønderby, UIrich Paquet, and Ole Winther. Sequential neural models with stochastic layers, 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, 2019.
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A. Pires, and Rémi Munos. Neural predictive belief representations, 2018
David Ha and Jürgen Schmidhuber. World Models. mar 2018. doi: 1/zenoo 1207631. URL /atta////eaau.ccraaaaa//////0000 http//xxxx.ooo...o/oooo/../..tot orw...oo..rr/1000 org/10
5281 6281/zenodo. zenodo. 1207631
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint arXiv: 18/11105511 2018.
Olivier J. Hénaff, Ali Razavi, Carl Doersch, S. M. Ali Eslami, and Aaron den Oord. Data-Eificient Image Recognition with Contrastive Predictive Coding, may 2019. URL /tLp:t////aa /20/5.905..0 09272
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, Alexander Lerchner, and Google Deepmind. beta-vae: Learning basic visual concepts with a constrained variational framework. Technical report.
Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon. Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. aug 2018. URL http::tpp/axii.oooo///
1808 06670
Dinesh Jayaraman, Frederik Ebert, Alexei A. Efros, and Sergey Levine. Time-Agnosier Prediction: Predicting Predictable Video Frames. aug 2018. URL htttp::/aaxwv..rgg/////////////t /188
07784
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model Based Reinforcement Learning for Atari. mar 2019.
00374]
10

---

Under review as a conference paper at ICLR 2020
Diederik P Kingma and Max Welling. arXiv:1312.6114 2013.
Auto-encoding variational bayes.
arXiv preprint
Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic Latent Actor- Critic: Deep Reinforcement Learning with a Latent Variable Model. jul 2019. URL http:
00953
Ralph Linsker. Seff-Organization in Perceptual Network, Computer, 21(3):105-117, 1988. ISSN 00189162. doi: 10.1 109/2.6..
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Near-optimal representation learning for hierarchical reinforcement learning, 2018.
Chen Sun, Fabien Baradel, Kevin Murphy, and Cordelia Schmid. Contrastive Bidirectional Trans- former for Temporal Representation Learning, jun 2019. URL http::://axw...oo//////////////p
1906. 05743
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint 1 2018.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive Multiview Coding. jun 2019. URL /tttp://arivve httgg//bb/990// 05849
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In Intelligent Robots and Systems (IROS), 2012 IEEE/RSJ International Conference on, pp. 5026- 5033. IEEE, 2012.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural Discrete Representation Learning. nov 2017.
00937
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation Learning with Contrastive Predictive Coding. jul 2018.
03748
Niklas Wahlström, Thomas B. Schön, and Marc Peter Deisemroth. From Pixels to Torques: Policy Learning with Deep Dynamical Models. feb 2015. URL httt:://aaxxaa org. hatp///b///50
02251
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking Model-Based Reinforce- ment Learning. jul 2019.
02057
Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images. jun 2015. URL
http::/arxixa wwwpxiivorrrreee wtt.:///////......////// 07365
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. SOLAR: Deep Structured Representations for Moodel-Based Reinforcement Learning, aug 2018.
URL http: www.rrivv.orgggboo.
wwt.//s////a..oooo////o// 09105

---

Under review as a conference paper at ICLR 2020
8 APPENDIX
8.1 HYPERPARAMETERS
Environment: Dimension of the input observation is 32 x 32 3. Horizon of all the tasks H 1000 with the same action repeat parameters as (Hafner et al. (0118).
Network: The dimension of the observation encoding space Z is 1024, dimension of the latent space 30. Horizon of the CEM planner is 12. The horizon of the Mutual Information Objective is 3. Number of negative examples for NCE objective is N 10.
Loss Function: In equation the coefficient for dynamics model loss set to A = -0.002, the coefficient for reward prediction loss is set to )2 0.02.
8.2 ABLATION STUDIES
In the following we present the ablation studies.
8.2.1 MUTUAL INFORMATION OBJECTIVE
To study the effect of the mutual information objective in the loss function, we ablated the MI objective and compared its performance with our full model (MIRO). As shown in 7 using solely dynamics model loss and reward prediction loss, planning fails on all four environments. This shows that the MI objective is a central part of the algorithm.
1000
81
03
.22
400
300
402
150
100
(22
662
(a) Cartpole Balance
(b) Reacher
300
500
80
200
139
50
330
150
300
250
(c) Finger
(d) Half Cheetah
MIRO
no objective
Figure 7: Learning curve of MIRO and ablated version without MI objective.
8.2.2 MODEL ARCHITECTURE
In this section we study the effect two architecture choices of the recurrent neural network used the model. We compare the performance when the latent state is stochastic (as used in MIRO) and when the latent state is deterministic as shown in Figure 9
With deterministic latent space, the performance hardly improves. This could be due to the fact that deterministic latents are more prone to overfit. The observation is consistent with (Hafner et al. 2018).
12

---

Under review as a conference paper at ICLR 2020
500
300
-00
(a) Cartpole Balance
(b) Reacher
110
600
430
80
130
Finger
(d) Half Cheetah
Stochastic Latent
Deterministic Latent
Figure 8: Comparison of deterministic latent space and stochastic latent space.
8.2.3 MUTUAL INFORMATION CONDITIONED ACTIONS
In this section we compare two different ways of incorporating actions in the mutual information objective. When predicting ž++ from 8y, we can incorporate actions " Qc+A-1 in two ways: 1. stacking the actions and concatenating them with 8, when predicting Zt+k; execute an open-loop rollout from 8 given a,, ",,".. 0444-1 and get & and predict 2++4 from 814k. For the first variant, we also examined the effect of different horizon h.
200
300
600
700
100
(a) Cartpole Balance
(b) Reacher
400
80
X00
139
200
200
300
300
602
700
100
200
800
700
803
(c) Finger
(d) Half Cheetah
Stack, MI_ 3orizizzn=
Stack, MI Ml_horizon=8
Openloop, MI Ml_horzion=3
Figure Different ways of conditioning the context vector on actions.
13


Reviews:
[{"review_id": "r1eaHySCFr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a model-based reinforcement algorithm in which the model is a sequential latent variable model and the actions planned with a cross-entropy method (CEM) planner. The model is learnt by maximizing a lower bound on the mutual information between the latent states and their successor observations (instead of the classical sequential ELBO). The authors argue that the latter objective function yield robustness to distraction in visual scenes. The algorithm, named MIRO, is experimented on 4 simulated environments.\n---\nOverall I did not find the paper particularly clear and easy to read. The method is only introduced in the 5th page and no ablation study is conducted. \nIt is still not obvious to me why maximizing the MI in the objective function would reduce the influence of potential distractors.` \nFurthermore, the paper overlooks a good part of the related work on extending VAEs to sequence data, published in the last 3 years and does not draw links to similar architectures. \nThe experiments are in my opinion not convincing, as the approach is only experimented on 2 non trivial -yet not particularly challenging- environments (Finger and Half Cheetah). \n\nMinor: in the equation of the ELBO, page 3, the parameters \\theta and \\phi are swapped. \n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "Hyl7bVL0tr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a latent dynamics model that is learned by maximizing a bound on mutual information between image embeddings and the latent state h time steps later. The model is evaluated on four standard visual control tasks that are solved by online planning.\n\nStrengths:\n- The paper addresses an important open problem with latent dynamics models.\n- The paper is written clearly.\n\nWeaknesses:\n- The paper does not discuss why the mutual information is expected to pay less attention to distractor objects.\n- The paper repeatedly mentions \"task relevant\" information. However, I cannot find anything about the method that would make the learned features more task relevant than reconstruction. This should be clarified.\n- The distractor objects in the experiments randomly change locations in each frame. How would the model be expected to behave if they changed in a predictable way?\n- The paper is lacking detail about hyper parameters and model architecture. What value for h is used? Is the transition function a vanilla RNN?\n- The paper is missing an ablation study. It would be interesting which of the design choices about the model contribute to its success.\n\nQuestions:\n- Could you please explain the KL term that is weighted by lambda_2 in Eq 1? The KL notation on random variables rather than distributions seems non-standard. It is unclear why information about the reward should be penalized.\n- Is the objective summed across time steps? How is the data sampled that the model trains on?\n\nComments:\n- The paper claims in the introduction that reconstruction based approaches cannot discard low-level information. This claim should be rephrased, since the decoder variance allows to discard low-level information (the amount can be controlled by scaling the KL).\n- I would suggest to remove the word \"robust\" from the title or find a more descriptive term to replace it with.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "B1x4s2u0Fr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n########### Post-rebuttal summary ############\nThe proposed method relies on the fact that the distractors have highly unpredictable movements such that a mutual information objective between frames of a sequence learns to ignore them. The experimental evaluations are performed with distractors that randomly change positions / directions, which fits the requirements of the method but not necessarily the behavior of distractors in real environments. To properly evaluate this, experiments in environments with more realistic distractors (who's dynamics are not chosen by the authors) are necessary. Therefore, I do not recommend acceptance of this submission.\n########################################\n\n\n## Paper Summary\n\nThis paper combines a mutual information maximisation objective a la CPC with objectives for dynamics and reward prediction to learn a representation for downstream planning / skill learning that is more robust to visual distractors than comparable representations learned with reconstruction-based objectives a la VAE. In particular, the MI objective maximises the mutual information between the representation of the current state and future observations. The authors show improved robustness of their representation to visual distractors over a baseline with pixel-reconstruction-based representation learning (PlaNet). As a result they are able to achieve better performance when model-predictive control is used on top fo the learned representation to perform control on simulated DM Control Suite tasks with added simple visual distractors.\n\n\n## Strengths\n\n- the paper addresses a relevant problem with an intuitive approach\n\n- the paper is well written and easy to read\n\n- the analytic toy experiments at the end of Sec 4.2 and in 6.2 help understand the properties of the learned representation\n\n- the proposed method, when applied to the model-based algorithm PlaNet, shows improved robustness in settings with artificially introduced visual distractors in simulated DM Control Suite tasks\n\n\n## Weaknesses\n\n(1) unclear whether artificial distractors are indicative of behavior with real distractors: the fact that distractors do not follow coherent motions but instead randomly change position between two consecutive frames makes it hard to estimate how this would perform on more natural distractors. As there is no MI between the distractor's position in one image and any other image, a CPC-style representation learning objective will naturally be encouraged to ignore them. However, if they move with more natural, coherent motion this might not be the case. I would suggest to add an experiment with more natural distractor motion (see below).\n\n(2) it is unclear how much of the invariance to distractors is coming from the MI objective and how much simply from the fact that the representation is learned with jointly predicting the reward. In order to justify that the MI objective is helpful for learning the representation I would suggest to run an ablation experiment that trains using *only* the reward prediction objective and compare performance. (see also suggested experiment below)\n\n(3) parts of the model formulation require clarification: when reading Sec 4.2 that describes the model some parts were unclear to me (see concrete questions below). \n\n(4) the amount of detail provided in the paper is insufficient for reproducing the results: the paper lacks detailed information about the used architectures, hyperparameters and versions of the baselines employed (e.g. PlaNet with stochastic/deterministic prediction?), code is not provided.\n\n(5) lacks reference to recent work on CPC-style objectives for RL (see suggested references below)\n\n\n## Questions\n\n(A) why is the latent state variable s_t observed in the model depicted in Fig 3 (i.e. part of the input data)? Shouldn't this variable be latent? How can the KL constraint on it be computed if it is not observed in the input data?\n(B) should the first term on the right hand side of equation (2) have \\hat{s}_t instead of \\hat{s}_{t+h}? Otherwise, how is the I_{NCE} computation conditioned on the current state s_t?\n(C) what is the scale of the x-axis in Fig 5, i.e. does it show the number of environment rollouts / steps or the number of model re-training iterations? if the latter is the case, how does that translate to the number of environment interactions?\n\n\n## Suggestions to improve the paper\n\n(for 1) add an experiment where the distractor has more natural dynamics so that there is MI between the distractor positions in consecutive frames (e.g. ball bouncing in the image frame in the background instead of randomly jumping to new positions)\n(for 2) add an experiment with a reward-prediction only baseline, i.e. only action-conditioned reward prediction so that task-irrelevant parts are ignored by default (i.e. also no reconstruction objective, but also no MI objective) -> show how PlaNet performance compares to the so far reported numbers when using this representation for planning\n(for 4) add details about the architecture, hyperparameters and training schedule, for both the method and all comparisons to the appendix\n(for 5) add references to related works that use CPC-style MI objectives for representation learning in the context of RL/skill learning: \n\t- [1] Nachum et al., ICLR 2019 -- applies CPC-style objective to hierarchical RL setting\n\t- [2] Anand et al., NeurIPS 2019 -- investigates MI objectives for representation learning on a wide range of Atari games (don't apply to RL)\n\t- [3] Gregor et al., NeurIPS 2019 -- while the main proposed model is generative they compare to a contrastive version that uses CPC to learn predictive representations (don't use it for RL)\n\t- [4] Guo et al., Arxiv 2018 -- similar investigation to [3] of CPC-style objective for representation learning in RL environments (don't use it for RL)\n\t- it should also be mentioned that the original CPC paper already showed that adding CPC-style auxiliary loss to RL improves performance (even though they did not compare to other model-based methods)\n- add qualitative rollouts for predictions from the PlaNet predictive network both with and without distractor to the appendix\n\n\n## Minor Edit Suggestions\n- \"Learning latent dynamics from pixels\", Hafner et al. is cited twice in the reference section\n- it might help to add the reward prediction module to Fig 3 or mention in the caption that it is omitted, it is only described later in the text and was confusing for me on first sight\n\n\n[Novelty]: okay\n[technical novelty]: minor\n[Experimental Design]: okay\n[potential impact]: high\n\n\n#######################\n[overall recommendation]: weakReject - I am inclined to accept this paper but am not fully convinced that the random distractors provide a good intuition about how the proposed method would behave with more natural distractors. If the authors are able to report positive results on the two requested experiments I am willing to raise my score.\n[Confidence]: High\n\n\n[1] Near-Optimal Representation Learning for Hierarchical Reinforcement Learning, Nachum et al., ICLR 2019\n[2] Unsupervised State Representation Learning in Atari, Anand et al., NeurIPS 2019\n[3] Shaping Belief States with Generative Environment Models for RL, Gregor et al., NeurIPS 2019\n[4] Neural Predictive Belief Representations, Guo et al., Arxiv 2018\n\n\n\n### Rebuttal Comment (copied here so that it's visible to the authors) ###\n\nI appreciate the effort the reviewers put into the rebuttal! My main concern was that the mutual information objective only encourages the model to ignore the distractors because they change positions randomly between frames, i.e. there is no mutual information between distractor positions in consecutive frames. On the other hand, if the distractor's motion was perfectly deterministic there would be infinite mutual information between distractor positions in consecutive frames and therefore the representation might exclusively model the distractor. I.e. the more predictable the distractor, the worse the proposed method will perform.\n\nI asked the authors to perform an experiment with a more predictable distractor that bounces in the image frame to test this hypothesis. The authors instead chose a distractor that moves predictably for a few steps before randomly changing movement directions. While this sounds similar it can actually make a big difference depending on how frames are sampled for the CPC objective, i.e. if the required pair of frames is sampled across a random direction change there is again no mutual information between the distractor positions in both frames.\n\nIn real scenes the behavior of distractors likely lies somewhere in between these extremes: they will likely not be fully deterministic but certainly not have frequent moments of purely random direction changes. Therefore, to properly evaluate the merit of the proposed approach, experiments on more realistic / previously published environments are needed where the distractor dynamics are given and cannot be altered to better fit the proposed method.\n\nDue to this fundamental concern about the method I cannot recommend acceptance of the submission.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The manuscript concerns a mutual information maximization objective for dynamics model learning, with the aim of using this representation for planning / skill learning. The central claim is that this objective promotes robustness to visual distractors, compared with reconstruction-based objectives. The proposed method is evaluated on DeepMind Control Suite tasks from rendered pixel observations, modified to include simple visual distractors. 

Reviewers concurred that the problem under consideration is important, and (for the most part) that the presentation was clear, though one reviewer disagreed, remarking that the method is only introduced on the 5th page. A central sticking point was whether the method would reliably give rise to representations that ignore distractors and preferentially encode task information. (I would note that a very similar phenomenon to the behaviour they describe has been empirically demonstrated before in Warde-Farley et al 2018, also on DM Control Suite tasks, where the most predictable/controllable elements of a scene are reliably imitated by a goal-conditioned policy trained against a MI-based reward). The distractors evaluated were criticized as unrealistically stochastic, that fully deterministic distractors may confound the procedure; while a revised version of the manuscript experimented with *less* random distractors, these distractors were still unpredictable at the scale of more than a few frames.

While the manuscript has improved considerably in several ways based on reviewer feedback, reviewers remain unconvinced by the empirical investigation, particularly the choice of distractors. I therefore recommend rejection at this time, while encouraging the authors to incorporate criticisms to strengthen a resubmission.
Decision: Reject