MamKO: Mamba-based Koopman operator for modeling and predictive control
Authors: ZHAOYANG LI
Minghao Han
Xunyuan Yin
Venue: ICLR 2025 Spotlight (2025)
Keywords: Mamba; Koopman operator; model predictive control; nonlinear systems
Primary area: learning on time series and dynamical systems

Abstract: The Koopman theory, which enables the transformation of nonlinear systems into linear representations, is a powerful and efficient tool to model and control nonlinear systems. However, the ability of the Koopman operator to model complex systems, particularly time-varying systems, is limited by the fixed linear state-space representation. To address the limitation, the large language model, Mamba, is considered a promising strategy for enhancing modeling capabilities while preserving the linear state-space structure.
In this paper, we propose a new framework, the Mamba-based Koopman operator (MamKO), which provides enhanced model prediction capability and adaptability, as compared to Koopman models with constant Koopman operators. Inspired by the Mamba structure, MamKO generates Koopman operators from online data; this enables the model to effectively capture the dynamic behaviors of the nonlinear system over time. A model predictive control system is then developed based on the proposed MamKO model. The modeling and control performance of the proposed method is evaluated through experiments on benchmark time-invariant and time-varying systems. The experimental results demonstrate the superiority of the proposed approach. Additionally, we perform ablation experiments to test the effectiveness of individual components of MamKO. This approach unlocks new possibilities for integrating large language models with control frameworks, and it achieves a good balance between advanced modeling capabilities and real-time control implementation efficiency.

Published as a conference paper at ICLR 2025
MAMKO: MAMBA-BASEL KOOPMAN OPERATOR FOR MODELING AND PREDICTIVE CONTROL
Zhaoyang Li' Minghao Han?, Xunyuan !2n
School of Chemistry, Chemical Engineering and Biotechnology. Nanyang Technological University 2 Nanyang Environment and Water Research Institute, Nanyang Technological University
LI ZIZH0022Ee.... ntu..uu.. eew..g.......o mi minghao. hwn@ntu. hww@ntu.edu.u................... edu. sg; xunyuan, yinentu. edu.sg .g
ABSTRACT
The Koopman theory, which enables the transformation of nonlinear systems into linear representations, is a powerful and efficient tool to model and control nonlin- ear systems. However, the ability of the Koopman operator to model complex sys- tems, particularly time-varying systems, is limited by the fixed linear state-space representation. To address the limitation, the large language model, Mamba, is considered a promising strategy for enhancing modeling capabilities while pre- frame-
serving the linear state-space structure. In this paper, propose enhanced work, the Mamba-basee Koopman operator (MamKO), which provides model prediction capability and adaptability, compared to Koopman models with constant Koopman operators. Inspired by the Mamba structure, MamKO generates Koopman operators from online data; this enables the model ef- fectively capture the dynamic behaviors of the nonlinear system over time. A model predictive control system then developed based on the proposed MamKO model. The modeling and control performance of the proposed method is evalu- ated through experiments on benchmark time-invariant and varyingg systems. The experimental results demonstrate the superiority of the proposed approach. Additionally, we perform ablation experiments to test the effectiveness of indi- vidual components of MamKO. This approach unlocks new possibilities for in- tegrating large language models with control frameworks, and achieves a good balance between advanced modeling capabilities and real-time control implemen- tation efficiency.
INTRODUCTION
Deep learning methods have made promising achievements in modeling large and complex systems with collected datasets (Wang et al. 2024). The time-series prediction tasks can be well achieved with existing learning technologies and applied various fields, including weather forecasting (Verma et al. 2024) and traffic prediction (Li et al. 2023). These data-driven methods, especially for deep learning methods, have proven to be highly capable modeling dynamic systems and provide strong support for advanced control applications.
The Koopman operator theory (Koopman) 1931) has emerged as a promising data-driven method in recent years. The Koopman operator can represent the dynamics of nonlinear systems by linear state-space model (SSM) (Kalman 1960) by lifting the states from the original space to higher- dimensional space through series of observable functions. Based the theory, extended dynamic mode decomposition (EDMD) (Williams et al., (013)) with manually selected observable functions proposed. With the advent of deep learning, neural networks (NNs) have been employed serve as the observable functions in the Koopman operator framework. Advanced deep learning methods, such as autoencoders (Takeishi et al. | (0177). probabilistic NNs (Han et al. (022]) and graph NNs (Li et al. 2020)) have been integrated with the Koopman operator to enhance the prediction accuracy of the model. The linear controller based on the Koopman model, including linear quadratic regulator
(Han et al. 2023), can be efficiently designed from the concise linear model. Although the Koopman model mentioned above shows relatively accurate prediction results, due to the parameter capacity
*Corresponding author: Xunyuan Yin.

---

Published as a conference paper at ICLR 2025
of the Koopman framework, challenges still exist in modeling real-world systems, especially for vime-narying systems.
Large language models (LLMs) (Zhao et al. 2023 are deep learning models that leverage large- scale number of parameters to comprehend and generate human language. The internal architecture of the model enables it to link and reason through contextual information. The most typical structure inside the large language model is Transformer (Wolf et al. 2020). Although the nonlinear function and attention mechanism (Vaswani et al., 2017) primarily improve the approximation capacity, the intricate structures pose challenges for model-based controller design. The Mamba structure (Gu & Dao 2023) based on SSM is developed for LLMs. Within the structure of Mamba, the similar linear SSM from Koopman operator theory is interpreted as a particular form that integrates features from both recurrent neural networks (RNNs) and convolutional neural networks (CNNs) (Dao &
Gu 2024). Selection mechanisms from Mamba are implemented to improve the modeling ability of linear SSM by generating matrix sequences, which can also be a potential solution for improving the capacity of the Koopman model.
In this work, we propose a Mamba-based Koopman operator (MamKO) by integrating the matrices generation network from Mamba (Gu & Dao 2023) with the original Koopman model to enhance the modeling accuracy. The structure can generate the Koopman operators according to the historical data and construct the SSM accordingly. Using the linear SSM, the MPC problem is formulated as a convex optimization problem. The proposed methods are then evaluated against existing represen- tative approaches across various settings. The results demonstrate the superiority of the proposed framework in both modeling and control performance compared to state-of-the-ar methods. The contributions of this work are summarized as follows:
We propose novel Mamba-based Koopman operator (MamKO) modeling method, which leverages matrices generated from the Mamba structure to model complex nonlinear systems.
We further enhance the matrix generation structure of Mamha accommodate unstable and time-varying systems.
We develop an MPC scheme based on the MamKO model to achieve computationally efficient optimal control of nonlinear systems.
Experiments are conducted on benchmark time-invariant and time-varying systems to illustrate the efficacy and superiority of the proposed framework.
2 RELATED WORKS
Koopman Operators Koopman operator theory (Koopman 1931) has been regarded as a powerful tool for dynamic system analysis (Brunton al.) 2022). In (Schulze et al. 2022). the scope of Koopman operator theory is extended from autonomous systems to controlled systems, enabling the design of corresponding controllers. Methods with manually selected observable functions show promising results both modeling and control (Zhang et al. 2023). However, the time consumption to find appropriate observable functions is significant drawback of this approach. An alternative method using NNs to automatically construct the observable functions can bypass this problem. This kind of Koopman model has been effectively applied in areas such as fluid physies (Morton al.) 2018), robotics (Shi & Meng 2022), chemical engineering (Han et al. 2024a). and vehicle systems (Chen & Lv 2024), where the neural network-based model serves powerful predictor.
Although the Koopman model has received much attention and is well-studied both in the field of modeling and control, challenges still exist when facing complex systems. especially for time- varying systems. (Zhang et al.) 2019) has developed an online dynamic mode decomposition strat- egy for time-varyyl varying systems with slowly changing parameters. With the online collected data, (Hao) et al. 2022) designs an adjusting algorithm to tune the parameters in the observable functions. Inte- grating with the Fourier filter to disentangle and exploit time-invariant and time-varying dynamics, the Koopman operators are updated in (Liu et al. 2023) to accommodate the time-varying dynam- ics. While these methods have shown promising results, the optimization process within the online updating framework can be time-consuming and may lead to infeasible solutions. In contrast, approach directly generates the Koopman operators from historical data, enabling real-time produc- tion of the time-varying SSM.
2

---

Published as a conference paper at ICLR 2025
Mamba Structure Previous works Mamba are based the research of state-space model (Gu
et al.
2020 2022b) for continuous data processing and applied to the field of audio and vision (Goel
et al.
2022 Nguyen et al. 2022). The structured state-space sequence models (S4) (Gu et al. (022a)
are expanded to selective state-space models (S6) in (Gu & Dao 2023) by adding the selection mechanism. The mechanism inside the Mamba structure transfers the time-invariant model to the time-varying model by generating the matrices of SSMs using the input data, which can facilitate the modeling of time-varying tasks. The modeling ability has largely increased with this mechanism, and corresponding hardware-ane
algorithm has been proposed to boost computational efficiency further. The Mamba structure has been applied in the fields including language processing (Lieber et al. 2024). image segmentation (Zhu et al. 2024 Ruan & Xiang 2024), video processing (Li ctt al. 2024). Notably, while the Mamba structure based on a linear SSM, is not directly suitable for controller design due to nonlinear functions within the selection mechanism. Modifications have been introduced to adapt the system to fit within the control framework.
Model-based Learning Control With the development of data-driven methods, the modeling curacy has increased, facilitating the application of model-based learning control methods. In the field of model-based reinforcement learning, multiple algorithms have been proposed to train the control strategy inside the established model (Ha & Schmidhuber, (0188) Hafner et al.] 2019: [2020] 2023). From the perspective of MPC, the computational burden for optimization problems is an important issue for model-based learning control. Some machine learning methods are preferred for lightweight use. Stochastic MPC based on Gaussian regression are discussed in (Hewing al. 2019 2020). Gaussian regression is applied to learn unknown disturbances inside the systems, which im- proves the robustness of the controller. From the collected input-output trajectories, linear model can be constructed in the form of the Hankel matrix and helps to generate optimal control output (Berberich al., 2020). Neural network- based models have also been integrated into the MPC framework (Chen et al. 2019; Nubert et al., (2020). However, the computational burden becomes significant concern as the modeling accuracy improves from the intricate structures. In this work, we incorporate the Mamba structure and Koopman operator theory to balance modeling performance with control efficiency.
3 MODELING
In this section, the basic concepts of the Koopman operator will first be introduced. Then, the derivation of the Mamba-based Koopman operator (MamKO) is presented.
3.1 THE KOOPMAN OPERATOR
In typical shopman-baseed approaches, general time-invariant controlled nonlinear system is con- sidered presented as:
Tk+1 f(x (x)
(1)
where I& E X R" is the state vector time instant k; us E C Rm denotes the input vector of the system at time instant k; f is a nonlinear function describing the dynamic behavior of the nonlinear system. The Koopman theory proposed in Koopman (I93I) shows that there exists an infinite-dimen dimentiodd Koopman operator K. H + H. acting on the observable functions such that the dynamics of the nonlinear process in ① can be linearly represented, which can be formulated as:
(ab 4x) = 3 o f(x, 44) = (1(+++1)
(2)
where e denotes the observable functions and represents function composition. From an applica- tion point of view, it favorable to approximate Koopman operators
finite-dimentee function space H C H. Considering the Koopman operator design in (Korda & Mezić 2018), the system in (T) modeled by Koopman operator can be formulated as:
Zk+1 A + Buk
(3)
4t = C'zu
where 2 (xxx) denotes the shifted state vector in H. via the observable function v) : R R
The corresponding dimenssonli Koopman operator is split into A E NxN and B E R N The details for learning the Koopman operator are included in Appendix
3

---

Published as a conference paper at ICLR 2025
3.2 MAMBA-BASED KOOPMAN OPERATOR
For real-world applications, the formulation in ① could fail in modeling systems where the future states are not only dependent on the current state and inputs but also time. The parameters or structures of the systems can change over time. A more general form considering the influence of time can be formulated as:
Ik+1 fz uks k):
(4)
At a given time instant k, the time-varying nonlinear system in can be regarded as a time-invariant system in (I). Thus, the dynamics at each time instant can be represented by the Koopman model in (3). which implics that distinct set of matrices {A, B, C}} can capture the dynamics at specific time instant. However, as time proceeds, the matrices {A, B, C} for the previous instant will not be sufficient for describing the dynamics at the current instant. An effective approach to model this time-varying system involves identifying a set of matrices A2, By, Cs for each time instant k which adapts to the changing dynamics. The corresponding linear time-varying SSM can capture the dynamics of the system for each time instant, which can be formulated as:
Sk+1 A& + Bgue
(5)
i/ C2
where 2ạ = x(*x) represents the shifted state vector. Some parameters from observable function 2h and Koopman operators As. Bx, and C₁ should be updated online to accommodate time-varying systems. Some research has been dedicated to updating parameters in s/k. As. Bg. and C& using online-collected datasets. In (Hao et al.) 2022: Chen et al. 2024 Zhang et al.) 2019), the Koopman operators observable functions are optimized online approximate the dynamics reflected in the newly collected data. Although these methods can provide good performance, solving the corre- sponding online optimization problems can be time-consuming and even infeasible. Instead, in our approach, we resort to the Mamba structure, which contains generative framework for matrices, to enable timely real-time updates of the Koopman matrices.
Mamba Structure Mamba (Gu & Dao 2023), a large language model based on the state-space model, provides a promising frame for modeling time-varying systems. Compared with the S4 framework (Gu et al., 2022a) for time-invariant systems, the selection mechanism is included to generate matrices from the input sequence, which formulated the time-varying SSM.
Inspired by the SSM model from the Mamba framework, we aim to develop a generative approach to compute the time-varying Koopman operators over the prediction horizon. However, directly applying the Mamba framework for the modeling and predictive control task is impractical. When considering the system in ( the Mamba framework only focuses on the inputs u in the form of the word sequence, while states in the LLMs do not have exact meaning. However, for controlled system in (4). the states and inputs with physical meaning should all be considered. Meanwhile, in (5). as B& is generated from the input sequence containing uk. the bilinear term B& can lead to a non-convex optimization problem for in the control part. As a substitute, we generate the matrices from the historical data. The differences between the proposed framework and the Mamba framework are discussed in Appendix B
Matrices Generation Inspired by the Mamba framework, we utilize discretization generate ma- trices efficiently. Given the matrices A, B in the continuous system and the sampling period T, the discrete matrices. A, B can be calculated by the zero-order hold (ZOH) as:
AT
[ Atd B
(6)
B
The matrix A is set as a diagonal matrix, facilitating the discretization process. In the Mamba framework, the diagonal elements of A are constrained to be negative through negative exponential function to promote the stability of the SSM. However, not all real-world systems are inherently stable, such as the oscillator system (Elowitz& & Leibler 2000). This constraint can significantly compromise modeling accuracy. In our work, we address this problem by substituting the negative exponential function with a negative continuously differentiable exponential linear unit (CELU) [Barron] 2017). The formulation of CELU can be presented as:
CFLU(()) max(0. + min(O,o(ee - 1))
4

---

Published as a conference paper at ICLR 2025
Forward Prediction
3-
oooOo
Aylk
Function
Observable Function
E5+3
Agia Ryia + Byjan
Historical Inputs
8g Cs
JIk
One-dImensional convolution
Koopman Operator Sequence
Discretization,
A
Matrices Generation
Aaje
Baok
Cujk
A
Bets-
CC+H-IR
Figure| I: Structure of the MamKO. With the matrices generation network and the observable func- tion, the linear time-varying SSM can be established. By transforming the initial state the shifted state through the observable function v the linear time-varying SSM can produce a sequence of the future states.
where a = in this paper. The CELU is continuously diffcrentiable, which can lead to smooth gradient updates. As CELU can generate negative output, the negative CELU allows for the exis- tence positive eigenvalues, thereby facilitating the representation of unstable systems. The new diagonal matrix activated by CELU is denoted as Á, with the eigenvalues limited in the range of [].
At time instant k the historical state sequence - from the past steps is shifted to space H through observable function vA, obtaining the shifted state sequence H:-H: The com- bination of shifted state sequence and historical imput sequence forms the historical data sequence [ Hik-1 of H One-dimensional convolution is implemented in the historical data se- quence to extract the temporal information inside. The extracted feature e can be calculated by:
C w
T 1 - b,
H Hy..., k h
where w and are the weight and bias parameters in the convolution network, and h is the ker- nel size. The matrix sequences Ba*+H-1 and Ck H-1 and the sampling periods sequence H-- are encoded from the extracted features using fully connected NNs (FCNNs). Discretiz- ing A, B using sampling periods sequence - by (6). the obtained matrix sequences A 11 BRcK+ -1 and Ckck+
applied the Koopman framework to accomplish multi- step-ahead prediction tasks via the model in (5) recursively. The structure the proposed MamKO presented in Figure [
Optimization Procedure The observable function that projects the state vector to finte-diimensional function space denoted (-100)) with trainable parameters 0g. The matrices generation net- works can be denoted as ((-10%) with trainable parameters 0g. The networks take the data sequence
[2] [=]- of -=--11 " and trainable matrix as input, and produce matrix sequence A B -Ij and Cka++ H-K To effectively train the matrix A, parameters 6p and 0,s, the op- timization problem for multi-step-ahead prediction tasks should be designed. Given the dataset
5

---

Published as a conference paper at ICLR 2025
{æy.ust jz k- n. the optimization can be formulated as follows:
1 k+H
min A.8 .0
Σ ||a
22yl2
(9a)
H
j=
s.t. 3 & A + ByJ
(9b)
2 = (zx((0x)
(9c)
8 - &
(9d)
A = -CELU(A)
(9c)
2j (x),100),3
H....
(9f)
Akc----1jkk BAck+H- CRAN,CCKHHH111 = <(A, - Uk 11:k 11:k-100))
(9g)
where 2zlk represents the predicted shifted state vector for time instant j at time instant 4; 231k represents the state vector for time instant predicted at time instant k; H represents the length of the model's prediction horizon. A BB. and C l are the generated matrix sequences from the historical data sequence. To minimize the prediction loss in (9a), all parameters are updated with gradient descent using ADAM (Kingma 2014).
4 CONTROL
In this section, for the time-varying system in (4), we propose a model predictive control framework that is inspired by Korda & Mezić (2018) based on the MamKO model.
At each time instant k. the matrix sequences AR Back+ -li and Ckck# 1|K for predic- tion are generated previously with the historical data. With the matrix sequence, the optimization problem for MPC is formulated as follows:
h+H-1 Σ
(1Cgk&2jk Iellà + [|ug]k uy (x + ||C B1 xzll?
min
M we "
(10a)
s.t.
2/-1|k = Agıkāl1 + BylkUgik
(10b)
2nja (x(k)
(10c)
Cž5lk X
(10d)
UgJh U. = 1: 1, k H-
(10e)
where "kjk.--.. N denotes the trajectory of optimal control signals determined given the current time instant k; 11-112 denotes vector's squared weighted Euclidean norm, which is computed
by ||=112 Tx; Q, R and P are dositive-definite weighting matrices; ||Czıæ z21lå and [IC++HI2*++H HIA xx 113 used reduce tracking errors; ||ub U 11 2 is set to improve the smoothness of the control inputs; (TOb) serves as the model constraint; (10c) provides the initial condition in the shifted state-space; (IOd) and (IOe) are state and input constraints, respectively. At time instant & 2 O, the optimal control input trajectory is calculated and the first element of the trajectory, "b is applied to the nonlinear process in (4) as the control input for the current time instant. The convex optimization problems derived from linear SSMs can be efficiently solved.
EXPERIMENTS
In this section, we will evaluate the performance of the MamKO model terms of modeling and control. Specifically, we evaluate the following aspects: (a) Convergence of the proposed train- ing algorithm on different systems with random parameter initialization. (b) Model accuracy of the MamKO compared to other methods in time-invariant and time- varying systems. (c) Control performance of the MPC based on the MamKO compared to other methods in time-nnariaa and time-varying systems.

---

Published as a conference paper at ICLR 2025
Five benchmark systems included to evaluate the modeling and control performance of the MamKO. The CartPole balancing system is included as a benchmark system for the controller de- sign task, which has been widely used in deep reinforcement learning (DRL) research (Lillierap 2015 Haarnoja et al. 2018). Correspondingly. a time-varying CartPole balancing system described in (Hao al.) 2022) is considered to test the performance of MamKO on the time-varying system. Moreover, we apply the MamKO to the biological gene regulatory networks (GRN) (Elowitz & Leibler 2000). The complex nonlinear systems in the chemical engineering field are also consid- ered. The reactor-separatee chemical process (Yin & Liu 2019) comprising two continuous stirred tank reactors and flash tank separator is chosen as a benchmark chemical process, denoted as the reactor-separator chemical process (RSCP) system. We also simulate the reactor-separator process with time-varying parameters (Nikravesh et al. 2000,, denoted as the time-varying RSCP system.
For the modeling performance evaluation, we compare the proposed method with two competitive baseline methods. (1) The Deep Koopman Operator (DKO) (Lusch et al., 2018: Han et al.) 2020) is a Koopman method that employs deep NNs in constructing the observable functions. Different from the generative Koopman operators from MamKO, the Koopman operators in the DKO are invariant. (2) The Multilayer Perceptrons (MLP) (Chua et al. 2018), based on multilayer fully connected NNs, is implemented as a baseline for modeling the systems.
For the control performance evaluation, we compare the MamKO- based MPC with three state-of- the-art baseline methods: (1) The DKO-based MPC (Lusch et al.) 2018: Han et al.] 2020); (2) MLP-based MPC, which utilizes the NNs predictive models and solves nonlinear optimization with the help of Casadi (Andersson et al. 2019) and Interior Point Optimizer, pronounced eye-pea- Opt (IPOPT) (Wächter & Biegler 2006); (3) Soft actor-critic (SAC) (Haarnoja al. 2018), which is a state-of-the-an model-free reinforcement learning algorithm. Despite the higher sample com- plexity of model-free methods compared to model-based ones, SAC often achieves superior control performance. SAC updates the controller to minimize cumulative stage costs, thereby implicitly optimizing for stabilizing controller.
For each environment, trajectories of state and action samples are gathered, generating training set of 36, 000 samples, a validation set of 4, 000 samples, and a test set of 4, 000 samples. Specifically, actions are uniformly sampled from the action space for each time instant for the CartPole system, the time-varying CartPole system, and the GRN system. For the RSCP and time-varying RSCP systems, actions are determined using a step function with added random noise. Details of the experimental setup are provided in Appendix C The methods are trained to predict state sequences over 30-step horizon. Hyperparameters for each method are listed in Appendix D For the SAC control task, each environment undergoes training with 1000k steps of state-action-reward pairs.
5.1 MODELING EVALUATION
The losses for the modeling task are calculated by the average prediction error of each step on the test set. For each system, we run ten model training trials with different datasets and randomly initial- ized parameters to test the convergence of the algorithms. The results for each system are presented in Figure 2 (a-e). In comparison to DKO, MamKO demonstrates superior modeling accuracy across all environments. Especially for the two time-varying systems, the MamKO outperforms the other two methods. Notably, with the nonlinear structures from the NNs, the MLP has fewer prediction errors in GRN and RSCP. Nonetheless, the inherent nonlinearities of the MLP model may compli- cate optimization tasks, leading reduced computational efficiency and suboptimal behavior. The experiments in the next section will show the possible limitations of MLP in controller design.
Performance on the rapidly changing dynamics For the time-varying CartPole system in (Hao et al.) 2022), a time-varying coefficient of the friction of the cart is added to the system in the form of sine waves. Based on the example, we compare the modeling performance of the three methods with different frequencies of the sine waves. The time-varying Cartpole system described before has the angular frequency of 1 rad/s. Two time-varying Cartpole systems with the angular frequencies of 0.1 rad/s and 10 rad/s are added. The results are shown in Figure 2 (f), illustrating that as the angular frequency of the time- varying parameters increases, the advantages of using MamKO become more pronounced. The generative matrices provide more accurate representation of the system with rapidly changing parameters, as different matrices approximate the dynamics of systems at each time instant.
7

---

Published as a conference paper at ICLR 2025
0.5
MLP DKO MamKO
MLP DKO
-0.5
-0.5
TestLoss
Test
MamKO
Test Loss
MamKo
1..5
-1.5
-2.5
-2.5
-3.5
-3.5
100
200
400
Epoch
Epoch
pooch
(a) CariPole
(b) GRN
(c) RSCP
0.5
MLP
MLP
0.10
DKO
0.08
DKO
Tost
-0.5
MamKO
Test
MamKO
Mean loss
MamkO
0.06
-1.5
0.04
-2.5
-3
10 radis Angular frequencies of the sine waves
Epoch
Epoch
(d) Time-varying CartPole
(c) Time-varying RSCP
(f) Time-varying CartPole with dif- ferent parameters
Figure 2: Average prediction enror on the test set (a-e) and average prediction error for the time-
varying CartPole with different coefficients (f). In (a-e), the Y-axis indicates the average mean- squared prediction error log space for the 30-step-ahead prediction task, and the X-axis indicates the training epochs. The shaded area represents the confidence interval (0.5 times the standard deviation)
across ten training trials. In (f), the Y-axis indicates the average loss after training for 400 epochs.
5.2 CONTROL EVALUATION
The control performance of the four methods is evaluated and compared in the five systems. The time-varying CartPole with the angular frequency of 10 rad/s is included to evaluate the control performance for the system with rapidly changing dynamics. The control task for each environment set-point tracking, and the details can be found in Appendix C The weighting parameters Q, Rt and for the MPCs are carefully adjusted to reach the best performance each baseline method. The prediction horizon is set 30 for all MPCs.
log space
costin space
0
Average castin
6
1.0
Avemge
-2.0
Time (s)
Time (a.u.)
Time (h)
(a) CartPole
(b) GRN
(c) RSCP
Average costin log space
Averege cnstin log space
Avergg costin log spare
1.0
0.5
0.5
1.0
1.5
10
15
20
0.5
1.0
1.5
2.0
5
10
15
20
Time
Time
Time (s)
(d) Time-varying CartPole
(c) Timc-varying RSCP
Time-varying CartPole with rapidly changing parameters NamKO-based MPC
(f)
SAC MLP-based MPC
DKO-based MPC
Figure 3: Average based track error. The Y-axis indicates the average cost in log space for ten experiments with random initial states. The shaded area represents the confidence interval (0.3 times the standard deviation) over the experiments.
8

---

Published as a conference paper at ICLR 2025
The cost trajectories shown in Figure 3 The state trajectories can be found in Appendix F As shown in Figure 3 the MamKO-basee MPC achieves the best control performance for all five systems. In companisn with DKO-based MPC, the MamKO-based MPC reduces the cost of the five systems by 5.05%, 3.70% 92.10%, 6.56%, 14.19%, 84.74%, respectively. A significant im- provement is observed in the time-varying CartPole with rapidly changing parameters. Although MLP achieves better modeling performance on GRN and RSCP systems, the corresponding control performance not as good
modeling performance, which can be attributed to suboptimal lutions from the non-convex optimization problems. Compared with the SAC, the two Koopman model-based MPCs can all achieve: a relatively stable control performance for the time-varying sys- tem. Although it can be hard to construct time-vare vinevia model, the Koopman operator in DKO can construct an average Koopman operator that accommodates various time-varying parameters in the systems, which can provide relatively good results.
5.3 EVALUATION ON THE ACTIVATION FUNCTION
As we simplify as a diagonal matrix, the cigenvalues of the matrix become the trainable param- eters directly. Compared to the Mamba framework, which sets the eigenvalues to be negative, replace the original negative exponential function with the negative CELU function. In this section, will compare the modeling performance of MamKO using the negative exponential function, the negative CELU function, and no activation function on the CartPole system, the GRN system, and the RSCP system. As observed in Figure the MamKO with CELU function has the best modeling performance.
1.5
CELU
1.5
CELU
CELU
0.5
0.5
Exp
Exp
None
-0.5
0.5
Test
-.5
-25
2.5
.3.5
100
200
300
400
-3.5
100
200
300
400
200
300
Epoch
Epoch
Epoch
(a) CartPole
(b) GRN
(c) RSCP
Figure 4: Evaluation on the activation function for the eigenvalues. For the CartPole system, the MamKO with the negative CELU function performs slightly better than the without the activa- tion function and largely better than the one with the negative exponential function. For the GRN and RSCP systems, the MamKO with the negative CELU has the smallest mean loss and lowest variance. Compared with the MamKO with no activation function, the improvement of the MamKO with the negative CELU is not significant. However, the initiation of A of the MamKO with no activation function should be carefully adjusted since no constraint is added to the eigenvalues. Oth- erwise, too large eigenvalues may result in unstable gradient descent.
5.4 ABLATION ON THE DISCRETIZATION
We also evaluate the impact of discretization on modeling performance. The benefit of generating the sequence of the Koopman operator through the method in (6) is demonstrated in this subsection. We compare the performance of this modeling approach with another method that applies the same discretization technique to matrix B as to matrix A. which is B = e "T, The pipeline is named as M Multiplication The modeling results for each environment are presented Table From Table T we can find that with the discretization on the matrices following (6), the modeling accuracy can be significantly improved, which validates the effect of the discretization.
Table 1: Results the ablation on the discretization.
Method
CartPole
GRN
RSCP
Time-varying CartPole
Time-varying RSCP
Discretization
6.50 x 10-4
1.91 10-3
2.92 x 10-3
4.31 x 10-4
9.67 10-3
Multiplication
8.30 x 10-4
3.52 10-3
7.60 x 10-3
9.11 x 10-4
1.19 10-2

---

Published as a conference paper at ICLR 2025
5.5 EVALUATION ON THE COMPUTATION TIME
In this subsection, evaluate the computation time of the control methods on different benchmark systems. The online control implementation of the control methods is conducted on a computer equipped with Intel Core 19-13900K CPU and 128 GB DDR4 RAM. The computation times of the MamKO-based MPC and baselines during control on the benchmark systems are presented in Table 2 Compared with the MLP-baased MPC, the MamKO-based MPC reduces the computation time by 98.38%, 83.39%, 99.21%, 90.34%, 99.17% for the five benchmark systems, respectively. To further demonstrate the computational efficiency of the proposed method, the sampling periods of the systems are presented in Table 3 A comparison of the sampling periods and computation times demonstrates that the proposed MamKO- MPC method can reliably ensure online implementa- tion for each of the considered systems. The efficient online implementation of the proposed method stems from the use of linear state-space model within the proposed framework, which facilitates the formulation of an optimal control problem that requires solving convex optimization despite the nonlinearity in the dynamics of the considered systems.
Table 2: Results of the computational time of each method in the benchmark systems.
Method
CartPole
GRN
RSCP
Time-varying CartPole
Time-varying RSCP
MLP-based MPC
7.43 x 10
1.09 x 10-1
3.31
1.7x x 10-1s
3.15 1,25 x 10-2s
DKO-based MPC
7.35 10-3
1.41 x 10-2s
1.06
10-2
7.63 x 10-3s
MamKo-bused MPC
1,02 10
1.81 10-2
2.62
10-2s
1.68 10-2
2.95 x 10-2 5
SAC
2.67 10
3.23 x 10-1s
3.41 x 10-1s
2.68 x 10-1s
3.32 x 10
Table 3: Sampling periods of the benchmark systems.
Method
CartPole
GRN
RSCP
Time-varying CartPole
Time-varying RSCP
Sampling period
0.02
1
18
0.02
18
The results also indicate that the extended computation time required by the MLP-based MPC, which is due to the need to solve nonlinear optimization problems, poses challenges for its online implementation. This limitation is particularly critical for systems with fast dynamics and short sampling periods, and for larger-scale systems with numerous control inputs and state variables to optimize.
CONCLUDING REMARKS
In this paper. a new modeling and control framework called Mamba-based Koopman Operator (MamKO) proposed by seamlessly integrating the Koopman operator with the large language model Mamba. The matrices generation network from Mamba was adapted to construct the linear time-varying state-space model based the Koopman modeling concept. The MamKO model can effectively predict the future states of nonlinear systems with time-varying parameters. A model pre- dictive controller was formulated based on the established MamKO model. Both time-invariant and time-varying benchmark systems were leveraged to evaluate the modeling and control performance of the proposed method. The experiments demonstrated the superior performance of the MamKO model in both modeling and control tasks. Future work will focus applying the MamKO model to large-scale systems to further evaluate its modeling capabilities.
Future research directions include the formal analysis of the convergence and stability of the pro- posed MamKO-based control approach and the utilization of the MamKO-based modeling frame- work for energy optimization of energy-intensive industrial systems, for example, through inte- grating the proposed modeling framework with Koopman-bar based economic model predictive control (Han et al., 2024b). We also plan to investigate the application of the MamKO based control meth- ods to industrial systems of larger scales and more complex dynamics for efficient and robust system operations.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This research is supported by the National Research Foundation, Singapore, and PUB, Singapore's National Water Agency under its RIE2025 Urban Solutions and Sustainability (USS) (Water) Centre of Excellence (CoE) Programme, awarded to Nanyang Environment Water Research Institute (NEWRI), Nanyang Technological University, Singapore (NTU). This research is also supported by Ministry of Education, Singapore, under Academic Research Fund Tier 1 (RS15/21 RG63/22). Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of the National Research Foundation, Singapore, and PUB, Singapore's s National Water Agency.
REFERENCES
Jens Alex, Lorenzo Benedetti, JB Copp. KV Gernacy, Ulf Jeppsson, Ingmar Nopens. MN Pons, Leiv Rieger, Christian Rosen, JP Steyer, et al. Benchmark simulation model no. (BSMI). Report by the IWA Taskgroup on benchmarking of control strategies for WWTPs, 1, 2008.
Joel A. E. Andersson, Joris Gillis, Greg Horn, James B. Rawlings, and Moritz Diehl. CasADi - software framework for nonlinear optimization and optimal control. Mathematical Programming Computation, 11()) 2019.
Jonathan T Barron. Continuously differentiable exponential linear units. arXiv:1704.07488 2017.
arXiv preprint
Julian Berberich, Johannes Köhler, Matthias A Müller, and Frank Allgöwer. Data-driven model predictive control with stability and robustness guarantees. IEEE Transactions on Automatic Control, 1702-1777 2020.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAl Gym. arXiv preprint
2016.
Steven I. Brunton, Marko Budišić, Eurika Kaiser, and J Nathan Kutz. Modern Koopman theory for dynamical systems. SIAM Review, 64(2):229-34)) 2022.
Hao Chen and Chen Lv. Incorporating ESO into deep Koopman operator modeling for control of autonomous vehicles. IEEE Transactions on Control Systems Technology, 38544:1854-1864, 2024.
Tianze Chen, Yongqiang Huang, and Yu Sun. Accurate pouring using model predictive control enabled by recurrent neural network. In IEEE/RSI International Conference on Intelligent Robots and Systems, pp. 7688-7694. IEEE, 2019.
Zhong Chen, Xiaofang Chen, Jinping Liu, Lihui Cen, and Weihua Gui. Learning model predictive control of nonlinear systems with time-varying parameters using Koopman operator. Applied Mathematics and Computation, 470i0228777777777 128577, 2024.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, volume 31, 2018.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060, 2024.
Michael B Elowitz and Stanislas Leibler. A synthetic oscillatory network of transcriptional regula- tors. Nature, 403(6767) 2000.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Re. It's Raw! Audio generation with state- space models. In International Conference on Machine Learning, volume 162, pp. 7616-7633, PMLR, 2022.
Albert Gu and Tri Dao. Mamba: Linear- sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00755 2023.
11

---

Published as a conference paper at ICLR 2025
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Advances in Neural Information Processing Systems, volume 33, pp. 1474-1487, 2020.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022a.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your HiPPO: State space models with generalized orthogonal basis projections. arXiv preprint arXiv:2206. /2037, 2022b.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint
2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 19Xiv: 2019.
Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.02199: 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04100 2023.
Minghao Han, Jacob Fuler-Rolle, and Robert K Katzschmann. Desko: Stability robust control with a deep stochastic Koopman operator. In International Conference on Learning Rep- resentations, 2021.
Minghao Han, Kiwan Wong. Jacob Euler-Rolle, Lixian Zhang, and Robert K. Katzschmann. Robust learning-based control for uncertain nonlinear systems with validation soft robot. IEEE Transactions on Neural Networks and Learning Systems, pp. 1-15, 2023.
Minghao Han, Zhaojian Li, Xiang Yin, and Xunyuan Yin. Robust learning and control of time- delay nonlinear systems with deep recurrent Koopman operators. IEEE Transactions on Industrial Informatics, 20(3):4675-4688 2024a.
Minghao Han, Jingshi Yao, Adrian predictive control of water treatment process with learning- based Koopman operator. Control Engineering Practice, 149:105975, 2024b.
Law, and Xunyuan Yin. Efficient economic model
Yiqiang Han, Wenjian Hao, and Umesh Vaidya. Deep learning of Koopman representation for control. In IEEE Conference on Decision and Control, pp. 1890-1895. IEEE, 2020.
Wenjian Hao, Bowen Huang, Wei Pan, Di Wu, and Shaoshuai Mou. Deep Koopman learning of nonlinear time-varying systems. arXiv e-prints, pp. arXiv-2210, 2022.
Lukas Hewing, Juraj Kabzan, and Melanic N Zeilinger. Cautious model predictive control using
Gaussian process regression. IEEE Transactions Control Systems Technology, 28(6):2736 2019.
Lukas Hewing, Kim P Wabersich, Marcel Menner, and Melanie N Zeilinger. Learning-bssed model predictive control: Toward safe learning in control. Annual Review of Control, Robotics, and Autonomous Systems, 3(1):269-299 2020.
Rudolph Emil Kalman. new approach linear filtering and prediction problems. 1960.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Bernard O Koopman. Hamiltonian systems and transformation in Hilbert space. Proceedings of the National Academy of Sciences, 17(5)3315-318, 1931.
12

---

Published as a conference paper at ICLR 2025
Milan Korda and Igor Mezić. Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control. Automatica, 93:149-166 1609 2018.
Fuxian Li, Jie Feng, Huan Yan, Guangyin Jin, Fan Yang, Funing Sun, Depeng Jin, and Yong Li. Dynamic graph convolutional recurrent network for traffic prediction: Benchmark and solution. ACM Transactions on Knowledge Discovery from Data, 17(1):1-21, 2023.
Kunchang Li, Xinhao Li, Yi Wang, Yinan He, Yali Wang, Limin Wang, and Yu Qiao. VideoMamba: State space model for efficient video understanding. arXiv preprint arXiv:2403. 06977, 2024.
Yunzhu Li, Hao He, Jiajun Wu, Dina Katabi, and Antonio Torralba. Learning compositional Koop- man operators for model-based control. In International Conference Learning Representa- tions, 2020.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Frez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, al. Jamha: A hybrid Transformer- Mamba language model. arXiv preprint arXiv:2403. 19887, 2024.
TP Lillicrap. Continuous control with deep reinforcement learning.
arXiv preprint
arXiv:1509.02971 02971, 2015.
Jinfeng Liu, David Muñoz de la Peña, Benjamin Ohran, Panagiotis Christofides, and James Davis. A two-tier architecture for networked process control. Chemical Engineering Science, 63 (22):5394-5409, 2008.
Yong Liu, Chenyu Li, Jianmin Wang, and Mingsheng Long. Koopa: Learning non-stationary time series dynamics with Koopman predictors. In Advances in Neural Information Processing Sys- lems, volume 36, pp. 12271-12290, 2023.
Bethany Lusch, J Nathan Kutz, and Steven L Brunton. Deep learning for universal linear embeddings of nonlinear dynamies. Nature Communications, 9()):4950, 2018
Jeremy Morton, Antony Jameson, Mykel J Kochenderfer, and Freddie Witherden. Deep dynamical modeling and control of unsteady fluid flows. In Advances in Neural Information Processing Systems, volume 31, 2018.
Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and vidcos multidimensional signals with state spaces. In Advances in Neural Information Processing Systems, volume 35, pp. 2846-2861, 2022.
M Nikravesh, AE Farell, and TG Stanford. Control of nonisothermal CSTR with time varying parameters via dynamic neural network control (DNNC). Chemical Engineering Journal, 76(1):
1-16, 2000.
Julian Nubert, Johannes Köhler, Vincent Berenz, Frank Allgöwer, and Sebastian Trimpe. Safe and fast tracking on a robot manipulator: Robust MPC and neural network control. IEEE Robotics and Automation Letters, 5(2):3050-3057, 2020.
Jiacheng Ruan and Suncheng Xiang. VM-UNet: Vision Mamba UNet for medical image segmen- tation. arXiv preprint arXiv:2402.0249,, 2024.
Jan C. Schulze, Danimir Doncevic, and Alexander Mitsos. Identification of MIMO Wiener-type Koopman models for data-driven model reduction using deep learning. Computers & Chemical Engineering, 161:107781, 2022.
Haojie Shi and Max Q-H Meng. Deep Koopman operator with control for nonlinear systems. IEEE Roboties and Automation Letters, 733):770--7707, 2022.
Aivar Sootla, Natalja Strelkowa, Damien Ernst, Mauricio Barahona, and Guy-Bart Stan. On periodic reference tracking using batch-mode reinforcement learning with application to gene regulatory network control. In IEEE Conference on Decision and Control, pp. 4086-4091, 2013.
Natalja Strelkowa and Mauricio Barahona. Switchable genetic oscillator operating in quasi-stable mode. Journal of The Royal Society Interface, 7(48):107111082, 2010.
13

---

Published as a conference paper at ICLR 2025
Naoya Takeishi, Yoshinobu Kawahara, and Takehisa Yairi. Learning Koopman invariant subspaces for dynamic mode decomposition. In Advances in Neural Information Processing Systems, vol- ume 30, 2017.
Peter Van Overschee and BLO888 De Moor. Subspace identification for linear systems: The- hnn-lmlleenntaionnpppiiiiiiiiiii Springer Science & Business Media, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances Neural Infor- mation Processing Systems, volume 30. Curran Associates, Inc., 2017.
Yogesh Verma, Markus Heinonen, and Vikas Garg. Climode: Climate and weather forecasting with physics-informed neural odes. In International Conference on Learning Representations, 2024.
Andreas Wächter and Lorenz T Biegler. On the implementation of an interior-pond point filter line- search algorithm for large-scale nonlinear programming. Mathematical Programming, 106(1): 25-57, 2006.
Kun Wang, Hao Wu, Yifan Duan, Guibin Zhang, Kai Wang, Xiaojiang Peng, Yu Zheng, Yuxuan Liang, and Yang Wang. Nuwadynamics: Discovering and updating causal spatio-temporal modeling. In International Conference Learning Representations, 2024.
Matthew O Williams, loannis G Kevrekidis, and Clarence W Rowley. A data driven approxima- tion of the Koopman operator: Extending dynamic mode decomposition. Journal of Nonlinear Science, 25:1307-1346, 2015.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-45, Online, 2020.
Xunyuan Yin and Jinfeng Liu. Subsystem decomposition of process networks for simultaneous distributed state estimation and control. AIChE Journal, 65(3)9904-114, 2019.
Jing Zeng and Jinfeng Liu. Economic model predictive control of wastewater treatment processes. Industrial Engineering Chemistry Research, 54(21):5710-57) ---11 2015.
Hao Zhang, Clarence Rowley, Eric Deem, and Louis N Cattafesta. Online dynamic mode decomposition for time-varying systems. SIAM Journal on Applied Dynamical Systems, 18(3): 1686-1609, 2019.
Xuewen Zhang, Minghao Han, and Xunyuan Yin. Reduced-order Koopman modeling and predictive control of nonlinear processes. Computers Chemical Engineering, 179:108440, 2023.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al. A survey of large language models. arXiv preprint arXiv:2303. 18223, 2023.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang. Wenyu Liu, and Xinggang Wang. Vi- sion Mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint arXiv:2401.094 00917, 2024.
14

---

Published as a conference paper at ICLR 2025
A KOOPMAN OPERATOR
According to the approach discussed in (Korda & Mezić 2018). a state vector incorporating both
state and input can be defined as Nx [ u
Defining a left shift operator S such that Sug Wk+1, the dynamics of this augmented state vector can be expressed as follows:
f ux)) Suk
[ ug)
Xx+1 F(Xx)
(11)
Uk+1
The Koopman operator K : H.5 - H.5 for (11) governs the dynamics of the augmented vector X linear fashion:
Ky = V F(XE) = 4 (A)++))
(12)
The observable functions inside transfer the original states into the higher-dimensional space.
In (Korda & Mezić 2018), only the states are lifted to nigher-dimensional space, while trans- formations are applied to the inputs. Specifically, the observable functions are defined as follows:
v (Xx) (z ux) [v (xx), [
(13)
where vi IR RN is the observable function that transforms the original states.
For real-world applications, a finite-dimensional approximation of the Koopman operator, denoted as KN,,, needs to be identified. Since future inputs not considered, only the elements in the first N rows the KN,, need to be identified. The formulation NN,,, can expressed as:
A
B
KN, =
(14)
where E NNN, BE Nmmm To capture the dynamics of 2x. it is sufficient to identify only the matrices and B in (14). while the blocks denoted by can be ignored.
B МАMBA FRAMEWORK
The matrices generation network in MamKO is inspired by the Mamba. It is worth noting that there are differences between our methods and the Mamba structure.
Firstly, the Mamba structure is designed for LLM with word sequences as inputs and generates cor- responding word sequences. Instead, our method attempts to model real-world systems containing both states and inputs. Compared with LLM,, MamKO model is closer the concept (Ha & Schmidhuber 2018). which contains the interaction between states and actions.
Secondly, the matrices generation network of Mamba utilizes the information in the whole input sequence. Given the sequence Up = {uj} y=k))= H-1 the matrices Ag. B& and C in (5) are obtained from the nonlinear neural networks, which can be denoted as A SAn (U+), B3 fB. (Ux). Cs = (Ux). The SSM from Mamba can be formulated as:
2k+1 fA. (Uk)2 fB. (Ux)uu
(15)
i& = fcr (Ux)
The nonlinear functions fAL JB., fc, and the bilinear terms fA, (U))zz) fB. (U,)un make the relationship between and Ug no longer linear. Applying MPC for this nonlinear system can cause non-convex optimization problems, which may reduce the control efficiency. In our work, the matrices are generated from historical data, which have no direct relationship between zμ and ug- The linear relationship between žμ and C allows for the construction of convex optimization problems.
Lastly, the SSMs the Mamba structure are separately generated for each channel of the input sequence. For the input E Rm with m channels, m pairs of SSMs with A E N B, E Nx1, C E IR N matrices are generated at time instant k. The multiple pairs of SSMs are not directly applicable to MPC. In our framework, the SSMs designed to describe the dynamics of all the input channels (features), which can be formulated as A₁ E N N Bg E R. xm C₁ E RaxN N for each time instant.
15

---

Published as a conference paper at ICLR 2025
Figure 5: CartPole system.
Fr
F2o
F
F10
T1
F1
T2
F2
T3
OO
OO
CSTR
CSTR
Separator
A + B C
A
Q1
Q2
Q3
3
Figure 6: Time-invarian RSCP system (Liu et al.
2008
C EXPERIMENTAL SETUP
The experiments are based OpenAl Gym (Brockman et al. are presented Figure and Figure
2016). Part of the environments
C.I CARTPOLE INVERTED PENDULUM ON A CART
We modified the CartPole system described in (Brockman et al. 2016) using continuous action space instead of a discrete one. The system consists inverted pendulum attached to While the cart fully actuated, the pendulum remains unactuated. In this experiment, the controller aims to keep the pendulum upright and vertical. The control input is the horizontal force applied to the cart (a [-20, 20)). Tihreshold and Oireshold represents the maximum of position and angle, respectively, Tihreshold
cart that moves horizontally with an
10 and Parresood 20 The episode ends (0) > Otheehoodd and the episodes end in advance. The episodes for control evaluation are of length 1000. For time-varying CartPole, the model is based on a modified system in (Hao et al. 2022). which can be presented as:
Fi + ml
(0 sin O. - Ōr cas 0.) μi sgu (dre)
Ft
me + m
(16)
0. cos 0e [-Fi ml0² sind, ; + ue sgn (ix) (me + m)
gsin 0g 11p0mml
1 [ - (mcoss 6.) / (me + m)]
1 1
(m cos? 6.) (m. + m)]
where n² 0.0005 cos(wt) is time-varying coefficient of friction of cart the track; is the angular frequency of sine wave; x; is the distance to the initial position; 0, offset angle; fr. 8, is the velocity and the angular velocity respectively: F, is the continuous control input. The parameters are set the same as the time-invariant CartPole, which can be presented as:
m. = 1.0 kg, m = 0.1 kg, / 0.5 m
For the lime-ngrying Cartpole system, the angular frequency is set as w 1; for the rapidly changing dynamics, the angular frequency is set 0.1, 1, 10, respectively.
16

---

Published as a conference paper at ICLR 2025
Notably, the control tasks for the two systems include position control and angular control. The controller should stabilize the CartPole in the state that 22 = and 8, = The cost designed for the two CartPole systems is defined as:
cost (F2, = It
+ 180 02
(17)
100
C.2 SYNTHETIC BIOLOGY GENE REGULATORY NETWORKS
The gene regulatory networks (GRNs) discussed in this paper function at the nanoscale, demon- strating unique physical properties that distinguish them from other examples. Notably, these GRNs exhibit intriguing oscillatory behavior.
In this example, we focus on classical dynamical system from systems and synthetic biology, specifically addressing a reference tracking task. The GRN under consideration is a synthetic three- gene regulatory network, where the dynamics of mRNAs and proteins inside the network exhibit oscillatory behavior. A detailed description of the network's 's mechanisms can be found in (Elowitz Leibler 2000). The discrete-time mathematical model of the GRNs can be represented by the following equations:
a1
21 (t =11=)11(1 + dt
(t) +
+ 41
+ (t)),
K1 + +g](t)
x2 + 1) = x1) + dt
(2) +
02
+ 1/2
+ Et(),,
K₂ + x2(()
03
(18)
21(t [ 1) = 21(1) + dt
(t) +
K3 + xz(() + U3y
+ Et)),
ma(1 + 1) =4(1) + dl
(1)
25(1 + 1) =z(() + di
[-cg0s (1)3) (1) + Et())
x6(1+ + = =6(()) + di
In this context, 21, 23 represent the concentrations of the mRNA transcripts, and 24, 25, ve rep- resent the concentrations of the corresponding proteins for genes 1, 2, and 3, respectively. The terms &,, for all i, are independent and identically distributed. Uniform noise variables are drawn from the range [- 5). For the simulations described in Section is set to O, while for the robustness tracking task in F.6 d is strengths for the respective genes. The symbols 71, 72, and "/3 represent the degradation rates of mRNA, whereas C1. C2, and 03 correspond the degradation rates of proteins. Moreover, 31, 82, and B3 denote the protein synthesis rates, and K1, K2, and K3 are the associated dissociation con- stants. The equations in (18) present a network structure in which gene I is inhibited by gene 2, gene 2 is suppressed by gene 3, and gene 3, in turn, is repressed by gene 1. The discretization time step denoted by dt. The protein concentrations are observable and are typically measured using fluorescent markers, such as green fluorescent protein (GFP) or red fluorescent protein (mCherry). The control inputs u; are applied via light control signals, which activate photo-sensitive promoters to induce gene expression. In line with the standard practice of GRN models in (Elowitz & Leibler 2000), we simplify the system dynamics by assuming that the parameters regulating mRNA and protein dynamics are uniform across all genes. Further details the mathematical modeling and control of synthetic gene regulatory networks can be found in (Strelkowa & Barahonal 2010, (Sootla et al. 2013). In the simulations mentioned in this paper, the parameters are chosen
0.5. The parameters 01, 02, and ag indicate the highest promoter
follows:
Vi: K; = 1, a; = 1.6, x = 0.16, 3 = 0.16, G = 0.06, dl
For the tracking task for 24, the cost function for the GRN system is defined as:
costGRN(Ta)) = (x.a 6)²
(19)
C.3
PROCESS
A reactor-separator process, consisting of two continuous stirred tank reactors (CSTRs) and flash tank separator (Liu et al. 2008), is considered in this paper. A schematic of the process is shown in
17

---

Published as a conference paper at ICLR 2025
Figure 6 The process consists of two chemical reactions: firstly, the reactant A is transformed into the desired product B, followed by secondary reaction where fraction of B is further converted The system starts with the pure reactant being fed into the first reactor (CSTR 1) at a flow rate of Pio. The output stream from CSTR 1. with a flow rate of Fi, together with an additional fresh supply of pure A at rate of Fzo. enters CSTR 2. The outflow from CSTR flow rate of Fz. Within the separator, recycle stream with flow rate of F,. redirected back to the first reactor for further conversion. All three vessels (CSTR 1, CSTR 2, and the separator) are equipped with a jacket to provide or remove heat, with a heating input rate Q, where = 2. 3. The system's 's ordinary differential equations can be expressed as
into the byproduct
2 is then directed to the separator
follows (Liu et al. 2008 Yin & Liu 2019):
A
F10
Fr Vi
IA1) + (æ
(I IA1) Ookie PT ZA1
dt V/2
dæ Fio
-Ei (5B10 x81) + (zB. - 2B1) Pokic rii FAI dokiz rTi 3811
Fr
dt
Vi
Vi
Fio dTi
Fr
AH1
AH2
Q
(T10 (T)- +
(Tg Ti) - -o
kie 12 TA1 &e
kze r3 JB +
dt
VI
Vi
pcp Vi
dz Fi
F20
(EAI - xA2) +
A20 -AA2)
Dokie r2 2A
dt
V/
da B
F20
R2
(TBL 2B2) +
(B 2B2) + óckie
T ZA2 dckiz
7T2 2
(20)
dt
V/2
1/2
dT:
FI
F20
AHl
-E1
Alla
BBg
Q2 pcp V2
(7)
-Tb)+
+ Tb)
(T'0 T2)
de
bie 722 IA
de
kzz The
+
dt
V/2
V/2
Cp
Cp
dz
F2
(F, + Fp)
(MA2 2AS)
(EAr xAA3)
di
V3
V/3
dz
F
(F, + Fp)
(x 2B3)
(I TBS)
dt
V3
V/3
dT'
F4
Qs
(F. (F++Fo) (r 2B, ICr Hmps)
(T2
Ts)
dt
V/3
PepVa
PcpVs
For more detailed explanation of this process, refer (Liu et 2008 Yin & Liu 2019) Addi-
tional process disturbances are introduced into the systems to evaluate the robustness of the proposed control strategy. These disturbances are sampled from a multivariate normal distribution (0, of), where a, = [0.01, 0.01, 0.50, 0.01, 0.01, 0.50, 0.01, 0.01, 0.50] corresponds to the standard devia- tions associated with the nine state variables. For time-invariant RSCP, parameter de is set as d. 1; For time-vane varying, parameter &e. is set as the e -0.011, The control objectives for both processes are to regulate the system states and ensure the states remain a steady-state set point:
[0. 18188, 0.67, 480.32 K, 0.20, 0.65, 472.79 79 K, 0.07, 0.67, 474.89 K)"
Given set scaling coefficients Escale [0.36, 0.18, 361.94 0.21, 0.18, 342.88 K, 0.26, 0.21, 361.59 KJ". the cost function for the two RSCP systems can be defined as:
I' Esoale
costRSCP(z)
(21)
D HYPERPARAMETERS
The hyperparameters for the MamKO, DKO, and MLP are listed in Table respectively.
Table
Table 6
18

---

Published as a conference paper at ICLR 2025
Table 4: Hyperparameters of MamKO
Hyperparameters
Value
Size of training set Size of test set Batch Size
40000
4000
256
Learning rate
le-3
Prediction horizon H
30
Structure of d (CartPole Time-varying CartPole)
(4, 64, 8)
Structure of d (GRN)
(6, 64, 10)
Structure of d (Time-invariant RSCP & Time-varying RSCP)
(9, 64, 15)
Activation function in dé
ReLU
Kernel size of (CartPole)
10
Kernel size h of (Time-varying CartPole)
15
Kernel size h of 6 (GRN)
10
Kernel size of (Time-invariant RSCP)
10
Kernel size h of á (Time-varying RSCP)
5
Dimension observables (CartPole & Time-varying CartPole)
8
Dimension of observables (GRN)
10
Dimension observables (time-invariant RSCP & Time-varying RSCP)
15
Table 5: Hyperparameters of DKO
Hyperparameters
Value
Size of training set
40000
Size of test
4000
Batch Size
256
Learning rate
le-3
Prediction horizon H
30
Structure eb (CartPole & Time-varying CartPole)
(4, 64, 8)
Structure of è (GRN)
(6, 64, 10)
Structure of v (Time-invariant RSCP Time-varying RSCP)
(9, 64, 15)
Activation function
ReLU
Dimension of observables (CartPole & Time-varying CartPole)
Dimension of observables (GRN)
10
Dimension of observables (Time-invariant RSCP & Time-varying RSCP)
15
Table 6: Hyperparameters of MLP
Hyperparameters
Value
Size of training set
40000
Size of test set Batch Size
4000
256
Learning rate Prediction horizon H
le-3
30
Structure of ó (CartPole & Time-varying CartPole)
(4, 24, 16, 8)
Structure of ó (GRN)
(6, 30, 20, 10)
Structure of á (Time-invariant RSCP & Time-varying RSCP)
(9, 45, 30, 15)
Activation function
ReLU
19

---

Published as a conference paper at ICLR 2025
E MODELING RESULTS
E COMPARISION WITH CLASSIC SYSTEM IDENTIFICATION METHOD
classical system identification method subspace identification (Vnn Overschee & De Moor 2012) is also applied to build data-based dynamic model for the dynamics of the CartPole system. For the 30-step-ahead prediction task, the model built based on subspace identification can only provide satisfactory predictions of the position and velocity of the cart, while the predictions of the state variables related the pole diverge. contrast, the proposed method can accurately forecast the future behaviors of all the states of the CartPole system.
Cart positon (m)
1.10
Caroiii (m/s)
0.95
0
0.8
0
0.1
0.2
0.3
0.4
0.5
0.6
0.1
0.2
0.3
0.4
0.5
0.6
Time (s)
Time (s)
Pole velocity (rad/s)
Pole angle (rad)
0.3
0.1
0
-0.1
0.1
0.2
0.3
0.4
0.5
0.6
0.1
0.2
0.3
0.4
0.5
0.6
Time (s)
Time (s)
Actual state
Subspace identification
MamKO
Figure 7: Modeling results of the CartPole system by MamKO and subspace identification.
E.2 THE SELECTION OF THE DIMENSION THE STATE-SPACE MODEL
The dimension of the linear state-space model is determined through trial and error. A case study the CartPole system conducted illustrate the selection of the dimension of the state-space model. The results, which include the test loss across various lifting dimensions, As the dimension of the lifted space increases from 5 to 8, the test loss decreases significantly to 6.93 10-4 Further increasing the dimension does not lead to significant improvement in the modeling performance.
presented in Table 7
Table 7: Test loss of different lifting dimensions on CartPole system modeling by MamKO.
Lifting dimension
10
Test loss
8.95 x 10 6.93 x 10
6.94 10
7.04 x 10
20

---

Published as a conference paper at ICLR 2025
E.3 PREDICTION ERROR ON EACH PREDICTION STEP
The modeling performance of the proposed method at each prediction step is evaluated. The results for the CartPole system, GRN system, and time-invariant RSCP system are presented in Figure g Figure [ and Figure TO From the results, it can be demonstrated that while the prediction error tends increase as the number of prediction steps progresses. there is notable reduction in error observed mid-way way through the prediction horizon. This phenomenon underscores the effectiveness of our method in managing error propagation.
2.5 x10
Prediction
1.5
1.0
5
10
20
25
30
Prediction Figure 8: Prediction error of the CartPole system at each prediction step.
x10
Prediction
1.0
0.5
5
10
15
20
25
30
Prediction step
Figure 9: Prediction error of the GRN system each prediction step.
x1
8.0
Prediction enrr
6.0
5
10
20
25
30
Prediction
Figure 10: Prediction error of the time-invariant RSCP system at each prediction step.
21

---

Published as a conference paper at ICLR 2025
F CONTROL EXPERIMENT RESULTS
F.1 CARTPOLE AND TIME-VARYING CARTPOLE SYSTEMS
The trajectories of the and for the Cartpole systems presented in Figure 11
0.02
SAC
MamKo
NamKo-hased based MPC
DKC based
0.01
-01
SAC
NamKO based MPC
-0.01
MLP-hased MPC
Reforence
DKO-based MPC
-0.2
-000
12
20
12
16
Time (5)
Time
(a) Z, of the CartPole system
(b) 0g of the CartPole system
1.0
0.02
MamKO-baased MP
MLP- -aaed MPC DEG based MPC
Reference
0.5
0.01
(rad)
Wh
-0.5
-0.01
pnlyy
MamkO-based
MLP ---eeeed MPC
Reference
DIO based MPC
-1.0
-0.22
12
16
20
12
26
20
Time
(c) 24 of the varying CartPole system
(d) 0g of the time-varying CartPole system
1.0
SAC
MamK based MPC
SAC
MamKC based MPC
MLP-based MPC
RReferecce
MPP -aaddaadddd MPC DKO-based MPC
Referencc
based MPC
0.01
(m))
(mdd)
0
16
Time (4)
Time (s)
(e) of the time-varying CartPole system with (1) Oc of the time-varying CartPole system with rapidly changing parameters
rapidly changing parameters
Figure 11: The trajectories of the Cartpole systems. The Y-axis indicates the average state trajectory of each time instant. The shaded area represents the confidence interval (one standard deviation) over the ten experiments.
The staad errors of the three CartPole systems are listed in Table g From the results, the HlamKO-basee MPC demonstrates the smallest steady-state errors in all the states of each of the three CartPole systems. Particularly, for the time-varying CartPole system with rapidly changing parameters, the MamKO-based MPC achieves substantial reduction in the steady-state error in
22

---

Published as a conference paper at ICLR 2025
cart position, outperforming the MLP-based MPC, DKO-based MPC and SAC by 98.38%, 92.92%, and 93.86%, respectively.
Table 8: Steady-state errors of the three CartPole systems by different control methods.
Method
Cat position
Carlole system
Fode ungle trad)
Cut postiion
Time-wrying CariPode sysaem Pole angle (ta))
CurPele system CanPulk systm with capht changing grransen Cur postiio (m))
Pole sugle (ral)
MamkO based MPC
6.40 10
2.15 10
3,98
4.74 10
584 x 10
2,00 10
MLP-aasedd MPC
1.28
5.5 10-4
662 10-1
1/02 x 10-3
3,06
1 10-3
D80-keo MPC
1.51 x 10-1
2.55 10-1
1.30 10-1
5.14 10-3
5:30 10
2.08 16-1
SAC
6.94 x 10-1
1,92 10-2
733 10
1.87 10-2
9.7 10
2.05 16-2
23

---

Published as a conference paper at ICLR 2025
F.2 GRN SYSTEM
The trajectories of I4 of the GRN system obtained based on different control methods are presented in Figure 12 The steady-state errors the GRN system are listed in Table Table99.The The proposed MamKO- based MPC provides the smallest steady-state erro, reducing it by 52.23% %0 92. 10% and 90.68% as compared MLP-based MPC, DKO-basee MPC, and SAC, respectively.
Table Steady-state error of the GRN system by different control methods.
Method
MamKO-based MPC
MLP--asedd MPC
DKO-based MPC
SAC
Steady-state enror
5.35 10
1.12 x 10
6.77 10
5.74 x 10
8
6
x
4
SAC
MamK O-based MP Reference
MLP -based MPC
OKO-based MP
50
100
150
200
250
300
350
400
Time (a.u.)
Figure 12: The trajectories of the GRN system by different control methods. The Y-axis indicates the average state trajectory of each time instant. The shaded area represents the confidence interval (one standard deviation) over the ten experiments
24

---

Published as a conference paper at ICLR 2025
F.3 TIME-INVARINNT RSCP AND
RSCP
One set of trajectories obtained from ten experiments for the time-invariant RSCP system is pre- sented in Figure 13 One of trajectories obtained from the ten experiments for the time-varying RSCP system is presented in Figure 14] Since the nine states the RSCP system have significantly different magnitudes. the tracking errors for the RSCP system are calculated based on the states after normalization. The stcady-state errors of the time-invariant RSCP system and time-varying RSCP system are listed in Table TO From the results, the proposed MamKO-based MPC demonstrates the smallest steady-state error. For the invariant RSCP system, the proposed MamKO-base steady-state error by 57.33%, 17.42%, and 61. 61. as compared to MLP-based MPC, DKO-based MPC, and SAC, respectively. For the time-varying RSCP system, the proposed MamKO based MPC reduces stcady-state error by 49.78%, 24.83%, and 22.25% as compared to MLP-based MPC, DKO-based MPC, and SAC, respectively.
MPC achieves reductions
Table 10: Steady-state errors of the time- innerinvaaiant RSCP system and the time-varying RSCP system.
Method
MamKO-based MPC
MI.P-based MPC
DKO-based MPC
SAC
Time-invariant RSCP system
1.28
10
3.00 10
1.55 10
3.34 x 10
Time-varying RSCP system
6.78 10
1.35 10-1
9.02 10
8.72 x 10
0.34
0.75
500
0.14
0.55
450
0.5
1.0
0.5
1,0
0.5
1.0
1.5
0.71
500
0.28
0.62
470
0.18
0.53
440
1.0
1.0
0.5
0.15
0.80
500
0.10
0.73
475
0.05
450
0.5
1.0
1.5
0.5
1.0
1.5
0.00
0.50
1.00
1.50
Time (h)
Time (h)
Time (h)
SAC MLP-based MPC
DKO-based MPC MMaMMM based MPC Refe-Refeeecce
Figure 13: The state trajectories for the time-invariant RSCP systems.
25

---

Published as a conference paper at ICLR 2025
0.76
500
0.24
0.68
475
0.14
0.60
450
0.5
1.0
0.5
1.0
0.5
1.0
0.72
490
0.26
0.66
470
0.18
0.60
450
0.5
1.0
1.5
0.5
1.0
1.5
0.5
1.0
1.5
0.14
0.80
500
0.09
0.72
475
0.04
0.64
450
0.5
1.0
1.5
0.5
1.0
1.5
0.5
1.0
1.5
Time (h)
Time (h)
Time (h)
SSAC -LLP-based MPC
DKO-based MPC MamkO -based MPC Reefeerenc
Figure 14: The state trajectories for the time-varying RSCP systems.
26

---

Published as a conference paper at ICLR 2025
F.4 WATEWATER TREATMENT SYSTEM
To further demonstrate the applicability of the proposed method on highly nonlinear systems, we evaluate the performance of MamKO large-scale nonlinear water treatment system (Alex al.) 2008), which contains 145 states, two control inputs and 14 known disturbances. The dynamic behaviors of this process are simulated using a high-fidelitty simulator that was built based on 145 ordinary differential equations (ODEs) Zeng & Liu 2013). A schematic of the water treatment plant is presented in Figure
15 invariant states are excluded from Koopman- based modeling and control tasks. The modeling task aims to build Koopman model that describes the evolution of the remaining 130 state vari- ables over 10-step prediction horizon. For the water treatment system, trajectories of state and action samples are gathered, generating training set of 20, 000 samples, a validation set 2, 000 samples, and a test set of 2, 000 samples. The modeling performance of the MamKO method and two baseline methods are presented Figure [6 From the experimental results, the proposed method outperforms the deep Koopman operator (DKO) and multilayer perception (MLP) methods in modeling performance.
A set-point tracking task is considered to evaluate the control performance of the proposed method. The objective of the control task to drive two selected states to a desired set point under the influence of the disturbances. In the control task, PID controller and MPC based the first- principles model denoted as NMPC are included as the baselines. From the experimental results, the proposed MamKO-basce MPC reduces tracking error by 55.60%, 45.95% and 10.07% compared to MLP-based MPC, DKO-based MPC, and SAC, respectively. The MPC based on the exact first- principles model (the same model as the simulator) exhibits the smallest tracking error. However, as emphasized in the introduction, obtaining exact first-principles model can be challenging. Additionally, the nonlinear optimization required for first fist-principles model systems often results in suboptimal solutions, leading to large deviation points in the process. Our method achieves reduction tracking error of 1.56% as compared to the PID controller. Furthermore, integrating state constraints with the PID controller for this system be challenging, and switching different set point may require retuning the parameters.
2e.
Biological reactor
Wastewater
K Kyag
K
K
Settler
Qu, Zn
Layer
Qg.Z,
Layer
Chamber
Chamber 2
Chamber
Chamber
Chamber
Layer 10
Non-aerated section
Aerated section
Qu
Qu Zu
Qui Zew
Zr
Figure 15: A schematic of the water treatment plant (Alex et al.
2008
27

---

Published as a conference paper at ICLR 2025
MLP DKO
MamKO
Test
-1
-2
0
50
100
150
200
250
300
350
400
Epoch
Figure 16: Test loss of water treatment process in log space from different methods. The shaded
area represents the confidence interval (0.5 times the standard deviation) across ten training trials.
IAAC MLP-based MPC -KO-baased -based MPC
MamkK-bbaa MPC PD 1NMPC
3.0
2.5
Tracking enro
2.0
1.5
1.0
Figure 17: Tracking errors the control task in the water treatment process from different methods. The red dots inside the boxes indicate the mean tracking errors in the process. The box chart shows the distribution of the tracking error. The middle line inside the box represents the median value of the collected data. The box represents the interquartile range (IQR), which encompasses the range between the first quartile (QI) and the third quartile (Q3). The whiskers extend the smallest and largest values within 1.5 times the IQR from the lower and upper quartiles. The outliers are the individual data points outside the whiskers.
28

---

Published as a conference paper at ICLR 2025
F.5 DEMAY RSCP SYSTEM
From the application perspective of the RSCP system, the transportation lag of reactant A and de- sired product B may introduce time delays in measuring the mass fractions z and Jni- In this system, a 0.025 time delay is included. At each new sampling instant k; the controller optimizes control output based on the measurements of zAi and ž at sampling instant measurements of T. and the historical trajectory. The modeling and control tasks of the time-delay RSCP system are same as the time-invariant RSCP and time-varying RSCP system Section 5 The modeling performance of each method is presented in Figure [8] and the control performance of each method is presented Figure 19 From the results of modeling tasks, the proposed method reduces the test loss by 55.56% and 66.67% as compared with the MLP and DKO methods, respec- tively. For the control task, the proposed LamKO-based MPC reduces the cost by 52.52%, 53.59% and 84.25% as compared to MLP-basce MPC, DKO-basee MPC, and SAC, respectively.
5, instantancous
0.5
MLP
DKO MamKO
0.5
Loss
-1.5
2.5
3355
0
50
100
150
200
250
300
350
400
Epoch
Figure 18: Test loss of the time-delay RSCP system. The shaded area represents the confidence
interval (0.5 times the standard deviation) across ten training trials.
2.0
SAC MLP-based MPC DKO-based MPC
1.5
Average cost
1.0
0.5
0.25
0.50
0.75
1.00
1.25
Time (h)
Figure 19: Average cost the time-delay RSCP system in log space. The shaded area represents the confidence interval (0.3 times the standard deviation) over the ten experiments.
29

---

Published as a conference paper at ICLR 2025
F.6 TRAJECTORY TRACKING TASK FOR GRN SYSTEM
A sinusoidal function with periods of 400 seconds is set as the tracking target of the task, which can be represented as:
xxt) sin
1) + +0
(22)
200
In the tracking process, extra uniform disturbance ranging from [-0.5, 0.5], i... & U((0.5, 0.5) added in the tracking process. The tracking trajectories of the GRN system are shown Figure 20 The mean tracking enrors of the different methods presented in Table [T] From the results, the MamKO-based MPC exhibits the best tracking performance among the control methods, which reduces the tracking based MPC, and SAC, respectively. The experimental results demonstrate that the proposed method can be applied to trajectory-trackin tasks.
by 9.84%, 20.29%, 32.93% compared to MLP-based MPC, DKO-
16
SAC
MLP-based
14
DKO -asebaaddd MPC
MamKO based MPC Reference
12
State value
4
2
0
50
100
150
200
250
300
350
400
Time (a.u.)
Figure 20: Tracking trajectories of the GRN system of different methods. The Y-axis indicates the average state trajectory of each time instant. The shaded area represents the confidence interval (one standard deviation) over the ten experiments.
Table 11: Mean tracking errors of GRN system of different methods.
System
MamKO-based MPC
MLP-based MPC
DKO-based MPC
SAC
Mcan tracking error
0.55
0.61
0.69
0.82
30


Reviews:
[{"review_id": "458Xmn7FeE", "reviewer": "Reviewer_9kPH", "summary": "In this paper, the authors propose MamKo, combining Koopman operator theory (projecting nonlinear systems into a high-dimensional linear space) with Mamba (a LLM based on state-space models) for accurate and efficient modeling of time-varying systems. Furthermore, based on the derived model, a model-based control approach is proposed. Simulations with different, time-varying systems demonstrate that MamKo can be superior to other SOTA methods. \n\nThe main contribution is combining these two methods (with some minor modifications), creating a powerful framework for time-varying systems.", "questions": "See weaknesses", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well written, the ideas are clearly formulated, and the results are well presented.\n- To the best of my knowledge, it is the first time that these two methods have been combined and applied to model time-varying systems. Thus, it's a novel and original approach.\n- The simulation results are quite convincing that the proposed method can significantly improve the modeling and control of time-varying systems.", "weaknesses": "- The simulation is limited to learning method but ignores \"classic\" system identification. There is a rich literature on identifying time-varying parameters in dynamical systems. In particular, for the cart pendulum example, a comparison to classic system identification would strengthen the paper.\n- Maybe I missed the part but how is the optimal dimension of the linear space, i.e. the dimension of A_k in (5) determined?\n- In the introduction, the computational time of other koopman based approaches is discussed. However, there is no proper discussion/evaluation/theory on the training/inference time with the proposed framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LtsGkdYmuz", "reviewer": "Reviewer_A4oR", "summary": "the paper MamKO leverages Mamba architecture to generate koopman operators  with projected benefits towards model predication capability and adaptability relative to koopman models with constant koopman operators.  MamKO generates the Koopman Operator from realtime data making the model learn nonlinear behavior over time. To demonstrate capability the MamKO model is applied to model predictive control (MPC).  The modeling and control performance of the proposed method is evaluated through experiments on benchmark time-invariant and time-varying systems.\nFour claims are made as advancing the state of the art\n1. Matrix generation using Mamba structure\n2. approach can handle unstable and time varying systems\n3. Computationally efficient MPC is developed based on the MamKO approach\n4. Experiments are provided in support. \nOverall, the paper is well written and the developments are clear and error free.  While the authors have not demonstrated the results against any real world practical system, they do provide adequate development for a practitioner to adopt into real world problems and study applicability.  In that regard I think this work is worth publishing so the SOTA in Kooman based approaches   can continue to be explored for real world use.", "questions": "Recommended fixes:\nwhich provides enhanced model predictability and adaptability, as compared to Koopman models with constant Koopman operators -- you mean prediction capability  ?\nSince time scales are important , figures should include time , instead of time steps ( or epochs) , Fig 8 just has  \"t\" with no units,  Please provide unit son all figure axes where possible   \n\nClaims requiring more evidence:  \nMamKO achieves good balance between advanced modeling capabilities and real-time control implementation efficiency --> \nThis claim in RT implementation needs more support hence actual time units are needed.  The RSCP problem has a scale of hours and is therefore generally very slow.  It will be good to see clearly the range of problems in time scale and frequency scale to understanding capability against system dynamics. \n\nTheoretical:  The paper could use more rigorous theoretical analysis or proofs to back the claims of stability and effectiveness in capturing time-varying dynamics. For example, formal proofs of the convergence properties of the Mamba-based Koopman operators or bounds on the error in approximating nonlinear systems would greatly increase the paper’s credibility.\n\nEquation (5): The transition from a time-invariant to a time-varying state-space model is theoretically justified but lacks a rigorous mathematical proof. The paper relies heavily on empirical evidence, perhaps some theoretical discussion may be help support eth generalizability of the method.\n\nCELU Activation: The use of the negative CELU function is an interesting choice to handle instability. The mathematical justification provided for this choice seems reasonable, but more in-depth analysis of how it impacts model stability across different systems would strengthen the claim.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "the paper has several strengths\n1. Leveraging Mamba for generating time-varying Koopman operators to address the limitations of fixed linear state-space models is innovative.\n2. The authors  effectively incorporate MamKO into an MPC framework and demonstrate practical applications.  ability to formulate the control problem as a convex optimization task adds to the framework's appeal for real-time applications.\n3. The experimental validation (although all simulations)  is based on multiple benchmark systems (e.g., CartPole, Gene Regulatory Networks, and CSTRs). Results demonstrate the advantages of MamKO over traditional methods like Deep Koopman Operator (DKO) and MLP-based models. \n4. Real-Time Adaptability: By generating matrices in real-time, MamKO addresses the challenge of modeling time-varying systems efficiently. The proposed method's performance on rapidly changing systems highlights its practical utility.\n5. Supplemental material allows reproducibility", "weaknesses": "The framework assumes that the Mamba-generated matrices adequately capture the dynamics of the underlying system. This assumption may not hold for  highly nonlinear systems where matrix representations could become unstable or inaccurate.\n\nExperimental results do support the claims of improved performance. However, the paper could enhance its rigor by including statistical analyses (e.g., confidence intervals or hypothesis tests) to demonstrate the significance of the results.\n\nAuthors do acknowledge that their method may not generalize to all real-world systems, but a more explicit discussion on the boundaries of applicability (e.g., systems with extremely high-frequency dynamics) would add value.\n\nComparison: The benchmark methods used have all been solved using classical/modern control methods, yet eth author did not include any non data-driven solutions for comparison.  This may be a useful addition to the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L4gPOmLDJh", "reviewer": "Reviewer_41k8", "summary": "The paper investigates complex dynamical systems and aims to address limitations in modeling these systems using Koopman operators. It explores the application of large language model (LLM) technology to manage the nonlinear and time-varying characteristics of dynamic systems. The work introduces a framework combining Mamba with Koopman operators, leveraging Mamba’s matrix generation capabilities to predict the future states of time-varying systems. This approach preserves a linear state-space structure while enhancing the model’s predictive accuracy and adaptability, creating a linear invariant state-space model. The paper is well-structured and logically organized, presenting innovative ideas with strong experimental support and highlighting the potential impact of this approach on complex dynamical system modeling and control.", "questions": "1. The paper chooses the Mamba model as the basis for the overall approach but needs to provide quantifiable evidence supporting this choice. Could the authors provide relevant data to clarify the rationale for choosing the Mamba model?\n2. All validation experiments were conducted on a simulation platform, with no hardware experiments presented. Simulation results are not a substitute for results from real-world experiments. Could the authors conduct hardware experiments to verify the feasibility and effectiveness of the approach?\n3. The proposed framework aims to solve time-varying issues in time-varying systems but needs more discussion of delay errors and accumulated multi-step prediction errors, as well as the effects these errors may have on system stability and robustness. Did the authors consider these issues in the initial design phase? If not, could they now provide evidence to show the approach can address these issues?\n4. The paper indicates that MamKO relies on multi-step prediction to achieve scalability for complex systems. Could you provide a detailed comparison of the computational overhead between MamKO and other methods? For instance, do quantitative data on computation time or hardware resource consumption for equivalent tasks exist to demonstrate the efficiency of MamKO?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides comprehensive experimental steps, data, and ablation analyses to validate the superiority of the MamKO model.\n2. Through experiments on multiple benchmark systems, the paper demonstrates the advantages of MamKO in handling time-varying systems.\n3. The paper shows that MamKO modeling relies on multi-step prediction, enabling the model to scale to complex systems with multi-step dependencies effectively.\n4. Based on the MamKO model, the paper develops an MPC (Model Predictive Control) scheme, emphasizing its computationally efficient optimal control for nonlinear systems.", "weaknesses": "1. The paper does not provide quantitative metrics to justify why the Mamba model was chosen over other models, nor does it include comparative experiments with other LLM models.\n2. All experiments on time-varying systems were conducted in simulated environments, lacking consideration of real-world physical factors. Thus, there are questions regarding the actual performance of the MamKO model.\n3. Although the paper emphasizes solving time-varying problems in control systems, it lacks an analysis or discussion of the control scheme’s stability and robustness.\n4. The proposed approach uses multi-step historical data but needs to address the accumulation of errors that this may cause, which significantly affects the rigor of the model.\n5. While the paper aims to apply the MamKO framework to time-varying systems to address related issues, it does not explicitly discuss the delay errors present in time-varying systems or propose corresponding solutions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yEvUJmiBkf", "reviewer": "Reviewer_PG54", "summary": "The paper innovatively combines the Mamba architecture with Koopman operators, introducing the MamKO framework - the first application of a large model architecture to Koopman modeling and control. The authors effectively leverage Mamba's capabilities in time-varying modeling and improve the matrix generation mechanism by replacing the negative exponential function with CELU activation function, enabling the model to handle unstable systems. Furthermore, based on the generated time-varying Koopman operator, they designed an MPC controller that significantly enhances modeling and control performance while maintaining computational efficiency. In the experimental section, the authors demonstrate MamKO's practical performance across multiple application scenarios (such as CartPole and chemical processes) and compare it with other control methods (like SAC) to validate its superiority.", "questions": "To enhance the impact and persuasiveness of this work, the following questions merit consideration:\n\n1. A detailed analysis or proof regarding the theoretical stability and convergence properties of the MamKO framework would be valuable.\n\n2. Validation of selecting the Mamba structure over other SSM-based models (such as S4, H3) would strengthen the paper, along with elaboration on Mamba's mathematical advantages and rationale in control tasks.\n\n3. Could more evidence about the MamKO framework's generalization ability and robustness when handling systems with different dynamic characteristics be supported by experimental data?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- Innovation: First-time integration of Mamba architecture with Koopman operators, introducing the MamKO framework, pioneering a new application of large models in control theory.\n\n- Performance: Demonstrated superior performance of MamKO-MPC across multiple practical application scenarios, with effective comparisons against existing control methods.\n\n- Model Enhancement: Improved the matrix generation mechanism in Mamba through CELU activation function, enabling the handling of unstable systems and enhancing model robustness.\n\n- Computational Efficiency: The designed MPC controller maintains high computational efficiency while improving modeling and control performance, reducing the computational burden associated with traditional online Koopman methods.", "weaknesses": "The paper has a clear overall structure, but lacks sufficient depth in key areas, particularly in justifying the choice of Mamba over other sequence model architectures. Furthermore, the limited explanation of theoretical background makes it challenging for readers to fully understand the design rationale and advantages of the MamKO framework. Specifically:\n\nTheoretical Analysis Needs Strengthening:\n- Lacks in-depth theoretical analysis of system stability and convergence. As a control method incorporating deep learning, it should provide more comprehensive theoretical guarantees for system stability\n\nInsufficient Model Selection Justification:\n- While the similarity between Mamba's SSM structure and Koopman operators is mentioned, a detailed mathematical comparison between the Mamba SSM structure and Koopman operators lacks and reasons for choosing Mamba over other sequence model architectures based on SSM (like S4, H3) inadequate\n\nLimited Control Task Types:\n- Control tasks primarily focus on setpoint tracking, missing evaluation of other control tasks (trajectory tracking, disturbance rejection, energy optimization)\n\nExperimental Evaluation Metrics Need Expansion:\n- Lacks systematic evaluation of control performance metrics (overshoot, steady-state error), recommend including more control performance indicators to comprehensively demonstrate MamKO-MPC's advantages", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)