Non-Markovian Predictive Coding For Planning In Latent Space
Authors: Tung Nguyen
Rui Shu
Tuan Pham
Hung Bui
Stefano Ermon
Venue:  (2021)
Keywords: representation learning
reinforcement learning
information theory

Abstract: High-dimensional observations are a major challenge in the application of model-based reinforcement learning (MBRL) to real-world environments. In order to handle high-dimensional sensory inputs, existing MBRL approaches use representation learning to map high-dimensional observations into a lower-dimensional latent space that is more amenable to dynamics estimation and planning. Crucially, the task-relevance and predictability of the learned representations play critical roles in the success of planning in latent space. In this work, we present Non-Markovian Predictive Coding (NMPC), an information-theoretic approach for planning from high-dimensional observations with two key properties: 1) it formulates a mutual information objective that prioritizes the encoding of task-relevant components of the environment; and 2) it employs a recurrent neural network capable of modeling non-Markovian latent dynamics. To demonstrate NMPC’s ability to prioritize task-relevant information, we evaluate our new model on a challenging modification of standard DMControl tasks where the DMControl background is replaced with natural videos, containing complex but irrelevant information to the planning task. Our experiments show that NMPC is superior to existing methods in the challenging complex-background setting while remaining competitive with current state-of-the-art MBRL models in the standard setting.

Under review as a conference paper at ICLR 2021
NON-MARKOVIAL PREDICTIVE CODING FOR PLANNING IN LATENT SPACE
Anonymous authors Paper under double-blind review
ABSTRACT
High-dimensional observations are major challenge in the application of model- based reinforcement learning (MBRL) to real-world environments. In order to handle high-dimensional sensory inputs, existing MBRL approaches use repre- sentation learning to map high- figh-dimensional observations into liwer-dimensional latent space that is more amenable dynamics estimation and planning. Cru- cially, the task-relevance and predictability of the learned representations play critical roles in the success of planning in latent space. this work, we present Non-Markovian Predictive Coding (NMPC), information-theorriie approach for planning from figh-dimensional observations with two key properties: 1) it formulates mutual information objective that prioritizes the encoding of task- relevant components of the environment; and 2) it employs a recurrent neural network capable of modeling Mon-Morkovi latent dynamics. To demonstrate NMPC's ability to prioritize task-relevant information, we evaluate our new model a challenging modification of standard DMControl tasks where the DMCon- trol background is replaced with natural vidcos, containing complex but irrelevant information to the planning task. Our experiments show that NMPC is superior existing methods in the challenging complex- belllx-bbakgroundd setting while remaining competitive with current state-of-the MBRL models the standard setting.
INTRODUCTION
Learning to control from high dimensional observations has been made possible due to the advance- ments in reinforcement learning (RL) and deep learning. These advancements have enabled notable successes such solving video games (Mnih al., 2015; Lample & Chaplot, 2017) and continuous control problems (Lillicrap et al., 2016) from pixels. However, it is well known that performing RL directly in the high-dimensional observation space is sample-inefficient and may require large amount of training data (Lake et al., 2017). This is a critical problem, especially for real-world applications. Recent model-based RL works (Kaiser et al., 2020; Ha Schmidhuber, 2018; Hafner et al., 2019; Zhang et al., 2019; Hafner et al., 2020) proposed to tackle this problem by learning a world model in the latent space, and then applying RL algorithms in the latent world model.
The existing MBRL methods that Iearn a latent world model typically do so via reconstruction-base objectives, which are likely encode task-irrelevant information, such as of the background. In this work, we take inspiration from the success of contrastive learning and propose Non-Markovian Predictive Coding (NMPC), a novel cotommtiion-thoreee approach for planning from pixels. In contrast to reconstruction, NMPC formulates mutual information (MI) objective to learn the latent space for control. This objective circumvents the need to reconstruct and prioritizes the encoding of task-relevant components of the environment, thus make NMPC more robust when dealing with complicated observations. Our primary contributions are as follows:
We propose Non-Markovian Predictive Coding (NMPC), novel information-theoree ap- proach to learn latent world models for planning from high-dimensional observations and theoretically analyze its ability to prioritize the encoding of task-relevant information.
We show experimentally that NMPC outperforms the state-of-the-as model when deal- ing with complex environments dominated by task-irrelevant information, while remaining competitive on standard DeepMind control (DMControl) tasks. Additionally, we conduct detailed ablation analyses to study the empirical importance of the components in NMPC.

---

Under review as a conference paper at ICLR 2021
2 BACKGROUND
The motivation and design of our model are largely based on two previous works (Shu et al., 2020; Hafner al., 2020). In this section, we briefly go over the key concepts in each work. Shu et al. (2020) proposed PC3, an information-thhoreli approach that uses contrastive predictive coding (CPC) to learn latent space amenable to locally-linear control. Specifically, they present the theory of predictive suboptimality to motivate CPC objective between the latent states of two consecutive time steps, instead of CPC between the frame and corresponding state. Moreover, they use the latent dynamics F as the variational device in the lower bound
F
E
E
O0
(8)
(4)
(Eo(E,F)=
(I)
K k
F
E
O++1
E
This particular design of the critic has two benefits, where it circumvents the instantiation of an auxiliary critic, and also takes advantage of the property that an optimal critic the true latent dynamics. However, the author also shows that this objective does not ensure the learning of a latent dynamics F that consistent with the true latent dynamics, therefore introduces the consistency loss to ensure the latent dynamics model F indeed approximates the true latent dynamics.
Lcons EE F) Eptor 1,,),s)) In F (E (0141) E (or),
Since PC3 only tackles the problem from an optimal control perspective, it is not readily applicable RL problems. Indeed, PC3 requires a depiction of the goal image in order to perform control, and also the ability to teleport to random locations of the state space to collect data, which are impractical in many problems. On the other hand, Dreamer (Hafner al., 2020) achieves state- of-the-art performance objective. While the authors included a demonstration of contrastive approach that yielded inferior performance to their reconstruction-bet based approach, their contrastive model applied CPC between the frame and its corresponding state, as opposed to between latent states across time steps.
many RL tasks, but learns the latent space using a peconstruction-base.
In this paper, present Non-Markovian Predictive Coding (NMPC), novel latent world model that leverages the concepts in PC3 and apply them to Dream paradigm and RL setting. Motivated by PC3, we formulate mutual information objective between historical and future latent states to learn the latent space. We additionally take advantage of the recurrent model in Dreamer to model more complicated dynamics than what was considered in PC3. The use of recurrent dynamics also allows us to extend Eqs. (1) and (2) to the non-Markovian setting, which was not considered in PC3.
3 NON-MARKOVIAN PREDICTIVE CODING FOR PLANNING FROM PIXELS
To plan in unknown environment, need to model the environment dynamics from experience. We do so by iteratively collecting new data and using those data to train the world model. In this section, focus presenting the proposed latent world model, its components and objective functions, and provide practical considerations when implementing the method.
3.1 MOR-MARKOVIAN PREDICTIVE CODING
We aim to learn a latent dynamics model for planning. To do that, we define an encoder E to embed high-dimensional observations into a latent space, latent dynamics F model the world in this space, and reward function, as follows
Encoder:
E(or) = St
Latent dynamics:
F(st 8cti act) p(s: sce, A(r, si) = p(r: st)
Reward function:
in which is the discrete time step. for, a,, r.r )T are data sequences with image observations or, continuous action vectors a,, scalar rewards rt, and s, denotes the latent state at time t. To handle po- tentially non-Markovian environment dynamics, we model the transition dynamics using recurrent
2

---

Under review as a conference paper at ICLR 2021
&L
à5
Av
Az
les
Figure 1: The graphical model of NMPC, in which we employ recurrent neural network model the dynamics. We omit the action and reward for simplicity. The approximation of & with 8, is done by using a combination of contrastive predictive coding and maximum likelihood.
neural network with deterministic state tion about the past, followed by the stochastic state model p(sr|S<t, deee) p(s, p(s,//)). In practice, we use a deterministic encoder, and Gaussian distribution for dynamics and reward functions. The graphical model is presented in Figure 1. We now introduce the three key components of our model: recurrent contrastive predictive coding, non-markovian consistency, and reward prediction.
which summarizes informa-
Recurrent contrastive predictive coding Instead of performing pixel prediction learn E and F as Hafner et al. (2019; 2020), we take a different approach, where we maximize the mu- tual information (MI) between the past latent codes and actions against the future latent code (O,); E(Oxt), Act). This objective prioritizes the encoding of predictable components from the environment, which potentially helps avoid encoding nuance information when dealing with complex image observations. A similar MI objective was previously used in Shu et al. (2020). However, Shu et al. (2020) only considered the Markovian setting and thus focused on the mutual information over consecutive time steps. In contrast, we consider non- Markovan setting and wish maximize the mutual information between the latent code any time step and the entire histori- cal trajectory latent codes and actions prior to f. Given trajectory over T steps, our non-Markovian MI objective thus sums over every possible choice of t E (2..... T}.
To estimate this quantity, we employ contrastive predictive coding (CPC) proposed by Oord et al. (2018). We perform CPC by introducing critic function construct the lower bound at partic- ular time step 1,
T(F((O.); F(Oc), A&d E Σ In exp ((E())) B(o 22)
p(t)
(5)
exp ((((0 (1) ((O2)),a2)
Pepe:
where the expectation is over K i.i.d. samples of (ot, 0e,, azt). Note that Eq. (5) uses past (F(o2)) (i) from unrelated trajectories as an efficient source of negative samples for the con- trastive prediction of the future latent code E(o Following Shu al. (2020), we choose tie to our recurrent latent dynamics model F.
exp fSt, Sitss a<i) F((eelscta
(6)
There are two favorable properties of this particular design. First, it is parameter-efficient since we circumvent the instantiation of a separate critic f. Moreover, takes advantage of the fact that an optimal critic is the true latent dynamics induced by the encoder E (Poole et al., 2019; Ma & Collins, 2018). We denote our overall CPC objective ( = e(t) E, F).
Non-Markovian consistency Although the true dynamics is an optimal critic for the CPC bound, maximizing this objective only does not ensure the learning of latent dynamics model F that
3

---

Under review as a conference paper at ICLR 2021
is consistent with the true latent dynamics, due to the non-uniqueness of the optimal critic (Shu et al., 2020). Since an accurate dynamics is crucial for planning in the latent space, we additionally introduce consistency objective, which encourages the latent dynamics model to maintain good prediction of the future latent code given the past latent codes and actions. Similar to the recurrent CPC objective, optimize for consistency every time step in the trajectory,
cons F) Σ Bploreenee In (E(oo))((oc)))
Reward prediction Finally, we train the reward function by maximizing the likelihood of the true reward value conditioned on the latent state s, = E(oo),
T
Cowward E, R) =
Σ 1=1
Eo(or) In R(r:)
(8)
3.2 THEORETICAL ANALYSIS OF NMPC AND VASK-RELEVANCE
In contrast to a reconstruction-bart based objective, which explicitly encourages the encoder to behave injectively the space of observations, our choice of mutual information objective specified in Eq. (5) may discard information from the observed scene. In this section, we wish to formally characterize the information discarded by our MI objective and argue that any information discarded by an optimal encoder under our MI objective is provably task-irrelevant.
Lemma 1. Consider an optimal encoder and reward predictor pair (E", ,R)) where
mag I (E(O.): E(0xz), E
(9)
DxL(p(r, oi) R [T, E*)))) 0.
(10)
Let T(a (o(o), a21) denote an -restiricteed policy whose access to the observations Oct is restricted by E*. Let waux E (oct), (oss), a<t) denote an (E*, "'',,'))re -'sstrestiette policy which has access to auxiliary information about Oct via some encoder E'. Let n(T) denote the expected cumulative reward achieved by a policy æ over a finite horizon T. Then there exists no encoder E' where the optimal E rrestriccee policy underperforms the optimal (E", FE) -eesttrcteteee policy,
EE' s.t. 17(TT) < (7( (Txx))
(11)
Intuitively, since E" optimizes our MI objective, any excess information contained in E' about Or (not already accounted for by E*) must be rmpporaly-umpredcctiv from the past nor predictive of the future. The excess information conveyed by E' effectively nuisance information and thus cannot be exploited to improve the agent's 's performance. is therefore permissible to dismiss the excess information in E" as being task-irrelevant. We provide proof formalizing this intuition in Appendix B.
is neither predictable
is worth noting that our MI objective does not actively penalize the encoding of nuisance informa- tion, nor does temporally- predictive information necessarily mean it will be task-relevant. However, the representation space has limited capacity to encode information about or (e.g.. due to dimen- sionality reduction, stochastic encoder, cemporally-predietivv
eererrally-reedictive
explicit information bottleneck regularizer), MI objective will favor temporally- predictive information over and in this sense, thus favor potentially task-relevant information over provably task-irrelevant infor- mation. This is in sharp contrast to a reconstruction objective, which makes no distinction between these two categories of information contained in the observation o,.
information
3.3 PRACTICAL IMPLEMENTATION OF NMPC
Avoiding map collapse Naively optimizing frepe and foons can lead to a trivial solution, where the latent map collapses into a single point to increase the consistency objective arbitrarily. In the pre- vious work, Shu et al. (2020) resolved this by adding Gaussian noise the future encoding, which balances the latent space retraction encouraged by foons with the latent space expansion encouraged by frape- However, found this trick insufficient, and instead introduce CPC objective between
4

---

Under review as a conference paper at ICLR 2021
the current observation and its corresponding latent code, which we call the instantaneous CPC. Our instantaneous CPC objective defines a lower bound for [(((O)) O.) where we employ a Gaussian distribution (EE(),)) E(O,), of) as choice of variational critic for f(E(O,) O,). Crucially, we keep the variance a fixed so that the instantaneous CPC objective encourages map expansion.
Smoothing the dynamics model Since our encoder is deterministic, the dynamics always receives clean latent codes as inputs during training. However, in behavior learning, we roll out multiple steps towards the future from stochastic dynamics. These roll outs susceptible cascading error problem, which hurts the value estimation and policy learning. To resolve this issue, we smooth the dynamics by adding Gaussian noise to the inputs of the recurrent dynamics during training. The
loise-adding procedure is as follows: assume the dynamics outputs 8, ~ N (J4,, o() o as the prediction time step f, we then add 4 ~ NO0 of) to 8 E(o,) and feed it to the latent dynamies, and repeat for every time step < t S T. We call this dynamics-associatd noise, which ensures that the latent dynamics can handle the amount noise that it produces when rolling out.
The overall objective of our model is thus
max A1 Expe (E, (EF)) + Aglcons (E, F) + Axgl-ap((E) Ag (E, R).
(12)
F.F.R
BEHAVIOR LEARNING
Following Hafner et al. (2020), For seff-containedness, in this section we give summary of this approach. Given the latent state 8, (o.), roll out multiple steps into the future using the learned dynamics F, estimate the return and perform backpropagation through the dynamics to maximize this return, which in turn improves the policy.
latent imagination to learn a parameterized policy for control.
Action and value models We define two components needed for behavior learning, the action model and the value model, which both operate the latent space. The value model estimates the expected imagined return when following the action model from particular state 8r, and the action model implements policy, conditioned on 8r, that aims to maximize this return. With imagine horizon H, we have
C+H
Action model:
a, π (ar(87)
Value model: v(s+) Exg. 8-) Σ 'r (13)
Value estimation
To learn the action and value model, we need to estimate the state values of 14+ where s, and a, are sampled according to the dynamics and imagined trajectories {s,, ar, T the policy. In this work, we use value estimation presented in Sutton Barto (2018).
V (s-) = -Egu
In + 2 n (Sh)
with h = min(n ( + k,t H)
(14)
VA (sx) A)
Σ A (sr) H (s7).
which allows us estimate the return beyond the imagine horizon. The value model is then opti- mized to regress this estimation, while the action model is trained to maximize the value estimation of all states 8, along the imagined trajectories.
(+
H Σ
max E
Σ Vx (8+)
min E
Im (87) - Va (8+)
(15)
EXPERIMENTS
In this section, empirically evaluate the proposed NMPC model in various settings. First, design experiments to compare the relative performance our model with the current best model- based method in several standard control tasks. Second, evaluate its ability to handle more
5

---

Under review as a conference paper at ICLR 2021
realistic but also more complicated scenario, in which replace the background of the environ- ment with a natural video. We discuss how our model superior in the latter case, while remaining competitive in the standard setting. Finally, conduct ablation studies to demonstrate the impor- tance of the components of our model.
Control tasks For the standard setting, we test our model on 6 DeepMind Control (DMC) tasks (Tassa et al., 2018y: Cartpole Swingup, Cheetah Run, Walker Run, Acrobot Swingup. Hopper Hop and Cup Catch. In the natural background setting, we replace the background of each data trajectory with video taken from the kinetics dataset (Kay et al., 2017). We use a different set of videos for training and testing also test the generalization of each method. In both settings, we use images of shape 64 x 64 x 3, episodes with 1000 steps and action repeat 2 across tasks.
Baseline methods We compare NMPC with Dreamer (Hafner et al., 2020), current state of the art model-based method for planning from pixels. Each method is evaluated by the environment return in 1000 steps. For Dreamer, we use the best set of hyperparameters as reported in their paper. We each task with different seeds.
5.1
COMPARISONS IN STANDARD SETTING
Acrobot Swingup
Cartpole Swingup
Cheetah Run
400
750
600
500
200
400
250
200
0
0.5
1.0
1.5
0.5
1.0
1.5
1e6
1e6
1e6
Cup Catch
Hopper Hop
Walker Run
800
300
750
600
500
200
400
250
100
200
1.5
0.5
1.0
1.5
1e6
1e6
1e6
Dreamer
NMPC
Figure 2: NMPC Dreamer in the standard DMControl tasks. The return is computed in 1000 environment steps. Each task is run with 3 seeds.
We demonstrate the performance both methods on standard control tasks in Figure 2. NMPC competitive in all the tasks, which is in contrast what previously observed in Dreamer (Hafner et al., 2020), where they showed the inferiority of contrastive learning approach (which applies CPC between an image observation and its corresponding latent code) compared to their reconstruction based approach. Our results thus show that the use of recurrent contrastive predictive coding is critical for learning latent space that suitable behavior learning.
5.2 COMPARISONS IN NATURAL BACKGROUND SETTING
In this setting, we evaluate the robustness of NMPC versus Dreamer in dealing with complicated, natural backgrounds. The performance of both models the natural backgrounds setting is shown in Figure 4. Dreamer fails to achieve meaningful performance across all six tasks. NMPC performs significantly better on four the six tasks, only failing to work on Acrobot Swingup and Hopper Hop. We note that Acrobot Swingup and Hopper HOp are the two most challenging tasks in the standard setting. Furthermore, the agents in these two tasks are also tiny compared to the complex backgrounds, thus making it difficult for the model to distill task-relevant information.

---

Under review as a conference paper at ICLR 2021
Figure 3: DeepMind Control tasks in the natural backgrounds setting. The images show sample training trajectory Cheetah Run task, where the background taken from the kinetics dataset.
Acrobot Swingup
Cartpole Swingup
Cheetah Run
800
400
400
200
200
0.5
1.0
1.5
0.5
1.0
1.5
0.5
1.0
1.5
1e6
1e6
Cup Catch
Hopper Hop
Walker Run
150
800
600
100
2
0.5
1.0
1.5
0.5
1.0
1.5
0.5
1.0
1.5
1e6
1e6
1e6
Dreamer
NMPC
Figure 4: NMPC vs Dreamer in the natural backgrounds setting. The return is computed in 1000 evironment steps. Each task is with seeds.
5.3
ABLATION ANALYSIS
Cartpole Swingup
Cheetah Run
Walker Run
800
750
600
600
500
400
400
250
200
200
0.5
1.0
1.5
1.5
1e6
1e6
1e6
NMPC
NMPC-No-Inst
NMPC-No-Rer
Figure 5: NMPC two different variants, where omit the recurrent CPC in NMPC-No-Rer and instantaneous CPC in NMPC-No-Inst. The return is computed in 1000 evironment steps. Each task with 3 seeds.
We conduct an ablation analysis to evaluate the importance of each CPC objective employed in NMPC. To do so, we compare the original model with two variants, where we omit the recurrent CPC or instantancous CPC objective respectively. The relative performance of these models on three control tasks is shown in Figure 5. The original NMPC achieves the best performance, while the variant with only instantaneous CPC (denoted NPMC-No-Ren fails across all three tasks. NMPC without instantancous CPC is unstable, since faces the map collapsing problem, which leads to poor performance.

---

Under review as a conference paper at ICLR 2021
6 RELATED WORK
Learning latent space for model-based RL via reconstruction Latent world models can be learned by jointly training the encoder and the latent dynamics model with observation reconstruc- tion loss. Learning Controllable Embedding (LCH) approach, including E2C (Watter et al., 2015), RCE (Banijamali et al., 2018) and PCC (Levine et al., 2020), uses randomly collected data to pre- train a Markovian latent dynamic model that is specifically designed for locally-linear control, then run offline optimal control on top of the learned latent space. CARL (Cui et al., 2020) extends these works for Soft Actor-Critic (Haarnoja al., 2018) and also proposes an online version, in which they iteratively learn the model and parameterized policy. World Models (Ha & Schmidhuber, 2018) learn recurrent latent dynamic model in a two-stage process to evolve their linear controllers in imagination. PlaNer (Hafner et al., 2019) jointly learns a recurrent state space model (RSSM) and plans in latent space using the cross entropy method, while Dreamer (Hafner et al., 2020) uses RSSM to iteratively learn the model and the policy by backpropagating through the dynamics. SO- LAR (Zhang et al., 2019) models the dynamics as time-varying Linear-Gaussin with quadratic costs and controls using guided policy scarch. However, training world models with reconstruction loss has several drawbacks: it requires a decoder as an auxiliary network for predicting images, and by reconstructing every single pixel, those methods are potentially vulnerable task-irrelevant infor- mation such as an irrelevant background.
Learning latent space for model-based RL via constrastive learning An alternative framework for learning latent world models is contrastive learning. Contrastive learning is a self- seff-super learning technique that aims learn representations by contrasting positive samples against nega- tive samples without having to reconstruct images (Oord et al., 2018: Chen et al., 2020a). Recently proposed contrastive learning methods have achieved significant successes in learning representa- tions purely from unlabeled data, which include works by (Chen et al., 2020a;b; Bachman et al., 2019; Hénaff et al., 2019; He et al., 2020; Tian et al., 2019). Poole et al. (2019) has also established close connection between contrastive learning and mutual information maximization. In the con- text of RL, recent works have proposed to use this framework to accelerate RL from pixels in two distinct directions: 1) cast contrastive learning
auxiliary representation learning task, and use model-free RL methods on top of the learned latent space (Oord et al., 2018: Srinivas et al., 2020); and 2) use contrastive learning in conjunction with learning latent dynamics for planning in the latent space (Shu et al., 2020; Ding et al., 2020; Hafner et al., 2020).
Our proposed method follows the latter direction, and can be seen as an extension of both PC3 (Shu et al., 2020) and Dreamer (Hafner et al., 2020). While having similar contrastive objective to PC3, we design a Man-Markovian version of it and also use a stronger controller, which enables our model to work on more complicated control tasks. We use the same controller proposed in Dreamer, however, while their contrastive variant maximizes the MI between the observation its latent code, our loss maximizes the MI between the historical latent codes and actions versus the future latent code. Experimental results show that NMPC is competitive to Dreamer with reconstruction in stan- dard DMControl tasks, which means outperform their contrastive variant.
7 CONCLUSION
In this work, we propose NMPC, a novel information-theoreeic approach that learns a non- Markovian world model for planning in the latent space. We employ recurrent contrastive predic- tive coding objective that specifically estimates the mutual information across temporal states. We show theoretically that our objective does not encourage the encoding of provably task-irrelevant information. This is critically different from reconstruction-bass learning objectives that only measure the mutual information between the current observation and its latent code, which indiscriminately favor the encoding of both task-relevant and irrelevant infor- mation. Our experiments show that NMPC outperforms the state-of-the tf- model when controls in environments dominated by task-irrelevant information while remaining competitive on standard control tasks from DMControl.
objectives as well as contrastive
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Philip Bachman, Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In Advances in Neural Information Processing Systems, pp. 15535-15545, 2019.
Ershad Banijamali, Rui Shu, Mohammad Ghavamzadeh, Hung Bui, and Ali Ghodsi. Robust locally- linear controllable embedding. volume 84 of Proceedings of Machine Learning Research, pp. 1751-1759, Playa Blanca, Lanzarote, Canary Islands, 09-11 Apr 2018. PMLR. URL http: press v /bani jama 11111.1110 html.
//proceedings
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint arXiv:200..05709 2020a.
Ting Chen, Simon Kormblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton. Big self- supervised models are strong cami-supervised learners. arXiv preprint wxXi:22006.1002, 2020b.
Brandon Cui, Yinlam Chow, and Mohammad Ghavamzadeh. Control-aware representations for
model-based reinforcement learning. arXiv preprint arXiv:2006. /3408, 2020.
Yiming Ding, Ignasi Clavera, and Pieter Abbeel. Mutual information maximization for robust plannable representations. arXiv preprint arXiv:2005.081 14, 2020.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolu- tion. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and
R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 2450- 2462. Curran Associates, Inc., 2018. URL http: //ap:///pppraaaa nips. cc/paper/ 7512-recurrent -ww.::/wwoeoo--oooaaaaaeeeooooo-
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. volume 80 of Proceedings of Machine Learning Research, pp. 1861-1870, Stockholmsmässann Stockholm Sweden, 10- 15 Jul 2018. PMLR. URL http: htttp:/pprocee: ngs..m.. ml press gress/v800 1080/haaroojaa 8b. html.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. volume 97 of Proceedings of Machine Learning Research, pp. 2555-2565, Long Beach, Californi,, USA, 09-15 Jun 2019. PMLR. URL http: http:///oocceeiieeeee
mir mir.a press v97 /twffeeeeeeeeeee html.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
URL https::///// / opwnrevieww net n forum? 1d=$110TC44
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9729-9738, 2020.
Olivier Hénaff, Aravind Srinivas, Jeffrey De Fauw, Ali Razavi, Carl Doersch, SM Eslami, and Aaron van den Oord. Data-efficient image recognition with contrastive predictive coding. arXiv preprint arXiv:1905.000
2019.
Lukasz. Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Camphell, 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. In International Conference on Learning Representations, 2020. URL https:// openreview.ner net forrm? www.ooommmmiiccoeecee 1=S11dC$
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 TxXXv:/705.066 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 44/266980. 2014.

---

Under review as a conference paper at ICLR 2021
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017.
Guillaume Lample and Devendra Singh Chaplot. Playing fps games with deep reinforcement learn- ing. In Proceedings of the Thirty-First AAAI Conference on Artificial Intelligence, AAAI'17, pp. 2140-2146. AAAI Press, 2017.
Nir Levine, Yinlam Chow, Rui Shu, Ang Li, Mohammad Ghavamzadeh, and Hung Bui. Prediction, consistency, curvature: Representation learning for locally-linear control. In International Con- ference on Learning Representations, 2020. URL https: /oppnreeeeeww..e//// net forum?
Ad=BJxx_OEttS.
Timothy Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In ICLR (Poster), 2016. URL
org/ abs /1599. 02971.
Zhuang Ma and Michael Collins. Noise contrastive estimation and negative sampling for conditional Consistency and statistical efficiency. Proceedings of the 2018 Conference on Em- Brussels, Belgium, October- November 2018. Association for Computational Linguistics. doi: 10.18653/v1/D18-14005 URL
models:
pirical Methods in Natural Language Processing, pp.
https:/ /www//ww....oo////oo wcwwwb..w.............. web. org/ /aa//antooolllo -
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 5187550):529-533 2015.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint arXiv: 1007003/8 2018.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On varia- tional bounds of mutual information. volume 97 of Proceedings of Machine Learning Re- search, pp. 5171-5180, Long Beach, California, USA, 09-15 Jun 2019. PMLR. URL http: / //proceedings mlr. press/ http////p//////e/////o 9a. html.
Rui Shu, Tung Nguyen, Yinlam Chow, Tuan Pham, Khoat Than, Mohammad Ghavamzadeh, Ste- fano Ermon, and Hung Bui. Predictive coding for locally-linear control. arXiv preprint arXiv:2003.01086, 2020.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. arXiv preprint arXiv:2004.00 04136, 2020.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
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 preprint arXiv: 00690, 2018.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. arXiv preprint arXiv: 05849. 2019.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in neural information processing systems, pp. 2746-2754, 2015.
2015.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.101 10742, 2020.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. SOLAR: Deep structured representations for model-hased reinforcement learning. volume 97 of Proceedings of Machine Learning Research, pp. 7444-7453, Long Beach, California. USA, 09- 15 Jun 2019. PMLR. URL http:/ http:://poceeeeeiiiiiiee /ptpppeeeingg...pppeepppp.p...se mlr wwe.spressssie / /twaagagggl 9m, html.
10

---

Under review as a conference paper at ICLR 2021
A HYPER PARAMETERS
A.1 STANDARD SETTING
NMPC and Dreamer share the following hyperparameters:
Model components
Latent state dimension: 30
Recurrent state dimension: 200
Activation function: ELU
The action model outputs tanh mean scaled by factor of and a softplus standard deviation for the Normal distribution that is then transformed using tanh (Haarnoja et al., 2018)
Learning updates
Batch size: 50 for Dreamer and 250 for NMPC
Trajectories length: 50
Optimizer: Adam (Kingma & Ba, 2014) with learning rates 6 10 for world model, 8 x 10-5 for value and action model.
Gradient update rate: 100 gradient updates every 1000 environment steps.
Gradient clipping norm: 100
Imagination horizon: 15
7 0.99 and À 0.95 for value estimation
Environment interaction
The dataset is initialized with S = 5 episodes collected using random actions.
We iterate between 100 training steps and collecting episode by executing the predicted mode action with (0, 0.3) exploration noise.
Action repeat: 2
Environment steps: x 106
Additionally.
Dreamer clips the KL below 3 nats
NMPC has fixed set of coefficient in the overall objective for all control tasks: A1
1, A2 = 0.1, A3 = 1, A4 = We use a fixed Gaussian noise e ~ (0, 0.22) to add to the future latent code when computing CPC, as suggested in (Shu et al., 2020), and also use 0.2 as the fixed variance in instantaneous CPC.
In NMPC, we also use a target network for the value model and update this network every 100 gradient steps. Note that we also tried to use target value network for Dreamer, but it does not improve the results, as suggested by their original paper (Hafner et al., 2020).
A.2 NATURAL BACKGROUND SETTING
To further encourage the model to focus task-relevant information from observations, increase the weight \ of the reward loss in the training objective for both Dreamer and NMPC. In each control task they share the same reward coefficient, which is specified in the table below. All of the remaining hyperparameters are kept the same.
11

---

Under review as a conference paper at ICLR 2021
Table 1: Reward coefficients for different tasks in the natural backgrounds setting
Task
Reward coefficient
Cartpole Swingup
1000
Cheetah Run
100
Walker Run
100
Acrobot Swingup
1000
Hopper Hop
100
Cup Catch
1000
B PROOF OF LEMMA
Our goal is to show that, under the conditions in Lemma I,
m1(=²) 2 (T(Tux)
(16)
for any choice of auxiliary encoder E".
We start by denoting = E (or) and s/c = (or). Note that the performance of π can be written
(Txux) = ((11)))())
(17)
= -Er r(ouT)
Il ploz oce, act)p(ss 56 oc) Taxx(a 8ge, s' act);
(18)
where Taux denotes the full trajectory of (0,8,8', a)ir and r(o:) evaluates the reward or (for simplicity, we shall assume p(re st) is deterministic. Since (K or) R (r, (o))) we can rewrite
(TTnx = Σ RC
I oca, (c))(((s de |0 (o) 85 &gg ac)),
(19)
Toux
where, with slight abuse of notation, note that Rs (o:) r(o:). We now further rewrite Taux 8ge, s'za, act) as
p(a: 8gg, 8ga, act, Taux),
(20)
and subsequently collapse the expression of the performance as
77 (axx))
Σ
(*1TT) ((o1: s1.7, si air Taux)
(21)
(0.8,5'.0) ,0hr
(11r)) Sim, our Taux),
(22)
(s,s',a)m
where the last step arises from marginalization of o1.. Note by chain rule that p(si.T, air Taux) becomes
[(( see, see, ace, (ps/ sge, s'z ace, Tpax s5n, s'gs aiee, Waux).
(23)
By analyzing the Markov blankets in p(81:7, 1Trr 01.r Taux), we can simplify the above expres- sion to
ppss sct, (ce)p()) sg1, s' act)p(a sga Seg acts Taux).
(24)
Note that omit the dependency Taux in the first two terms since, given only the history of past actions and observations, the next observation does not depend on our choice of policy but only on the environment dynamics.
12

---

Under review as a conference paper at ICLR 2021
Since E is optimal under the MI objective, note that
[S See Act St St) 1(5ct, Act St). Eq. (25) implies that s'zr is independent of & given (s<e, act), and that (#<es s'cns act) is indepen-
(25)
dent of 84 given 8t. This allow further simplify Eq. (24) to
Il p(s, 8cc,, a<))p(s) 8t) Ta)x (at &g
(26)
Thus, the performance expression equates to
(Tuxx)
II p(s scn act)p(s | st) sger
(27)
Tamx
Note by way of similar reasoning (up and including Eq. (24)) that
n(TT) Σ Rd (s1:)) II p(s: s<),a<))((a sgr,
(28)
By comparing Eq. (27) and Eq. (28), we that s'. effectively serves source of noise that makes Faux behave like a stochastic policy depending on the seed choice for si. To take advantage of this, we introduce reparameterization of s' as such that
77 (Tau))
= Σ r(SiIT)
II p(se sic,,
sgn, Ege, acx)
(29)
Taux
Ep(ere))
Σ r(SITT) II p(se 8ee, act) Taux
(30)
(e,0)1.r
max
Σ r(s1r)
Il p(s: 8ce aaiaTaaaa (gt)
(31)
61:T
(s,0)ur
<mxx)(())
(32)
where the last inequality comes from defining a policy
n' := Taux ogt, Eggo a2t))
(33)
and noting that the performance of xuxx must be bounded by the performance x²
13

---

Under review as a conference paper at ICLR 2021
C ADDITIONAL RESULTS
In Figure 6, we show the comparison NMPC with more baselines in the standard setting. They include Contrastive-Dreamer the contrastive version of Dreamer reported in Hafner et al. (2020), and Markovian Predictive Coding (MPC) in which we stack multiple frames and use a feed- forward dynamics instead of recurrent one. The figure shows the results with two frames stacked. MPC does not work in all of the control tasks. During training we observed that the CPC and Consistency objectives of MPC were much worse than those of NMPC. Therefore, we hypothesize that the ferw-forward dynamics with stacked frames is not capable of predicting the far future, which leads to poor performance in tasks that require long-term planning.
Acrobot Swingup
Cartpole Swingup
Cheetah Run
400
750
600
500
200
400
250
200
0-
0.5
1.0
1.5
1.5
1e6
1e6
1e6
Cup Catch
Hopper Hop
Walker Run
1000
800
300
600
200
400
250
100
200
0
0.5
1e6
1e6
1e6
Contrastive-Dreamrt
Dreamer
MPC
NMPC
Figure 6: NMPC vs the baselines in the standard DMControl tasks. The return computed in 1000environment steps. Each task is run with 3 seeds.
Figure demonstrates the ablation analysis which compares NMPC with other contrastive variants, including Pnatrastive-rraamee It casy to notice that all other models perform significantly worse than NMPC, verifying the necessary of different components that we proposed.
Cartpole Swingup
Cheetah Run
Walker Run
800
750
600
500
400
400
250
200
200
0
0
0.5
1.0
1.5
0.5
1.0
1.5
0.5
1.0
1.5
1e6
1e6
1e6
Contrastive-Dreamer
NMPC
NMPC-No-Inst
NMPC-No-Rer
Figure 7: Abaltion study with more contrastive variants in the standard setting
We also tried stacking three four frames and observed similar performances.
14

---

Under review as a conference paper at ICLR 2021
control tasks to make results comparable to Deep Bisimulation for Control (DBC) Zhang et al. (2020). We note that the gradient update rates of NMPC and DBC are different. In NMPC, update 100 gradient steps after collecting 1000 environment steps, while DBC performs one gradient update after each environment step collected. Even with setup differences that put NMPC at disadvantage, method either is competitive vastly outperforms DBC.
In Figure and 9, NMPC
Cartpole Swingup
Cheetah Run
Hopper Hop
800
600
400
400
200
0.5
1.0
1.5
0.5
1.0
1.5
1e6
1e6
1e6
Reacher Easy
Walker Run
Walker Walk
750
600
600
500
400
400
250
200
200
0.5
1.0
1.5
0.5
1.0
1.5
0.5
1.0
1.5
1e6
NMPC
Figure 8: NMPC on more control tasks in the standard setting
Cartpole Swingup
Cheetah Run
Hopper Hop
800
3
600
400
200
1
200
0
0
0.5
1.0
1.5
1.0
1.5
1e6
1e6
1e6
Reacher Easy
Walker Run
Walker Walk
800
150
500
600
400
w
100
200
200
50
0.5
1.0
1.5
0.5
1.0
1.5
le6
le6
le6
NMPC
Figure 9: NMPC on more control tasks in the natural background setting
15


Reviews:
[{"review_id": "XxE5d9Isotw", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "Problem Setup: The paper proposes a mutual information objective to learn a latent representation which\ncan be used for planning. The paper note that most of the existing model based RL methods learn a \nmodel of the world via reconstruction objective, which requires to predict each and every detail of the visual \ninput, and hence can be detrimental in case of noisy inputs or in the presence of distractors. \n\nProposed idea: In order to tackle this problem, the paper proposes a mutual information objective to maximize \nthe mutual information between the latent codes at distinct time steps.  In order to capture the history of the past, \nthe authors utilize a recurrent model (from Dreamer Model) to encode information about the history of the trajectory.\nThe paper also uses 2 different objectives in order to prevent the representation from collapsing. The paper proposes\nto use a mutual information objective between the observation and the encoding of the observation (as in Dreamer), \nas well as consistency objective in the latent space (already used before). Essentially the underlying idea behind the proposed method is not new per se, but as far as I know this is the first paper, which has shown to make it work on DeepMind Control tasks.\n\nExperiments: The authors compare the proposed method to DREAMER model on 6 DeepMind Control (DMC) tasks. \nThe authors also evaluate the robustness of the proposed method by evaluating the capability of the proposed method \nin dealing with complicated backgrounds (given the scenario, when the entity of interest occupies a small region in the input).\n\nClarity: The paper is clearly written. \n\nReferences: Their are bunch of references that could be cited. Shaping belief paper [1] also uses a CPC style objective for learning a model of the environment. [2] also learns a model of the environment by predicting only the relevant information by constructing a temporal information bottleneck but still within the framework of maximum likelihood prediction. [3] also uses a mutual information based objective and without any explicit reconstruction.\n\n- [1] Shaping Belief States with Generative Environment Models for RL https://arxiv.org/abs/1906.09237\n- [2] Learning dynamics model in reinforcement learning by incorporating the long term future https://arxiv.org/abs/1903.01599\n- [3] Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction\n https://arxiv.org/abs/2007.14535\n\nScalability: It would be interesting to see how the proposed method evaluates on more challenging tasks just as on atari or on continuous control tasks such as \"box\" stacking which requires some relational reasoning. Since the underlying idea has been tried in some other context and in this work the contribution is to make it work for deep RL problems, it becomes important to evaluate on more challenging problems and tasks. \n\n======\n\nAfter Rebuttal: I have read the rebuttal, as well as reviews by other reviewers. I keep my original score. Hope to see a better version of the paper soon.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "RCXPwwy5aMt", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe paper introduces a new method for model based RL that learns a dynamical latent representation from pixel data (images) using a maximum mutual information criterion together with a predictability loss. The core idea is that maximizing mutual information between states bias the encoder toward learning predictable (and therefore potentially task relevant) latent features while discarding unpredictable features.\n\nRelevance:\nThe paper addresses the very relevant problem of learning representations usable in a control task.\n\nOriginality:\nThe core novelty of the paper is to combine the mutual information approach introduced with PC3 in the context of control theory with the \"dream to control\" approach for model-based reinforcement learning in pixel space.\n\nScientific quality:\n- The proposed approach is in general well motivated. However. I am not convinced by the emphasis that the authors put in the proposition that the maximum mutual information loss helps to learn task relevant features. While it is true that the proposed approach is biased towards temporally predictable features, most distinctive features both in the real wold and in most games and simulations have as much temporal predictability as the task relevant ones. For example, the video backgrounds in the experiments are completely task irrelevant and at the same time highly predictable. In general, the encoder cannot trulely promote task-relevant features without having access to the reward structure.\n- The experiment section offers a decently wide range of experiments. However, the authors should include more baselines, possibly including other model-based methods such as [1] and model-free methods such as some variant of DQNs.\n\nPros:\n- Very relevant research area\n-Rather original combination of methods\n-Clear and well-written paper\n\nCons:\n- The main claim that the method is biased towards learning task-relevant features is questionable. \n-  The experiments should contain more baselines including other model based approaches and some model free approach.\n\nReferences:\n[1] Hafner, Danijar, et al. \"Learning latent dynamics for planning from pixels.\" International Conference on Machine Learning. PMLR, 2019.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "PCmD6cMTTXp", "reviewer": "AnonReviewer4", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThis paper proposes an information-theoretic framework for learning a world model that encodes task-relevant information of the world. It shows that the learned encoder and dynamics model can be used to train the policy and fitting the value function to agent to perform comparibly well to Dreamer on standard tasks and outperform them when there are distractions in the scene. The paper also provides a theoretical anylisis of the task-relevant information in the encoding.\n\nPros:\n\nThis paper provides a clean framework that learns the embeddings that contain task-relevant information.\n\nI appreciate the theoretical formulation to show that the optimal embedding will containis sufficient to train the policy on. I believe this theoretical analysis is novel.\n \nThe experiments provide an insightful ablation and show that the proposed contrastive model is less agnostic to the background distraction.\n\nCons:\n\nMy main conern is regarding the novelty. This paper misses a few recent relevant work [1, 2, 3]. These work also use contrastive losses to learn a latent world model and outperforms Dreamer on DMControl tasks and natural background. I'd like to see the comparison against [1,2] or clarify how NMPC is different and/or better in some ways. \nHow do you think NMPC performs against this model-free version with similar task-relevant objective [4]? \n\nIn the orignal paper dreamer also proposes a contrastive loss for training the latent model although it is not temporal. Is the ablation NMPC-No-Rcr in figure 2 comparable to that or different? Also, how important is having a recurrent dynamics model as supposed to a feedforward dynamics model?\n\nHow senstive the algorithm is to the hyperparameter combinations? The loss seems to have a few terms to be tuned. It would be helpful if the ablation can show the importance of other components such as Non-Markovian consistency and dynamics smoothing.\n\nConclusion:\n\nOverall, the paper is well written and easy to understand. It proposes a simple framework with experimental and theoretical support. The main downsight lies in its novelty as a few other works that aim to tackle this issue in a similar way. If this can be addressed, I think it is a really good paper.\n\nReferences:\n\n[1] Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction (https://arxiv.org/abs/2007.14535)\n\n[2] Contrastive Variational Model-Based Reinforcement Learning for Complex Observations (https://arxiv.org/abs/2008.02430)\n\n[3] Learning Predictive Representations for Deformable Objects Using Contrastive Estimation (https://arxiv.org/abs/2003.05436)\n\n[4] Learning Invariant Representations for Reinforcement Learning without Reconstruction (https://openreview.net/pdf?id=-2FCwDKRREu_", "extra_scores": {}, "extra_text": {}}, {"review_id": "C7VoDvO8XkC", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "---- Summary ----\n\nThe paper proposes a method for visual model-based reinforcement learning that relies on contrastive learning to learn the predictive model. Building on Hafner’20, the paper replaces the image reconstruction objective with a noise contrastive estimation (NCE) objective for the latent dynamics model, an NCE objective between the images and representations, and an additional maximum likelihood objective for the latent dynamics. It is shown that the method is competitive to Hafner’20 on the DM Control benchmark, and outperforms Hafner’20 on DM Control tasks with natural images used as background. The paper also theoretically analyzes one of the used objectives, arguing that it may lead to discarding irrelevant information.\n\n---- Decision ----\n\nThe paper tackles a relevant and important problem of building predictive models that do not rely on image reconstruction, and proposes a promising method to this end. However, there are several technical weaknesses of the presented approach, and further important baselines and ablations are omitted. Due to lacking experimental evaluation, I lean toward rejecting the paper, but would be happy to update my score if the experimental evaluation were improved.\n\n---- Strengths ----\n\nThe paper investigates a relevant and promising direction. The experiments show the effectiveness of the proposed method in simple continuous control scenarios, further supported by some theoretical analysis.\n\n---- Weaknesses ----\n\nThe major weakness of the paper is a lack of representative baselines. The paper cites several similar papers, including Ding’20, and Shu’20, which the model is called “an extension of”. How is the proposed method better than these prior methods? Further, the central claim of the paper is the importance of non-markovian predictive coding, also claimed to be the main technical novelty (although this was already used in the foundational paper by van den Oord). It is never evaluated whether the non-markovian part improves performance. The contrastive learning version of Dreamer should be added to the plots as well.\n\nA harder to fix issue is that the proposed method is rather unprincipled. Three different competing objectives are used (e.g. maximum likelihood of the latent dynamics reduces I(z;z’) while the other objectives increase it), and the interactions between them are not discussed. The paper presents only intuitive justification for the three objectives. This issue also causes the need to tune balance terms between different losses. For instance, the method requires tuning of the weight on the reward prediction, and requires different weights for different environments (Table 1 in the appendix), while in Dreamer this weight is always set to 1 and no hyperparameters are environment-specific. The weight for the recurrent CPC objective $\\lambda_1$ does not seem to be specified in the paper.\n\n---- Additional comments ----\n\nEq 1 - none of the symbols in the equation are defined. \n\n### ---- Update ----\n\nThe authors' response does not satisfactorily address my concerns. My main concern is that the paper does not properly evaluate alternative choices, even though a large literature on contrastive learning exists.  While in the revision one baseline was added to the appendix, the main experiment still only contains a comparison to Dreamer and ablations. Further, it appears that the proposed method fails completely when the non-markovian part is removed. This is rather concerning since learning markovian latent dynamics is important and also possible with other methods (e.g. PlaNet-RNN). As far as I can tell, the paper does not discuss this issue and does not explain why learning non-markovian dynamics is crucial.  \n\nOverall, the paper proposes an interesting method but fails to provide any insight into how the method compares with possible alternatives. I, therefore, maintain my borderline score. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents Non-Markovian Predictive Coding (NPMC), a method for learning state representations in visual RL domains that can be used for planning. This work builds on recent work on PC3 (Shu et al. 2020) and PlaNet (Hafner et al. 2020). Concretely, NPMC replaces the image reconstruction objective in PlaNet with a noise contrastive estimation (NCE) objective for the latent dynamics model, an NCE objective between the images and representations, and an additional maximum likelihood objective for the latent dynamics.

Reviewers were in agreement that this paper tackles an important problem and appreciated the writing quality, the experiments that demonstrate effectiveness of NPMC in continuous control scenarios, and accompanying theoretical analysis. However, reviewers were on balance in consensus that this paper needs another iteration before it can appear. Aside from discussion of related work, the main weakness noted by reviewers is that the manuscript in its current form makes it clear how NPMC differs from closely related methods from a technical point of view, but does not make it sufficiently clear to what extent non-Markovian predictive coding leads to improved planner performance. In particular, the paper lacks detailed comparisons to baselines, and reviewers were not sufficiently convinced by experiments that were added to the appendix after discussion.

The authors indicate that their contribution is that NPMC extends PC3 to RL tasks. The metareviewer appreciates that experimental comparisons can require creative thinking when baselines are not directly applicable to the tasks of interest, but would nonetheless like to encourage the authors to consider how they can improve their experiments.
Decision: Reject