Learning to Solve
Optimization Problems Constrained
with Partial Differential Equations

The first objective is to establish a neural network-based mapping that transforms the system state $y(t,x)$ of a PDE-constrained optimization problem 1 to a control action $u(t,x)$. Practically, this mapping is implemented as a neural network $\mathcal{U}_{\omega}$, which takes the current system state $\hat{y}(t,x)$ and desired terminal state $y(T,x)$ as input, and outputs an array of coefficients $\hat{\bm{c}}(t)$ that are then used to estimate the corresponding optimal decision variables. These coefficients are the weights used to combine a set of basis functions to represent the control action, and are described next.

Basis function representation of the control action. To approximate the optimal control action using a neural network, we represent the control trajectory ${u}(t,{x})$ as a linear combination of basis functions. This representation has also been proposed in basis_cite_1 ; basis_cite_2 . This design choice is motivated by the fact that, although the output of a neural network lies in a finite-dimensional space, combining it with continuous basis functions ${\phi_{i}(x)}$ allows us to construct a continuous approximation of the control trajectory ${u}(t,{x})$. In this way, the controller model $\mathcal{U}_{\omega}$ operates in a low-dimensional space, while the basis functions map these finite number of outputs into a continuous-time signal. As shown in the additional experiments provided in Appendix 9.6, direct modeling of continuous-time control actions, either with a neural network or classical method, leads to qualitatively similar performance to our basis expansion approach, but while the computational cost of our method remains extremely low, the running time of the classical methods grows significantly. These results show that our method is suitable for real-time tasks.
Consider the control action ${u}(t_{k},{x})$, where $t_{k}$ denotes a collocation point; at each time instant $t_{k}$, the control ${u}(t_{k},{x})$ is expressed as a weighted sum of $M$ continuous basis functions $\{\phi_{i}(x)\}_{i=1}^{M}$:

$${u}(t_{k},{x})=\sum_{i=1}^{M}c_{i}(t_{k})\,\phi_{i}(x),\qquad t_{k}=0,\ldots,T,$$

(3)

where $\bm{c}(t_{k})=[c_{1}(t_{k})\;\ldots\;c_{M}(t_{k})]^{\top}$ denotes the coefficients of the basis functions at time instant $t_{k}$. For example, in one of our experiments, we adopt sinusoidal basis functions, e.g. $\phi_{1}(x)=\sin(\pi x/X)$, $\phi_{2}(x)=\cos(\pi x/X)$, $\dots$, $\phi_{2M-1}(x)=\sin(M\pi x/X)$, $\phi_{2M}(x)=\cos(M\pi x/X)$, which provide a Fourier-like expansion of the control signal. Here $M$ denotes the number of sinusoidal modes included in the expansion, with higher values of $M$ allowing for better approximations of the control variable.

Surrogate controller model training. The role of the controller surrogate $\mathcal{U}_{\omega}$ is to predict the array of coefficients $\bm{c}(t_{k})=[c_{1}(t_{k})\;\dots\;c_{M}(t_{k})]^{T}$. Formally, the network $\mathcal{U}_{\omega}$ approximates the function

$${U}:\big({y}(t_{k},{x}),y(T,x)\big)\mapsto\bm{c}(t_{k}),$$

(4)

which maps the state variable ${y}(t_{k},x)$ and desired terminal state ${y}(T,x)$ to the corresponding controller weights $\bm{c}(t_{k})$, subsequently used to construct the control action $u(t_{k},x)$ according to (3). The surrogate controller model $\mathcal{U}_{\omega}$ is trained on-the-fly, using a dataset $\mathcal{D}=\{(\hat{{\bm{y}}}^{i}(t_{k},{x}),\,{y}(T,{x})\}_{i=1}^{N}$, where $\hat{{\bm{y}}}(t_{k},{x})$ denotes the state variables estimated by the dynamic predictor $\mathcal{Y}_{\theta}$, ${y}(T,{x})$ is the desired terminal state and $N$ is the number of samples. Here, bold symbols $\bm{y}(t_{k},x)$ denote the space-discrete version of the corresponding vector $y(t_{k},x)$, with each collocation point corresponding to a specific output of $\mathcal{Y}_{\theta}$, e.g. $\bm{y}(t_{k},x)=\big[\hat{y}(t_{k},x_{1})\;\ldots\;\hat{y}(t_{k},x_{n})\big]^{T}$. Note that this setting is unsupervised: optimal weights $\bm{c}(t_{k})$ are unknown, and candidate weights $\hat{\bm{c}}(t_{k})=\mathcal{U}_{\omega}(\hat{\bm{y}}(t_{k},x),y(T,x))$ are evaluated solely by the objective value $\mathcal{J}$ in (1a) they induce. Figure 2 illustrates the interaction between PDE-OP’s proxy optimizer $\mathcal{U}_{\omega}$ and dynamic predictor $\mathcal{Y}_{\theta}$ at time instant $t=t_{k}$. Given the predicted state $\hat{{y}}(t_{k},{x})$ (for $k=0$ we have $\hat{{y}}(t_{0},{x})={y}_{0}(x)$), the surrogate controller outputs an estimate of $\bm{c}(t_{k})$ which are then used to construct the control action ${u}(t_{k},{x})$; this is subsequently used by $\mathcal{Y}_{\theta}$ to predict the next state $\hat{{y}}(t_{k+1},{x})$. This process is repeated at each collocation point, until the terminal state ${y}(T,{x})$ is estimated.
To build our surrogate controller, we build on the proxy optimizer method proposed in (park2023self, ), which allows self-supervised training without access to precomputed optimal solutions, while directly encouraging feasibility through primal–dual updates. Once trained, the model ${\cal U}_{\omega}$ can be used to generate near-optimal solutions at low cost. The next section discusses how the state variables are estimated using a discrete-time Neural Operator architecture.

The second component of PDE-OP is a discrete-time Neural Operator architecture, here denoted as $\mathcal{Y}_{\theta}$, which models the PDE solution operator. The proposed dynamic predictor employs two subnetworks: the first, which we denote branch network, encodes input functions (e.g., control signals), and the second, called trunk network, encodes the coordinates where the solution is evaluated (e.g., spatial locations). Their outputs are then combined to approximate the target operator. This architecture is inspired by (Lu_2021, ), a neural operator framework designed to approximate mappings between infinite-dimensional function spaces.

Recall that the PDE-constrained problem couples system dynamics and decision variables. To take account for this coupling, the proposed architecture incorporates both the system state and the controller output. Given the current estimated state $\hat{\bm{y}}(t_{k},x)$ and the control weights $\hat{\bm{c}}(t_{k})$ estimated by $\mathcal{U}_{\omega}$, the network $\mathcal{Y}_{\theta}$ outputs an estimate of the state at the subsequent collocation point in time:

$$\hat{\bm{y}}(t_{k+1},x)=\mathcal{Y}_{\theta}\!\left(t,x,\hat{\bm{y}}(t_{k},x),\hat{\bm{c}}(t_{k})\right).$$

(5)

Rolling out (5) across $k=0,\ldots,T-1$ yields the full estimated state trajectory, $\{\hat{\bm{y}}(t_{k},x)\}_{k=0}^{T}$.

Branch network. The branch net encodes the estimated system state sampled at $n$ fixed sensor locations, $\hat{\bm{y}}(t_{k},x)$, together with estimated control weights $\hat{\bm{c}}(t_{k})$, and maps this concatenated input to a $d$-dimensional vector of coefficients $\bm{b}\!\left(\hat{\bm{y}}(t_{k},x),\hat{\bm{c}}(t_{k})\right)\in\mathbb{R}^{d}$.
Trunk network. The trunk net encodes the spatial query coordinate $x$, producing a latent feature vector $\bm{\gamma}(x)\in\mathbb{R}^{d}$.
The estimated subsequent (in time) state $\hat{\bm{y}}(t_{k+1},x)$ at location $x_{i}$ is obtained by combining the branch coefficients and trunk features through an inner product:

$$\hat{y}(t_{k+1},x_{i})\;=\;\sum_{j=1}^{d}b_{j}\!\left(\hat{{y}}(t_{k},x_{i}),\hat{\bm{c}}(t_{k})\right)\,\gamma_{j}(x_{i}),\;\;i=1,\dots,n.$$

(6)

The remainder of this section details the method by which the state variables are precisely estimated using the proposed discrete-time Neural Operator.

Model initialization and training. Given the proposed time-discrete neural operator model $\mathcal{Y}_{\theta}$ takes as input the weights of the basis functions estimated by the surrogate controller $\mathcal{U}_{\theta}$, it is practical to initialize the dynamic predictor model. To achieve this, we initialize $\mathcal{Y}_{\theta}$ in a supervised fashion; its training data are generated by sampling time-varying control weight sequences $\bm{c}(t_{k})$ from a gaussian distribution, e.g., $c_{i}(t_{k})\sim\mathcal{N}(0,\sigma_{i})$, and by specifying the terminal state $y(T,x)$. The corresponding solution $y(t,x)$ satisfying (1b)–(1f) is computed using a numerical PDE solver. Each trajectory is discretized into one-step training samples of the form $(y(t_{k},x),\bm{c}(t_{k}))\mapsto y(t_{k+1},x)$. Formally, the dataset of training pairs is given by $\mathcal{D}=\Big\{\big(y^{i}(t_{k},x),\bm{c}^{i}(t_{k}),y^{i}(t_{k+1},x)\big)\;\Big|\;t_{k}\in\{0,\dots,T-1\}\Big\}_{i=1}^{N}$. The network $\mathcal{Y}_{\theta}$ is trained to minimize the prediction error between its outputs and the ground-truth values. Formally, this model is trained by minimizing the loss:

$$\min_{\theta}\;\mathbb{E}_{(\bm{c},y)\sim\mathcal{D}}\Big[\big\|\hat{y}(t,x)-y(t,x)\big\|^{2}\Big],$$

(7)

where $\mathcal{D}$ denotes the dataset of pairs $(\bm{c},y)$. This training strategy ensures that $\mathcal{Y}_{\theta}$ learns to approximate the PDE solution operator given a distribution of weights $\bm{c}(t)$ approximating the distribution of estimated weights generated by the surrogate controller $\mathcal{U}_{\omega}$, so that it can provide accurate PDE-solutions given the predicted controller weights $\hat{\bm{c}}(t)$.

To integrate the proposed time-discrete neural perator within the decision process, this paper proposes a Primal Dual (LD) learning approach, which is inspired by the augmented Lagrangian relaxation technique (Hestenes1969MultiplierAG, ) adopted in classic optimization. In Lagrangian relaxation, some or all the problem constraints are relaxed into the objective function using Lagrangian multipliers to capture the penalties induced by violating them. The proposed formulation leverages Lagrangian duality to integrate trainable regularization terms that encapsulates violations of the constraint functions. When all the constraints are relaxed, the violation-based Lagrangian of problem (1) is

$$\operatorname*{Minimize}_{{u}}\mathcal{J}({{u(t,x)}},{{y}}(t,x))+\bm{\lambda}^{\top}_{h^{\prime}}|\bm{h}^{\prime}({{u(t,x)}},{{y}}(t,x))|+\bm{\lambda}^{\top}_{g}\max(0,\bm{g}({{u(t,x)}},{{y}}(t,x))),$$

where $\mathcal{J}$, $\bm{g}$ are defined in (1a), and (1f), respectively, and $\bm{h}^{\prime}$ is defined as follows,

$$\bm{h}^{\prime}(u,y)\;=\;\begin{bmatrix}\partial_{t}y(t,x)-\mathcal{I}_{t}\!\big(y(t,x),u(t,x),t,x\big)\\[2.0pt] \partial_{x}y(t,x)-\mathcal{I}_{x}\!\big(y(t,x),u(t,x),t,x\big)\\[2.0pt] y(0,x)-y_{0}(x)\\[2.0pt] y(t,x)-y_{\partial}(t,x)\\[2.0pt] \bm{h}(u,y)\end{bmatrix}=\bm{0}.$$

(8)

It denotes the set of all equality constraints of problem (1), thus extending the constraints $\bm{h}$ in (1e), with the system dynamics (1b)-(1c) and the initial and boundary conditions equations (1d)-(1e) written in an implicit form as above. Therein, $\bm{\lambda}_{h^{\prime}}$ and $\bm{\lambda}_{g}$ are the vectors of Lagrange multipliers associated with functions $\bm{h}^{\prime}$ and $\bm{g}$, e.g. ${\lambda}^{i}_{h^{\prime}},{\lambda}^{j}_{g}$ are associated with the $i$-th equality $h^{\prime}_{i}$ in $\bm{h}^{\prime}$ and $j$-th inequality $g_{j}$ in $\bm{g}$, respectively. The key advantage of expressing the system dynamics (1b)-(1c) and initial and boundary conditions (1d)-(1e) in the same implicit form as the equality constraints $\bm{h}$, (as shown in (8)), is treating the system dynamics in the same manner as the constraint functions $\bm{h}$. This enables us to satisfy the system dynamics and the static set of constraints ensuring that they are incorporated seamlessly into the optimization process.

The proposed primal-dual learning method uses an iterative approach to find good values of the primal $\omega$, $\theta$ and dual $\bm{\lambda}_{h^{\prime}},\bm{\lambda}_{g}$ variables; it uses an augmented modified Lagrangian as a loss function to train the prediction $\hat{{u}}$, $\hat{{y}}(t)$ as employed

$$\begin{split}\mathcal{L}^{\text{PDE-OP}}(\hat{{u}},\hat{{y}})&\!\!=\!\!\mathcal{J}({{\hat{u}(t,x)}},{{\hat{y}}}(t,x))\!+\!\bm{\lambda}^{\top}_{h^{\prime}}|\bm{h}^{\prime}(\hat{{u}}(t,x),\hat{{y}}(t,x))|\!+\!\bm{\lambda}^{\top}_{g}\max(0,\bm{g}(\hat{{u}}(t,x),\hat{{y}}(t,x))),\!\!\!\!\end{split}$$

(9)

where $\mathcal{J}({{\hat{u}(t,x)}},{{\hat{y}}}(t,x))$ represents the total objective cost, while $\bm{\lambda}^{\top}_{h^{\prime}}|\bm{h}^{\prime}(\hat{{u}}(t,x),\hat{{y}}(t,x))|$ and $\bm{\lambda}^{\top}_{g}\max(0,\bm{g}(\hat{{u}}(t,x),\hat{{y}}(t,x)))$ measures the constraint violations incurred by prediction $\hat{{u}}(t,x)$ and $\hat{{y}}(t,x)$. The loss function in (9) combines an objective term to ensure minimization of the objective function (1a) with a weighted penalty on violations of the constraint functions $\bm{h}^{\prime},\bm{g}$. This accounts for the contribution of both networks $\mathcal{U},\mathcal{Y}$ during training, which is balanced via iterative updates of the multipliers in (9) based on the amount of violation of the associated constraint function.

At iteration $j+1$, finding the optimal parameters $\omega,\theta$ requires solving

$$\omega^{j+1},\theta^{j+1}=\arg\min_{\omega,\theta}\mathbb{E}_{y(t,x)\sim\mathcal{D}}\left[\mathcal{L}^{\text{PDE-OP}}\left(\mathcal{U}_{\omega^{j}}^{\bm{\lambda}^{j}}(t,x,\hat{y}),\mathcal{Y}_{\theta^{j}}^{\bm{\lambda}^{j}}\left(t,x,\hat{u}\right)\right)\right],$$

(10)

where $\mathcal{U}_{\omega}^{\lambda^{j}}$ and $\mathcal{Y}_{\theta}^{\lambda^{j}}$ denote the DE-OP’s optimization and predictor models $\mathcal{U}_{\omega}$ and $\mathcal{Y}_{\theta}$, at iteration $j$, with $\bm{\lambda}^{j}=[\bm{\lambda}^{j}_{h^{\prime}}\ \bm{\lambda}^{j}_{g}]^{\top}$. This step is approximated using a stochastic gradient descent method

$$\begin{split}\omega^{j+1}&=\omega^{j}-\eta\nabla_{\omega}\mathcal{L}^{\text{PDE-OP}}\left(\mathcal{U}_{\omega^{j}}^{\bm{\lambda}^{j}}(t,x,\hat{y}),\mathcal{Y}_{\theta^{j}}^{\bm{\lambda}^{j}}\left(t,x,\hat{u}\right)\right)\\ \theta^{j+1}&=\theta^{j}-\eta\nabla_{\theta}\mathcal{L}^{\text{PDE-OP}}\left(\mathcal{U}_{\omega^{j}}^{\bm{\lambda}^{j}}(t,x,\hat{y}),\mathcal{Y}_{\theta^{j}}^{\bm{\lambda}^{j}}\left(t,x,\hat{u}\right)\right)\end{split}$$

(11)

where $\eta$ denotes the learning rate and $\nabla_{\omega}\mathcal{L}$ and $\nabla_{\theta}\mathcal{L}$ represent the gradients of the loss function $\mathcal{L}$ with respect to the parameters $\omega$ and $\theta$, respectively, at the current iteration $k$. Importantly, this step does not recompute the training parameters from scratch, but updates the weights $\omega$, $\theta$ based on their value at the previous iteration. Finally, the Lagrange multipliers are updated as

$$\begin{split}\bm{\lambda}^{j+1}_{h^{\prime}}&=\bm{\lambda}^{j}_{h^{\prime}}+\rho|\bm{h}^{\prime}(\hat{{u}}(t,x),\hat{{y}}(t,x))|\\ \bm{\lambda}^{j+1}_{g}&=\bm{\lambda}^{j}_{g}+\rho\max\left(\bm{0},\bm{g}(\hat{{u}}(t,x),\hat{{y}}(t,x))\right)\end{split}$$

(12)

where $\rho$ denotes the Lagrange step size and the max operation is performed element-wise.

Next, we evaluate our method on a range of PDE-constrained optimization tasks with varying levels of complexity, including cases space-time varying control actions and nonlinear PDE dynamics.

The problem of regulating the voltage magnitude in a leaky electrical transmission line, subject to a reaction-diffusion PDE with Neumann conditions is described as follows:

	$\displaystyle\operatorname*{Minimize}_{u(x)}\quad$	$\displaystyle\int_{\Omega}(y(T,x)-y_{\text{target}}(x))^{2}dx+\gamma\int_{\Omega}u(x)^{2}dx$			(13a)
	s.t.	$\displaystyle\frac{\partial y}{\partial t}=D\frac{\partial^{2}y}{\partial x^{2}}-\beta(y-y_{\text{ref}}(x))+\alpha u(x),\quad$	$\displaystyle\forall(x,t)\in\Omega\times[0,T]$		(13b)
		$\displaystyle y(0,x)=y_{0}(x),\quad$	$\displaystyle\forall x\in\Omega,\{\ t=0\}$		(13c)
		$\displaystyle\frac{\partial y}{\partial x}(t,x)=0,\quad$	$\displaystyle\forall x\in\partial\Omega,\ t\in[0,T]$		(13d)
		$\displaystyle u_{\text{min}}\leq u(x)\leq u_{\text{max}}.\quad$	$\displaystyle\forall x\in\Omega.$		(13e)

Here, $y(t,x)\in\mathbb{R}$ represents the voltage profile across the space-time domain defined by the cartesian product between $\Omega$ and time horizon $[0,T]$. The system dynamics, expressed by (13b), model the physical phenomena of diffusion ($D\frac{\partial^{2}y}{\partial x^{2}}$), voltage leakage towards a reference $y_{\text{ref}}(x)$ ($-\beta(y-y_{\text{ref}}(x))$) and the applied control action ($\alpha u(x)$). $D$, $\beta$, and $\alpha$ are all constants representing the diffusion rate, leakage rate, and control gain, respectively. The system is initialized from a state $y_{0}(x)$ according to (13c) and is constrained by zero-flux Neumann boundary conditions (13d), implying no voltage diffusion across the boundaries of the domain. The control decision $u(x)$ is a space-dependent function representing the externally applied voltage, constrained by actuator limits in (13e). The objective (13a) minimizes a combination of terminal cost and control energy.

Datasets and methods. In the system dynamics (13b), we set the diffusion rate $D=0.1$, leakage rate $\beta=1.0$, and control gain $\alpha=2.0$, with a uniform reference voltage $y_{\text{ref}}(x)=1.0$. The ground-truth dynamics are simulated using a second-order implicit Crank-Nicolson finite-difference method(Crank_Nicolson_1947, ). To initialize the proposed time-discrete Neural Operator $\cal Y_{\theta}$, we generate a dataset by sampling a diverse set of control functions $u(x)$ from a zero-mean Gaussian Random Field (GRF) with a squared-exponential kernel, with mean $m(x)$ and covariance function $k_{\ell}(x_{1},x_{2})$ as:

$$m(x)=0,\qquad k_{\ell}(x_{1},x_{2})\;=\;\sigma^{2}\exp\!\left(-\frac{\|x_{1}-x_{2}\|_{2}^{2}}{2\ell^{2}}\right),$$

(14)

where $\ell>0$ is the length–scale and $\sigma^{2}>0$ the marginal variance. The corresponding solution $y(t,x)$ is computed using a numerical solver, and then split into $80\%$ training, $10\%$ validation, and $10\%$ test set.

PDE-OP’s model $\mathcal{Y}_{\theta}$ estimates the voltage profile $y(t,x)$; for $t=T$ this prediction yield $y(T,x)$, which is compared against the target profile $y_{\text{target}}(x)$. PDE-OP’s surrogate controller $\mathcal{U}_{\omega}$ is trained to produce optimal control action $\hat{u}(x)$ while minimizing objective (13a). Our approach is benchmarked against Direct and Adjoint-Method, which are described in Appendix 8.1- 8.2. At inference time, PDE-OP’s surrogate controller $\mathcal{U}_{\omega}$ is tested on three unseen target profiles of different level of complexities: constant ($y_{\text{target}}(x)=p_{1}$), linear ramp ($y_{\text{target}}(x)=p_{1}x+p_{2}$), and sine wave ($y_{\text{target}}(x)=p_{1}+p_{2}(\sin f_{0}x+p_{3})$). For each target, the control $\hat{u}(x)$ is generated in a single forward pass.

Results. Figure 4 shows the predicted voltage (central sub-figure of Fig. 4) of PDE-OP’s dynamic component $\mathcal{Y}_{\theta}$ at test time, which is compared to the solution (left sub-figure) obtained with a numerical PDE-solver. The figure shows that the predicted voltage closely matches the numerical PDE-solver solution, with an absolute error (right sub-figure) on the order of $10^{-2}$ across the space–time domain.

Figure 3 presents a comparison between PDE-OP’s predicted terminal state $\hat{y}(T,x)$ and that obtained using the classical methods, over $3$ target profiles $y_{\text{target}}(x)$: sinusoidal, linear and constant. This figure shows that the voltage distribution produced by PDE-OP is well aligned with those produced by the Direct and the Adjoint-sensitivity method. As shown in the first row of Table 1, the direct method is the slowest, requiring $26.097s$ and yielding an MSE of $7.46\times 10^{-3}$. The Adjoint method runs in $0.755s$ with a similar accuracy, $7.13\times 10^{-3}$. Notably, PDE-OP is the fastest, by a significant margin: it runs in only $0.035$ s ($\sim\!21.6\times$ faster than Adjoint; $\sim\!746\times$ faster than Direct) and yield the lowest MSE, with $6.95\times 10^{-3}$. These results show the enormous potential of the proposed approach, which produces solutions that are qualitatively similar (or even slightly better) than the classical methods, while being between $20$ and $750$ times faster than them.

This task involves controlling the evolution of temperature in a one-dimensional domain described by a reaction–diffusion PDE. Given an initial profile, the controller aims to reach a prescribed target state with minimal control effort. The problem is formulated as:

	$\displaystyle\operatorname*{Minimize}_{\bm{c}(t)_{t\in[0,T]}}\quad$	$\displaystyle\mathcal{L}=\mathcal{L}_{\text{terminal}}+\lambda\mathcal{L}_{\text{running}}+\gamma\mathcal{L}_{\text{effort}}$			(15a)
	s.t.	$\displaystyle\frac{\partial y}{\partial t}=D\frac{\partial^{2}y}{\partial x^{2}}-\beta(y-y_{\text{ref}}(x))+\alpha u(t,x),\quad$	$\displaystyle\forall(x,t)\in\Omega\times[0,T]$		(15b)
		$\displaystyle u(t,x)=\bm{c}(t)^{\top}\bm{\phi}(x),\quad$	$\displaystyle\forall(x,t)\in\Omega\times[0,T]$		(15c)
		$\displaystyle y(0,x)=y_{0}(x),\quad$	$\displaystyle\forall x\in\Omega$		(15d)
		$\displaystyle\frac{\partial y}{\partial x}(t,x)=0,\quad$	$\displaystyle\forall x\in\partial\Omega,\ t\in[0,T]$		(15e)
		$\displaystyle c_{\text{min}}\leq\bm{c}(t)\leq c_{\text{max}},\quad$	$\displaystyle\forall\ t\in[0,T].$		(15f)

The loss components in (15a) defines as follows.

	$\displaystyle\mathcal{L}_{\text{terminal}}$	$\displaystyle:=\int_{\Omega}\left|y(T,x)-y_{\text{target}}(x)\right|^{2}\,dx,$		(16a)
	$\displaystyle\mathcal{L}_{\text{running}}$	$\displaystyle:=\int_{0}^{T}\int_{\Omega}\left|y(t,x)-y_{\text{target}}(x)\right|^{2}\,dxdt,$		(16b)
	$\displaystyle\mathcal{L}_{\text{effort}}$	$\displaystyle:=\int_{0}^{T}||\bm{c}(t)||_{2}^{2}\,dt.$		(16c)

Here $y(t,x)$ represents the temperature profile. The objective (15a) balances three terms: a terminal loss for final state accuracy, a running loss for temperature tracking, and an effort loss for control efficiency, described by (16a)–(16c).

Datasets and methods. The solutions of (15b) are generated using a Crank-Nicolson method, which deal with time-varying source term by averaging the control input over each time step. To train our models, we generate a dataset of diverse trajectories. First, smooth, time-varying control weights ($\bm{c}(t)$) are generated by sampling each of the $M$ weight trajectories from a GRF over the time domain. The full spatial-temporal control field $u(t,x)$ is then reconstructed from these weights constrained by (15f) using a set of fixed basis functions, as defined in (15c). For each generated control field, the full PDE is solved to generate a ground-truth temperature evolution $y(t,x)$. Each complete simulation yields a training sample, consisting of the control input weight trajectory $\{\bm{c}(t_{k})\}$, and the corresponding state trajectory $\{y(t_{k},x)\}$ evaluated at a set of fixed sensor locations. The dataset is then split into 80% training, 10% validation and 10% test.

Our framework uses two neural networks. First discrete-time neural operator, $\mathcal{Y}_{\theta}$, is trained as a surrogate model described in Section (4.2). It act as a one-step time-advancement operator, learning the mapping $(y(t_{k},x),\bm{c}(t_{k}))\rightarrow y(t_{k+1},x)$. Second, PDE-OP’s surrogate controller $\mathcal{U}_{\omega}$ is a recurrent neural network that acts as sequential decision-maker. At each time step $t_{k}$, it takes the current state $y(t_{k},x)$ and the final target $y_{\text{target}}(x)$ to produce the control weights $(\bm{c}(t_{k}))$. The controller is trained end-to-end by unrolling the trajectory using the surrogate $\mathcal{Y}_{\theta}$. We benchmark our approach against a highly-tuned Adjoint-Method and Linear MPC(LMPC). Further implementation details are in Appendix 8.3 - 8.4.

At inference, we evaluate PDE-OP model on three unseen target profiles: constant($y_{\text{target}}(x)=p_{1}$), linear ramp($y_{\text{target}}(x)=p_{1}x+p_{2}$), and sine wave($y_{\text{target}}(x)=p_{1}+p_{2}(\sin f_{0}x+p_{3})$).

Results. Figure 6 shows the temperature (central sub-figure of Fig. 6) estimated by PDE-OP’s dynamic component $\mathcal{Y}_{\theta}$ at test time, which is compared to the solution (left sub-figure) obtained with a numerical PDE-solver. The figure shows that also in this setting the predicted temperature aligns closely with the numerical PDE-solver solution, with an absolute error on the order of $10^{-3}$ across the space–time domain.

Figure 5 illustrates the comparison of PDE-OP ’s predicted terminal state $\hat{y}(T,x)$ and those obtained using the classical methods, over the same target profiles introduced in the previous experiment. As also shown in Table 1, all three methods achieve the same accuracy, reporting a MSE of $3.5\times 10^{-4}$. However, PDE-OP is by far the fastest method, running in $0.124\,\text{s}$. This corresponds to a $3.1\times$ speedup over LMPC (running in $0.384\,\text{s}$) and a $5.3\times$ speedup over the Adjoint method (which runs in $0.662\,\text{s}$). These results demonstrate that even in a more complex setting, with time-variant control actions and linear PDE-dynamics, our method produces solutions that are qualitatively comparable to the classical methods, but at a significantly reduced computation cost.

To test the ability of PDE-OP to handle nonlinear dynamics, we address the control of the 1D viscous Burgers’ equation with Dirichlet boundary condition. The control input and external force $u(t,x)$, is parameterized identically to the heat equation task, using time-varying weights $(\bm{c}(t_{k}))$ for a set of basis functions. The optimization goal is to find the optimal trajectory of these weights subject to Burger’s equation (17b), initial (17d) and Dirchlet’s boundary condition (17e) and box constraints (17f).

	$\displaystyle\operatorname*{Minimize}_{\bm{c}(t)_{t\in[0,T]}}\quad$	$\displaystyle\mathcal{L}=\mathcal{L}_{\text{terminal}}+\lambda\mathcal{L}_{\text{running}}+\gamma\mathcal{L}_{\text{effort}}$			(17a)
	s.t.	$\displaystyle\frac{\partial y}{\partial t}+y\frac{\partial y}{\partial x}=\nu\frac{\partial^{2}y}{\partial x^{2}}+u(t,x),$	$\displaystyle\forall(x,t)\in\Omega\times[0,T]$		(17b)
		$\displaystyle u(t,x)=\bm{c}(t)^{\top}\bm{\phi}(x),$	$\displaystyle\forall(x,t)\in\Omega\times[0,T]$		(17c)
		$\displaystyle y(0,x)=y_{0}(x),$	$\displaystyle\forall x\in\Omega$		(17d)
		$\displaystyle y(t,x)=0,$	$\displaystyle\forall x\in\partial\Omega,\ t\in[0,T]$		(17e)
		$\displaystyle c_{\text{min}}\leq\bm{c}(t)\leq c_{\text{max}},\quad$	$\displaystyle\forall\ t\in[0,T].$		(17f)

$y(t,x)$ is velocity field, and $\nu$ represents the viscosity parameter. The PDE in (17b) contains convective nonlinearity $(y\frac{\partial y}{\partial x})$ which has a significant control challenge not present in the linear heat equation. The objective retrains the same structure, balancing final state accuracy, path tracking, and control efficiency.

Datasets and methods. The ground-truth dynamics for the nonlinear Burgers’ equation are simulated using a high-fidelity numerical solver. The dataset generation process mirrors that of the heat equation: time-varying weights $\bm{c}(t)$ are sampled from a GRF, and the force field $u(t,x)$ is reconstructed using spatial basis functions, and the PDE is solved to generate trajectories. Due to the system’s high nonlinearity, we benchmark our method against a much more general and computationally expensive Nonlinear MPC (NMPC) and Adjoint-Sensitivity Method. Further implementation details are in Appendix 8.3 - 8.4.

At inference, we test our PDE-OP model on three unseen target profiles: constant ($y_{\text{target}}(x)=p_{1}$), parabolic ($y_{\text{target}}(x)=p_{1}x(1-x)$) and sine wave ($y_{\text{target}}(x)=p_{1}(\sin f_{0}x+p_{2})$). We don’t use a a linear target profile, because of the Dirichlet BCs.

Results. Figure 8 shows the system dynamics (central sub-figure of Fig. 8) estimated by PDE-OP’s dynamic component $\mathcal{Y}_{\theta}$ at test time, which is compared to the solution (left sub-figure) obtained with a numerical PDE-solver. The figure shows that the predicted temperature aligns closely with the numerical PDE-solver solution; with absolute error on the order of $10^{-3}$ across the space–time domain, demonstrating the ability of PDE-OP’s dynamic predictor to accurately capture system dynamics even under nonlinearities.

Figure 7 illustrates the comparison of PDE-OP ’s predicted terminal state $\hat{y}(T,x)$ and those obtained using the classical methods, over the same target profiles introduced in the previous experiments. As shown in Figure 7 and Table 1, each method produces highly accurate solutions, with NMPC achieving the lowest error ($\text{MSE}=7.0\times 10^{-5}$) but requiring $878.667s$. The Adjoint method is slightly less accurate, yielding a $\text{MSE}=2.3\times 10^{-4}$ in $120.696s$. In contrast, PDE-OP produces qualitatively comparable predictions while being the fastest $0.062s$ by a big margin - approximately $1.42\times 10^{4}$ faster than NMPC and $1.95\times 10^{3}$ faster than the Adjoint method!

The direct method treats the simulator as a black box mapping the bounded control vector $\bm{u}$ to the discrete objective $J(\bm{u})$ and employs a standard optimizer to minimize it. Here the control variable is discretized at each collocation point: $\bm{u}\in\mathbb{R}^{N_{x}}$.

Mechanism. We use L-BFGS-B[30] with per–coordinate box bounds $u_{\min}\leq u_{j}\leq u_{\max}$ to minimize (31). No analytic gradients are provided; instead, numerical derivatives are queried via finite differences. With a forward-difference method,

$$\frac{\partial J}{\partial u_{j}}\;\approx\;\frac{J(\bm{u}+\epsilon_{j}\bm{e}_{j})-J(\bm{u})}{\epsilon_{j}},\qquad\epsilon_{j}=\eta\,\bigl(1+|u_{j}|\bigr),$$

(18)

where $\bm{e}_{j}$ is the $j$-th unit vector and $\eta$ is a small scalar (e.g., $\mu=10^{-6}$). For higher accuracy (to which corresponds a double computational cost), the central-difference method is adopted:

$$\frac{\partial J_{h}}{\partial u_{j}}\;\approx\;\frac{J_{h}(\bm{u}+\epsilon_{j}\bm{e}_{j})-J_{h}(\bm{u}-\epsilon_{j}\bm{e}_{j})}{2\epsilon_{j}}.$$

(19)

Computational cost. Each gradient evaluation requires $N_{x}\!+\!1$ simulator calls with forward differences (or $2N_{x}$ with central differences), where $N_{x}$ denotes the number of spatial collocation points. Since one simulator call is a complete rollout (27) over $N_{t}$ time steps, the complexity per iteration is $\mathcal{O}\big((N_{x}\!+\!1)\,N_{t}\,N_{x}\big)$. This method is robust and simple, but its cost scales linearly with $N_{x}$, which becomes prohibitive for dense spatial discretizations.

The Adjoint-Sensitivity Method efficiently computes exact gradients of PDE-constrained objectives w.r.t. high-dimensional control variables. Crucially, the cost of one gradient evaluation is one forward solve + one backward adjoint solve, regardless of the number of control parameters. We derive the continuous adjoint for the open-loop, time-invariant Voltage Control Optimization (5.1) and then we present the discrete-time adjoint state used in our Crank–Nicolson (CN) implementation.

Function spaces & BCs. Consider $u\in L^{2}(\Omega)$, $y\in L^{2}([0,T]\!\times\!\Omega)$, and impose homogeneous Neumann BCs on $V$ so that the boundary terms vanish in the variational calculus; the adjoint inherits the same Neumann BCs.

The problem consists of finding $u(x)$ minimizing

$$J(u)=\tfrac{1}{2}\!\int_{\Omega}\!\!(y(T,x)-y_{\text{target}}(x))^{2}\,dx\;+\;\tfrac{\gamma}{2}\!\int_{\Omega}\!u(x)^{2}\,dx,$$

(20)

subject to

	$\displaystyle\frac{\partial y}{\partial t}-D\frac{\partial^{2}y}{\partial x^{2}}+\beta\big(y-y_{\text{ref}}(x)\big)-\alpha u(x)=0,$	$\displaystyle\forall(t,x)\in[0,T]\times\Omega,$		(21a)
	$\displaystyle y(0,x)=y_{0}(x),$	$\displaystyle\forall x\in\Omega,$		(21b)
	$\displaystyle\partial_{x}y(t,x)=0,$	$\displaystyle\forall(t,x)\in[0,T]\times\partial\Omega,$		(21c)

where $\gamma>0$ is the regularization weight in the control effort.

Lagrangian. We define

$$\mathcal{L}(u,y,p)=J(u)+\int_{0}^{T}\!\!\int_{\Omega}p(t,x)\!\left[\frac{\partial y}{\partial t}-D\frac{\partial^{2}y}{\partial x^{2}}+\beta\big(y-y_{\text{ref}}(x)\big)-\alpha u(x)\right]dx\,dt,$$

(22)

where $p(t,x)$ is the adjoint state. By imposing $\nabla_{y}J=0$ and integrating by parts in $t$ and $x$ (with Neumann BCs) yields the adjoint system.

Adjoint PDE.

$$-\frac{\partial p}{\partial t}-D\frac{\partial^{2}p}{\partial x^{2}}+\beta\,p=0,\qquad\partial_{x}p|_{\partial\Omega}=0.$$

(23)

Terminal condition.

$$p(T,x)=y(T,x)-y_{\text{target}}(x).$$

(24)