Learning Latent State Spaces for Planning through Reward Prediction

1 INTRODUCTION . Deep reinforcement learning ( DRL ) has demonstrated some of the most impressive results on several complex , high-dimensional sequential decision-making tasks such as video games ( Mnih et al. , 2015 ) and board games ( Silver et al. , 2018 ) . Such algorithms are often model-free , i.e . they do not learn or exploit knowledge of the environment ’ s dynamics . Although these model-free schemes can achieve state-of-the-art 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 a model for the environment . The learned model can then be used for planning , thereby minimizing the number of environment interactions . However , model-based methods are typically limited to low-dimensional tasks or settings with certain assumptions imposed on the dynamics ( Li & Todorov , 2004 ; Deisenroth & Rasmussen , 2011 ; Levine & Koltun , 2013 ; Chebotar et al. , 2017 ) Several recent works have learned dynamics models parameterized by deep neural networks ( Lenz et al. , 2015 ; Finn & Levine , 2017 ; Williams et al. , 2017 ) , but such methods generally result in high-dimensional models which are unsuitable for planning . In order to efficiently apply model-based planning in high-dimensional spaces , a 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 a good policy ( Watter et al. , 2015 ; Banijamali 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-dimensional latent state must retain a sufficient amount of information from the original state to accurately predict the state in the high-dimensional space . Although the ability to perform full-state prediction is sufficient for planning in latent space , it is not strictly necessary . Predicting the full-state may result in a 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 to 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 a maze , it may be more fruitful to instead optimize a latent model to be able to perform reward prediction . In some sense , value-based modelfree 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 , modelbased 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 a planning method requires is a latent reward prediction model : a model which predicts current and future rewards over a horizon from a 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-space . Learning the latent reward prediction model benefits from the sample-efficiency of model-based methods , and the learned latent model is more useful for planning than its full-state prediction counterparts as 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 us 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 multi-pendulum and multi-cheetah environments where the agent must ignore a large amount of irrelevant information . The remainder of this paper is structured as follows . In Section 2 , we present related work in modelbased DRL . In Section 3 , we provide the requisite background on RL 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 a handful of episodes , however they usually have difficulty scaling to high-dimensional observations . Guided-policy search methods ( Levine & Koltun , 2013 ; Chebotar et al. , 2017 ) make use of local models to update a global deep policy , however these local assumptions could fail to capture global properties and suffer at points of discontinuity . For high-dimensional observations such as images , deep neural networks have been trained to 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 a 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. , 2015 ; Banijamali et al. , 2017 ; Finn et al. , 2016 ; Ichter & Pavone , 2019 ; Zhang et al. , 2019 ) . However , these methods construct latent spaces based on state reconstruction/prediction 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 state-of-the-art model-free methods , but its latent model might also suffer from the aforementioned issues with irrelevant observation predictions . In this work we use a similar planning algorithm to PlaNet , however , removing the dependency of observation prediction provides a more concise latent representation for planning . Reward-based representations learning for RL The work of VPN ( Oh et al. , 2017 ) proposed a RL framework which learns a model to predict the future values rather than the full-states of the environment . Although VPN shares a very similar idea as our framework to not predict irrelevant information in full-state , it 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 as auxiliary objectives . However , with only the single-step prediction loss , DeepMDP may not be suitable for making long term reward predictions necessary for 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 of DeepMDP by learning latent dynamics only from the reward sequences , but focusing on multi-step reward prediction allows us to achieve higher performance in model-based planning . 3 PRELIMINARIES . 3.1 PROBLEM SETUP . In this paper we consider a discrete time nonlinear dynamical system f : S×A → S with continuous state and action spaces S ⊆ Rdn , A ⊆ Rdm , and a reward function R : S ×A → R. Then , given an admissible action at ∈ A at time t , the system state st evolves according to the dynamics st+1 = f ( st , at ) ( 1 ) where st+1 is the next state at time t+ 1 . Our goal is to find a policy π : S → A that selects actions { at = π ( st ) , t = 0 , 1 , . . . } so as to maximize the cumulative discounted rewards ∞∑ t=0 γtR ( st , at ) ( 2 ) where the initial state s0 is assumed to be fixed , and γ is the discount factor . Note that Equations ( 1 ) and ( 2 ) formulate an 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 a dynamics model is available , model predictive control ( MPC ) is a 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 a H-horizon trajectory optimization problem : max ât : ( t+H−1 ) J ( st , ât : ( t+H−1 ) ) ( 3a ) subject to J ( st , ât : ( t+H−1 ) ) = t+H−1∑ τ=t γτ−tR ( ŝτ , âτ ) ( 3b ) ŝt = st , and ŝτ+1 = f ( ŝτ , âτ ) for τ = t , t+ 1 , . . . , ( t+H − 2 ) ( 3c ) Suppose the optimal control sequence is â∗t : ( t+H−1 ) , then the MPC agent will select â ∗ t as the control action at time t. This method is termed model predictive control due to its use of the model f to predict the future trajectory ŝ ( t+1 ) : ( t+H−1 ) from state st and the intended action sequence ât : ( t+H−1 ) . Let πMPC ( st|f , R ) denote the MPC policy using the dynamics model f and reward R. Then the MPC agent selects its action by πMPC ( st|f , R ) = â∗t . 3.3 SAMPLE-BASED TRAJECTORY OPTIMIZATION . From the description in Section 3.2 , an MPC agent ’ s need to solve the trajectory optimization problem specified by Equation equation 3 . In this work , we use the cross entropy method ( CEM ) ( Rubinstein , 1997 ) to solve problem equation 3 . CEM is a very simple sampling-based planning method which only relies on a forward model and trajectory returns which may make direct use of our learned latent mappings somewhat like a simulator . The advantage to using a method like CEM , is that it ’ s sampling procedure can be easily parallelized directly in the learned latent space .

Summary: 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.