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

Abstract

We present a data-driven control framework for partial differential equations (PDEs). Our approach integrates time-integrated deep operator networks (TI-DeepONet) as differentiable PDE surrogate models within 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 inviscid 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 to five orders of magnitude acceleration at inference time compared to nonlinear model predictive control benchmarks. These results highlight the promise of neural operator learning for scalable model-based control of PDEs.