Model-Based Closed-Loop Control Algorithm for Stochastic Partial Differential Equation Control

TL;DR

SPDE control is difficult due to stochastic forcing and degraded state regularity. The authors propose Model-Based Closed-Loop Control (MB-CC), combining a Regularity Structure Features (RF) block for robust forward modeling with an operator-encoded policy net for real-time closed-loop control. Across a 1-D stochastic reaction-diffusion equation and a 2-D stochastic Navier–Stokes equation, MB-CC demonstrates improved tracking accuracy and reduced inference time, with ablations confirming the contributions of both RF and the policy net. MB-CC offers a practical, plug-and-play framework for robust SPDE control that leverages physical priors and data-driven learning, outperforming open-loop baselines and RL methods like SAC.

Abstract

Neural operators have demonstrated promise in modeling and controlling systems governed by Partial Differential Equations (PDEs). Beyond PDEs, Stochastic Partial Differential Equations (SPDEs) play a critical role in modeling systems influenced by randomness, with applications in finance, physics, and beyond. However, controlling SPDE-governed systems remains a significant challenge. On the one hand, the regularity of the system's state (which can be intuitively understood as smoothness) deteriorates, making modeling and generalization more challenging. On the other hand, this stochasticity also renders control more unstable and thus less accurate. To address this gap, we propose the Model-Based Closed-Loop Control Algorithm (MB-CC), the first model-based closed-loop control method for SPDEs. MB-CC introduces two key innovations to enhance control robustness and efficiency: a Regularity Feature (RF) block and a closed-loop strategy with an operator-encoded policy network. The RF block, inspired by the regularity structure theory of SPDEs, addresses noise-induced irregularities by transforming the network's input, including the system state and noise-perturbed external forces, into a refined feature space for improved forward prediction. Compared to previous works using regularity features, we introduce a new parameterization, data augmentation, and extend the RF block as a plug-and-play component. Additionally, to achieve closed-loop control, we introduce an operator-encoded policy network to map the current state to optimal control, which integrates physical priors and swiftly makes decisions based on states returned by the environment. We conduct a systematic evaluation of MB-CC on two notable SPDEs, showcasing its effectiveness and efficiency. The ablation studies show its ability to handle stochasticity more effectively.

Figures (6)

• Figure 1: Overview of the Model-Based Closed-Loop Control Algorithm (MB-CC).
• Figure 2: Overall architecture of the combination of the RF block and the base model.
• Figure 3: Visualization of results on the stochastic reaction-diffusion equation. The figure shows the visualized results of controlling three samples using MB-CC, results using CNN & OpenLoop, and the control targets. It is obvious that the results controlled by MB-CC are significantly closer to the target.
• Figure 4: Training and testing the forward models with different scales of the space-time white noise.
• Figure 5: Base Model Architecture.