Data-driven Koopman Operator-based Prediction and Control
Using Model Averaging

While conventional systems modeling methods describe the dynamics on a state-space, operator theoretic approaches offer an alternative view of the system behaviors through the lens of functions, often called feature maps or observables. The Koopman operator, a linear operator acting on a space of feature maps, has been widely used in the context of data-driven modeling, analysis, and control of nonlinear dynamical systems[1, 2]. It allows linear evolution of the embedded systems on a function space even if the original dynamics is nonlinear in the state-space. Also, it can be approximated numerically from data using either linear regression or neural network training. The Koopman operator framework is especially appealing when applied to controller synthesis since linear controller designs such as Linear Quadratic Regulator (LQR) and linear Model Predictive Control (MPC) can be utilized with the obtained linear embedding model.

However, obtaining accurate and reliable Koopman operator-based models that can be successfully applied to decision-making tasks is still challenging. Specifically, it is known that the convergence property of EDMD does not hold for control systems[3], which implies that Koopman operator-based models may fail to accurately describe control systems even if the training data is sufficiently rich and unbiased. To overcome the modeling error of such Koopman operator-based models for control applications, several data-driven controller designs have been proposed that take model uncertainty into account[4, 5, 6, 7]. Also, if one restricts the attention to control affine systems instead of general nonlinear dynamics, finite-data error bounds are available and stability guarantees may be established on the basis of robust control theories[8, 9]. While many efforts have been made on the model uncertainty of Koopman operator-based models from a controller design perspective, it is also of great importance to devise a modeling method that can yield accurate and generalizable models for control systems. In [10], a model refinement technique is proposed to incorporate data from closed-loop dynamics into model learning so that the modeling error when applied to controller design problems can be directly mitigated. To improve the generalizability of the Koopman operator-based models for a wide range of applications, [11] develops a two-stage learning method utilizing the oblique projection in the context of linear operator learning in a Hilbert space.

In this paper, we adopt a Bayesian approach to learn unknown control systems with the use of the Koopman operator, which takes into consideration that a single Koopman operator-based model is not likely to be accurate and reliable for a wide range of operating points and it will be beneficial to aggregate an ensemble of models and utilize them to find the most useful output. Ensemble approaches are a popular class of methods that aim to improve the accuracy in learning problems by combining predictions of multiple models[12, 13]. We show that the Bayesian Model Averaging (BMA)[14], a Bayesian inference-based model averaging technique, yields a Koopman operator-based model whose parameters are represented as sums of corresponding parameters of individual models weighted by the posterior model evidence. The proposed Koopman Model Averaging (KMA) first populates multiple models, starting from a base model whose feature maps are parameterized by a neural network. Computing point estimates of the original state and the embedded state then results in the same type of linear embedding model while being model-uncertainty aware. The outputs of the individual models are incorporated into this uncertainty-aware model and the overall performance is expected to be improved.

In Section II, the Koopman operator framework for control systems is reviewed. Model learning with neural networks is formulated in Section III, and the proposed modeling approach is derived in Section IV. Numerical examples are provided in Section V for evaluation of the proposed method.

Consider a discrete-time, non-autonomous dynamical system:

where $x_{k}\in\mathcal{X}\subseteq\mathbb{R}^{n}$, $u_{k}\in\mathcal{U}\subseteq\mathbb{R}^{p}$, and $f:\mathcal{X}\times\mathcal{U}\rightarrow\mathcal{X}$ are the state, the input, and the (possibly nonlinear) state-transition mapping, respectively. It is assumed throughout the paper that the dynamics (1) is unknown while we have access to data of the form $\{(x_{k},u_{k},y_{k})\mid y_{k}=f(x_{k},u_{k})\}$.

We embed the state $x_{k}$ into a latent space by applying feature maps, or observables, $g:\mathcal{X}\rightarrow\mathbb{R}$. Let $\mathcal{G}$ denote the function space to which the feature maps $g$ belong. If the input is constant s.t. $u_{k}\equiv\bar{u}$, $\forall k$, (1) induces autonomous dynamics:

The state transition through $g$ is then represented by

where the composition operator $\mathcal{K}_{\bar{u}}:\mathcal{G}\rightarrow\mathcal{G}:g\mapsto g\circ f_{\bar{u}}$ is called the Koopman operator associated with the autonomous dynamics (2). It is obvious from (3) that $\mathcal{K}_{\bar{u}}$ is a linear operator and it describes the possibly nonlinear dynamics (2) linearly in the latent space $\mathcal{G}$. Since $\mathcal{K}_{\bar{u}}$ is infinite dimensional acting on the function space $\mathcal{G}$, its finite dimensional approximation is often considered for the purpose of dynamical systems modeling. Given $N_{x}$ feature maps $g_{i}\in\mathcal{G}$ ($i=1,\cdots,N_{x}$), there exists a matrix $K\in\mathbb{R}^{{N_{x}}\times{N_{x}}}$ s.t.

if and only if $\text{span}(g_{1},\cdots,g_{N_{x}})$ is an invariant subspace, i.e., $\mathcal{K}_{\bar{u}}g\in\text{span}(g_{1},\cdots,g_{N_{x}})$ for $\forall g\in\text{span}(g_{1},\cdots,g_{N_{x}})$. Note that finding feature maps $g_{i}$ that form an invariant subspace may be challenging in practice. In the data-driven setting, one can approximately obtain $K$ by solving the following linear regression problem:

where $\bm{g}(x_{k}):=[g_{1}(x_{k})\cdots g_{N_{x}}(x_{k})]^{\mathsf{T}}$ and the training data is given as $\{(x_{i},y_{i})\mid y_{i}=f_{\bar{u}}(x_{i})\}$. It admits the unique analytical solution with the use of pseudo inverse and this procedure is called Extended Dynamic Mode Decomposition (EDMD)[15]. Note that the design of the feature maps $g_{i}$ needs to be pre-specified in EDMD. The finite dimensional approximation $K_{\text{approx}}$ then yields a linear embedding model of the form:

where $x_{k+1}^{\text{pred}}$ is the predicted state at the next time step and the decoder $W\in\mathbb{R}^{n\times{N_{x}}}$ is similarly obtained by solving linear regression on the given data set. The decoder can be also introduced as a nonlinear mapping depending on the problem.

For general non-autonomous dynamics (1) with control inputs $u_{k}$, the corresponding Koopman operator can be defined as follows[3]. For the space of input sequences: $l(\mathcal{U}):=\{(u_{0},u_{1},\cdots)\mid u_{k}\in\mathcal{U},\forall k\}$, consider a mapping $\hat{f}:\mathcal{X}\times l(\mathcal{U})\rightarrow\mathcal{X}\times l(\mathcal{U}):(x,(u_{0},u_{1},\cdots))\mapsto(f(x,u_{0}),(u_{1},u_{2},\cdots))$. Also, let $\hat{g}:\mathcal{X}\times l(\mathcal{U})\rightarrow\mathbb{R}$ be an embedding feature map from an extended space $\mathcal{X}\times l(\mathcal{U})$ to $\mathbb{R}$. Then, the Koopman operator associated with the non-autonomous dynamics (1) is defined as a linear operator $\mathcal{K}:\hat{\mathcal{G}}\rightarrow\hat{\mathcal{G}}:\hat{g}\mapsto\hat{g}\circ\hat{f}$ s.t. $\hat{\mathcal{G}}$ is a function space to which feature maps $\hat{g}:\mathcal{X}\times l(\mathcal{U})\rightarrow\mathbb{R}$ belong and the dynamics (1) along with a sequence $(u_{k},u_{k+1},\cdots)$ of future inputs is represented by

The argument on the subspace invariance of the Koopman operator (eq. (4)) also holds for the non-autonomous case, which is written as

where $\hat{K}\in\mathbb{R}^{\hat{N}\times\hat{N}}$ and ${\hat{N}}$ denotes the number of feature maps. In analogous to (8), if we consider $\hat{N}=N_{x}+p$ feature maps $\hat{g}_{i}$ ($i=1,\cdots,N_{x}+p$) of the form:

the first $N_{x}$ rows of (9) reads $\bm{g}(x_{k+1})=\hat{A}\bm{g}(x_{k})+\hat{B}u_{k}$, where $[\hat{A}\ \hat{B}]\in\mathbb{R}^{{N_{x}}\times(N_{x}+p)}$ denotes the first $N_{x}$ rows of $\hat{K}$. Similar to (8), this yields a linear embedding model:

in which the parameters $A,B$, and $C$ may be obtained by EDMD with training data $\{(x_{i},u_{i},y_{i})\mid y_{i}=f(x_{i},u_{i})\}$, i.e.,

A notable feature of the model (13) is that the model dynamics in the embedded space is given as a linear time-invariant system and linear controller designs can be utilized to control the possibly nonlinear dynamics (1). For instance, with the parameters $A$ and $B$ in (13), one can compute an LQR gain that stabilizes the following virtual system:

Note that if $\bm{g}$ span an invariant subspace, we have the exact relation: $\bm{g}(x_{k+1})=A\bm{g}(x_{k})+Bu_{k}$. Since the model (13) has a linear decoder, the quadratic loss $\bm{g}(x_{k})^{\mathsf{T}}Q\bm{g}(x_{k})$ of the LQR problem ($Q$ is a weight matrix) can be associated with the original state $x_{k}$ by $\bm{g}(x_{k})^{\mathsf{T}}Q\bm{g}(x_{k})\approx x_{k}^{\mathsf{T}}(C^{\mathsf{T}}QC)x_{k}$.

While we adopt linear embedding models represented in the form (13), which is a common choice in the literature, more expressive ones such as bilinear models are also available if the strict linearity w.r.t. $u_{k}$ in (13) is not enough to reconstruct the target dynamics[16, 17]. Also, it is noted that in addition to the extension of the Koopman operator framework to control systems reviewed in this section[3], there are other formalisms to utilize the Koopman operator for modeling non-autonomous systems[18].

The accuracy of the linear embedding model (13) depends on the design of the feature maps $\bm{g}$ as well as how the model dynamics parameters $A$, $B$, and $C$ are obtained. While EDMD provides a simple model learning procedure with the analytic solution, the design of feature maps needs to be user-specified such as monomials and Fourier basis functions and it may not be sufficiently expressive for complex nonlinear dynamics. Neural networks, on the other hand, are a popular choice of the feature map design since they allow greater expressivity of the model compared to EDMD (e.g., [19, 20, 21]). With the feature maps $\bm{g}$ characterized by a neural network, both the model dynamics parameters and the feature maps can be learned simultaneously, which is formulated as the following problem:

Although neural network-based models are expected to be more expressive and accurate than EDMD-based ones, Problem 1 is typically a high-dimensional non-convex problem and the resulting models may suffer from overfitting or poor learning. For instance, solving Problem 1 can lead to inaccurate models if the optimization is terminated at a local minimum with a high loss value, or if the quantity and/or quality of training data are not sufficient to reconstruct the target dynamics for a wide range of operating points. Therefore, one may need to repeat the data-collection and learning processes multiple times to obtain a satisfying model in practice, which often takes a long time and consumes a large amount of computational resources.

As the first step of the proposed method, we execute Step 1 to obtain a base model. Considering that the base model may not be perfect, we aim to improve its accuracy further by combining multiple models in the second step. Specifically, an ensemble of linear embedding models (13) is generated using additional data points with the design of feature maps fixed to that of the base model. These models, including the base model, share the same model structure but are trained on different data sets and their predictive capabilities vary. Therefore, even if some model is most accurate for certain unseen data among the model ensemble, others may show better accuracy for different data. To handle this lack of generalizability of individual models, we use all the models so that the accuracy for unknown regimes of dynamics will be improved. Specifically, a Bayesian inference-based model averaging technique is utilized to merge the individual models into a new linear embedding model.

In this section, we first outline the Bayesian Model Averaging (BMA). It is then followed by the formulation of the proposed Koopman operator-based model averaging method.

We propose a model averaging method for learning Koopman operator-based models for prediction and control utilizing the Bayesian model averaging. While the Koopman operator framework allows to obtain linear embedding models from data so that linear systems theories can be applied to control possibly nonlinear dynamics, it is challenging to accurately reconstruct the dynamics and deploy the learned model in several applications. Considering that training a single model only may not be sufficient to obtain model predictions that are accurate enough for a wide range of operating points even if the model possesses high accuracy w.r.t. a certain regime of dynamics such as training data, the proposed method first trains a base model with the use of neural networks and then populates an ensemble of models on different data points. Based on the ideas of the Bayesian model averaging, these models are merged into a new weighted linear embedding model, in which individual outputs of the model ensemble are weighted according to the posterior model evidence. This model explicitly takes uncertainty into account and the overall performance is expected to be improved. Using numerical simulations, it is shown that the proposed weighted model achieves better state predictive accuracy as well as greater generalizability to different control applications compared to other Koopman operator-based models.