Learning to Control PDEs with Differentiable Predictive Control and Time-Integrated Neural Operators

TL;DR

This paper introduces a novel data-driven control framework for PDEs that combines differentiable surrogate models with neural control policies, enabling efficient offline optimization and real-time inference.

Contribution

The authors develop a framework integrating TI-DeepONets with Differentiable Predictive Control, allowing scalable, offline neural control of PDEs without online optimization.

Findings

Achieved target tracking and constraint satisfaction in PDE control tasks.

Demonstrated four orders of magnitude faster inference compared to benchmarks.

Generalized policies across initial conditions and parameters.

Abstract

We present a data-driven control framework for partial differential equations (PDEs). Our approach integrates Time-Integrated Deep Operator Networks (TI-DeepONets) as differentiable PDE surrogate models within the Differentiable Predictive Control (DPC)-a self-supervised learning framework for constrained neural control policies. The TI-DeepONet architecture learns temporal derivatives and couples them with numerical integrators, while the DPC algorithm uses automatic differentiation to compute policy gradients by backpropagating the expectations of the optimal control loss through the learned TI-DeepONet. This approach enables efficient offline optimization of neural policies without the need for online optimization or supervisory controllers. We empirically demonstrate the proposed method across diverse PDE systems, including the heat, the nonlinear Burgers', and the reaction-diffusion equations. The learned policies achieve target tracking, constraint satisfaction, and curvature minimization objectives, while generalizing across distributions of initial conditions and parameters. Moreover, we demonstrate four orders of magnitude acceleration at inference time compared to nonlinear model predictive control benchmarks. These results highlight the promise of operator learning for scalable model-based control of PDEs.