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 exploits the duality between the Perron-Frobenius and Koopman operators to formulate a Deep Stochastic Koopman Operator (DeSKO) to control an unknown system.  In particular, in contrast to prior work in this space that assumes a deterministic system, the paper proposes to learn a distribution over observables (encoded via a probabilistic neural network) and linear dynamics that propagate this distribution forward.  Theoretically, it is shown that under certain assumptions, a model-predictive-controller (MPC) consisting of a nominal policy and a stabilizing feedback gain around the nominal trajectory can guarantee robust stability.  Empirically, extensive experiments are conducted across 8 (4 nominal + 4 perturbed) environments, showing that the proposed approach leads to favorable prediction error (Fig. 2), control performance (Fig.3), and robustness (Fig. 4) when compared to other Koopman and Model-Free RL algorithms.  Further an ablation study on the effect of an entropy constraint that is introduced is also performed (Fig. 5).