Learning Latent State Spaces for Planning through Reward Prediction
Authors: Aaron Havens
Yi Ouyang
Prabhat Nagarajan
Yasuhiro Fujita
Venue:  (2020)
Keywords: Deep Reinforcement Learning
Representation Learning
Model Based Reinforcement Learning

Abstract: Model-based reinforcement learning methods typically learn models for high-dimensional state spaces by aiming to reconstruct and predict the original observations. However, drawing inspiration from model-free reinforcement learning, we propose learning a latent dynamics model directly from rewards. In this work, we introduce a model-based planning framework which learns a latent reward prediction model and then plan in the latent state-space. The latent representation is learned exclusively from multi-step reward prediction which we show to be the only necessary information for successful planning.  With this framework, we are able to benefit from the concise model-free representation, while still enjoying the data-efficiency of model-based algorithms.  We demonstrate our framework in multi-pendulum and multi-cheetah environments where several pendulums or cheetahs are shown to the agent but only one of them produces rewards. In these environments, it is important for the agent to construct a concise latent representation to filter out irrelevant observations. We find that our method can successfully learn an accurate latent reward prediction model in the presence of the irrelevant information while existing model-based methods fail. Planning in the learned latent state-space shows strong performance and high sample efficiency over model-free and model-based baselines.

Under review as a conference paper at ICLR 2020
LEARNING LATENT STATE SPACES FOR PLANNING THROUGH REWARD PREDICTION
Anonymous authors Paper under double-blind review
ABSTRACT
Model-based reinforcement learning methods typically learn models for high- dimensional state spaces by aiming to reconstruct and predict the original obser- vations. However, drawing inspiration from model-free reinforcement learning, we propose learning a latent dynamics model directly from rewards. In this work, we introduce a model-based planning framework which learns latent reward prediction model and then plan in the latent state-space. The latent representation is learned exclusively from multi-step reward prediction which we show to be the only necessary information for successful planning. With this framework, we are able to benefit from the concise model-free representation, while still enjoying the data-efficiency of model-based algorithms. We demonstrate in multi- multi--endulum and multi-cheete environments where several pendulums or cheetahs are shown the agent but only one of them produces rewards. In these environments, it is important for the agent to construct a concise latent representa- tion to filter out irrelevant observations. We find that our method can successfully learn an accurate latent reward prediction model in the presence of the irrelevant information while existing model based methods fail. Planning in the learned latent state-space shows strong performance and high sample efficiency over model-free and model-based baselines.
framework
INTRODUCTION
Deep reinforcement learning (DRL) has demonstrated some of the most impressive results on several complex, high nigh-dimensional sequential decision-making tasks such as video games (Mnih et al. (0015) and board games (Silver al. 2018). Such algorithms are often model-free, i.e. they do not learn or exploit knowledge of the environment's 's dynamics. Although these model-free schemes can achieve state-of of-te-of-th performance, they often require millions of interactions with the environment. The high sample complexity restricts their applicability to real environments, where obtaining the required amount of interactions is prohibitively time-consuming. On the other hand, model-based approaches can achieve low sample complexity by first learning model for the environment. The learned model can then be used for planning, thereby minimizing the number of environment interactions. However, model-basee methods typically limited to low-dimensional tasks or settings with certain assumptions imposed on the dynamics (Li & Todoroy 2004 Deisenroth & Rasmussen 2011 Levine
& Koltun 2013 Chebotar et al. 2017)
Several recent works have learned dynamics models parameterized by deep neural networks (Lenz ct al. 2015 Finn & Levine 2017: Williams al., 2017), but such methods generally result in high-dimensional models which are unsuitable for planning. In order to efficiently apply model-based planning tigh-dimensional spaces. promising approach is to learn a low-dimensional latent state representation of the original high-dimensional space. A compact latent representation is often more conducive for planning, which can further reduce the number of environment interactions required to learn good policy (Watter et al. (2155) (aanjaaaal et al.) 2017]] Finn et al., 2016 Ichter & Pavone 2019 Zhang et al. 2019 Hafner et al. 2019). Typically, such a model has a variational autoencoder (VAE) (Kingma & Welling 2013 like structure that learns a latent representation. In particular, it learns this representation through reconstruction and prediction of the full-state from the latent state. In other words, the learned low- -iw-iineniional latent state must retain sufficient amount of information from the original state to accurately predict the state in the high- aigh-dimensional space. Although the ability to perform full-state prediction is sufficient for planning in latent space, it is not

---

Under review as a conference paper at ICLR 2020
strictly necessary. Predicting the full-state may result in latent state representation which contains information irrelevant to planning. For example, consider a task of maze navigation, with a TV placed in the maze, displaying images (Burda et al. 2018a). The TV content is irrelevant to the task, but the latent representation may learn features designed to predict the state of the TV.
In RL, the objective is maximize the cumulative sum of rewards. Given this objective, it may be excessive to learn to perform full-state prediction, if aspects of the full state have no influence on the reward. In examples like the aforementioned TV in maze, it may be more fruitful to instead optimize a latent model to be able to perform reward prediction. In some sense, value-based model- free methods can be seen as reward prediction models. However, they attempt to predict the optimal value which is policy dependent and requires a large number of samples to learn. Fortunately, model- based planning methods such as model predictive control (MPC) only maximize the cumulative sum of future rewards over the choices of action sequences. This implies that what planning method requires a latent reward prediction model: model which predicts current and future rewards over a horizon from latent state under different action sequences. In this work, we introduce a model-based DRL framework where we learn a latent dynamics model exclusively from the multi-step reward prediction criterion and then use MPC to plan directly in the latent state- spate. Learning the latent reward prediction model benefits from the camplee-fficiency of model-based methods, and the learned latent model is more useful for planning than its full-state prediction counterparts discussed above.
The contributions of this paper are as follows. The proposed latent model is learned only from multi-step reward prediction. Optimizing exclusively for reward prediction allows to circumvent the inefficiency stemming from learning irrelevant parts of the state space, while maintaining the model-based sample efficiency. We provide a performance guarantee on planning in the learned latent state-space without assumptions on the underlying task dynamics. Empirically, we demonstrate that our method outperforms model-free and model-based baselines on mulli-penduulle and multi-cheete environments where the agent must ignore large amount of irrelevant information.
The remainder of this paper is structured as follows. In Section 2 we present related work in model- based DRL. In Section 3 we provide the requisite background on RI. and MPC. In Section 4 we introduce our method. In Section [5] we present our experimental results, and conclude by indicating potential future directions of research.
2 RELATED WORK
Model-based DRL Model-based works have aimed to increase sample efficiency and stability by utilizing a model during policy optimization. Methods like PILCO (Deisenroth & Rasmussen 2011). DeepMPC (Lenz et al. 2015), and MPPI (Williams et al.. 2017) have exhibited promising performance on some control tasks using only handful of episodes, however they usually have difficulty scaling to high-dimensional observations. Guided-policy search methods (Levine & Koltun et al.) 2017) make use of local models to update global deep policy, however these local assumptions could fail to capture global properties and suffer at points of discontinuity. For high-dimensional observations such images, deep neural networks have been trained directly predict the observations (Finn & Levine 2017). but the resulting high-dimensional models limits their applicability in efficient planning.
Learning latent state spaces for planning To make model-based methods more scalable, several works have adopted an auto-encoder scheme to learn reduced-dimensional latent state representation.
Then one may plan directly in the latent space using some conventional planning methods such as linear quadratic regulator (LQR), or rapidly exploring random tree (RRT) (Watter et al. (0155] Banijamali et al.) [017]] Finn et al., 2016 Ichter & Pavone 2019 Zhang et al.) 2019). However, these methods construct latent spaces based on state reconstruction/ preeicccoooo which may not be entirely relevant to the planning problem as discussed in the previous section. A recent work PlaNet (Hafner et al., 2019) offered a latent space planning approach by predicting multi-step observations and rewards. The PlaNet has shown to achieve good sample complexity and performance compared to the ofate-of-thh-an model-free methods, but its latent model might also suffer from the aforementioned issues with irrelevant observation predictions. In this work we use similar planning algorithm to PlaNet, however, removing the dependency of observation prediction provides a more concise latent representation for planning.
2

---

Under review as a conference paper at ICLR 2020
Reward-based representations learning for RL The work of VPN (Oh et al. 2017) proposed RL framework which learns a model to predict the future values rather than the ftal-states of the environment. Although VPN shares very similar idea as framework to not predict irrelevant information in full-state, suffers the sample complexity issue of learning the optimal value as in model-free methods. Performance guarantees of latent models were analyzed under the DeepMDP framework in a recent paper (Gelada et al., 2019) without full-state reconstruction or prediction. With Lipschitz assumptions, their work provides performance bounds for latent representations when the two prediction losses of the current reward and the next latent state are optimized. They also show improvements in model-free DRL by adding the DeepMDP losses auxiliary objectives. However, with only the single-step prediction loss, DeepMDP may not be suitable for making long term reward predictions necessary
planning algorithms. In addition, the latent state prediction loss may be unnecessary and could result in local minima observed in their paper. Our framework shares the same spirit DeepMDP by learning latent dynamics only from the reward sequences, but focusing on multi-step reward prediction allows to achieve higher performance in model-bases planning.
3
PRELIMINARIES
3.1 PROBLEM SETUP
In this paper we consider discrete time nonlinear dynamical system f f -> S witt continuous state and action spaces S S RM A S Rdm, and reward function R: 8 x A + R. Then, given an admissible action a; A at time the system state 8, evolves according to the dynamics
St41 az)
where 81+11 is the next state time +
Our goal is to find a policy π S + A that selects actions far = ((s)). t 0, 0......) so as to maximize the cumulative discounted rewards
Σ r² R(st. at) t=0
(2)
where the initial state so assumed to be fixed, and y is the discount factor.
Note that Equations ( and (2) formulate infinite horizon optimal control problem. This problem can also be viewed as a deterministic Markov decision process (MDP) by the tuple (S,A, f, R,
3.2 MPC PLANNING
When dynamics model is available, model predictive control (MPC) is powerful planning framework for optimal control problems. A MPC agent chooses its action by online optimization over a finite planning horizon H. More specifically, at each time t, the agent solves H--horizon trajectory optimization problem:
max (St, (=((+-1)))
(3a)
8-1-1
subject
J(s, (
Σ T- (-((6),,
(3b)
s sc, =,, and 8,41 = (8,, à,) for = 1,, (1+
(3c)
Suppose the optimal control sequence is (* then the MPC agent will select à as the control action time t. This method is termed model predictive control due to its use of the model f to predict the future trajectory $te+ ():ft1):ft+-/// --1) from state s, and the intended action sequence ar.(t+ Let (s
(s: If R) denote the MPC policy using the dynamics model ff and reward R. Then the MPC agent selects its action by (MPc(s:]] f., R) 05-
3.3 SAMPLE-BANE TRAJECTORY OPTIMIZATION
From the description Section 3.2
an MPC agent's 's need to solve the trajectory optimization problem In this work, the cross entropy method (CEM) (Rubinstein)
specified by Equation equation
3

---

Under review as a conference paper at ICLR 2020
1997) to solve problem equation CEM is very simple sampling- basedd planning method which only relies on a forward model and trajectory returns which may make direct use of our learned latent mappings somewhat like simulator. The advantage using method like CEM, is that it's sampling procedure can he easily parallelized directly in the learned latent space.
LATENT REWARD PREDICTION MODEL FOR PLANNING
In order to perform model-based planning directly from observations without a model, we need to learn dynamics model. For nigh-dimensional problems, common model-based DRI. approach is learn latent representation and then conduct planning in the low-dimensional latent space. this work, rather than predict the full state/observation, we learn latent state-space model which predicts only the current and future rewards conditioned on action sequences. We will discuss that observation reconstruction is not necessary if predict the rewards well over the planning horizon.
4.1 LEARNING A LATENT REWARD PREDICTION MODEL
Our model consist of three distinct components which can be trained an end-to-end fashion. To do we first embed the state observation s, the latent state with the parameterized function dg S + Z C Idd,, In order to propagate the latent state forward in latent space, we define the x
forward dynamics function which maps f3 dynamical system in latent state space. Finally, require reward function R 2 A - R which provides us an estimate of the reward given a latent state 2₁ and control action dĄ. The full latent reward prediction model is depicted in Figure
2, corresponding the discrete-time
2x
(st), 2441 fz(2, at), = Rē ac)
Note that we do not require a "decoding" function since planning algorithms often only need to evaluate the rewards along a trajectory given an action sequence.
ds
1
f3
å
8L
R
aL
R
0++1
R
O--I-1
f141
f.. H-
Figure 1: The latent reward prediction model (de. fū. R2). Blue and green circles represent the input and output variables. The squares indicate components of the prediction model and white circles represent the latent variables. The model takes in an initial state 8, and sequence of actions (n(+++-11 to return a sequence of H-step predicted rewards fe:(t+
To achieve the aforemmentioned reward prediction ability, we learn these functions Φe, fū. R² over multi-step reward prediction losses. Formally, given H H-step sequence of states and actions {(sk, ox), k = t....,t- H 1}, the training objective is defined to be the mean- squa--s error between the original rewards and the multi-step latent reward predictions.
6+H-1
Σ
11
R(s,, - - (R a7)11
(5)
H
T=1
where
f a.),
&,
(6)
We later show that this -see reward prediction loss sufficient to bound the planning performance over the same horizon.
Note that we do not impose other losses such as the observation prediction loss or latent state prediction loss as in previous latent model frameworks (Gelada al. 2019 Hafner al.) 2019).
4

---

Under review as a conference paper at ICLR 2020
The reason is that other loss functions not necessary for planning performance. Having those additional losses would increase the training complexity and may induce undesirable local minima.
4.2 MPC WITH THE LATENT REWARD PREDICTION MODEL.
Once a latent state-space model is learned, we can perform MPC using the latent embedding, Let TMPc($ (ho, fū RR)) be the latent MPC policy with the latent reward model (do. J RR). Similar to
Equation 33. the latent MPC agent each time solves H-step trajectory planning problem.
max (s, (c)++H-1))
(7a)
-1)
2+H-1
subject
(
Σ ---
(2r,
(7b)
& (0 and 2r+1 ( 2,) for t,t t 1...., (t
(7c)
After solving the optimization problem, the latent MPC agent will choose its action by
(M Jū: R2) = 0z where (1:(++H-1) argmaxx Js dc((+H-1)) (1:(c+H11)
Similar to other model-based methods, the latent reward prediction model can be used in offline setting with a pre-collected dataset. Together with an exploration policy, e.g. e-greedy, our model can also be used in online iterative scheme. See appendix for the descriptions of Algorithm (offline) and Algorithm (online iterative).
4.3 PLANNING PERFORMANCE GUARANTEE
To evaluate the performance of MPC planning, define recursively the optimal n-step Q-function by
Qi(8, a) R(s, a); (x, R(s. + max Q Qf for 2.3...... a'EAA
Then the Bellman optimality shows that (7 a) gives the optimal value for the MPC planning problem B) i.., maxa [ss, (1((++-11) max, Q ($;, a).
Similarly, define the optimal n-step Q-functions for the latent state-space model (do. fú: Rš):
0 = a), 0 (z,a) Re(2,a) a) + max (((,))()) for 2....... H
d'EA
We have the following result connecting the H H-step reward prediction loss and the H H-step Q-functions. Theorem 1. Suppose the H-step prediction losses LH S 6² for any H H-step trajectory. Then the H-step Q-functions satisfy
O (0o(s))) 2 (s) - 1/2
for all E S, A. Moreover, if Túpc(s 60, fū. Rt) is optimal for the latent MPC planning policy for the MPC planning problem 3)
problem equation then is also a
This result provide a planning performance guarantee for the latent reward prediction model. also confirms that don't it need other unnecessary losses achieve the desired performance from the latent model. Note that this result does not require additional assumptions such as Lipschitz continuity of the environment. See the appendix for proof of the theorem.
5

---

Under review as a conference paper at ICLR 2020
5 EXPERIMENTS
In this section we introduce several experiments and their results which are aimed to investigate the representation ability, long-term prediction ability and sample efficiency. We implement the mappings Φ0, f R by deterministic feed-forward neural networks. The extensions to stochastic and/or recurrent neural networks are potential future directions. In order to do MPC, we choose to use CEM with a planning horizon of H 12 and K 1000 trajectory samples. We use the 100 best trajectories to compute the control, re-planning after every action. CEM will use model somewhat like simulator to collect samples in latent space (in case this is batch forward pass in a deep network). We consider two sets of environments with sigh-dimensional irrelevant observations. Results for additional image-based based environment is also available in the appendix.
5.1 BASELINES
Before introducing the experiments, briefly introduce the baseline algorithms used for comparison. Throughout our experiments, we have several baselines, each serving to demonstrate the attributes of reward model: concise representation, long-term prediction and sample efficiency.
To demonstrate the representation ability of our reward prediction model, we consider latent model where the latent space is learned from state prediction only. We train reward predictor separately for the model. This model works well for the concise representations, however, like many model-based methods it fails to scale to high-dimensional observations.
State prediction model
DeepMDP We use the DeepMDP model (Gelada al. 2019), which is essentially one-step variant of the reward prediction model with an additional latent state prediction loss. We aim to show that our multi-step step reward prediction-only loss improves the planning performance.
SAC Finally to demonstrate sample efficiency, compare against the model free Soft Actor-Critic algorithm (SAC) (Haarnoja et al. 2018). SAC currently a popular DRL algorithm known for its sample efficiency and performance.
5.2
MUL.TI-PENDULLMM
In several RL tasks, there may be irrelevant aspects of the state, e.g. the noisy TV problem (Burda et al. 2018b), so reward prediction might be extremely crucial to learn a more conducive and useful model for planning. To verify this claim, develop multi nulti-pendulum environment, of the classic pendulum control task (Brockman et al. 2016), where the agent must swing up and stabilize a pendulum at the up-right state with torque constraints. Although this task is simple, encompasses nonlinear dynamics and long planning horizons in order to accumulate enough energy task, the agent operates in an environment with N pendulums,
adaptation
to swing up. In the multi-pendulle but is only rewarded for swinging up a single pendulum. The other N pendulums serving as determined observational noise being driven by uniform random actions. The state observation for the N -pedduluu environment is given
8 (0i), cos(0,) (0), sin (fz), cos((2), 02, sin cos(O R3N
Since reward prediction model can be used either completely offline iterative online setting, present two sets of experiments. First we show the average final performances of Algorithm and baselines after training completely offline from 20000 steps of the environment under time-correlated noised control. We vary the number of pendulums N E (1, 2, 5, 10, 15) where we maintain the latent space to 3-dimensional throughout. In the second experiment we train the model and baselines under the iterative scheme described in Algorithm [ consuming approximately the same number of samples. We also compare against the SAC algorithm as a state-of-the -hee---a----rtt model-free baseline.
As shown in Figure 2 the reward prediction loss does not degrade quickly for our reward prediction model as the number irrelevant pendulums increases, but degrades almost immediately for the state prediction model. DeepMDP struggles to predict well in all cases, due to only training single-step latent plus reward reconstruction loss. In Table we can further see that performance of the reward prediction model (Algorithm T) is not very sensitive to the number of excess pendulum

---

Under review as a conference paper at ICLR 2020
environments, where the state prediction model immediately fails to solve the task. Note that the SAC method requires a significantly greater number of samples to achieve similar performance.
For the iterative scheme (Algorithm 2). the reward prediction model clearly outperforms all baselines as can be seen in Figure 3 Note that the DeepMDP model has a much higher variance than other methods. This high variance could be due the local optima from the latent state prediction loss in DeepMDP. On the other hand, our reward prediction model provides a stable low-variance result.
10sstep rewerd model
10-step stace peediction
Deepadp model
Peudduuns
Peadulams
racan
10: Penéibume
4
15 Pendilluss
5000
30000
15000
20000
5000
16000
18000
20000
5000
10000 15000
20003
Ireaation (barch of 256)
Ineraction (barch of 256)
Itenation (barch of 256)
Figure 2: Training curves for the 10-step reward prediction loss. All experiments are ran in the offline setting using 20000 steps from the true environment under a random policy and trained for 300 epochs. The results of our reward prediction model is shown in the left, the state prediction model the middle, and DeepMDP in the right.
Reward Model
State Model mean(std)
Dee
SAC (10' samples) mean(std)
# pendulums
mean( std)
mean(std)
138. 238.28(115 103.83)
137. 137.65 (79.01))1)
3374.2(1113.2)
145. 74..26(92. (0)
2
163.99 79.76)
451. 151.72(188.01)
337. (46.48)
150.31 (90.53)
5
166.178(79.708
858. (7((15.000)))
665509((99)) (196.32)
151.85(
(51.85(100.82 (100.82)
10
(91.46)
911263(160.33
-539. -539.85(131) (1300.8)
-151.41(83.21
15
-185.36(128.84) -922.45(151.76) -578.89(172.26)
-14..92994.21)
Table 1: Statistics over the final performance of the multi multi-pendulum environment trained under the offline setting from 20000 environment steps under random-policy (excluding SAC). For each of the latent model, we use CEM for MPC planning, and each method is evaluated from 100 episodes.
5.3 MULTI-CHEBTAH
In order to demonstrate that method scales to higher dimensional environment, use the Deepmind control suite "cheetah-run" environment (Tassa et al. 2018) to construct "multi-cheetah" The cheetah-run environment has continuous 17-dimensional observation space and 6-dimensional action space. Similar to multi-pendulum, we concatenate each cheetah observation vector together into one observation, where only the first cheetah environment return the reward signal. Because random policy is not sufficient to explore the state space in cheetah, the remaining N - cheetahs act according to an expert SAC policy. In multi-cheetah, we mainly compare our method with the model-free baseline SAC to highlight the strong performance and sample efficiency. Shown in figure the reward prediction model able to learn and plan near-optimally using far less samples than SAC in the 1-cheetah environment (i.e., the regular cheetah-run). Within 500 episodes, the reward model is able to achieve a near-optimal running behavior (503 and 501 final return evaluation for 1-cheetah and 5-cheetah respectively). which takes SAC 3000 episodes to achieve. After 1000 episodes of 1-cheetah (not shown the figure), the reward model achieves 570 return, approaching SAC's maximum performance of 590 after 4000 episodes. When there several irrelevant running cheetahs, the performance SAC further drops. In comparison, our reward prediction model is rather agnostic to the irrelevant observations and maintains high performance and sample efficiency.
7

---

Under review as a conference paper at ICLR 2020
5 Feusullan Itrraiie Training
-230
step reward mode.
10-ttpp state modal
-400
ruvan inwaadd aang
Dergimip madel
6.00
SAC
-830
-1030
-1230
1400
100 200
300
400
500 600
episodes
Figure 3: Training curves for the iterative scheme in the 5-pendulum environment. We initialize the each model with 2500 steps, and then perform 100 training iterations batch-size 256 after each episode collected with the current policy with 0.7-greedy exploration. Each method is run with different seeds. The solid lines and shaded areas are the average training returns and their one standard deviation regions under the exploration policy. The corresponding dashed lines mark the average final evaluation performance after x 23800 environment samples.
Lecintth Buriiie Taining
caretali Ieeratvee Training
300
10-step
Deepond
600
S.C.
3000
100
200
2011
100
100
optocobss
epiocdess
Figure 4: The average return training curves for SAC, Deepmdp and our reward prediction model in the 1-cheetah and 5-cheetah environments (85-dimensional). The reward model is trained iteratively online with zero-mean Gaussian exploration noise. Solid lines indicate training return (with seeds and shaded one standard deviation regions), and the dashed lines are the final evaluation returns without exploration noise.
CONCLUSION
In this paper we introduced method for learning latent state representation from rewards. By constructing a latent space based on multi-step reward predictions, we obtained a concise represen- tation similar that of model-free DRL algorithms while maintaining high sample efficiency by planning with model-bate base algorithms. In this work we have used sample-based CEM algorithm for MPC planning, but the use of other planning algorithms also plausible. We demonstrated in the multi-pendulum and multi-cheetah environments that the latent reward prediction model is able to succeed in the presence of high-dimensional irrelevant information with both offline and iterative online schemes.
Although, there are several aspects which still need be investigated such as more complex environ- ments and sparse reward settings, where typically a finite-horizon planner would struggle. We would also like to investigate how to incorporate model uncertainty through either stochastic dynamics or ensemble methods. One interesting question have is what is the optimal observer what the optimal combination reward and state reconstruction to do well in the task. We believe our minimal formulation good starting point for this question.
8

---

Under review as a conference paper at ICLR 2020
REFERENCES
Ershad Banijamali, Rui Shu, Mohammad Ghavamzadeh, Hung Bui, and Ali Ghodsi. Robust locally- linear controllable embedding. arXiv preprint arXiv:/710.053373, 2017.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. CoRR, abs/1600.015500, 2016. URL http:://axxii.orrt
abs/1606 01540
Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A Efros. Large-scale study of curiosity-driven learning. arXiv preprint arXiv: 180,00550 2018a.
Yuri Burda, Harrison Edwards, Deepak Pathak, Amos J. Storkey, Trevor Darrell, and Alexei A. Efros. Large-scale study of curiosity-driven learning. CoRR, abs/1808.00 808.04355, 2018b. URL
http://arxlv.org//
04355
Yevgen Chebotar, Mrinal Kalakrishnan, Ali Yahya, Adrian Li, Stefan Schaal, and Sergey Levine. Path integral guided policy search. In 20/7 IEEE international conference robotics and automation (ICRA), pp. 3381-3388. IEEE, 2017.
Mare Deisenroth and Carl E Rasmussen. Pilco: A model-basee and data-efficient approach policy search. In Proceedings of the 28th International Conference on machine learning (ICML-11), pp- 465-472, 2011.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. 20/7 IEEE International Conference Robotics and Automation (ICRA), pp. 2786-2793. IEEE, 2017.
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pp. 512-519. IEEE, 2016.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pp. 2170-2179, 2019.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Schoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv: 1812.05905, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2565, 2019.
Brian Ichter and Marco Pavone. Robot motion planning in learned latent spaces. IEEE Robotics and Automation Letters, 4(3):2407-2414, 2019.
Diederik P. Kingma and Max Welling, Auto-encoding variational bayes. CoRR, 1bs/1312.6114, 2013.
Ian Lenz, Ross A Knepper, and Ashutosh Saxena. Deepmpe: Learning deep latent features for model predictive control. In Robotics: Science and Systems. Rome, Italy, 2015.
Sergey Levine and Vladlen Koltun. Guided policy search. International Conference on Machine Learning, pp. 1-9, 2013.
Weiwei Li and Emanuel Todorov. Iterative linear quadratic regulator design for nonlinear biological movement systems. In ICINCO (1), pp. 222-229, 2004.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529, 2015.
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In Advances in Neural Information Processing Systems, pp. 6118 18-1188,888881 2017.

---

Under review as a conference paper at ICLR 2020
Reuven Y Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1):89-112, 1997.
David Silver, Thomas Hubert, Julian Schrittwisser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140-1144, 2018.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. CoRR, 0bss/1801.00699 2018. URL htttp://arxi.. www.xiv.orrr..oooooo org/ abs/
1801. 00690
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.
Grady Williams, Nolan Wagener, Brian Goldfain, Paul Drews, James M Rehg, Byron Boots, and Evangelos A Theodorou. Information theoretic mpe for model--ased reinforcement learning. In 2017 IEEE International Conference on Robotics and Automation (ICRA), pp. 1714-1721. IEEE, 2017.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for moodel-based reinforcement learning. In International Conference Machine Learning, pp. 7444-7453, 2019.
10

---

Under review as a conference paper at ICLR 2020
A ALGORITHM DESCRIPTIONS
Algorithm 1 Offline Training
Inputs: planning horizon H, dataset D, latent model (0o, J R
Train (0o, f RR) using the loss function Cи in equation and data D.
Output: policy ( f2. R()
Algorithm 2 Online Iterative Scheme
Inputs: planning horizon H, number of iterations latent model (øe. f Rt).
Initialize dataset 1) with random trajectories.
for i = 1 to N do
Train (do, fā. 1Rt) using the loss function Ls in equation 5 and data D.
Run latent MPC ( Φn. f RR) in the environment with e-greedy and collect data D;.
D-DUD, end for
Output: policy ( l0. f3. R
B PROOF OF THEOREM
Proof. Note that the H-step Q-function can be written
#-1
Qn((s,)) = Σ
R(s az)
1=
where 5g = s, oñ = a, and for
81 ( f((s, a; = n"" (st)
Here n² denotes the optimal policy for the MPC problem equation Because Q (00(8), a) is the optimal Q-function for the latent MPC problem equation we have
Y (oo(s))
Σ ? (
Adding and subtracting Qú (8, a) obtain
()(()
7' (GG() (s,
H-1
(s, +
Σ
(10(( -)(()))))))
H-1
(8)((,
Σ
1I' (o) R(s,, af))le
1=0
(8, (H
Q2 (s, 1/2
(10)
Now let (2) be the optimal policy for the latent MPC problem in the latent space, then from the Q-function definition have
(0o(s)))
Σ
o' (
T=
11

---

Under review as a conference paper at ICLR 2020
where 20 #0(ss, ap a, and for t
"t = (+)-1, (o,(()0))) (
a , = =
The H- Htsp Q-function of policy ग- Φ0 in problem equation 3 is then given by
11
QH 0 (s,
=
y R[ff(s, 00:11);
1=
QN (6()).))
Σ
> (olfe(s,00-11).00 -
#-1
(R(S on (0
(0 - H 1/2
(11)
Now combining equation 10 and equation we get
QH
(s, a) 2 (s, 1/
which means that n o 6 is 2H /2 e-optimal policy.
C EXPERIMENTS WITH IMAGE OBSERVATIONS
To demonstrate that this method can solve tasks with rather high-dimensional state observations, we solve the single Pendulum task from two-gray-scale images (2 100 x 100 pixels). Similar to the Kutli-Pendulum experiments, we let the latent space 3-dimensional. In provide two images per state to insure that velocity can be inferred and fully-observable assumption in maintained. We find that the prediction results are fairly similar to the kinematic state representation, but more noisy due to the course image observations. We train the model using only mostly offline, using only 1 iteration. We initialize the model with 10000 samples under random control trained for 100 epochs and then trained once more with another 10000 samples using the new control policy. This policy is able to solve the task efficiently (142 return) similarly to the kinematic state representation, producing similar prediction results. Its interesting to note that the open-looy model predicts that the pendulum will continue to oscillate and not stabilize. This could be evidence that the model over fits to the random training data, however it would also require very accurate model to predict stabilizing behavior far into the future.
Open-kop xeedictio Inage Pundulun civiormni
20
15
rewarl
-5
10
-15
time stop
Figure 5: Depicted is the open loop prediction result of the Image- mnage-Pendulum environment successfully stabilized by the learned rewadd CEM policy. The model is first trained 10000 samples from a random policy, then another 10000 sample under the new control policy.
12


Reviews:
[{"review_id": "SJxkccvy5H", "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": "Summary:\nThis paper proposes a novel algorithm for planning on specific domains through latent reward prediction. The proposed model uses an encoder to learn embedding the state to the latent state, a forward dynamics function to learn dynamical system in latent state space, and a reward function to estimate the reward given a latent state and an action. Using these functions, the authors define the objective using the mean-squared error between true and multi-step prediction of rewards. To justify the proposed method, the authors provide a theoretical analysis and experimental results on specific RL domains, multi-pendulum and multi-cheetah, which contain irrelevant aspects of the state.\n\nComments:\nThis paper is well-written and easy to understand.\n- In this paper, the authors assume deterministic transition and use deterministic function for latent transition. It seems to be the authors want to use MPC, which is a powerful planning algorithm. However, many RL tasks are modeled with stochastic transition. In stochastic transition cases, is the proposed algorithm still valid?\n- As shown in Figure 3, even proposed method shows better performance than SAC in early episode but table 1 says that SAC shows the best convergence results in any number of pendulums except the single pendulum case. It seems to be different results from intuition, because the authors emphasize that the strength of the proposed method is efficiency of learning in RL tasks with irrelevant information. \n\nQuestions and minor comments:\n- What objective is used to learn the latent model of the state-prediction model algorithm? \n- Providing detailed experimental settings, like detailed settings for three deterministic feed-forward neural networks, and results such as consumed CPU time will help the comparison algorithms.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HJlQrTscYr", "reviewer": "AnonReviewer2", "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": "This paper presents a technique for model based RL/planning with latent dynamics models, which learns the latent model only using reward prediction. This is in contrast to existing work which generally use a combination of reward prediction and state reconstruction to learn the latent model. The paper suggests that by removing the state reconstruction loss, the agent can learn to ignore irrelevant parts of the state, which should enable better performance in settings where state reconstruction is challenging. \n\nOverall the motivation for this work is good, and the idea is promising. Difficulty in reconstructing high dimensional states is a challenge for learning latent dynamics models. The paper is also very well written and easy to follow.\n\nMy concerns are centered around the experimental evaluation. Specifically, I see the following issues: (1) the experimental environments seem artificial, and hand tailored for this method, (2) given that the proposed method is a minor modification to the PlaNet paper, it seems that PlaNet should be included as a comparison (especially because it has been shown to work on high dimensional states), and (3) the proposed method seems very prone to overfitting to the given task, and there should be an analysis of how the proposed change affects generalization and robustness.\n\n(1): The testing environments contain many distractor pendulums/cheetahs, which makes state reconstruction especially challenging. While this does seem to be the point the authors are trying to show, the environments are an extreme, almost artificial, case of difficult state reconstruction. Would the same results hold in more realistic settings, for example, visual robot manipulation in a cluttered scene? Model based RL with video prediction models has been shown to work in such real cluttered robot manipulation environments. Showing that the proposed method can outperform such approaches in robot manipulation settings would be a powerful result. The results on images in the appendix seem to show a delta between the true and predicted reward, suggesting that the proposed method does not yet work on images. Why might this be the case?\n\n(2): From what I can see, the proposed method is very similar to the PlaNet algorithm with state reconstruction loss removed. Given the similarity, PlaNet should be included as a comparison in both the pendulum and cheetah environments. Similarly why was DeepMDP performance not shown in the Cheetah environment?\n\n(3): One of the strengths of model based reinforcement learning is the ability to plan to reach unseen goals with a model trained via self-supervision or different goals. Does the proposed approach lose some of this, by overfitting to only the task reward? I suspect that in generalizing to unseen tasks, a model trained with state prediction would potentially perform much better. If trained on many tasks, could this method achieve similar generalization? \n\nDue to some of these questions which remain unanswered by the experimental evaluation my current rating is Weak Reject. If the authors are able to clarify some of the questions above I may adjust my score.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}, {"review_id": "BJgKdT6JFS", "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": "This paper claims that one only needs a reward prediction model to learn a good latent representation for model-based reinforcement learning. They introduce a method that learns a latent dynamics model exclusively from multi-step reward prediction, then use MPC to plan directly in the latent space. They claim this is sample efficient in the model-based way, and is more useful than predicting full states. They learn a model that predicts only current and future rewards conditioned on action sequences, and that observation reconstruction is unnecessary to learn a good latent space. They provide planning performance guarantees for approximate latent reward prediction models. \n\nI tend to reject this work, because although I support the premise and believe it is very important, and like the style of experiments run with the use of distractors, I believe it is not impactful if only looking at the dense reward setting. The type of environment they describe that requires using lossy representations is also likely to only have sparse reward, so to only note in the conclusion as future work is not enough. The contributions consist only of learning a multi-step reward model for planning, and only provide results in two dense reward environments. In the second experiment with more difficult, high-dimensional observation and action space setting, two of the 3 baselines are left out, namely the state model and DeepMDP. I think it is crucial to include DeepMDP, as it is the one most likely to perform competitively with the proposed method. \n\nThe justification for Table 1 vs. Figure 3 are also very unclear, as to why SAC is trained with 10^6 samples while DeepMDP is trained under a random policy, whereas the original paper utilizes samples collected by the policy as its trains. SAC is off policy, and can therefore be evaluated with random data, rather than being used as an \"upper bound baseline\". The final evaluation performance in dashed line in Figure 3 also doesn't include standard deviation across the 5 seeds, which it should. The final results for SAC also do not match the performance in Figure 3, although it is hard to tell since the final performance in Table 1 is written in terms of number of environment steps while Figure 3 the axis is in terms of episodes. \n\nIncluding sparse reward experiments would vastly help support the claims in the paper, as well as including the DeepMDP results for HalfCheetah and additional explanation of the difference in performance of SAC in Table 1 and Figure 3.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose a model-based RL algorithm, consisting of learning a                                                           
deterministic multi-step reward prediction model and a vanilla CEM-based MPC                                                       
actor.                                                                                                                             
In contrast to prior work, the model does not attempt to learn from observations                                                   
nor is a value function learned.                                                                                                   
The approach is tested on task from the mujoco control suit.                                                                       
                                                                                                                                   
The paper is below acceptance threshold.                                                                                           
It is a variation on previous work form Hafner et al.                                                                              
Furthermore, I think the approach is fundamentally limited: All the learning                                                       
derives from the immediate, dense reward signal, whereas the main challenges in RL                                                 
are found in sparse reward settings that require planning over long horizons, where value                                          
functions or similar methods to assign credit over long time windows are                                                           
absolutely essential.
Decision: Reject