Prediction, Consistency, Curvature: Representation Learning for Locally-Linear Control

1 INTRODUCTION . Decomposing the problem of decision-making in an unknown environment into estimating dynamics followed by planning provides a powerful framework for building intelligent agents . This decomposition confers several notable benefits . First , it enables the handling of sparse-reward environments by leveraging the dense signal of dynamics prediction . Second , once a dynamics model is learned , it can be shared across multiple tasks within the same environment . While the merits of this decomposition have been demonstrated in low-dimensional environments ( Deisenroth & Rasmussen , 2011 ; Gal et al. , 2016 ) , scaling these methods to high-dimensional environments remains an open challenge . The recent advancements in generative models have enabled the successful dynamics estimation of high-dimensional decision processes ( Watter et al. , 2015 ; Ha & Schmidhuber , 2018 ; Kurutach et al. , 2018 ) . This procedure of learning dynamics can then be used in conjunction with a plethora of decision-making techniques , ranging from optimal control to reinforcement learning ( RL ) ( Watter et al. , 2015 ; Banijamali et al. , 2018 ; Finn et al. , 2016 ; Chua et al. , 2018 ; Ha & Schmidhuber , 2018 ; Kaiser et al. , 2019 ; Hafner et al. , 2018 ; Zhang et al. , 2019 ) . One particularly promising line of work in this area focuses on learning the dynamics and conducting control in a low-dimensional latent embedding of the observation space , where the embedding itself is learned through this process ( Watter et al. , 2015 ; Banijamali et al. , 2018 ; Hafner et al. , 2018 ; Zhang et al. , 2019 ) . We refer to this approach as learning controllable embedding ( LCE ) . There have been two main approaches to this problem : 1 ) to start by defining a cost function in the high-dimensional observation space and learn the embedding space , its dynamics , and reward function , by interacting with the environment in a RL fashion ( Hafner et al. , 2018 ; Zhang et al. , 2019 ) , and 2 ) to first learn the embedding space and its dynamics , and then define a cost function in this low-dimensional space and conduct the control ( Watter et al. , 2015 ; Banijamali et al. , 2018 ) . This can be later combined with RL for extra fine-tuning of the model and control . In this paper , we take the second approach and particularly focus on the important question of what desirable traits should the latent embedding exhibit for it to be amenable to a specific class of control/learning algorithms , namely the widely used class of locally-linear control ( LLC ) algorithms ? We argue from an optimal control standpoint that our latent space should exhibit three properties . The first is prediction : given the ability to encode to and decode from the latent space , we expect ∗Equal contribution . Correspondence to nirlevine @ google.com the process of encoding , transitioning via the latent dynamics , and then decoding , to adhere to the true observation dynamics . The second is consistency : given the ability to encode a observation trajectory sampled from the true environment , we expect the latent dynamics to be consistent with the encoded trajectory . Finally , curvature : in order to learn a latent space that is specifically amenable to LLC algorithms , we expect the ( learned ) latent dynamics to exhibit low curvature in order to minimize the approximation error of its first-order Taylor expansion employed by LLC algorithms . Our contributions are thus as follows : ( 1 ) We propose the Prediction , Consistency , and Curvature ( PCC ) framework for learning a latent space that is amenable to LLC algorithms and show that the elements of PCC arise systematically from bounding the suboptimality of the solution of the LLC algorithm in the latent space . ( 2 ) We design a latent variable model that adheres to the PCC framework and derive a tractable variational bound for training the model . ( 3 ) To the best of our knowledge , our proposed curvature loss for the transition dynamics ( in the latent space ) is novel . We also propose a direct amortization of the Jacobian calculation in the curvature loss to help training with curvature loss more efficiently . ( 4 ) Through extensive experimental comparison , we show that the PCC model consistently outperforms E2C ( Watter et al. , 2015 ) and RCE ( Banijamali et al. , 2018 ) on a number of control-from-images tasks , and verify via ablation , the importance of regularizing the model to have consistency and low-curvature . 2 PROBLEM FORMULATION . We are interested in controlling the non-linear dynamical systems of the form st+1 = fS ( st , ut ) +w , over the horizon T . In this definition , st ∈ S ⊆ Rns and ut ∈ U ⊆ Rnu are the state and action of the system at time step t ∈ { 0 , . . . , T − 1 } , w is the Gaussian system noise , and fS is a smooth non-linear system dynamics . We are particularly interested in the scenario in which we only have access to the high-dimensional observation xt ∈ X ⊆ Rnx of each state st ( nx ns ) . This scenario has application in many real-world problems , such as visual-servoing ( Espiau et al. , 1992 ) , in which we only observe high-dimensional images of the environment and not its underlying state . We further assume that the high-dimensional observations x have been selected such that for any arbitrary control sequence U = { ut } T−1t=0 , the observation sequence { xt } Tt=0 is generated by a stationary Markov process , i.e. , xt+1 ∼ P ( ·|xt , ut ) , ∀t ∈ { 0 , . . . , T − 1 } .1 A common approach to control the above dynamical system is to solve the following stochastic optimal control ( SOC ) problem ( Shapiro et al. , 2009 ) that minimizes expected cumulative cost : min U L ( U , P , c , x0 ) : = E [ cT ( xT ) + T−1∑ t=0 ct ( xt , ut ) | P , x0 ] , 2 ( SOC1 ) where ct : X ×U → R≥0 is the immediate cost function at time t , cT ∈ R≥0 is the terminal cost , and x0 is the observation at the initial state s0 . Note that all immediate costs are defined in the observation space X , and are bounded by cmax > 0 and Lipschitz with constant clip > 0 . For example , in visualservoing , ( SOC1 ) can be formulated as a goal tracking problem ( Ebert et al. , 2018 ) , where we control the robot to reach the goal observation xgoal , and the objective is to compute a sequence of optimal open-loop actions U that minimizes the cumulative tracking error E [ ∑ t ‖xt − xgoal‖2 | P , x0 ] . Since the observations x are high dimensional and the dynamics in the observation space P ( ·|xt , ut ) is unknown , solving ( SOC1 ) is often intractable . To address this issue , a class of algorithms has been recently developed that is based on learning a low-dimensional latent ( embedding ) space Z ⊆ Rnz ( nz nx ) and latent state dynamics , and performing optimal control there . This class that we refer to as learning controllable embedding ( LCE ) throughout the paper , include recently developed algorithms , such as E2C ( Watter et al. , 2015 ) , RCE ( Banijamali et al. , 2018 ) , and SOLAR ( Zhang et al. , 2019 ) . The main idea behind the LCE approach is to learn a triplet , ( i ) an encoderE : X → P ( Z ) ; ( ii ) a dynamics in the latent space F : Z ×U → P ( Z ) ; and ( iii ) a decoder D : Z → P ( X ) . These in turn can be thought of as defining a ( stochastic ) mapping P̂ : X ×U → P ( X ) of the form P̂ = D ◦F ◦E . We then wish to solve the SOC in latent space Z : min U , P̂ E [ L ( U , F , c , z0 ) | E , x0 ] + λ2 √ R2 ( P̂ ) , ( SOC2 ) such that the solution of ( SOC2 ) , U∗2 , has similar performance to that of ( SOC1 ) , U ∗ 1 , i.e. , L ( U∗1 , P , c , x0 ) ≈ L ( U∗2 , P , c , x0 ) . In ( SOC2 ) , z0 is the initial latent state sampled from the encoder E ( ·|x0 ) ; c̄ : Z × U → R≥0 is the latent cost function defined as c̄t ( zt , ut ) =∫ ct ( xt , ut ) dD ( xt|zt ) ; R2 ( P̂ ) is a regularizer over the mapping P̂ ; and λ2 is the corresponding 1A method to ensure this Markovian assumption is by buffering observations ( Mnih et al. , 2013 ) for a number of time steps . 2See Appendix B.3 for the extension to the closed-loop MDP problem . tion SOC2 under dynamics F , and ( c ) ( red ) in equation SOC3 under dynamics P̂ . regularization parameter . We will define R2 and λ2 more precisely in Section 3 . Note that the expectation in ( SOC2 ) is over the randomness generated by the ( stochastic ) encoder E . 3 PCC MODEL : A CONTROL PERSPECTIVE . As described in Section 2 , we are primarily interested in solving ( SOC1 ) , whose states evolve under dynamics P , as shown at the bottom row of Figure 1 ( a ) in ( blue ) . However , because of the difficulties in solving ( SOC1 ) , mainly due to the high dimension of observations x , LCE proposes to learn a mapping P̂ by solving ( SOC2 ) that consists of a loss function , whose states evolve under dynamics F ( after an initial transition by encoder E ) , as depicted in Figure 1 ( b ) , and a regularization term . The role of the regularizer R2 is to account for the performance gap between ( SOC1 ) and the loss function of ( SOC2 ) , due to the discrepancy between their evolution paths , shown in Figures 1 ( a ) ( blue ) and 1 ( b ) ( green ) . The goal of LCE is to learn P̂ of the particular form P̂ = D ◦ F ◦ E , described in Section 2 , such that the solution of ( SOC2 ) has similar performance to that of ( SOC1 ) . In this section , we propose a principled way to select the regularizer R2 to achieve this goal . Since the exact form of ( SOC2 ) has a direct effect on learning P̂ , designing this regularization term , in turn , provides us with a recipe ( loss function ) to learn the latent ( embedded ) space Z . In the following subsections , we show that this loss function consists of three terms that correspond to prediction , consistency , and curvature , the three ingredients of our PCC model . Note that these two SOCs evolve in two different spaces , one in the observation space X under dynamics P , and the other one in the latent space Z ( after an initial transition from X to Z ) under dynamics F . Unlike P and F that only operate in a single space , X and Z , respectively , P̂ can govern the evolution of the system in both X and Z ( see Figure 1 ( c ) ) . Therefore , any recipe to learn P̂ , and as a result the latent space Z , should have at least two terms , to guarantee that the evolution paths resulted from P̂ in X and Z are consistent with those generated by P and F . We derive these two terms , that are the prediction and consistency terms in the loss function used by our PCC model , in Sections 3.1 and 3.2 , respectively . While these two terms are the result of learning P̂ in general SOC problems , in Section 3.3 , we concentrate on the particular class of LLC algorithms ( e.g. , iLQR ( Li & Todorov , 2004 ) ) to solve SOC , and add the third term , curvature , to our recipe for learning P̂ .

Summary: This paper considers learning low-dimensional representations from high-dimensional observations for control purposes. The authors extend the E2C framework by introducing the new PCC-Loss function. This new loss function aims to reflect the prediction in the observation space, the consistency between latent and observation dynamics, and the low curvature in the latent dynamics. The low curvature term is used to bias the latent dynamics towards models that can be better approximated as locally linear models. The authors provide theory (error bounds) to justify their proposed PCC-Loss function. Then variational PCC is developed to make the algorithm tractable. The proposed method is evaluated in 5 different simulated tasks and compared with the original E2C method and the RCE method.