DeSKO: Stability-Assured Robust Control with a Deep Stochastic Koopman Operator

1 INTRODUCTION . Modeling is crucial to the development of intelligent machines and robots . It is needed to predict the behavior of a dynamical system , analyze a system ’ s interior properties , realize a reliable planning and control ( Müller et al. , 2007 ) . For many systems of interest , such as a soft robotic arm that can interact with an unknown object ( George Thuruthel et al. , 2018 ) , the governing equations are either unknown or too complex to be useful ( Bakker et al. , 2020 ) . Thus , it is essential to use an accurate but concise model . The need for a fast model that is able to represent complex systems has led to recent advancements in data-driven modeling of complex dynamics , and learning world models . These advancements have been achieved by constructing surrogate neural network ( NN ) models ( Chua et al. , 2018 ; Ha & Schmidhuber , 2018 ; Zhang et al. , 2019 ) . By formulating the modeling problem as a supervised learning task , the neural network can be trained to infer the systems ’ dynamics based on the dense state and action input ( Nagabandi et al. , 2018b ) or high-dimensional image observations ( Hafner et al. , 2019 ) . However , as nonlinear systems are often approximated by complex and over-parameterized neural networks ( Morton et al. , 2019 ) , the learned models are restrictive when they are used for dynamical analysis and controller design . Optimal control problems based on NN models are typically non-convex , and they need to be solved by using inefficient population-based optimization methods ( Bharadhwaj et al. , 2020 ) , such as random shooting ( Zhang et al. , 2019 ; Nagabandi et al. , 2018a ) or cross entropy methods ( Chua et al. , 2018 ; Hafner et al. , 2019 ; Wang & Ba , 2019 ) . Recently , the Koopman operator theory has attracted much attention . The Koopman operator can model complex nonlinear systems by linearly propagating observables in an infinite dimensional functional space ( Koopman , 1931b ) . This propagation allows us to apply the well-established linear control and analysis theory to nonlinear systems . Based on a system ’ s runtime data , a linear operator , which propagates the observables in time , is then learned . Based on this idea , practical algorithms , like dynamic mode decomposition ( DMD ) and extended dynamic mode decomposition ( EDMD ) , were developed by selecting a finite set of parameterized linear and nonlinear observable functions . These methods have been applied to the modeling of soft robotic arms ( Bruder et al. , 2020 ; Haggerty et al. , 2020 ) , multirotors ( Folkestad et al. , 2020 ) , and mobile robots ( Shi & Karydis , 2021 ) . More recently , Deep-DMD exploits deep learning techniques to automate the design of observable functions ( Lusch et al. , 2018 ; Otto & Rowley , 2019 ; Han et al. , 2020 ; Yeung et al. , 2018 ; Morton et al. , 2018 ) . DeepDMD automatically learns a richer set of observable functions than DMD or EDMD , and it can also be more accurate , even though it uses fewer observables ( Bakker et al. , 2020 ) . However , these methods focus on modeling nominal systems with clean data sets , while noise , either in the process or observation , has not yet been accounted for . Furthermore , it is challenging to ensure closed-loop stability in a controller design that is based on a learned Koopman operator model ( Bakker et al. , 2020 ) . In this work , we learn and control uncertain nonlinear dynamics using the deep stochastic Koopman operator ( DeSKO ) model . Instead of encoding the observables deterministically , DeSKO infers the probabilistic distribution of observables and learns a linear operator that can propagate the inferred distribution forward in time . Furthermore , we propose an efficient and robust model predictive controller based on the DeSKO model . In terms of robotic control benchmarks , we show that our proposed method outperforms state-of-the-art deep Koopman operators ( Lusch et al. , 2018 ; Han et al. , 2020 ) in modeling and in stabilizing systems that have been corrupted by process and observation noises . Furthermore , we show that the proposed framework is scalable ; it can handle high-dimensional and complex systems , including legged robots , soft robotic arms , and biological gene networks . The outline of this paper is as follows . In Section 2 , we introduce and discuss the related works . In Section 3 , we introduce the basic notations and concepts of the Koopman operator ( KO ) , and we present our proposed DeSKO model . In Section 4 , we discuss our robust control framework and prove its stability . In Section 5 , we evaluate the DeSKO model in terms of modeling , control , and robustness in eight different environments , and we compare it to baseline methods . Finally , in Section 6 , we conclude the paper and discuss future challenges . 2 RELATED WORKS . Koopman operators Following the Koopman operator theory formalized in Koopman ( 1931a ) , various practical algorithms have been proposed to extract the Koopman spectral properties from state-transition data . DMD and EDMD adopt a hand-crafted set of linear and nonlinear functions to encode the observables and obtain the linear operator through least square optimization ( Schmid , 2010 ; Tu et al. , 2014 ; Williams et al. , 2015 ) . A straightforward approach to automate the design of observable functions is to parameterize the observable functions as neural networks and use deep learning techniques for training ( Lusch et al. , 2018 ; Otto & Rowley , 2019 ; Yeung et al. , 2018 ; Morton et al. , 2018 ; Takeishi et al. , 2017 ) . Unlike DMD and EDMD , Deep-DMD approaches do not suffer from the curse of dimensionality ; that is , the number of observables does not grow exponentially with the number of states . Furthermore , Deep-DMD can also express a richer class of observable functions than DMD or EDMD . Graph neural networks and block-wise Koopman operators are introduced in ( Li et al. , 2020 ) to model objects composed of similar building blocks , such as ropes . While the described works have greatly advanced the field of modeling and control by applying the Koopman operator , several challenges still remain . Azencot et al . ( 2020 ) proposed a Koopman autoencoder framework that can learn a Koopman model from noisy data for long horizon prediction , but only considered autonomous systems without control input . Morton et al . ( 2019 ) proposed to measure uncertainty with an ensemble of Koopman models sampled from the encoded distribution . However , their work is focused on deterministic systems with clean datasets . Furthermore , how to guarantee stability with a learned Koopman model remains an open question . Mamakoukas et al . ( 2020 ) ; Kolter & Manek ( 2019 ) ; Pan & Duraisamy ( 2020 ) studied how to learn a stable dynamics of embeddings on the latent space , while we are focused on stabilizing the original system with the learned model . Our method differs from previous works by learning a stochastic model from noisy data , and providing a robust control framework with a stability guarantee . We show that this framework can easily generalize to systems with process and observation noises . Model-based Learning Control In recent years , model-based reinforcement learning has dominated a large portion of model-based learning control studies ( Pascanu et al. , 2017 ; Hamrick et al. , 2017 ; Racanière et al. , 2017 ) . Several researchers have combined learned neural network models with model predictive control ( Chua et al. , 2018 ; Zhang et al. , 2019 ; Hafner et al. , 2019 ; Nagabandi et al. , 2018a ; Wang & Ba , 2019 ) . In this work , we utilize the Koopman operator ’ s linear propagation property , and we conduct model predictive control efficiently . Furthermore , we also show that the Koopman operator substantially facilitates analyzing and assuring stability of the closed-loop system . Gaussian process-based methods are another popular class of learning control methods ( Kocijan et al. , 2004 ; Berkenkamp et al. , 2017 ; Vinogradska et al. , 2016 ; Berkenkamp & Schoellig , 2015 ) . As a type of non-parametric method , Gaussian process-based methods are efficient in modeling with small data sets . Indeed , Gaussian process models have been applied to various robotic systems ( Chang et al. , 2017 ; Cao et al. , 2017 ; Kabzan et al. , 2019 ; Chang et al. , 2020 ) . Nonetheless , Gaussian process-based approaches are typically limited to low-dimensional systems and small data sets ( Wang et al. , 2020 ) . In this work , we present a framework that is scalable to high-dimensional robotic control tasks . 3 MODELING . In this section , we will first introduce the basic notations and concepts of the Koopman operator ( KO ) ( Koopman , 1931a ) . Then the derivation of the DeSKO model is established . 3.1 THE KOOPMAN OPERATOR . For the sake of clarity , we introduce the KO using notations that are similar to the ones in Bakker et al . ( 2020 ) . Consider the nonlinear discrete time dynamical system xt+1 = f ( xt , ut ) where f is a nonlinear differentiable function . The system can be lifted to an infinite-dimensional function space F composed of all square-integrable real-valued functions within the compact domain X × U ⊂ Rn+m . Elements of F are called observables . On this space , the flow of the systems is characterized by the Koopman operator K : F → F , which is an infinite-dimensional linear operator that satisfies Kψ ( xt , ut ) = ψ ◦ f ( xt , ut ) where ψ denotes the observable function . However , the infinite-dimensional function space F and the KO K are impractical . Therefore , a finite dimensional function space F ⊂ F , spanned by a linearly independent basis function ψ : Rn → Rh is used instead , where h is the dimensionality of the observable function specified by the designer . Here , i denotes the index of each observable function . The linear property of the KO enables the use of linear control techniques for efficient control . The Koopman equation can be written as Aψ ( xt ) +But = ψ ◦ f ( xt , ut ) ( 1 ) where the KO is split into the Koopman matrix A ∈ Rh×h and the control matrix B ∈ Rh×m , . For the reader ’ s convenience , we also have also included the details for learning Koopman matrices with a set of observables in the Appendix . The Koopman operator conducts a deterministic inference of observables and ignores the existence of any uncertainty that may stem from the system and the lack of data about the system . A controller design based on such a defective model would inevitably perform suboptimally , and even result in failure . In this work , we are interested in learning and controlling a nonlinear system that has been corrupted by additive noise : xt+1 = f ( xt , ut ) + wt ( 2 ) where wt ∼ pw denotes random noises subject to the distribution pw . In this case , the transition of states is probabilistic , and it is denoted by p ( xt+1|xt , ut ) . In the next section , we will discuss how to overcome this challenge by encoding and propagating the distribution of observables with Koopman operators .

Summary: This paper leverage the Koopman operator theory to represent a nonlinear dynamics as a high dimensional linear dynamics. The map from the state space to the high-dimensional space as well as the Koopman Matrix (linear transition matrix) are learned from data. Similar ideas has been explored in existing works. The contribution of this work is bringing stochasticity into the learning of Koopman models. The paper assumes the data are noisy and try to uncover a Koopman model with noisy data. Another contribution is the joint learning of a control matrix in the Koopman setting and a robust control framework for the proposed method. The proposed robust control framework is based on model predictive control and is provably stable (thanks to the linear transition of the observables). The proposed method is tested on rigid body systems (CartPole and HalfCheetah), a soft robotic arm (SoPrA) and cell biology (GRN). The proposed method outperform existing Koopman-based method (DKO) and RL method (SAC).