Learning Noise-Robust Stable Koopman Operator for Control with Hankel DMD

Consider a nonlinear discrete-time dynamical system:

where $f(\cdot)$ represents a generally nonlinear mapping, and $x\in\mathbb{R}^{n}$ and $u\in\mathbb{R}^{m}$ denote the state and input of the system, respectively. The objective is to predict the trajectory of the system, given an initial condition $x_{0}$ at time $t_{0}$, and a sequence of inputs $\{u_{i}\}_{i=0}^{N_{p}}$, where $N_{p}$ is the prediction horizon length of interest. To achieve this, we define the lifted linear predictor as a linear dynamical system determined by a triplet of matrices $(A,B,C)$:

where $z\in\mathbb{R}^{N_{\phi}}$ is the lifted state, $\Phi:\mathbb{R}^{n}\rightarrow\mathbb{R}^{N_{\phi}}$ is the lifting mapping, and $\hat{x}$ is the prediction of the original system’s state $x$. The lifting mapping $\Phi(\cdot)$ is defined as:

where $\phi_{i}:\mathbb{R}^{n}\rightarrow\mathbb{R}$ is a scalar function referred to as an observable.

It is important to note that under mild assumptions, the predictor could also take a bilinear form:

as noted in several works [36, 37, 38, 39]. While it is natural to extend our work to accommodate the bilinear form, we avoid this for simplicity and maintain readiness for application in linear control systems.

To provide a rigorous justification for constructing a linear predictor via state-space lifting, we briefly outline the Koopman operator approach for analyzing controlled dynamical systems. We consider the following discrete-time nonlinear system with control:

where, $x\in\mathbb{R}^{n}$ is the state of the system and $u\in\mathcal{U}\subset\mathbb{R}^{m}$ is the control input. $\mathcal{U}$ is the space of the available control inputs. $f:\mathbb{R}^{n}\times\mathbb{R}^{m}\rightarrow\mathbb{R}^{n}$ is the transition mapping. $x^{+}$ is the system state of the next time step. We define the extended state comprising of the state $x$ of system and the control sequence $\mathbf{u}\in\ell(\mathcal{U})$, where $\ell(\mathcal{U})$ is the space of all sequences $(u_{i})_{i=0}^{\infty}$ with $(u_{i})_{i=0}^{\infty}\in\mathcal{U}$,

The dynamics of the extended state $\chi$ is described by

where $\mathcal{S}$ is the left shift operator, i.e., $(\mathcal{S}\mathbf{u})(i)=u(i+1)$, and $u(i)$ denotes the $i$th element of the sequence $\mathbf{u}$. The Koopman operator $\mathcal{K}:\mathcal{H}\to\mathcal{H}$ associated to eq. 6 is defined by

for any observable $\psi:\mathbb{R}^{n}\times\ell(\mathcal{U})\to\mathbb{R}$ belonging to some space of observables $\mathcal{H}$. Despite Koopman operator $\mathcal{K}$ being a linear operator fully describing the time evolution of any observable function on the non-linear dynamical system, it is inherently infinite dimensional that makes learning Koopman operator challenging.

For practical purposes, we adapt to the Extended Dynamic Mode Decomposition (EDMD) [40] approach followed by Korda and Mezić [18] for control. The EDMD provides a straightforward method to approximate the Koopman operator from data, thereby obtaining the lifted linear predictor in eq. 2. Since our goal is not to predict the control sequence, the observable function will depend only on the system states. Thus, the process begins by selecting $N_{\phi}$ observables $\phi_{1},\ldots,\phi_{N_{\phi}}$ and forming the lifted state $z$ from the measured state $x$ of the system (1) as follows:

The one-step prediction loss based EDMD algorithm assumes that a collection of data $(x_{j},u_{j},x_{j}^{+}),j=1,\ldots,K$ satisfying $x_{j}^{+}=f(x_{j},u_{j})$ is available and seeks matrices $A$ and $B$ minimizing

To select suitable basis functions $\Phi$ that span an invariant subspace, we are motivated by Polyflow basis functions from Wang and Jungers [28] and by prior work such as Krylov Subspace methods discussed by Mezic [27]. When the dynamics $f(x,u)$ is known, we can use a special set of basis functions $\Phi(x)=T_{N}(x)$, known as Polyflow basis functions of order $N\in\mathbb{Z}^{+}$ which is defined,

Using these Polyflow basis functions to lift the system state, we can derive the following system, which aligns with equation eq. 2

Note that the matrix $C$ is typically known from the selection of $T_{N}(x)$.

Despite the similarity, time delay embedding [41] and Polyflow offer slightly distinct approaches to state lifting. Time delay embedding relies on observing and recording the physical system’s state over time to embed the state vector. In contrast, Polyflow directly utilizes the governing equations of the system, enabling an instantaneous embedding of the state without accessing the history of the physical system. Our framework assumes a known governing equation, but for systems with unknown dynamics, methods such as SINDy [42] or Neural ODE [43] can be employed as a surrogate of the governing equation.

To enhance long-term prediction accuracy, we implement a recurrent roll-out loss instead of using one-step prediction loss eq. 8. For this, we consider a set of $M$ trajectories, each consisting of $R$ steps. Let $x_{i,j}$ denote the state at the $i$-th step of the $j$-th trajectory, where $i=0,1,\ldots,R$ and $j=1,2,\ldots,M$. The following optimization problem is then formulated:

where, $A$ and $B$ are the matrices to be determined, and $u_{i}^{j}$ denotes the input at the $i$-th timestep of the $j$-th trajectory. To determine an appropriate basis size $N$, we solve the optimization problem (13) iteratively, increasing $N$ until the approximation error is sufficiently small.

Curriculum learning [44] is a training strategy where a model is first trained on simpler data points and then gradually introduced to more challenging ones. This approach helps the model to progressively learn and adapt to increasing complexity.

In our framework, we implement a form of curriculum learning by progressively increasing the roll-out length $R$ during the training process. Initially, the model is trained with shorter roll-out lengths, focusing on immediate predictions. As the training progresses and the model stabilizes, the roll-out length $R$ is increased geometrically after certain epochs. This gradual increase is expected to allow the model to learn from easier short-term predictions before tackling the more complex task of long-term prediction, thereby aligning with the curriculum learning paradigm. By doing so, the model is expected to build its capacity to handle extended sequences, improving overall prediction accuracy over long time horizons.

To promote stability of the lifted system, we enforce stable eigenvalues of $A$ by employing the fact that a skew symmetric matrix with negative real valued diagonal elements have stable or neutral stable eigenvalues in continuous time dynamical system [32]. For any $Q\in\mathbb{R}^{\tilde{n}\times\tilde{n}}$ matrix and $C\in\mathbb{R}^{\tilde{n}\times\tilde{n}}$ diagonal matrix with negative real diagonal elements, $\tilde{A}=Q-Q^{T}+C$ has eigenvalues in the negative real part of the complex plane. We compute the discrete-time $\hat{A}\in\mathbb{R}^{\tilde{n}\times\tilde{n}}$ by taking the matrix exponentiation of $\tilde{A}=Q-Q^{T}+C$ multiplied with the sampling timestep $\Delta t$.

Note that a similarity transform of $\hat{A}$ via the coordinate transform matrix $P\in\mathbb{R}^{\tilde{n}\times\tilde{n}}$ is required to recover $A$.

This formulation ensures the stability of the Koopman Operator throughout the training, which we refer as Dissipative Koopman Operator. In contrast, we refer Standard Koopman Operator as lifted system with an unconstrained matrix $A$.

One of the most important applications of lifted linear predictor is to provide a surrogate model for MPC, which is referred as Koopman-MPC (KMPC) [18]. At each discrete time step, the MPC solves the following optimization problem over the prediction horizon $N_{p}$:

where $u_{\min/\max}$ are the input bounds, $z_{\min/\max}$ are the bounds on the lifted states, and $r_{k}$ is the reference signal. Only the first control input $u_{0}$ from the optimal sequence is applied, and the optimization problem is resolved at the next time step. The cost function $J$ is given by:

Models trained on trajectories that are generated by randomly sampling the phase space and control input could perform poorly when subjected to the optimal feedback control inputs. This is primarily because these trained models lack exposure to the desired scenarios during the training process, especially for relatively high-dimensional system. To address this issue, we propose a straightforward iterative data augmentation technique inspired by the work of Uchida and Duraisamy [45]. The idea is to iteratively collect response of the true system under control input synthesized with the previous data-driven model. Such response together with control input is used to augment the original dataset to further refine and improve the data-driven model. Our rationale is that such experience from the true system can be used to better explore the phase space of nonlinear system than merely ad-hoc random sampling the phase space, which is subject to the curse of dimensionality. Hence, our proposed framework is named as STIRK (Stable Training with Iterative Roll-out Koopman Operator). A summary of the proposed framework is illustrated in the figure 1.

We compare the proposed framework with the following versions of DMD:

1. 1.

   DMD [46],

2. 2.

   eDMD  [40]with Polyflow basis function and Radial Basis function,

3. 3.

   OptDMD [47] with Polyflow basis function, Radial Basis functions and Identity Observables,

4. 4.

   TLSDMD [48] with Polyflow basis function, Radial Basis functions and Identity Observables,

5. 5.

   FBDMD [49] with Polyflow basis function, Radial Basis functions and Identity Observables.

Figure 9a compares the performance of the Dissipative Koopman Model and the Standard Koopman Model across different noise levels. The Dissipative Koopman Model consistently shows lower mean errors and mean test errors, underscoring its superior prediction accuracy and robustness. However, the training time for the Dissipative Koopman Model is considerably longer than that of the Standard Koopman Model, due to the increased number of learning parameters and the computationally intensive operations, such as matrix exponentiation and pseudo-inverse computation. Despite its relatively short training time, a significant drawback of the Standard Koopman Model is its sensitivity to the initial conditions of the Koopman matrix. If these initial conditions are not carefully tuned, the model can produce unstable eigenvalues, leading to poor performance. Figure 8 illustrates this issue by showing the discrete-time eigenvalues of the standard and dissipative Koopman operators for the Van der Pol Oscillator System at a noise level of 0.0599, using a different weight initialization than in previous cases. For fair comparison, both standard and dissipative Koopman models are initialized with stable Koopman matrices. However, after 1,000 training epochs, some of the discrete-time eigenvalues in the Standard Koopman matrix move outside the unit circle, leading to a loss of stability while dissipative Koopman model never loses stability during training.

Figure 9b compares the performance of the model using constant learning rate (LR) scheduling and cyclic learning rate scheduling under different noise levels. For low noise, the cyclic learning rate scheduling achieves a 20% reduction in training error and a 16.67% reduction in test error, with a 9.09% reduction in time taken compared to the constant learning rate. In medium noise scenarios, the cyclic learning rate scheduling shows a 16.67% improvement in training error, 14.29% in test error, and an 8.70% improvement in time taken. However, at high noise levels, cyclic learning rate scheduling underperforms compared to the constant learning rate, with higher training and test errors, as well as slightly longer training times. These results indicate that while the cyclic learning rate scheduler is beneficial at low and medium noise levels, it is less effective in high-noise environments.

Figure 9c examines the effect of switching optimizers on the model’s performance. Different combinations of Adam and LBFGS optimizers were used, and their impact on mean errors, mean test errors, and training times was analyzed. Switching from the Adam optimizer to LBFGS after 4000 epochs in a 5000-epoch training cycle resulted in about a 25% error reduction compared to the Adam-only optimizer at low and medium noise levels, demonstrating the efficacy of LBFGS during fine-tuning. However, this advantage diminished at higher noise levels, where the differences between the optimizers became less pronounced. In experiments where optimizers were switched periodically every 500 epochs, results varied, but overall, periodic switching helped minimize the error in certain cases, though it did not consistently outperform the Adam-LBFGS approach. Notably, using only LBFGS led to significantly higher errors, suggesting that while LBFGS is beneficial for fine-tuning, it tends to settle in local minima and does not further improve the error.

Figure 9d compares the performance of the model using progressive roll-out length scheduling and constant roll-out length scheduling under different noise levels. Progressive roll-out scheduling shows a 20% improvement in mean training errors and a 16.67% improvement in mean test errors compared to the constant roll-out length, indicating better learning and generalization. Although the time taken for training with progressive roll-out scheduling is 6.94% lower, the overall training duration remains slightly higher due to the incremental increase in roll-out length. These results highlight the effectiveness of progressive roll-out scheduling in enhancing model performance.

Finally, we compare the overall performance of our optimized configuration with the baseline configuration, as shown in 10 and 11. 10 shows the prediction performance comparison for Van der Pol system and 11 shows the Model Predictive Costs comparison for Cartpole system.

For the Van der Pol system 10, the baseline setup utilized a standard Koopman operator with a fixed maximum rollout length of 99, without any learning rate scheduler, and maintained a constant learning rate of 0.01. In contrast, our optimized configuration employed a dissipative Koopman operator alongside a progressive rollout strategy, starting with a rollout length of 1 and geometrically increasing by a factor of 2 every 200 epochs until reaching the maximum rollout length. Additionally, we implemented a cyclic learning rate scheduler in triangular mode, with the maximum learning rate set to 0.1, halving every 500 epochs, and a base learning rate of 0.01. Both configurations used the Adam optimizer. Overall, the optimized configuration resulted in a 39% improvement in training error and a 37.6% improvement in testing error compared to the baseline. This significant reduction in error highlights how the combined components of the optimized configuration effectively enhance the model’s performance.

In the case of the Cartpole system (11, the baseline strategy used the Adam optimizer with a constant rollout length of 16 and a constant learning rate. The standard Koopman operator was utilized throughout the training. Conversely, the optimized strategy employed the Adam optimizer for the initial 1500 epochs, followed by fine-tuning using the LBFGS optimizer. The rollout length was progressively increased from 1 to 16, a cyclic learning rate was implemented, and a dissipative Koopman operator was used. The optimized training strategy outperforms the baseline training strategy across different initial conditions over multiple iterative training data augmentations.