Many real-world sequential decision-making problems can be formulated as optimal control with high-dimensional observations and unknown dynamics. A promising approach is to embed the high-dimensional observations into a lower-dimensional latent representation space, estimate the latent dynamics model, then utilize this model for control in the latent space. An important open question is how to learn a representation that is amenable to existing control algorithms? In this paper, we focus on learning representations for locally-linear control algorithms, such as iterative LQR (iLQR). By formulating and analyzing the representation learning problem from an optimal control perspective, we establish three underlying principles that the learned representation should comprise: 1) accurate prediction in the observation space, 2) consistency between latent and observation space dynamics, and 3) low curvature in the latent space transitions. These principles naturally correspond to a loss function that consists of three terms: prediction, consistency, and curvature (PCC). Crucially, to make PCC tractable, we derive an amortized variational bound for the PCC loss function. Extensive experiments on benchmark domains demonstrate that the new variational-PCC learning algorithm benefits from significantly more stable and reproducible training, and leads to superior control performance.   Further ablation studies give support to the importance of all three PCC components for learning a good latent space for control. 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 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. In this paper, we argue from first principles that learning a latent representation for control should be guided by good prediction in the observation space and consistency between latent transition and the embedded observations. Furthermore, if variants of iterative LQR are used as the controller, the low-curvature dynamics is desirable. All three elements of our PCC models are critical to the stability of model training and the performance of the in-latent-space controller. We hypothesize that each particular choice of controller will exert different requirement for the learned dynamics. A future direction is to identify and investigate the additional bias for learning an effective embedding and latent dynamics for other type of model-based control and planning methods. where D TV is the total variation distance of two distributions. The first inequality is based on the result of the above lemma, the second inequality is based on Pinsker's inequality (Ordentlich & Weinberger, 2005) , and the third inequality is based on Jensen's inequality (Boyd & Vandenberghe, 2004) of (·) function. Now consider the expected cumulative KL cost: t=0 KL(P (·|x t , u t )|| P (·|x t , u t )) | P, x 0 with respect to some arbitrary control action sequence {u t } T −1 t=0 . Notice that this arbitrary action sequence can always be expressed in form of deterministic policy u t = π (x t , t) with some nonstationary state-action mapping π . Therefore, this KL cost can be written as: where the expectation is taken over the state-action occupation measure t=0 P(x t = x, u t = u|x 0 , U ) of the finite-horizon problem that is induced by data-sampling policy U . The last inequality is due to change of measures in policy, and the last inequality is due to the facts that (i) π is a deterministic policy, (ii) dU (u t ) is a sampling policy with lebesgue measure 1/U over all control actions, (iii) the following bounds for importance sampling factor holds: To conclude the first part of the proof, combining all the above arguments we have the following inequality for any model P and control sequence U : For the second part of the proof, consider the solution of (SOC3), namely (U * 3 , P * 3 ). Using the optimality condition of this problem one obtains the following inequality: Using the results in (11) and (12), one can then show the following chain of inequalities: where U * 1 is the optimizer of (SOC1) and (U * 3 , P * 3 ) is the optimizer of (SOC3). Therefore by letting λ 3 = √ 2T 2 · c max U and R 3 ( P ) = E x,u KL(P (·|x, u)|| P (·|x, u)) and by combining all of the above arguments, the proof of the above lemma is completed. A.2 PROOF OF LEMMA 2 For the first part of the proof, at any time-step t ≥ 1, for any arbitrary control action sequence {u t } T −1 t=0 , and any model P , consider the following decomposition of the expected cost : . Now consider the following cost function: E[c(x t−1 , u t−1 ) + c(x t , u t ) | P , x 0 ] for t > 2. Using the above arguments, one can express this cost as By continuing the above expansion, one can show that where the last inequality is based on Jensen's inequality of (·) function. For the second part of the proof, following similar arguments as in the second part of the proof of Lemma 1, one can show the following chain of inequalities for solution of (SOC3) and (SOC2): where the first and third inequalities are based on the first part of this Lemma, and the second inequality is based on the optimality condition of problem (SOC2). This completes the proof.
Learning embedding for control with high-dimensional observations
{"src": "scitldr-AIC_train", "paper_id": "BJxG_0EtDS"}