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

TL;DR

The paper addresses PDE-constrained optimal control in infinite-dimensional spaces by replacing online PDE solvers with differentiable neural operator surrogates. It introduces Time-Integrated Deep Operator Networks (TI-DON) to learn the instantaneous temporal derivative $\partial u/\partial t$ and couples them with classical integrators to preserve causal evolution, enabling stable long-horizon predictions. These surrogates are integrated with Differentiable Predictive Control (DPC) to train a parametric neural policy offline via backpropagation through the closed-loop dynamics, eliminating the need for supervisory controllers. Empirical results on the heat, Burgers', and Fisher-KPP equations demonstrate accurate target tracking, shock mitigation, and population-density control, with policies generalizing across initial conditions and problem parameters and transferring to high-fidelity finite-difference solvers. Open-source code supports reproducibility and further research in PDE-constrained, model-based self-supervised control.

Abstract

We present an end-to-end learning to 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, thus preserving the temporal causality of infinite-dimensional PDEs while reducing error accumulation in long-horizon predictions. Within DPC, we leverage automatic differentiation to compute policy gradients by backpropagating the expectations of optimal control loss through the learned TI-DeepONet, enabling efficient offline optimization of neural policies without the need for online optimization or supervisory controllers. We empirically demonstrate that the proposed method learns feasible parametric policies 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 problem parameters. These results highlight the promise of combining operator learning with DPC for scalable, model-based self-supervised learning in PDE-constrained optimal control.

Figures (4)

• Figure 1: Schematic of the proposed Differentiable Predictive Control with Neural Operators. Forward propagation (green dashed arrows) computes control actions via a neural policy and evolves the system dynamics through a time-integrated neural operator. Backward propagation (dashed red arrows) computes gradients by differentiating through the closed-loop system, enabling end-to-end learning of constrained control policies for PDEs.
• Figure 2: HE control performance. Each scenario shows: (left) uncontrolled evolution from initial state (blue) to final state (red) versus target (black dotted); (middle) controlled trajectory achieving target; (right) applied control signals $f_i(t)$.
• Figure 3: BE shock control. Each row: (left) uncontrolled shock development; (middle) controlled smooth evolution; (right) control signals $f_i(t)$; (far right) curvature loss reduction.
• Figure 4: RDE density control. Each row shows: (left) uncontrolled evolution from initial state (blue) to final state (red) versus target (black dotted); (middle) controlled trajectory achieving target; (right) applied control signals $f_i(t)$.