Polynomial Parametric Koopman Operators for Stochastic MPC

Abstract: This paper develops a parametric Koopman operator framework for Stochastic Model Predictive Control (SMPC), where the Koopman operator is parametrized by Polynomial Chaos Expansions (PCEs). The model is learned from data using the Extended Dynamic Mode Decomposition -- Dictionary Learning (EDMD-DL) method, which preserves the convex least-squares structure for the PCE coefficients of the EDMD matrix. Unlike conventional stochastic Galerkin projection approaches, we derive a condensed deterministic reformulation of the SMPC problem whose dimension scales only with the control horizon and input dimension, and is independent of both the lifted state dimension and the number of retained PCE terms. Our framework, therefore, enables efficient nonlinear SMPC problems with expectation and second-order moment constraints with standard convex optimization solvers. Numerical examples demonstrate the efficacy of our framework for uncertainty-aware SMPC of nonlinear systems.

• The paper presents a PPKO framework that integrates Koopman operator theory with polynomial chaos expansions to effectively capture parametric uncertainty in nonlinear systems.
• It utilizes an EDMD-DL approach to learn a polynomial-parameterized operator, enabling efficient propagation of state means and variances with dramatic computational speed-ups.
• The methodology demonstrates robust closed-loop control in diverse applications, notably in Duffing oscillators and CSTRs, achieving up to 50× acceleration over traditional SMPC methods.

Polynomial Parametric Koopman Operators for Stochastic Model Predictive Control

Introduction and Problem Context

Stochastic Model Predictive Control (SMPC) for nonlinear dynamical systems with significant parametric uncertainty remains an outstanding challenge due to the computational complexity of propagating uncertainty and evaluating expectation-type objectives and constraints. Classical approaches leveraging polynomial chaos expansions (PCEs) have been effective for linear systems, but scale poorly with system size and are inapplicable to the nonlinear case without severe approximations or prohibitive computation. Recent Koopman operator approaches enable linearization in a lifted space, greatly facilitating deterministic MPC but neglecting direct accommodation of parametric uncertainty in the dynamic model. This work introduces a Polynomial Parametric Koopman Operator (PPKO) framework, which marries data-driven Koopman representations with PCE-based stochastic parameterization for efficient, uncertainty-aware SMPC.

Technical Framework and Algorithms

The proposed framework learns a parametric Koopman operator using Extended Dynamic Mode Decomposition–Dictionary Learning (EDMD-DL), where the operator matrices are explicitly parameterized using orthonormal polynomial bases defined by the probability distribution of uncertain parameters, i.e., PCEs. In contrast to neural network parameterizations of the operator over parameters [guo2025learning], this approach yields polynomial analytic dependence of the system matrices on the parameters, enabling highly efficient propagation of both state means and second moments.

More formally, for a discrete-time nonlinear system $x_{t+1} = f(x_t, u_t, \theta)$ with time-invariant parametric uncertainty $\theta$, the lifted linear system is learned as

$z_{t+1} = A(\theta) z_t + B(\theta) u_t$

where the Koopman matrices $A(\theta), B(\theta)$ admit expansions

$$A(\theta) = \sum_{k=0}^{N_K-1} A_k \phi_k(\theta), \quad B(\theta) = \sum_{k=0}^{N_K-1} B_k \phi_k(\theta)$$

with PCE basis functions $\phi_k(\theta)$. Koopman observables are learned by neural network dictionary over trajectory data indexed by parameter samples.

Expectation-type costs and constraints under parametric uncertainty are then exactly condensed (without stochastic Galerkin projection) into a low-dimensional deterministic quadratic program (QP) or quadratically constrained QP, whose decision variables span only the input horizon. The resulting SMPC problem dimension is independent of both the lifted state and PCE basis dimension, in sharp contrast to classical PCE-based Galerkin approaches.

Numerical Performance and Empirical Results

The framework is validated on two canonical nonlinear systems:

1. Duffing Oscillator with Parametric Uncertainty:

The Duffing oscillator with uniformly distributed damping and stiffness parameters is discretized and used to train the PPKO model. In open-loop prediction, the learned PPKO model accurately reproduces both the mean and variance trajectories compared to Monte Carlo simulations over thousands of parameter samples.

Figure 2: Open-loop Duffing oscillator dynamics under parametric uncertainty: PPKO-approximated mean and variance closely match ground truth Monte Carlo simulations.

In closed-loop SMPC, with different parameter realizations yielding distinct dynamical regimes (double-well, single-well, undamped), the PPKO-SMPC controller robustly stabilizes the system in all cases, successfully traversing the complex geometry of the state space arising from parametric variabilities.

Figure 3: Closed-loop Duffing oscillator regulation for three representative parameter samples, illustrating robust performance across fundamentally different open-loop regimes.

2. Series-Parallel Reaction Network in CSTR:

For a high-dimensional chemical reaction in a continuous stirred-tank reactor (CSTR) with parametric uncertainty in reaction rates, the PPKO is learned using deviation variables relative to parameter-dependent steady states. Open-loop predictions again closely match the true system mean and variance.

Figure 4: PPKO reproduces the true stochastic dynamics of CSTR concentrations under random parameterizations over time horizons relevant for control design.

Under receding horizon implementation for regulation, the PPKO-SMPC controller achieves tight concentration tracking with strong rejection of input disturbances, whereas the uncontrolled open-loop trajectories exhibit large deviations.

Figure 1: Closed-loop regulation of CSTR intermediate concentration using PPKO-SMPC, demonstrating significant improvement over uncontrolled operation across uncertainty ensemble.

Computational Scale and Runtime

A critical practical contribution is the demonstration of dramatic sample and runtime efficiency. Compared to scenario-based Quadrature SMPC and classic intrusive PCE-Galerkin SMPC, the PPKO approach's condensed QP formulation yields order-of-magnitude speed-ups without sacrificing accuracy in closed-loop performance.

Figure 5: Cumulative solve time comparison reveals that PPKO-SMPC offers up to $50\times$ acceleration over baseline SMPC approaches, especially pronounced in high-dimensional chemical process applications.

This robust scaling is attributed to the analytic structure of the cost and constraints after integrating out uncertainty, enabling all heavy expectations to be accurately precomputed offline.

Implications, Limitations, and Future Directions

The formal condensation of the SMPC problem dimension opens new possibilities for deploying nonlinear SMPC under uncertainty in resource-constrained environments such as embedded process control. The PPKO framework crucially allows data-driven learning of operator representations that capture both system nonlinearity and parametric uncertainty, moving beyond robust/tube-based designs that only account for model residuals.

However, accuracy and stability rely on the fitting fidelity of PCE parameterizations and convex neural network dictionary architectures. Extending the framework to time-varying or non-i.i.d. uncertainties, or to systems with more complex parameter dependencies, will require further theoretical advances. Potential directions include quantifying approximation errors for finite PCE truncation, integrating multi-step operator learning strategies for improved closed-loop guarantees [wu2026least], and establishing formal robust stability results as in [strasser2024koopman, zhang2022robust].

Conclusion

This work establishes a new paradigm for incorporating data-driven parametric Koopman models into convex, scalable stochastic MPC formulations. By leveraging PCE-structured operator representations, the framework enables real-time, uncertainty-aware control of nonlinear systems, with computational requirements suitable for embedded deployment. The condensation technique fundamentally shifts the complexity bottleneck from online solution to offline learning and quadrature evaluation, an attractive tradeoff for industrial practice. Future research will target generalization to broader uncertainty classes and rigorous guarantee derivation.

---

References:

• "Polynomial Parametric Koopman Operators for Stochastic MPC" (2604.00935)
• "Learning Parametric Koopman Decompositions for Prediction and Control" [guo2025learning]
• "Least-Squares Multi-Step Koopman Operator Learning for Model Predictive Control" (Wu et al., 17 Jan 2026)
• "Koopman-Based Feedback Design with Stability Guarantees" [strasser2024koopman]
• "Robust Tube-Based Model Predictive Control with Koopman Operators" [zhang2022robust]