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