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

Abstract: The Koopman operator theory linearly describes nonlinear dynamical systems in a high-dimensional functional space and it allows to apply linear control methods to highly nonlinear systems. However, the Koopman operator does not account for any uncertainty in dynamical systems, causing it to perform poorly in real-world applications.
Therefore, we propose a deep stochastic Koopman operator (DeSKO) model in a robust learning control framework to guarantee stability of nonlinear stochastic systems. The DeSKO model captures a dynamical system's uncertainty by inferring a distribution of observables. We use the inferred distribution to design a robust, stabilizing closed-loop controller for a dynamical system. Modeling and control experiments on several advanced control benchmarks show that our framework is more robust and scalable than state-of-the-art deep Koopman operators and reinforcement learning methods. Tested control benchmarks include a soft robotic arm, a legged robot, and a biological gene regulatory network. We also demonstrate that this robust control method resists previously unseen uncertainties, such as external disturbances, with a magnitude of up to five times the maximum control input. Our approach opens up new possibilities in learning control for high-dimensional nonlinear systems while robustly managing internal or external uncertainty.

Reviews and author responses:

 Hi,

Thank you. The changes made to the manuscript address the concerns I raised in my review. I would be happy to see this work published in its revised form. I prefer to keep my original score.

 We greatly appreciate the reviewer’s affirmative comments on the technical contributions of this work and detailed suggestions on the related works and text organization. We have revised the manuscript according to the reviewer’s comments and suggestions. We summarized our revisions of the manuscript in the following to address the reviewer’s concerns. 

Takeishi et al. 2017 investigated the data-driven modeling of deterministic autonomous dynamical systems, an approach that can be classified as the “Deep-DMD method”, which has been discussed in Sec.2-Koopman operators of the manuscript. On the other hand, DeSKO can model the stochastic controlled systems by encoding the distribution of observables and assure robust stability of the closed-loop system.

Pan et al. 2020 proposed a Deep-NN based framework with structural parameterization such that the learned Koopman operator is stable. This study is different from DeSKO in two aspects: i) it is focused on learning continuous-time deterministic autonomous systems while DeSKO focuses on stochastic controlled systems; ii) though the dynamic of the observables is stable, the control problem and stability of the original system were not addressed.

Morton et al. 2019 proposed a deep variational Koopman model to encode a distribution of observables and showed its effectiveness in terms of modeling and control. This study is different from DeSKO in two aspects: i) though it also encodes a distribution of observables, this work is focused on deterministic systems with clean datasets, whereas DeSKO is valid for stochastic systems with noisy datasets; ii) DeSKO provides robust closed-loop stability guarantee and is validated on various robotic tasks, while Morton et al. 2019 lacks such guarantees and only showed the results on cartpole and acrobot. 

Azencot et al. 2020 proposed a Koopman autoencoder framework and exploited both forward and backward dynamics, and showed its effectiveness in modeling high-dimensional systems. Their work is different from DeSKO in the following aspects: i) their work encourages the model to learn quasi-stable Koopman operator to prevent divergence of prediction, i.e., they are concerned with the stability of the observables, while DeSKO deals with the controlled system and assures the robust stability of the original system, i.e. stability of the states; ii) Azencot et al. 2020 is focused on the learning of autonomous systems, while DeSKO deals with the controlled systems.

The above discussions on the related works are included in the revised manuscript in Section 2.


References:

Takeishi N, Kawahara Y, Yairi T. Learning Koopman invariant subspaces for dynamic mode decomposition[C]//Proceedings of the 31st International Conference on Neural Information Processing Systems. 2017: 1130-1140.

Pan S, Duraisamy K. Physics-informed probabilistic learning of linear embeddings of nonlinear dynamics with guaranteed stability[J]. SIAM Journal on Applied Dynamical Systems, 2020, 19(1): 480-509.

Azencot O, Erichson N B, Lin V, et al. Forecasting sequential data using consistent Koopman autoencoders[C]//International Conference on Machine Learning. PMLR, 2020: 475-485.

Proctor J L, Brunton S L, Kutz J N. Generalizing Koopman theory to allow for inputs and control[J]. SIAM Journal on Applied Dynamical Systems, 2018, 17(1): 909-930.

Korda M, Mezić I. Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control[J]. Automatica, 2018, 93: 149-160.

 In response to the minor comments: 

- The reference noa 2016 in Sec. 2 seems to be wrong. 

Thanks for pointing out the typo, it has been fixed!

- Please specify the domain and range of A and B in Eq. (1) 

Thanks for the reviewer’s advice! The domain has been included in the revised version, the Koopman matrix $A\in \mathbb{R}^{h\times h}$, the control matrix $B\in \mathbb{R}^{h\times m}$, and the observation matrix $C\in\mathbb{R}^{n\times h}$, $h$ is the dimensionality of observables specified by the designer.

- Why K is separable in the dynamics and control? 

Thanks for the reviewer’s question. We separate the Koopman operator into matrices$A$ and $B$ to allow for a control input in the typical structure used in the DMD and EDMD literature. For controlled systems, the observable could take the form of $\psi(x,u)=[\phi(x)^T, u^T]^T$, i.e., taking a linear encoding on $u$ and nonlinear encoding on $x$, respectively. The resulting Koopman matrix would then be a $(n+m)\times (n+m)$ dimensional matrix. Since we are not interested in the prediction of future control input, we can disregard the last $m$ rows of $K$, and separate the remaining parts into two components $A$ and $B$, interacting with the state encoding and action encoding, respectively. 

- I would move Sec. 3.1 & 3.2 to the appendix; and - The contribution of the Perron-Frobenius to this work is unclear. Please elaborate.

We very much appreciate the reviewer’s kind suggestion, we have now moved Sec. 3.2 to the Appendix. Sec. 3.1 introduces the basic idea of Koopman and notations used in this paper, thus we think keeping it in the main text would help the readers better understand the paper.  The aim of Sec 3.2 is to justify the use of Koopman operator from the perspective of its relation to Perron-Frobenius operator. As it plays a less relevant role in the paper, we have moved it to Appendix. 

- The subscript/superscript notation in Eq. (5) is inconsistent with Sec. 3.1 

The subscript and superscript notation have been unified, thanks for pointing out!

- The Figs. 2, 4 & 5, and possibly 3 should use a log-scale on the y-axis. 

We have used log-scale in the figures to better compare the performance of different methods in the revised manuscript.

- Why SAC is not in Fig. 2? Similarly, why MLP is not in Figs. 3 & 4?

Thanks for the reviewer’s question. Fig. 2 shows comparisons in regards to modeling approaches. SAC is a model-free RL method that does not learn the dynamics explicitly, thus SAC can’t be compared against modeling approaches and can therefore not be included in Fig. 2. MLP is an ensemble of multi-layer perceptrons and used as a baseline in terms of modeling, but does not have an associated approach to conduct efficient control tasks, thus is not compared in terms of control nor included in Figs. 3&4. 

 We greatly appreciate Reviewer R938’s insightful question and affirmative comments on the model structure, empirical results and writing. We have revised the manuscript according to the reviewer’s suggestions, and summarized our responses in the following to address the reviewer’s concerns:

The reviewer is making a valid point that Proposition 1 considers only the case where the Koopman operator is exact, and we agree with the reviewer that stability in the case of sub-optimal/approximated Koopman operator needs to be addressed. Therefore, we added as a new theoretical result Proposition 2 to address this concern in Sec. 4.2. We show that the proposed controller (10) can stabilize the system even with an approximated Koopman operator, though the uniform ultimate bound is larger than the bound in the ideal case due to the existence of prediction residuals.

For the reviewer’s concern with the linear control input on latent space, we agree that a more general structure could be proposed where the input is also shifted with a nonlinear function. However, such a Koopman representation relying on an nonlinear control input makes it infeasible to exploit efficient linear control techniques. There is also a bilinear Koopman representation (Goswami et al.(2020), Peitz et al.(2020), Bruder et al.(2021)) that is more general than the linear input representation. However, the bilinear representation can’t cover all types of nonlinear systems and, most importantly, does not preserve the convexity of the resulting control problem either. Due to similar reasons,  Korda et al.(2018)only discussed the generality of bilinear representations and still adopted a linear representation. To recall, the goal of this work is to preserve the convexity of the resulting control problem and leverage linear control techniques to achieve efficient control. Thus, in this paper, we focused on systems that can be approximated by a Koopman representation with linear input and validated its ability to model and control various systems. Designing a more general nonlinear input model that preserves the convexity of the resulting control problem is an interesting research problem, and we would like to leave that for future work. 

References:

Korda M, Mezić I. Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control[J]. Automatica, 2018, 93: 149-160.

Bruder D, Fu X, Vasudevan R. Advantages of bilinear koopman realizations for the modeling and control of systems with unknown dynamics[J]. IEEE Robotics and Automation Letters, 2021, 6(3): 4369-4376.

Goswami D, Paley D A. Global Bilinearization and Reachability Analysis of Control-Affine Nonlinear[J]. The Koopman Operator in Systems and Control: Concepts, Methodologies, and Applications, 2020, 484: 81. 

Peitz S, Otto S E, Rowley C W. Data-driven model predictive control using interpolated Koopman generators[J]. SIAM Journal on Applied Dynamical Systems, 2020, 19(3): 2162-2193.


 We greatly appreciate the reviewer’s evaluation and affirmative comments on the manuscript! We will continue improving this work with the reviewers.


 We greatly appreciate the reviewer’s insightful questions and suggestions. We have revised the manuscript according to the reviewer’s suggestions, and summarized our response in the following to address your concerns:

(1) Can use more experiments to show a more convincing picture.

To further address the reviewer’s concern, we have now included two more simulated experiments and benchmarked the performance against baselines. In the experiments, we inject noise in the Halfcheetah and SoPrA models, and then compare our DeSKO against other baselines in terms of modelling and control. In addition to the simulated experiments, we expanded the experiments by offering a real-world demonstration on the soft robotic arm SoPrA. Videos of the real-world experiment are included in the updated supplementary materials. We have presented the detailed experimental results in the Appendix F and G due to space limitations. To summarize, we have compared DeSKO against three baselines on eleven benchmark problems in terms of modelling, control, and robustness to uncertainties. We hope the presented results are now sufficient to show a more convincing picture. 

(2) Compared with SAC, the proposed method does not show a clear edge in many cases.

We thank the reviewer for this question about clarifying the advantage DeSKO has over SAC. The comparison between DeSKO and SAC is two-folded: i) DeSKO needs 1 to 2 orders of magnitude fewer data than SAC (one twenty-fifth) and is thus more suitable for real-world implementation; ii) with a limited data set, DeSKO still outperforms SAC in terms of robustness and achieves better or comparative performance with SAC. The specifics on how DeSKO outperforms SAC are shown in Figure 4. The paper now adds this summary of performance improvements in the Discussion section.

(3) Explanation of how the Perron-Frobenius operator evolves the system distribution is not entirely clear, i.e. why f−1(A)?

Apologies for the typo in notation, f−1(A) should be noted as $f^{-1}(\mathbb{A})$ and we have corrected this typo in the revised version. In equation (15), the right-hand side denotes the distribution of states at the next time instant, i.e. $f(x)$. Denoting $g’(x)$ as the density of $f(x)$, one has $\int_\mathbb{A}f(x)g’(x) \text{d}\mu(x)= \int_f^{-1}(\mathbb{A})g(x)$. More detailed discussions on this can be found in Klus et al. (2016). As the introduction of Perron-Frobenius operators does not play a key role in the derivation of the DeSKO algorithm, we have moved this part to the Appendix according to the advice of Reviewer GCEV. 

(4) How does equation 3 change the training process compared with the DKO baseline? I assume it has to do with the expected distribution ED in equation 4 and 5? Please clarify, and also draw the similarity/difference with DKO side by side.

Thanks for the reviewer’s question. DeSKO encodes a distribution of observables $\psi_{t}$, while in DKO $\psi_{t}$ is a deterministic function of the state $x_t$. Thus in DeSKO, the loss function (3) includes the expectation over the observables $\mathbb{E}_{\psi_t}$ and this does not exist in DKO. The most outstanding difference between DeSKO and DKO is that DeSKO encodes a distribution of observables, enabling the model to learn the stochasticity contained in the dynamic, which is not supported by DKO.  Both DKO and DeSKO would need expectation over the dataset $\mathbb{E}_\mathcal{D}$.

(5) Is multi-step loss also included in DKO training? If not, I don’t think it is a fair comparison. Please clarify.

Yes, both DKO and MLP are trained to minimize the multi-step loss, with the same prediction horizon as DeSKO. We have clarified this in Section 5 of the revised manuscript.

References:

Klus S, Koltai P, Schütte C. On the numerical approximation of the Perron-Frobenius and Koopman operator[J]. Journal of Computational Dynamics, 2016, 3(1): 51-77.



In this paper, authors propose  "Deep Stochastic Koopman Operator" which can handle uncertainties in the system dynamics. The authors demonstrate that their approach can produce state loss similar to large capacity MLPs, yet maintains the linear structure in the projected space that is suitable to use linear MPC control methods such as LQR. The proposed method can achieve best tracking performance compared with previous non-stochastic version and SAC trained policies.
 The main strengths of the paper:

(1) Extends the Deep-DMD work by allowing uncertainties in the system dynamics.
(2) By projecting the system dynamics to linear spaces, one can explore existing optimal control methods such as LQR for linear systems.
(3) Achieve good control results (i.e. lowest tracking error or and in general more stable) in standard benchmarks. 

The  main  weaknesses of the paper includes:

(1) Can use more experiments to show a more convincing picture. 

(2) Compared with SAC, the proposed method does not show a clear edge in many cases. 

(3) Explanation of how the Perron-Frobenius operator evolves the system distribution is not entirely clear, i.e. why $f^{-1}(A)$ ?

(4) How does equation 3 change the training process compared with the DKO baseline? I assume it has to do with the expected distribution $E_D$ in equation 4 and 5? Please clarify, and also draw the similarity/difference with DKO side by side. 

(5) Is multi-step loss also included in DKO training? If not, I don’t think it is a fair comparison. Please clarify. 

 The authors proposed a novel method that allows Deep Koopman Operators to handle noisy systems, and demonstrate that their approach beats various baselines. However, there are not enough experiment varieties, and the results are close to the SAC baseline (not by large margin).


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). The learning and control of dynamical system from noisy data is not new, but this paper brings a novel angle into solving this problem - stochastic Koopman operator. The learning of the Koopman model and the proposed robust control provides a strong approach in terms of learning accuracy and reliabilily of controller. The proposed approach outperforms existing models and opens up a new perspective for the important problem of learning and control for dynamical systems. 

 Although the building blocks of this paper are not new, the paper leverage these ideas and empirically shows the advantage of the proposed framework. I recommend acceptance of this paper. 

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

+ The proposed method is elegant and intuitive, and the resulting algorithm practical and easy to implement.

+ The empirical evaluations are extensive and convincing, with the method performing favorably against strong baselines.

+ The paper is well written and very clear

Weaknesses

- The theoretical results implicitly assume that the learned Koopman embedding (encoded in equation (6)) is exact in the absence of noise across the entire state-space.  This is easily seen by setting $w_t=0$ in equation (10).  Of course, this is never the case in practice, as only prediction error bounds on the parameters $(A,B)$ can be obtained, and these are distribution dependent: i.e., once we switch from the exploratory policy used to identify the parameters (such as random noise, as was used in the experiments) to the closed-loop control policy, we have no guarantees that equations (6) and (10) are valid.  This is one of the main challenges that imitation and reinforcement learning must deal with both practically and theoretically.  Therefore while the main theory result is correct, it is only correct under impractical assumptions and implies nothing about stability of the actual closed-loop system.  Further, the assumption that the lifted system can be written to be linear in the control input (i.e., that the control input enters as $Bu_k$ oin equation (1)) is also very strong, and not supported by the theory, see for example, section 3.2.2 of 
@article{korda2018linear,
  title={Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control},
  author={Korda, Milan and Mezi{\'c}, Igor},
  journal={Automatica},
  volume={93},
  pages={149--160},
  year={2018},
  publisher={Elsevier}
}

 If the issues with the theoretical statements are adequately addressed, I find this to be a nice paper with convincing empirical results.

This paper introduces a new control framework that is based on Koopman theory. A key advantage of the proposed approach is the guaranteed stability, yielding a robust closed-loop controller. The method is evaluated on several challenging control benchmarks, and it is compared to several state-of-the-art approaches. The results highlight the advantages of the approach. In particular, in the setting of large external disurbances, the proposed method achieves the best results.  This work presents a new and elegant approach to modeling uncertainty in data-driven dynamical systems methods. The proposed method is simple and easy to code, yet it achieves impressive results on challenging benchmark control tasks. The guaranteed stability is a nice addition to this framework, which further facilitates the design and application of this approach to new tasks. The extensive evaluation of the approach shows its features in comparison to other methods, and independently. Finally, the versailty of framework and its simplicity opens up new possibilities for future research as partially listed by the authors in the discssion section.

My main concern regarding this work is the somewhat not detailed discussion of related work. Specifically, the following references are missing and should be added:

1. "Learning Koopman invariant subspaces for dynamic mode decomposition", Takeishi et al. 2017
2. "Physics-informed probabilistic learning of linear embeddings of nonlinear dynamics with guaranteed stability", Pan et al. 2020
3. "Forecasting Sequential Data Using Consistent Koopman Autoencoders", Azencot et al. 2020

Moreover, I would like to better understand the differences of the current approach with respect to Morton et al. 2019, and Pan et al. 2020. The former work considers distribution of observables, whereas the latter work has stability guarantees. Similarly, the Azencot et al. 2020 work shows several examples with sensor noise, and they also discuss empirical stability.


minor comments:
	- The reference noa 2016 in Sec. 2 seems to be wrong.
	- Please specify the domain and range of A and B in Eq. (1)
	- Why $\mathcal{K}$ is separable in the dynamics and control?
	- I believe you should change $L_\infty$ to $L_2$ below Eq. (3)
	- I would move Sec. 3.1 & 3.2 to the appendix. 
	- The contribution of the Perron-Frobenius to this work is unclear. Please elaborate.
	- The subscript/superscript notation in Eq. (5) is inconsistent with Sec. 3.1
	- The Figs. 2, 4 & 5, and possibly 3 should use a log-scale on the y-axis.
	- Why SAC is not in Fig. 2? Similarly, why MLP is not in Figs. 3 & 4?
 This paper presents an elegant solution to the modeling of uncertainty in dynamical systems for the puprpose of designing closed-loop controllers. The model is backed by control theory and has a guaranteed stability. The method is evaluated on several tasks showing its properties.


Meta-review: The paper was seen positively by all reviewers. The strength of the paper are:
- Intuitive and interesting combination of Koopman Operators and Optimal Control for Reinforcement Learning
- Convincing experiments on challenging benchmark tasks
- All of the issues of the reviewers (advantages to SAC, gaps in the theory and missing references) have been properly addressed in the rebuttal.

I therefore recommend acceptance of the paper.

Decision: Accept (Poster)