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.
########### Post-rebuttal summary ############
The 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.
########################################
## Paper Summary
This 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.
## Strengths
- the paper addresses a relevant problem with an intuitive approach
- the paper is well written and easy to read
- the analytic toy experiments at the end of Sec 4.2 and in 6.2 help understand the properties of the learned representation
- 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
## Weaknesses
(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).
(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)
(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). 
(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.
(5) lacks reference to recent work on CPC-style objectives for RL (see suggested references below)
## Questions
(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?
(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?
(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?
## Suggestions to improve the paper
(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)
(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
(for 4) add details about the architecture, hyperparameters and training schedule, for both the method and all comparisons to the appendix
(for 5) add references to related works that use CPC-style MI objectives for representation learning in the context of RL/skill learning: 
	- [1] Nachum et al., ICLR 2019 -- applies CPC-style objective to hierarchical RL setting
	- [2] Anand et al., NeurIPS 2019 -- investigates MI objectives for representation learning on a wide range of Atari games (don't apply to RL)
	- [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)
	- [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)
	- 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)
- add qualitative rollouts for predictions from the PlaNet predictive network both with and without distractor to the appendix
## Minor Edit Suggestions
- "Learning latent dynamics from pixels", Hafner et al. is cited twice in the reference section
- 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
[Novelty]: okay
[technical novelty]: minor
[Experimental Design]: okay
[potential impact]: high
#######################
[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.
[Confidence]: High
[1] Near-Optimal Representation Learning for Hierarchical Reinforcement Learning, Nachum et al., ICLR 2019
[2] Unsupervised State Representation Learning in Atari, Anand et al., NeurIPS 2019
[3] Shaping Belief States with Generative Environment Models for RL, Gregor et al., NeurIPS 2019
[4] Neural Predictive Belief Representations, Guo et al., Arxiv 2018
### Rebuttal Comment (copied here so that it's visible to the authors) ###
I 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.
I 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.
In 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.
Due to this fundamental concern about the method I cannot recommend acceptance of the submission.