AONN: An adjoint-oriented neural network method for all-at-once
  solutions of parametric optimal control problems
Parametric optimal control problems governed by partial differential
equations (PDEs) are widely found in scientific and engineering applications.
Traditional grid-based numerical methods for such problems generally require
repeated solutions of PDEs with different parameter settings, which is
computationally prohibitive especially for problems with high-dimensional
parameter spaces. Although recently proposed neural network methods make it
possible to obtain the optimal solutions simultaneously for different
parameters, challenges still remain when dealing with problems with complex
constraints. In this paper, we propose AONN, an adjoint-oriented neural network
method, to overcome the limitations of existing approaches in solving
parametric optimal control problems. In AONN, the neural networks are served as
parametric surrogate models for the control, adjoint and state functions to get
the optimal solutions all at once. In order to reduce the training difficulty
and handle complex constraints, we introduce an iterative training framework
inspired by the classical direct-adjoint looping (DAL) method so that penalty
terms arising from the Karush-Kuhn-Tucker (KKT) system can be avoided. Once the
training is done, parameter-specific optimal solutions can be quickly computed
through the forward propagation of the neural networks, which may be further
used for analyzing the parametric properties of the optimal solutions. The
validity and efficiency of AONN is demonstrated through a series of numerical
experiments with problems involving various types of parameters.
Pengfei Yin, Guangqiang Xiao, Kejun Tang, Chao Yang
AONN: an adjoint-oriented neural network method for all-at-once solutions of parametric optimal control problems+

###### Abstract

Parametric optimal control problems governed by partial differential equations (PDEs) are widely found in scientific and engineering applications. Traditional grid-based numerical methods for such problems generally require repeated solutions of PDEs with different parameter settings, which is computationally prohibitive especially for problems with high-dimensional parameter spaces. Although recently proposed neural network methods make it possible to obtain the optimal solutions simultaneously for different parameters, challenges still remain when dealing with problems with complex constraints. In this paper, we propose AONN, an adjoint-oriented neural network method, to overcome the limitations of existing approaches in solving parametric optimal control problems. In AONN, the neural networks are served as parametric surrogate models for the control, adjoint and state functions to get the optimal solutions all at once. In order to reduce the training difficulty and handle complex constraints, we introduce an iterative training framework inspired by the classical direct-adjoint looping (DAL) method so that penalty terms arising from the Karush-Kuhn-Tucker (KKT) system can be avoided. Once the training is done, parameter-specific optimal solutions can be quickly computed through the forward propagation of the neural networks, which may be further used for analyzing the parametric properties of the optimal solutions. The validity and efficiency of AONN is demonstrated through a series of numerical experiments with problems involving various types of parameters.

p -- 49M41, 49M05, 65N21

## 1 Introduction

Optimal control modeling has been playing an important role in a wide range of applications, such as aeronautics [56], mechanical engineering [52], haemodynamics [43], microelectronics [39], reservoir simulations [57], and environmental sciences [48]. Particularly, to solve a PDE-constrained optimal control problem, one needs to find an optimal control function that can minimize a given cost functional for systems governed by partial differential equations (PDEs). Popular approaches for solving PDE-constrained optimal control problems include the direct-adjoint looping (DAL) method [32, 21] that iteratively solves the adjoint systems, the Newton conjugate gradient method [47] that exploits the Hessian information, the semismooth Newton method [58, 10] that includes control and state constraints, and the alternating direction method of multipliers [8] designed for convex optimization. In practice, the cost functionals and PDE systems often entail different configurations of physical or geometrical parameters, leading to parametric optimal control modeling. These parameters usually arise from certain desired profiles such as material properties, boundary conditions, control constraints, and computational domains [22, 23, 48, 35, 46, 39, 31].

Most of the aforementioned methods cannot be directly applied to parametric optimal control problems. The main reason is that, in addition to the already costly process of solving the PDEs involved in the optimal control modeling, the presence of parameters introduces extra prominent complexity, making the parametric optimal control problems much more challenging than the nonparametric ones [18]. An efficient method for solving parametric optimal control problems is the reduced order model (ROM) [40, 43, 48, 36], which relies on surrogate models for the parametric model order reduction, and can provide both efficient and stable approximations if the solutions lie on a low-dimensional subspace [18]. However, because of the coupling of the spatial domain and the parametric domain, the discretization in ROM still suffers from the curse of dimensionality, thus is unable to obtain all-at-once solutions to parametric optimal control problems [22, 43, 40].

Numerical methods based on deep learning have been receiving increasingly more attentions in solving PDEs [41, 42, 12, 16, 59, 44, 45]. Recently, several successes have been made in solving PDE-constrained optimal control problems with deep-learning-based approaches. For example, a physics-informed neural network (PINN) method is designed to solve optimal control problems by adding the cost functional to the standard PINN loss [34, 29]. Meanwhile, deep-learning-based surrogate models [56, 30] and operator learning methods [55, 20] are proposed to achieve fast inference for the optimal control solution without intensive computations. Although these methods are successful for solving optimal control problems, few of which can be directly applied in parametric optimal control modeling. In a recent work [11], an extended PINN is proposed to augment neural network inputs with parameters, so that the Karush-Kuhn-Tucker (KKT) conditions and neural networks can be combined. In this
way, the optimal solution with a continuous range of parameters could be obtained for parametric optimal control problems with simple constraints. However, it is difficult for this method to generalize to solve more complex parametric optimal control problems, especially when the control function has additional inequality constraints [2, 3]. In such scenarios, too many penalty terms have to be introduced into the loss function to fit the complex KKT system, which is very hard to optimize [26]. A more detailed discussion of aforementioned methods can be found in Section 4.

To tackle the challenges in solving parametric optimal control problems and avoid the curse of dimensionality, we propose AONN, an adjoint-oriented neural network method that combines the advantages of both the classic DAL method and the deep learning technique. In AONN, we construct three neural networks with augmented parameter inputs, and integrate them into the framework of the DAL method to get an all-at-once approximation of the control function, the adjoint function, and the state function, respectively. On the one hand, neural networks enable the classic DAL framework to solve parametric problems simultaneously with the aid of random sampling rather than the discretization of the coupled spatial domain and parametric domain. On the other hand, unlike the PINN-based penalty methods [42, 29, 34, 11], the introduction of DAL avoids directly solving the complex KKT system with various penalty terms. Numerical results will show that, AONN can obtain high precision solutions to a series of parametric optimal control problems.

The remainder of the paper is organized as follows. In Section 2, the problem setting is introduced. After that, we will present the AONN framework in Section 3. Some further comparisons between AONN and several recently proposed methods are discussed in Section 4. Then, numerical results are presented in Section 5 to demonstrate the efficiency of the proposed AONN method. The paper is concluded in Section 6.

## 2 Problem setup

Let \(\boldsymbol{\mu}\in\mathcal{P}\subset\mathbb{R}^{D}\) denote a vector that collects a finite number of parameters. Let \(\Omega(\boldsymbol{\mu})\subset\mathbb{R}^{d}\) be a spatial domain depending on \(\boldsymbol{\mu}\), that is bounded, connected and with boundary \(\partial\Omega(\boldsymbol{\mu})\), and \(\mathbf{x}\in\Omega(\boldsymbol{\mu})\) denote a spatial variable. Consider the following parametric optimal control problem

\[\mathrm{OCP}(\boldsymbol{\mu}):\quad\left\{\begin{aligned} &\min_{(y(\mathbf{x}, \boldsymbol{\mu}),u(\mathbf{x},\boldsymbol{\mu}))\in Y\times U}J(y(\mathbf{x },\boldsymbol{\mu}),u(\mathbf{x},\boldsymbol{\mu});\boldsymbol{\mu}),\\ &\mathrm{s.t.}\ \ \mathbf{F}(y(\mathbf{x},\boldsymbol{\mu}),u( \mathbf{x},\boldsymbol{\mu});\boldsymbol{\mu})=0\ \ \mathrm{in}\ \Omega(\boldsymbol{\mu}),\ \mathrm{and}\ \ u(\mathbf{x},\boldsymbol{\mu})\in U_{ad}( \boldsymbol{\mu}),\end{aligned}\right. \tag{1}\]

where \(J:Y\times U\times\mathcal{P}\mapsto\mathbb{R}\) is a parameter-dependent objective functional, \(Y\) and \(U\) are two proper function spaces defined on \(\Omega(\boldsymbol{\mu})\), with \(y\in Y\) being the state function and \(u\in U\) the control function, respectively. Both \(y\) and \(u\) are dependent on \(\mathbf{x}\) and \(\boldsymbol{\mu}\). To simplify the notation, we denote \(y(\boldsymbol{\mu})=y(\mathbf{x},\boldsymbol{\mu})\) and \(u(\boldsymbol{\mu})=u(\mathbf{x},\boldsymbol{\mu})\). In \(\mathrm{OCP}(\boldsymbol{\mu})\) (1), \(\mathbf{F}\) represents the governing equation, such as, in our case, parameter-dependent PDEs, including the partial differential operator \(\mathbf{F}_{I}\) and the boundary operator \(\mathbf{F}_{B}\) (see Section 5 for examples). The admissible set \(U_{ad}(\boldsymbol{\mu})\) is a parameter-dependent bounded closed convex subset of \(U\), which provides an additional inequality constraint for \(u\), e.g., the box constraint \(U_{ad}(\boldsymbol{\mu})=\{u(\boldsymbol{\mu})\in U:u_{a}(\boldsymbol{\mu}) \leq u(\boldsymbol{\mu})\leq u_{b}(\boldsymbol{\mu})\}\).

Since the \(\mathrm{OCP}(\boldsymbol{\mu})\) (1) is a constrained minimization problem, the necessary condition for the minimizer \((y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu}))\) of (1) is the following KKT system [52, 9, 19]:

\[\begin{cases}J_{y}(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu}); \boldsymbol{\mu})-\mathbf{F}_{y}^{*}(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol {\mu});\boldsymbol{\mu})p^{*}(\boldsymbol{\mu})=0,\\ \mathbf{F}(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu});\boldsymbol{\mu})=0,\\ (\mathrm{d}_{u}J(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu});\boldsymbol{ \mu}),v(\boldsymbol{\mu})-u^{*}(\boldsymbol{\mu}))\geq 0,\ \forall v(\boldsymbol{\mu})\in U_{ad}( \boldsymbol{\mu}),\end{cases} \tag{2}\]

where \(p^{*}(\boldsymbol{\mu})\) is the adjoint function which is also known as the Lagrange multiplier, and \(\mathbf{F}_{y}^{*}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu});\boldsymbol{\mu})\) denotes the adjoint operator of \(\mathbf{F}_{y}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu});\boldsymbol{\mu})\). As \(y(\boldsymbol{\mu})\) can always be uniquely determined by \(u(\boldsymbol{\mu})\) through the state equation \(\mathbf{F}\), the total derivative of \(J\) with respect to \(u\) in (2) can be formulated as

\[\mathrm{d}_{u}J(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu}); \boldsymbol{\mu})=J_{u}(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu}); \boldsymbol{\mu})-\mathbf{F}_{u}^{*}(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol {\mu});\boldsymbol{\mu})p^{*}(\boldsymbol{\mu}). \tag{3}\]

The solution of \(\mathrm{OCP}(\boldsymbol{\mu})\) satisfies the system (2). So the key point is to solve this KKT system, based on which it is expected to find a minimizer for the \(\mathrm{OCP}(\boldsymbol{\mu})\). In general, it is not a trivial task to solve (2) directly, and solving the parametric PDE involved in the KKT system poses additional computational challenges (e.g. the discretization of parametric spaces). In this work, we focus on the deep learning method to solve (2). More specifically, we use three deep neural networks to approximate \(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu})\) and \(p^{*}(\boldsymbol{\mu})\) separately with an efficient training algorithm.

## 3 Methodology

Let \(\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y}\right), \hat{u}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{u}\right)\), and \(\hat{p}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{p}\right)\) be three independent deep neural networks parameterized with \(\boldsymbol{\theta}_{y},\boldsymbol{\theta}_{u}\) and \(\boldsymbol{\theta}_{p}\) respectively. Here, \(\mathbf{x}(\boldsymbol{\mu})\) is the augmented input of neural networks,
which is given by

\[\mathbf{x}(\boldsymbol{\mu})=\left[\begin{array}{cccccccc}x_{1},&\ldots,&x_{d},& \mu_{1},&\ldots,&\mu_{D}\end{array}\right].\]

We then use \(\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y}\right),\hat{u} \left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{u}\right)\), and \(\hat{p}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{p}\right)\) to approximate \(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu})\) and \(p^{*}(\boldsymbol{\mu})\) through minimizing three loss functions defined as

\[\mathcal{L}_{s}(\boldsymbol{\theta}_{y},\boldsymbol{\theta}_{u}) =\left(\frac{1}{N}\sum_{i=1}^{N}\left|r_{s}\left(\hat{y}(\mathbf{ x}(\boldsymbol{\mu})_{i};\boldsymbol{\theta}_{y}),\hat{u}(\mathbf{x}( \boldsymbol{\mu})_{i};\boldsymbol{\theta}_{u});\boldsymbol{\mu}_{i}\right) \right|^{2}\right)^{\frac{1}{2}}, \tag{1a}\] \[\mathcal{L}_{a}(\boldsymbol{\theta}_{y},\boldsymbol{\theta}_{u}, \boldsymbol{\theta}_{p}) =\left(\frac{1}{N}\sum_{i=1}^{N}\left|r_{a}\left(\hat{y}( \mathbf{x}(\boldsymbol{\mu})_{i};\boldsymbol{\theta}_{y}),\hat{u}(\mathbf{x}( \boldsymbol{\mu})_{i};\boldsymbol{\theta}_{u}),\hat{p}(\mathbf{x}(\boldsymbol {\mu})_{i};\boldsymbol{\theta}_{p});\boldsymbol{\mu}_{i}\right)\right|^{2} \right)^{\frac{1}{2}},\] (1b) \[\mathcal{L}_{u}(\boldsymbol{\theta}_{u},u_{\mathsf{step}}) =\left(\frac{1}{N}\sum_{i=1}^{N}\left|\hat{u}(\mathbf{x}( \boldsymbol{\mu})_{i};\boldsymbol{\theta}_{u})-u_{\mathsf{step}}(\mathbf{x}( \boldsymbol{\mu})_{i})\right|^{2}\right)^{\frac{1}{2}}, \tag{1c}\]

where \(\{\mathbf{x}(\boldsymbol{\mu})_{i}\}_{i=1}^{N}\) denote the collocation points. The functionals \(r_{s}\) and \(r_{a}\) represent the residuals for the state equation and the adjoint equation induced by the KKT conditions, i.e.,

\[r_{s}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu});\boldsymbol{\mu}) \triangleq\mathbf{F}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}); \boldsymbol{\mu}), \tag{2a}\] \[r_{a}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}),p(\boldsymbol{\mu} );\boldsymbol{\mu}) \triangleq J_{y}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}); \boldsymbol{\mu})-\mathbf{F}_{y}^{*}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}); \boldsymbol{\mu})p(\boldsymbol{\mu}), \tag{2b}\]

and \(u_{\mathsf{step}}(\mathbf{x}(\boldsymbol{\mu}))\) is an intermediate variable during the update procedure of the control function for the third variational inequality in the KKT conditions (2), which will be discussed in Section 3.2. These three loss functions try to fit the KKT conditions by adjusting the parameters of the three neural networks \(\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y}\right), \hat{u}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{u}\right)\), and \(\hat{p}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{p}\right)\), and the training procedure is performed in a sequential way. The derivatives involved in the loss functions can be computed efficiently by automatic differentiation in deep learning libraries such as TensorFlow [1] or PyTorch [38].

### Deep learning for parametric PDEs

The efficient solution of parametric PDEs is crucial for parametric optimal control modeling because of extra parameters involved in the physical system (1). To deal with parametric PDEs, we augment the input space of the neural networks by taking the parameter \(\boldsymbol{\mu}\) as additional inputs, along with the coordinates \(\mathbf{x}\) to handle the parameter-dependent PDEs. In addition, the penalty-free techniques [44, 27] are employed to enforce boundary conditions in solving parametric PDEs. Next, we illustrate how to apply penalty-free deep neural networks to solve the parametric state equation (2a), which can be directly generalized to the solution of the adjoint equation (2b).

The key point of the penalty-free method is to introduce two neural networks to approximate the solution, of which one neural network \(\hat{y}_{B}\) is used to approximate the essential boundary conditions and the other \(\hat{y}_{I}\) deals with the rest part of the computational domain. In this way, the training difficulties from the boundary conditions are eliminated, which improves the accuracy and robustness for complex geometries. For problems with simple geometries, we can also construct an analytical expression for \(\hat{y}_{B}\) to further reduce the training cost (see Section 5 for examples). The approximate solution of the state equation is constructed by

\[\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y}\right)=\hat{y }_{B}(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{B}})+\ell(\mathbf{x }(\boldsymbol{\mu}))\hat{y}_{I}\left(\mathbf{x}(\boldsymbol{\mu}); \boldsymbol{\theta}_{y_{I}}\right), \tag{3}\]

where \(\boldsymbol{\theta}_{y}=\{\boldsymbol{\theta}_{y_{B}},\boldsymbol{\theta}_{y_{I}}\}\) collects all parameters of two sub-neural networks \(\hat{y}_{B}\) and \(\hat{y}_{I}\), and \(\ell\) is a length factor function that builds the connection between \(\hat{y}_{B}\) and \(\hat{y}_{I}\), satisfying the following two conditions:

\[\left\{\begin{array}{l}\ell(\mathbf{x}(\boldsymbol{\mu}))>0,\quad\text{in} \;\Omega(\boldsymbol{\mu}),\\ \ell(\mathbf{x}(\boldsymbol{\mu}))=0,\quad\text{on}\;\partial\Omega( \boldsymbol{\mu}).\end{array}\right.\]

The details of constructing the length factor function \(\ell\) can be found in ref.[44]. With these settings, training \(\hat{y}_{B}\) and \(\hat{y}_{I}\) can be performed separately, i.e., one can first train \(\hat{y}_{B}\), and then fix \(\hat{y}_{B}\) to train \(\hat{y}_{I}\). For a fixed \(u(\boldsymbol{\mu})\), we have

\[\mathbf{F}(\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y} \right),u(\boldsymbol{\mu});\boldsymbol{\mu})=\left[\begin{array}{l} \mathbf{F}_{I}(\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y} \right),u(\boldsymbol{\mu});\boldsymbol{\mu})\\ \mathbf{F}_{B}(\hat{y}_{B}(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{B }}),u(\boldsymbol{\mu});\boldsymbol{\mu})\end{array}\right],\]

and the residual of the state equation can be rewritten as

\[r_{s}(\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y}\right),u (\boldsymbol{\mu});\boldsymbol{\mu})=\left[\begin{array}{l}r_{s_{I}}(\hat{y} \left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y}\right),u( \boldsymbol{\mu});\boldsymbol{\mu})\\ r_{s_{B}}(\hat{y}_{B}(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{B}}),u( \boldsymbol{\mu});\boldsymbol{\mu})\end{array}\right].\]
We then sample a set \(\{\mathbf{x}(\boldsymbol{\mu})_{i}\}_{i=1}^{N}\) of collocation points to optimize \(\boldsymbol{\theta}_{y}\) through minimizing the state loss function (1a) if \(\boldsymbol{\theta}_{u}\) is fixed.

For parametric problems, we take parameters as the additional inputs of neural networks. This approach is used to solve parametric forward problems [24] and control problems [49]. A typical way for sampling training points is to separately sample data in \(\Omega\) and \(\mathcal{P}\) to get \(\{\mathbf{x}_{i}\}\) and \(\{\boldsymbol{\mu}_{j}\}\), and then compose product data \(\{(\mathbf{x}_{i},\boldsymbol{\mu}_{j})\}\) in \(\Omega\times\mathcal{P}\). Rather than taken from each slice of \(\Omega(\boldsymbol{\mu})\) for a fixed \(\boldsymbol{\mu}\), collocation points are sampled in space \(\Omega_{\mathcal{P}}\) in this work, where

\[\Omega_{\mathcal{P}}=\{\mathbf{x}(\boldsymbol{\mu}):\mathbf{x}\in\Omega( \boldsymbol{\mu})\}\]

represents the joint spatio-parametric domain. The reason is that, for parametric geometry problems, as the spatial domain \(\Omega(\boldsymbol{\mu})\) is parameter-dependent, the sampling space cannot be expressed as the Cartesian product of \(\Omega\) and \(\mathcal{P}\).

### Projection gradient descent

Due to the additional inequality constraints \(u\in U_{ad}(\boldsymbol{\mu})\) for the control function, the zero gradient condition \(\mathrm{d}_{u}J(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{\mu});\boldsymbol{ \mu})=0\) cannot be directly applied to the optimal solution to get the update scheme for \(u\). One way to resolve this issue is to introduce additional Lagrange multipliers with some slack variables to handle the inequality constraints. However, this will bring additional penalty terms that could affect the procedure of optimization [7]. Furthermore, the inequality constraints also lead to the non-smoothness of the control function, making it more difficult to capture the singularity by penalty methods [29, 15]. To avoid these issues, we here use a simple iterative method to handle the variational inequality without utilizing a Lagrange multiplier, where a projection gradient descent method is employed, based on which we can obtain the update scheme for \(u\). The projection operator onto the admissible set \(U_{ad}(\boldsymbol{\mu})\) is defined as:

\[\mathbf{P}_{U_{ad}(\boldsymbol{\mu})}(u(\boldsymbol{\mu}))=\arg\min_{v( \boldsymbol{\mu})\in U_{ad}(\boldsymbol{\mu})}\|u(\boldsymbol{\mu})-v( \boldsymbol{\mu})\|_{2},\]

which performs the projection of \(u(\boldsymbol{\mu})\) onto the convex set \(U_{ad}(\boldsymbol{\mu})\). In practice, the above projection is implemented in a finite dimensional vector space, i.e., \(u(\boldsymbol{\mu})\) is discretized on a set of collocation points (e.g. grids on the domain \(\Omega\)). So it is straightforward to build this projection since \(U_{ad}(\boldsymbol{\mu})\) is a convex set. For example, if

\[U_{ad}(\boldsymbol{\mu})=\{u\in U:u_{a}(\mathbf{x}(\boldsymbol{\mu}))\leq u( \mathbf{x}(\boldsymbol{\mu}))\leq u_{b}(\mathbf{x}(\boldsymbol{\mu})), \forall\mathbf{x}\in\Omega(\boldsymbol{\mu})\} \tag{2}\]

provides a box constraint for \(u\), where \(u_{a}\) and \(u_{b}\) are the lower bound function and the upper bound function respectively, both of which are dependent on \(\boldsymbol{\mu}\), and \([u_{1},\ldots,u_{N}]^{\mathrm{T}}\) represents the control function values at \(N\) collocation points \(\{\mathbf{x}(\boldsymbol{\mu})_{i}\}_{i=1}^{N}\) in \(\Omega_{\mathcal{P}}\), then we can construct the projection \(\mathbf{P}_{U_{ad}(\boldsymbol{\mu})}\) in an entry-wise way [52, 19]:

\[\mathbf{P}_{U_{ad}(\boldsymbol{\mu})}(u_{i})=\begin{cases}u_{a}(\mathbf{x}( \boldsymbol{\mu})_{i}),\quad\text{if }u_{i}<u_{a}(\mathbf{x}(\boldsymbol{\mu})_{i}),\\ u_{i},\quad\text{if }u_{b}(\mathbf{x}(\boldsymbol{\mu})_{i})\geq u_{i}\geq u_{a}( \mathbf{x}(\boldsymbol{\mu})_{i}),\quad i=1,\ldots,N.\\ u_{b}(\mathbf{x}(\boldsymbol{\mu})_{i}),\quad\text{if }u_{i}>u_{b}(\mathbf{x}( \boldsymbol{\mu})_{i}).\end{cases} \tag{3}\]

The projection gradient step can be carried out according to the above formula, so as to obtain the update of the control function denoted by \(u_{\mathsf{step}}\), which is

\[u_{\mathsf{step}}(\boldsymbol{\mu})=\mathbf{P}_{U_{ad}(\boldsymbol{\mu})}\left( u(\boldsymbol{\mu})-\mathrm{cd}_{u}J(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}); \boldsymbol{\mu})\right), \tag{4}\]

and the loss for updating the control function is naturally defined as in (1c), making an approximation of \(\hat{u}\) obtained through minimizing (1c).

The optimal control function \(u^{*}(\boldsymbol{\mu})\) satisfies the following variational property:

\[u^{*}(\boldsymbol{\mu})-\mathbf{P}_{U_{ad}(\boldsymbol{\mu})}\left(u^{*}( \boldsymbol{\mu})-\mathrm{cd}_{u}J(y^{*}(\boldsymbol{\mu}),u^{*}(\boldsymbol{ \mu});\boldsymbol{\mu})\right)=0,\quad\forall c\geq 0.\]

Here \(\mathrm{d}_{u}J\) is associated with the adjoint function \(p^{*}(\boldsymbol{\mu})\) from total derivative expression (3), and thus we define the residual for the control function

\[r_{v}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}),p(\boldsymbol{\mu}),c;\boldsymbol {\mu})\triangleq u(\boldsymbol{\mu})-\mathbf{P}_{U_{ad}(\boldsymbol{\mu})} \left(u(\boldsymbol{\mu})-c\mathrm{d}_{u}J(y(\boldsymbol{\mu}),u(\boldsymbol{ \mu});\boldsymbol{\mu})\right), \tag{5}\]

and its corresponding variational loss is defined as

\[\mathcal{L}_{v}(\boldsymbol{\theta}_{y},\boldsymbol{\theta}_{u},\boldsymbol{ \theta}_{p},c)=\left(\frac{1}{N}\sum_{i=1}^{N}|r_{v}\left(\hat{y}(\mathbf{x}( \boldsymbol{\mu})_{i};\boldsymbol{\theta}_{y}),\hat{u}(\mathbf{x}(\boldsymbol{ \mu})_{i};\boldsymbol{\theta}_{u}),\hat{p}(\mathbf{x}(\boldsymbol{\mu})_{i}; \boldsymbol{\theta}_{p}),c;\boldsymbol{\mu}_{i}\right)|^{2}\right)^{\frac{1}{2}}. \tag{6}\]
The first two losses in (10) together with (11) reflect how well \(y(\mathbf{\mu}),u(\mathbf{\mu})\) and \(p(\mathbf{\mu})\) approximate the optimal solution governed by the KKT system (2). Note that \(r_{v}\) and \(\mathcal{L}_{v}\) are dependent on the constant \(c\), which is actually the step size for gradient descent. For verification, the variational loss is constructed to verify the convergence of algorithm, and \(c\) is often chosen as the last step size.

### AONN algorithm

Now putting all together, we are ready to present our algorithm. Our goal is to efficiently approximate the minimizer of (1) via adjoint-oriented neural networks (AONN). The overall training procedure of AONN consists of three steps: training \(\hat{y}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{y}\right)\), updating \(\hat{p}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{p}\right)\) and refining \(\hat{u}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}\right)\). The schematic of AONN for solving the parametric optimal control problems is shown in Figure 1. The three neural networks \(\left(\hat{y}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{y}\right),\hat{p}\left( \mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{p}\right),\hat{u}\left(\mathbf{x}(\mathbf{\mu} );\mathbf{\theta}_{u}\right)\right)\) with augmented parametric input, as illustrated in panels A and B, are optimized to iteratively minimizing the objective functional with respect to the corresponding variables once at a time. More specifically, according to the loss functions derived by the three equations shown in panel C, the training procedure is performed as in panel D.

Starting with three initial neural networks \(\hat{y}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{y}^{0}\right),\hat{p}\left( \mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{p}^{0}\right)\), and \(\hat{u}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}^{0}\right)\), we train and obtain the state function \(\hat{y}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{y}^{1}\right)\) through minimizing \(\mathcal{L}_{s}\left(\mathbf{\theta}_{y},\mathbf{\theta}_{u}^{0}\right)\) (see (11a)), which is equivalent to solving the parameter-dependent state equation. With \(\hat{y}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{y}^{1}\right)\), we minimize the loss \(\mathcal{L}_{a}\left(\mathbf{\theta}_{y}^{1},\mathbf{\theta}_{u}^{0},\mathbf{\theta}_{p}\right)\) (see (11b)) for the adjoint equation to get \(\hat{p}\left(\mathbf{x};\mathbf{\theta}_{p}^{1}\right)\), corresponding to solving the parameter-dependent adjoint equation. To update the control function \(\hat{u}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}\right)\), \(u_{\mathsf{step}}^{0}(\mathbf{x}(\mathbf{\mu}))\) is computed first by gradient descent followed by a projection step (see (12)), and then \(\hat{u}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}^{1}\right)\) is obtained by minimizing \(\mathcal{L}_{u}(\mathbf{\theta}_{u},u_{\mathsf{step}}^{0}(\mathbf{x}(\mathbf{\mu})))\) (see (11c)). Then another iteration starts using \(\mathbf{\theta}_{y}^{1},\mathbf{\theta}_{p}^{1},\mathbf{\theta}_{u}^{1}\) as the initial parameters. In general, the iterative scheme is specified as follows:

\[\text{training }\hat{y}: \mathbf{\theta}_{y}^{k}=\arg\min_{\mathbf{\theta}_{y}}\mathcal{L}_{s} \left(\mathbf{\theta}_{y},\mathbf{\theta}_{u}^{k-1}\right),\] \[\text{updating }\hat{p}: \mathbf{\theta}_{p}^{k}=\arg\min_{\mathbf{\theta}_{p}}\mathcal{L}_{a} \left(\mathbf{\theta}_{y}^{k},\mathbf{\theta}_{u}^{k-1},\mathbf{\theta}_{p}\right),\] \[\text{refining }\hat{u}: \mathbf{\theta}_{u}^{k}=\arg\min_{\mathbf{\theta}_{u}}\mathcal{L}_{u} \left(\mathbf{\theta}_{u},u_{\mathsf{step}}^{k-1}\right),\]

where

\[u_{\mathsf{step}}^{k-1}(\mathbf{x}(\mathbf{\mu}))=\mathbf{P}_{U_{ad}(\mathbf{\mu})} \left(\hat{u}(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}^{k-1})-c^{k}\mathrm{d}_{u}J (\hat{y}(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{y}^{k}),\hat{u}(\mathbf{x}(\mathbf{\mu} );\mathbf{\theta}_{u}^{k-1});\mathbf{\mu})\right), \tag{12}\]

and

\[\begin{split}\mathrm{d}_{u}J\left(\hat{y}(\mathbf{x}(\mathbf{\mu}); \mathbf{\theta}_{y}^{k}),\hat{u}(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}^{k-1});\mathbf{ \mu}\right)&=J_{u}\left(\hat{y}(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_ {y}^{k}),\hat{u}(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}^{k-1});\mathbf{\mu}\right)\\ &-\mathbf{F}_{u}^{*}\left(\hat{y}(\mathbf{x}(\mathbf{\mu});\mathbf{\theta} _{y}^{k}),\hat{u}(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}^{k-1});\mathbf{\mu}\right) \hat{p}(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{p}^{k}).\end{split} \tag{13}\]

The iteration of AONN revolves around the refinement of \(\hat{u}\) with the aid of \(\hat{y}\) and \(\hat{p}\), forming the direct-adjoint looping (DAL),which is indicated by red lines in Figure 1. This procedure shares similarities to the classical DAL framework, but there is a crucial difference between AONN and DAL. That is, a reliable solution of OCP(\(\mathbf{\mu}\)) for any parameter can be efficiently computed from the trained neural networks in our AONN framework, while DAL cannot achieve this. More details can be found in the discussions of Section 4.

The training process is summerized in Algorithm 1, where the loss function \(\mathcal{L}_{v}(\mathbf{\theta}_{y},\mathbf{\theta}_{u},\mathbf{\theta}_{p},c)\) (see (11)) is used for the verification. In our practical implementation, we employ the step size decay technique with a decay factor \(\gamma\) for robustness. The AONN method can be regarded as an inexact DAL to some extent since the state equation and the adjoint equation are not accurately solved but approximated with neural networks at each iteration. So the number of epochs is increased by \(n_{\mathrm{aug}}\) compared with the previous step (on line 9 of Algorithm 1) to ensure the accuracy and convergence. It is worth noting that the training of network \(\hat{u}\left(\mathbf{x}(\mathbf{\mu});\mathbf{\theta}_{u}^{*}\right)\) can be put after the while loop, if the collocation points are always fixed, since training the state function only uses the value of \(u\) at the collocation points (the calculation of line 5 in Algorithm 1).
**Remark 3.1**.: A post-processing step can be applied to continue training \(\mathbf{\theta}_{y}^{*}\) until a more accurate solution of the state function (or the adjoint function) is found. That is, we can fix \(\hat{u}\left(\mathbf{x};\mathbf{\theta}_{u}^{*}\right)\) and train \(\mathbf{\theta}_{y}^{*}\) by minimizing the state loss (3.1a). We can also fix \(\mathbf{\theta}_{y}^{*},\mathbf{\theta}_{u}^{*}\) and train \(\mathbf{\theta}_{p}^{*}\) using (3.1b). The initial step size \(c^{0}\) is crucial for the convergence of Algorithm 1. A large step size may lead to divergence of the algorithm, while a small one could result in slow convergence.

## 4 Comparison with other methods

Unlike solving the deterministic optimal control problems, the existence of parameters in \(\mathrm{OCP}(\mathbf{\mu})\) causes difficulties for traditional grid-based numerical methods. A straightforward way is to convert the \(\mathrm{OCP}(\mathbf{\mu})\) into the deterministic optimal control problem. For each realization of

Figure 1: The schematic of AONN for solving the parametric optimal control problems. (A) _Spatial coordinates and parameters form the input of neural networks._ (B) _AON consists of three separate neural networks \(\hat{y},\hat{p},\hat{u}\) and return the approximation of state, adjoint and control respectively._ (C) _The state equation, the adjoint equation and the projected gradient equation are derived to formulate the corresponding loss functions._ (D) _The gradients in the state PDE and the adjoint PDE are computed via automatic differentiation [38]. \(\hat{y},\hat{p},\hat{u}\) are then trained sequentially via the Adam [25] or the BFGS optimizer.
parameters, the OCP(\(\boldsymbol{\mu}\)) is reduced to the following

\[\text{OCP}:\quad\begin{cases}\min_{(y,u)\in Y\times U}J(y,u),\\ \text{s.t. }\mathbf{F}(y,u)=0\ \ \text{in}\ \Omega,\ \ \text{and}\ u\in U_{ad}. \end{cases} \tag{11}\]

The classical direct-adjoint looping (DAL) method [32, 21] is a popular approach for solving this problem, where an iterative scheme is adopted to converge toward the optimal solution by solving subproblems in the KKT system with numerical solvers (e.g. finite element methods). At each iteration in the direct-adjoint looping procedure, one first solves the governing PDE (10a) and then solves the adjoint PDE (10b) which formulates the total gradient (10c) for the update of the control function.

\[\mathbf{F}(y,u)=0, \tag{12a}\] \[J_{y}(y,u)-\mathbf{F}_{y}^{*}(y,u)p=0,\] (12b) \[\text{d}_{u}J(y,u)=J_{u}(y,u)-\mathbf{F}_{u}^{*}(y,u)p. \tag{12c}\]

Despite its effectiveness, DAL is not able to handle the OCP(\(\boldsymbol{\mu}\)) problem, directly due to the curse of dimensionality of the discretization of \(\Omega_{\mathcal{P}}\). An alternative strategy is the reduced order model (ROM) [36], which rely on surrogate models for parameter-dependent PDEs. The idea is that the solution of PDE for any parameter can be computed based on a few basis functions that are constructed from the solutions corresponding to some pre-selected parameters. However, it is still computationally unaffordable for ROM when the parameter-induced solution manifold does not lie on a low-dimensional subspace.

Recently, some deep learning algorithms are used to solve the optimal control problem for a fixed parameter [29, 34]. By introducing two deep neural networks, the state function \(y\) and the control function \(u\) can be approximated by minimizing the following objective functional:

\[\min_{(y,u)\in Y\times U}J(y,u)+\beta_{1}\mathbf{F}(y,u)^{2}+\beta_{2}\|u- \mathbf{P}_{U_{ad}}(u)\|_{U}, \tag{13}\]

where two penalty terms are added, and \(\beta=(\beta_{1},\beta_{2})\) are two parameters that need tuning. As the penalty parameters increase to \(+\infty\), the solution set of this unconstrained problem approaches to the solution set of the constrained one. However, this penalty approach has a serious drawback. On the one hand, as the penalty parameters increase, the optimal solution becomes increasingly difficult to obtain. On the other hand, the constraint is not satisfied well if the penalty parameter is small. To alleviate this difficulty, one can use hPINN [29] which employs the augmented Lagrangian method to solve (13). However, it is still challenging to directly extend this approach to OCP(\(\boldsymbol{\mu}\)) due to the presence of parameters. This is because it is extremely hard to optimize a series of objective functionals with a continuous range of parameters simultaneously.

### PINN for OCP(\(\boldsymbol{\mu}\))

For handling the parametric optimal control problems, an extended PINN method [11] with augmented inputs is used to obtain a more accurate parametric prediction. That is, the inputs of the neural networks consist of two parts: the spatial coordinates and the parameters. The optimal solution for any parameter is approximated by a deep neural network that is obtained from solving the parameter-dependent KKT system (2). In particular, when there is no restriction on the control function \(u(\boldsymbol{\mu})\) (e.g., \(U_{ad}(\boldsymbol{\mu})\) is the full Banach space), the KKT system is

\[\mathcal{F}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}),p(\boldsymbol{\mu}); \boldsymbol{\mu})=\begin{bmatrix}J_{y}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}) ;\boldsymbol{\mu})-\mathbf{F}_{y}^{*}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu}) ;\boldsymbol{\mu})p(\boldsymbol{\mu})\\ \text{F}(y(\boldsymbol{\mu}),u(\boldsymbol{\mu});\boldsymbol{\mu})\\ \text{d}_{u}J(y(\boldsymbol{\mu}),u(\boldsymbol{\mu});\boldsymbol{\mu})\end{bmatrix} =\mathbf{0}, \tag{14}\]

where the total gradient \(d_{u}J\) is given in (3). In such cases, one can use the PINN algorithm to obtain the optimal solution through minimizing the least-square loss derived from the KKT system. Nevertheless, to apply this method to the cases where there are some additional constraints on the control function \(u\), such as the box constraint (11), one may need to introduce the Lagrange multipliers \(\lambda(\boldsymbol{\mu})=(\lambda_{a}(\boldsymbol{\mu}),\lambda_{b}( \boldsymbol{\mu}))\) corresponding to the
constraints \(u(\mathbf{\mu})\geq u_{a}(\mathbf{\mu})\) and \(u(\mathbf{\mu})\leq u_{b}(\mathbf{\mu})\). For such cases, the KKT system is

\[\mathcal{F}(y(\mathbf{\mu}),u(\mathbf{\mu}),p(\mathbf{\mu}),\lambda(\mathbf{\mu}); \mathbf{\mu}) =\begin{bmatrix}J_{y}(y(\mathbf{\mu}),u(\mathbf{\mu});\mathbf{\mu})-\mathbf{F }_{y}^{*}(y(\mathbf{\mu}),u(\mathbf{\mu});\mathbf{\mu})p(\mathbf{\mu})\\ \mathbf{F}(y(\mathbf{\mu}),u(\mathbf{\mu});\mathbf{\mu})\\ \mathrm{d}_{u}J(y(\mathbf{\mu}),u(\mathbf{\mu});\mathbf{\mu})-\lambda_{a}(\mathbf{\mu})+ \lambda_{b}(\mathbf{\mu})\\ \lambda_{a}(\mathbf{\mu})(u_{a}(\mathbf{\mu})-u(\mathbf{\mu}))\\ \lambda_{b}(\mathbf{\mu})(u_{b}(\mathbf{\mu})-u(\mathbf{\mu}))\end{bmatrix}=\mathbf{0},\] (11a) and \[\quad\begin{cases}u_{a}(\mathbf{\mu})\leq u(\mathbf{\mu})\leq u_{b}(\mathbf{ \mu}),\\ \lambda_{a}(\mathbf{\mu})\geq 0,\lambda_{b}(\mathbf{\mu})\geq 0.\end{cases} \tag{11b}\]

Applying the framework of PINN to solve the system (11) needs to deal with several penalty terms in the loss function including the penalties of equality terms (11a) and inequality terms (11b), leading to an inaccurate solution even for the problem with fixed parameters, which will be presented in the next section. In addition, the extra constraint of \(U_{ad}(\mathbf{\mu})\) often introduces inequality terms and nonlinear terms and brings more singularity to the optimal control function [2, 3], which limits the application of PINN for solving OCP(\(\mathbf{\mu}\)) with control constraints.

### PINN+Projection for OCP(\(\mathbf{\mu}\))

To find a better baseline for comparison, we propose to improve the performance of PINN by introducing a projection operator. In this way, the KKT system (2) can be reformulated to a more compactly stated condition [37]:

\[\mathcal{F}(y(\mathbf{\mu}),u(\mathbf{\mu}),p(\mathbf{\mu}),c;\mathbf{\mu})=\begin{bmatrix}J_ {y}(y(\mathbf{\mu}),u(\mathbf{\mu});\mathbf{\mu})-\mathbf{F}_{y}^{*}(y(\mathbf{\mu}),u(\mathbf{ \mu});\mathbf{\mu})p(\mathbf{\mu})\\ \mathbf{F}(y(\mathbf{\mu}),u(\mathbf{\mu});\mathbf{\mu})\\ u(\mathbf{\mu})-\mathbf{P}_{U_{ad}(\mathbf{\mu})}\left(u(\mathbf{\mu})-c\mathrm{d}_{u}J(y( \mathbf{\mu}),u(\mathbf{\mu});\mathbf{\mu})\right)\end{bmatrix}=\mathbf{0}, \tag{12}\]

where \(c\) could be any positive number. Note that choosing an appropriate \(c\) can accelerate the convergence of the algorithm. For example, the classic way is to choose \(c=1/\alpha\) for canceling out the control function \(u\) inside the projection operator, where \(\alpha\) is the coefficient of the Tikhonov regularization term (see the experiment in Section 5.1). The complementary conditions and inequalities caused by the control constraints are avoided in (12), thus significantly reducing the difficulty of optimization. In this paper, we call the method PINN+Projection, which combines the projection strategy with the KKT system to formulate the PINN residual loss. Although PINN+Projection alleviates the solving difficulty brought by control constraints to PINN, it still has limitations on nonsmooth optimal control problems. For nonsmooth optimization such as sparse \(L_{1}\)-minimization, the KKT system can no longer be described by (12) because of the nondifferentiable property of the \(L_{1}\)-norm [9]. Instead, the dual multiplier for the \(L_{1}\)-cost term is required. The difficulty arises from the third nonsmooth variational equation of (12), which makes the neural network difficult to train. AONN reduces this difficulty by leveraging the update scheme in the DAL method without the implicit variational equation in (12). Numerical results also show that the KKT system (12) cannot be directly used to formulate the loss functions of neural networks. Such results of \(L_{1}\)-minimization involved in OCP(\(\mathbf{\mu}\)) (see Section 5.5) strongly suggest that AONN is a more reliable and efficient framework.

The proposed AONN method has all the advantages of the aforementioned approaches while avoiding their drawbacks. By inheriting the structure of DAL, the AONN method can obtain an accurate solution through solving the KKT system in an alternative minimization iterative manner. So it does not require the Lagrange multipliers corresponding to the additional control constraints and thus can reduce the storage cost as well as improve the accuracy. Moreover, AONN can accurately approximate the optimal solutions of parametric optimal control problems for any parameter and can be generalized to cases with high-dimensional parameters.

## 5 Numerical study

In this section, we present results of five numerical experiments to illustrate the effectiveness of AONN, where different types of PDE constraints, objective functionals and control constraints under different parametric settings are studied. In the following, AONN is first validated by solving OCP, and further applied to solving OCP(\(\mathbf{\mu}\)) with continuous parameters changing over a specific interval. For comparison purposes, we also use the PINN method and the PINN+Projecton method to solve OCP(\(\mathbf{\mu}\)). We employ the ResNet model [17] with sinusoid activation functions to build the neural networks for AONN and other neural network based algorithms. Unless otherwise specified, the quasi Monte-Carlo method is used to generate collocation points from \(\Omega_{\mathcal{P}}\) by calling the SciPy module [53]. Analytical length factor functions (see (11)) are
constructed for all test problems to make the approximate solution naturally satisfy Dirichlet boundary conditions. The training of neural networks is performed on a Geforce RTX 2080 GPU with PyTorch 1.8.1. The Broyden-Fletcher-Goldfarb-Shanno (BFGS) algorithm with a strong Wolfe line search strategy is used to update the neural network parameters to speed up the convergence, where the maximal number of iterations for BFGS is set to 100.

### Test 1: Optimal control for the semilinear elliptic equations

We start with the following nonparametric optimal control problem:

\[\left\{\begin{aligned} &\min_{y,u}J(y,u):=\frac{1}{2}\left\|y-y_{d} \right\|_{L_{2}(\Omega)}^{2}+\frac{\alpha}{2}\|u\|_{L_{2}(\Omega)}^{2},\\ &\text{subject to }\left\{\begin{aligned} -\Delta y+y^{3}=u+f& \text{ in }\Omega\\ y=0&\text{ on }\partial\Omega,\\ \end{aligned}\right.\\ &\text{and }\quad u_{a}\leq u\leq u_{b}\quad\text{ a.e. in }\Omega.\end{aligned}\right. \tag{10}\]

The total derivative of \(J\) with respect to \(u\) is \(\mathrm{d}_{u}J(y,u)=\alpha u+p\), where \(p\) is the solution of the corresponding adjoint equation:

\[\left\{\begin{aligned} -\Delta p+3py^{2}&=y-y_{d}& \text{ in }\Omega,\\ p&=0&\text{ on }\partial\Omega.\end{aligned}\right. \tag{11}\]

We take the same configuration as in ref.[14], where \(\Omega=(0,1)^{2}\), \(\alpha=0.01\), \(u_{a}=0\), and \(u_{b}=3\). The analytical optimal solution is given by

\[y^{*} =\sin\left(\pi x_{1}\right)\sin\left(\pi x_{2}\right), \tag{12}\] \[u^{*} =\mathbf{P}_{[u_{a},u_{b}]}(2\pi^{2}y^{*}),\] \[p^{*} =-2\alpha\pi^{2}y^{*},\]

where \(\mathbf{P}_{[u_{a},u_{b}]}\) is the pointwise projection operator onto the interval \([u_{a},u_{b}]\). The desired state \(y_{d}=(1+4\pi^{4}\alpha)y^{*}-3y^{*2}p^{*}\) and the source term \(f=2\pi^{2}y^{*}+y^{*3}-u^{*}\) are given to satisfy the state equation and the adjoint equation.

To solve the optimal control problem with AONN, we construct three networks \(\hat{y}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{I}} \right),\hat{p}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{ I}}\right)\) and \(\hat{u}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{u}\right)\), whose network structures are all comprised of two ResNet blocks, each of which contains two fully connected layers with 15 neurons and a residual connection. We randomly sample \(N=4096\) points inside \(\Omega\) to form the training set. A uniform meshgrid with size \(256\times 256\) in \(\Omega\) is generated for testing and visualization. We use fixed step size and training epochs in subproblems, i.e. \(c^{k}\equiv 1/\alpha=100,n^{k}\equiv 500\). The loss behavior and the relative error \(\|u-u^{*}\|/\|u^{*}\|\) with \(\ell_{2}\)-norm and \(\ell_{\infty}\)-norm are reported in Figure 1(a), while Figure 1(b) evaluates the difference between the AONN solution and the analytical solution. As reported in ref.[14], to achieve the error \(1\times 10^{-4}\) in the \(\ell_{2}\) sense requires 7733 degrees of freedom with the finite element method, while the AONN method needs only 781 neural network parameters to approximate the control function.

Test 2: Optimal control for the semilinear elliptic equations with control constraint parametrization

We then consider the same optimal control problem with control constraint parametrization. The control constraint upper bound \(u_{b}\) is set to be a continuous variable \(\boldsymbol{\mu}\) ranging from 3 to 20 instead of a fixed number. Thus (10) actually constructs a series of optimal control problems and the optimal solutions (12) are dependent on \(\boldsymbol{\mu}\). We now verify whether the all-at-once solutions can be obtained by AONN when \(\boldsymbol{\mu}\) changes continuously over the interval \(3\leq\boldsymbol{\mu}\leq 20\). We seek optimal \(y(\boldsymbol{\mu}),u(\boldsymbol{\mu})\) defined by the following problem:

\[\left\{\begin{aligned} &\min_{y(\boldsymbol{\mu}),u( \boldsymbol{\mu})}J(y(\boldsymbol{\mu}),u(\boldsymbol{\mu})):=\frac{1}{2}\left\| y(\boldsymbol{\mu})-y_{d}(\boldsymbol{\mu})\right\|_{L_{2}(\Omega)}^{2}+ \frac{\alpha}{2}\|u(\boldsymbol{\mu})\|_{L_{2}(\Omega)}^{2},\\ &\text{subject to }\left\{\begin{aligned} -\Delta y( \boldsymbol{\mu})+y(\boldsymbol{\mu})^{3}&=u(\boldsymbol{\mu})+f( \boldsymbol{\mu})&\text{ in }\Omega\\ y(\boldsymbol{\mu})&=0&\text{ on }\partial\Omega,\\ \end{aligned}\right.\\ &\text{and }\quad u_{a}\leq u(\boldsymbol{\mu})\leq \boldsymbol{\mu}&\text{ a.e. in }\Omega.\end{aligned}\right. \tag{13}\]

To naturally satisfy the homogeneous Dirichlet boundary conditions in the state equation and the adjoint equation, three neural networks for approximating the AONN solutions of OCP(\(\boldsymbol{\mu}\)) (13) are defined as follows:

\[\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_ {I}}\right) =\ell(\mathbf{x})\hat{y}_{I}\left(\mathbf{x}(\boldsymbol{\mu}); \boldsymbol{\theta}_{y_{I}}\right), \tag{14}\] \[\hat{p}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_ {I}}\right) =\ell(\mathbf{x})\hat{p}_{I}\left(\mathbf{x}(\boldsymbol{\mu}); \boldsymbol{\theta}_{y_{I}}\right),\] \[\hat{u}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_ {I}}\right) =\hat{u}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{u} \right),\]
where the length factor function is formed by

\[\ell(\mathbf{x})=x_{0}(1-x_{0})x_{1}(1-x_{1}). \tag{10}\]

The network structures of \(\hat{y}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{I}}\right), \hat{p}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{p_{I}}\right)\) and \(\hat{u}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{u}\right)\) are the same as those of the previous test except for the input dimension being \(3\) and the number of neurons in each hidden layer being \(20\), resulting in \(1361\) undecided parameters. To evaluate the loss, we sample \(N=20480\) points in the spatio-parametric space \(\Omega_{\mathcal{P}}\). We keep the same step size \(c^{k}\) and training epochs \(n^{k}\) as the previous test, and perform \(N_{\text{iter}}=20\) iterations until Algorithm 1 converges. The test errors are computed on the uniform meshgrid with size \(256\times 256\) for each realization of \(\boldsymbol{\mu}\).

In Figure 2, we plot the analytical solutions, the AONN solutions and the PINN solutions for eight equidistant realizations of \(\boldsymbol{\mu}\), where it can be seen that the AONN solutions are better than the PINN solutions in the sense of absolute error. Looking more closely, the large errors are concentrated around the location of singularity of \(u\), i.e., the curve of active constraints \(\left\{\mathbf{x}:u(\mathbf{x}(\boldsymbol{\mu}))=\boldsymbol{\mu}\right\}\), except for the case \(\boldsymbol{\mu}=20\) where the inequality constraint is nonactive, keeping the smoothness of the optimal control function. Note that adaptive sampling strategies [51, 50, 13] may be used to improve the accuracy in the singularity region, which will be left for future study.

### Test 3: Optimal control for the Navier-Stokes equations with physical parametrization

The next test case is the parametric optimal control problem

\[\min_{y(\boldsymbol{\mu}),u(\boldsymbol{\mu})}J(y(\boldsymbol{\mu}),u( \boldsymbol{\mu}))=\frac{1}{2}\left\|y(\boldsymbol{\mu})-y_{d}(\boldsymbol{ \mu})\right\|_{L_{2}(\Omega)}^{2}+\frac{1}{2}\|u(\boldsymbol{\mu})\|_{L_{2}( \Omega)}^{2}, \tag{11}\]

subject to the following steady-state incompressible Navier-Stokes (NS) equations:

\[\left\{\begin{aligned} -\boldsymbol{\mu}\Delta y(\boldsymbol{\mu})+(y( \boldsymbol{\mu})\cdot\nabla)y(\boldsymbol{\mu})+\nabla p(\boldsymbol{\mu})& =u(\boldsymbol{\mu})+f(\boldsymbol{\mu})&\text{ in }\Omega,\\ \text{div}\,y(\boldsymbol{\mu})&=0&\text{ in }\Omega,\\ y(\boldsymbol{\mu})&=0&\text{ on }\partial\Omega,\end{aligned}\right. \tag{12}\]

in \(\Omega=(0,1)^{2}\) with parameter \(\boldsymbol{\mu}\) representing the reciprocal of the Reynolds number. Note that the nonparametric problems without control constraint for \(\boldsymbol{\mu}=0.1\) and \(\boldsymbol{\mu}=1.0\) were studied in refs.[28, 54]. We set the physical parameter \(\boldsymbol{\mu}\in[0.1,100]\) and in addition, we consider the following constraint for \(u(\boldsymbol{\mu})=(u_{1}(\boldsymbol{\mu}),u_{2}(\boldsymbol{\mu}))\):

\[u_{1}(\boldsymbol{\mu})^{2}+u_{2}(\boldsymbol{\mu})^{2}\leq r^{2}, \tag{13}\]

Figure 1: Test 1: training loss and test error of problem 1. Test error is evaluated at 256\(\times\)256 grid points. (a) Loss behaviour measured in terms of (10a)-(10b), and test errors in both \(\ell_{2}\)-norm and \(\ell_{\infty}\)-norm during training process. (b) The AONN solution and its absolute errors compared with the analytical solution.
with \(r=0.2\), posing additional challenges to this problem. The desired state \(y_{d}(\mathbf{\mu})\) and the source term \(f(\mathbf{\mu})\) are given in advance to ensure that the analytical solution of the above \(\mathrm{OCP}(\mathbf{\mu})\) is given by

\[y^{*}(\mathbf{\mu}) =e^{-0.05\mathbf{\mu}}\left(\begin{array}{c}\sin^{2}\pi x_{1}\sin \pi x_{2}\cos\pi x_{2}\\ -\sin^{2}\pi x_{2}\sin\pi x_{1}\cos\pi x_{1}\end{array}\right),\] \[\lambda^{*}(\mathbf{\mu}) =\left(e^{-0.05\mathbf{\mu}}-e^{-\mathbf{\mu}}\right)\left(\begin{array} []{c}\sin^{2}\pi x_{1}\sin\pi x_{2}\cos\pi x_{2}\\ -\sin^{2}\pi x_{2}\sin\pi x_{1}\cos\pi x_{1},\end{array}\right).\]

The adjoint equation is specified as

\[\left\{\begin{aligned} -\mathbf{\mu}\Delta\lambda(\mathbf{\mu})-(y(\mathbf{\mu}) \cdot\nabla)\lambda(\mathbf{\mu})+(\nabla y(\mathbf{\mu}))^{T}\lambda(\mathbf{\mu})+ \nabla\nu(\mathbf{\mu})&=y(\mathbf{\mu})-y_{d}(\mathbf{\mu})&\text{ in }\Omega,\\ \mathrm{div}\,\lambda(\mathbf{\mu})&=0&\text{ in }\Omega,\\ \lambda(\mathbf{\mu})&=0&\text{ on }\partial\Omega,\end{aligned}\right. \tag{5.10}\]

where \(\lambda(\mathbf{\mu})\) denotes the adjoint velocity and \(\nu(\mathbf{\mu})\) denotes the adjoint pressure. The optimal pressure and adjoint pressure \(p^{*}(\mathbf{\mu}),\nu^{*}(\mathbf{\mu})\) are both zero. In order to satisfy the state equation and the adjoint equation, \(y_{d}(\mathbf{\mu})\) and \(f(\mathbf{\mu})\) are chosen as

\[f(\mathbf{\mu}) =-\mathbf{\mu}\Delta y^{*}(\mathbf{\mu})+(y^{*}(\mathbf{\mu})\cdot\nabla)y^{* }(\mathbf{\mu})-u^{*}(\mathbf{\mu}), \tag{5.11}\] \[y_{d}(\mathbf{\mu}) =y^{*}(\mathbf{\mu})-\left(-\mathbf{\mu}\Delta\lambda^{*}(\mathbf{\mu})+(y^{* }(\mathbf{\mu})\cdot\nabla)\lambda^{*}(\mathbf{\mu})-(\nabla y^{*}(\mathbf{\mu}))^{T} \lambda^{*}(\mathbf{\mu})\right).\]

It is easy to check that the optimal control is \(u^{*}(\mathbf{\mu})=\mathbf{P}_{B(0,r)}(\lambda^{*}(\mathbf{\mu}))\), where \(B(0,r)\) is a ball centered at the origin of radius \(r\). The state equation (5.8) and the adjoint equation (5.10) together with the variational inequality where \(\mathrm{d}_{u}J(\mathbf{\mu})=u(\mathbf{\mu})-\lambda(\mathbf{\mu})\) formulate the optimality system.

For AONN, we use a neural network to approximate \(y\), and it is constructed by two ResNet blocks, each of which contains two fully connected layers with 20 units and a residual connection, resulting in 1382 parameters. The neural network for approximating \(p\) has two ResNet blocks built by two fully connected layers with 10 units, resulting in 381 parameters. The architectures of the neural networks for \(\lambda\) and \(\nu\) are the same as those of \(y\) and \(p\) respectively. We select \(N=20000\) randomly sampled points in the spatio-parametric space \(\Omega_{\mathcal{P}}\). The maximum iteration number in Algorithm 3.1 is set to \(N_{\mathrm{iter}}=300\) and the step size is \(c^{k}\equiv c^{0}=1.0\). We choose an initial training epoch \(n^{0}=200\) and increase it by \(n_{\mathrm{aug}}=100\) after every 100 iterations. For the PINN method, the architectures of the neural networks are the same as those of AONN except for adding another neural network

Figure 5.2: Test 2: the control solutions \(u(\mathbf{\mu})\) of AONN and PINN with eight realizations of \(\mathbf{\mu}\in[3,20]\), and their absolute errors.
for \(\zeta(\mathbf{\mu})\) to satisfy the following KKT system:

\[\left\{\begin{array}{l}\text{state equation }(\ref{eq:KKT}),\\ \text{adjoint equation }(\ref{eq:KKT}),\\ u_{1}(\mathbf{\mu})-\lambda_{1}(\mathbf{\mu})+2u_{1}(\mathbf{\mu})\zeta(\mathbf{\mu})=0,\\ u_{2}(\mathbf{\mu})-\lambda_{2}(\mathbf{\mu})+2u_{2}(\mathbf{\mu})\zeta(\mathbf{\mu})=0,\\ (u_{1}(\mathbf{\mu})^{2}+u_{2}(\mathbf{\mu})^{2}-r^{2})\zeta(\mathbf{\mu})=0,\\ u_{1}(\mathbf{\mu})^{2}+u_{2}(\mathbf{\mu})^{2}\leq r^{2},\zeta(\mathbf{\mu})\geq 0, \end{array}\right. \tag{5.12}\]

where \(\zeta(\mathbf{\mu})\) is the Lagrange multiplier of the control constraint (5.9).

We compare the solutions of AONN with those obtained using PINN and plot their absolute errors in Figure 5, where it shows the control function \(u=(u_{1},u_{2})\) for a representative parameter \(\mathbf{\mu}=10\). From the figure, it can be seen that AONN can obtain a more accurate optimal control function than that of PINN, even when the training of PINN costs more epochs than that of AONN. Also, the quadratic constraint is not satisfied well for the PINN solution because there are more penalties from the KKT system (5.12) for the PINN loss. We compute the relative error \(\|u-u^{*}\|/\|u^{*}\|\) on a uniform \(256\times 256\) meshgrid for each parameter \(\mathbf{\mu}\) and plot the results in Figure 6. For most of the parameters, the relative errors of the AONN solutions are smaller than that of PINN, indicating that AONN is more effective and efficient than PINN in solving parametric optimal control problems. Note that this problem becomes harder when the parameter \(\mathbf{\mu}\) gets smaller [54]. In particular, the relative errors of AONN and PINN are both large as \(\mathbf{\mu}\) closes to \(0.1\).
### Test 4: Optimal control for the Laplace equation with geometrical parametrization

In this test case, we are going to solve the following parametric optimal control problem:

\[\left\{\begin{aligned} &\min_{y(\boldsymbol{\mu}),u(\boldsymbol{\mu})}J \left(y(\boldsymbol{\mu}),u(\boldsymbol{\mu})\right)=\frac{1}{2}\left\|y( \boldsymbol{\mu})-y_{d}(\boldsymbol{\mu})\right\|_{L_{2}(\Omega(\boldsymbol{ \mu}))}^{2}+\frac{\alpha}{2}\left\|u(\boldsymbol{\mu})\right\|_{L_{2}(\Omega( \boldsymbol{\mu}))}^{2},\\ &\text{subject to }\begin{cases}-\Delta y(\boldsymbol{\mu})=u( \boldsymbol{\mu})&\text{ in }\Omega(\boldsymbol{\mu}),\\ y(\boldsymbol{\mu})=1&\text{ on }\partial\Omega(\boldsymbol{\mu}),\\ \text{and }&u_{a}\leq u(\boldsymbol{\mu})\leq u_{b}&\text{ a.e. in }\Omega(\boldsymbol{\mu}),\end{cases}\right.\end{aligned}\right. \tag{5.13}\]

where \(\boldsymbol{\mu}=(\mu_{1},\mu_{2})\) represents the geometrical and desired state parameters. The parametric computational domain is \(\Omega(\boldsymbol{\mu})=([0,2]\times[0,1])\backslash B((1.5,0.5),\mu_{1})\) and the desired state is given by

\[y_{d}(\boldsymbol{\mu})=\begin{cases}1&\text{ in }\Omega_{1}=[0,1]\times[0,1], \\ \mu_{2}&\text{ in }\Omega_{2}(\boldsymbol{\mu})=([1,2]\times[0,1]) \backslash B((1.5,0.5),\mu_{1}),\end{cases} \tag{5.14}\]

where \(B((1.5,0.5),\mu_{1})\) is a ball of radius \(\mu_{1}\) with center \((1.5,0.5)\). We set \(\alpha=0.001\) and the parameter interval to be \(\boldsymbol{\mu}\in\mathcal{P}=[0.05,0.45]\times[0.5,2.5]\).

This test case is inspired by the literature [36, 23] that involve the application of local hyperthermia treatment of cancer. In such case, it is expected to achieve a certain temperature field in the tumor area and another temperature field in the non-lesion area through the heat source control. The circle cut out from the rectangular area represents a certain body organ, and by using AONN we aim to obtain all-at-once solutions of the optimal heat source control for different expected temperature fields and different organ shapes. In particular, we consider a two-dimensional model problem corresponding to the hyperthermia cancer treatment. One difficulty of this problem is the geometrical parameter \(\mu_{1}\) that leads to various computational domains, which causes difficulties in applying traditional mesh-based numerical methods. In the AONN framework, we can solve this problem by sampling in the spatio-parametric space:

\[\Omega_{\mathcal{P}}=\{(x_{0},x_{1},\mu_{1},\mu_{2})|0\leq x_{0}\leq 2,\,0\leq x _{1}\leq 1,\,0.05\leq\mu_{1}\leq 0.45,\,0.5\leq\mu_{2}\leq 2.5,\,(x_{0}-1.5)^{2}+(x _{1}-0.5)^{2}\geq\mu_{1}^{2}\}.\]

The computational domain \(\Omega(\boldsymbol{\mu})\) as well as the 40000 training points are given in Figure 5.5(a) and Figure 5.5(b).

The state neural network \(\hat{y}\) is constructed by \(\hat{y}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{I}}\right)= \ell(\mathbf{x},\boldsymbol{\mu})\hat{y}_{I}\left(\mathbf{x}(\boldsymbol{\mu}); \boldsymbol{\theta}_{y_{I}}\right)+1\) to naturally satisfy the Dirichlet boundary condition (5.13), where the length factor function is

\[\ell(\mathbf{x},\boldsymbol{\mu})=x_{0}(2-x_{0})x_{1}(1-x_{1})(\mu_{1}^{2}-(x_ {0}-1.5)^{2}-(x_{1}-0.5)^{2}).\]

The three neural networks \(\hat{y}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{I}} \right),\hat{p}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{p_{ I}}\right)\) and \(\hat{u}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{u}\right)\) are comprised of three ResNet blocks, each of which contains two fully connected layers with 25 units and a residual connection. The input dimension of these three neural networks is 4 and the total number of parameters of these three neural networks is \(3\times 3401=10203\). We take \(\gamma=0.985\) and the number of epochs for training the state function and the adjoint function increases from 200 to 700 during training. The configurations of the neural networks for the PINN and PINN+Projection methods are the same as those of AONN, and the number of training epoch is 50000. For this test problem, the AONN algorithm converges in 300 steps. Note that our AONN method can obtain all-at-once solution for any parameter \(\boldsymbol{\mu}\). To evaluate the performance of AONN, we employ the classical finite element method to solve the OCP\((\boldsymbol{\mu})\) with a fixed parameter. More specifically, a limited-memory BFGS algorithm implemented with bounded support is adopted in the dofini-adjoint [33] to solve the corresponding OCP. The solution obtained using the dofini-adjoint can be regarded as the ground truth. Among the four methods, AONN, PINN and PINN+Projection are able to solve parametric optimal control problems, while the dofini-adjoint solver can only solve the optimal control problem with a fixed parameter.

Figure 5.6 shows the optimal control solution obtained using AONN for the parametric optimal control problem (5.13). We choose several different parameters \(\boldsymbol{\mu}\) for visualization. The left column of Figure 5.6 corresponds to \(\mu_{2}=1\), in which case the optimal control is exactly zero because the desired state is achievable for \(y=y_{d}\equiv 1\). The middle and right column of Figure 5.6 indicate that the decrease of \(\mu_{1}\) and increase of \(\mu_{2}\) could increase the magnitude of \(u\). The results obtained by the dofini-adjoint solver, AONN, PINN and PINN+Projection with different values of \(c\) are displayed in Figure 5.7, where the control functions at \(\boldsymbol{\mu}=(0.3,2.5)\) are compared. The mesh with 138604 triangular elements are used in the dofini-adjoint solver, and after 16 steps, the final projected gradient norm achieves \(2.379\times 10^{-10}\). Figure 5.7 shows that AONN can
converge to the reference solution obtained by the dolfin-adjoint solver but PINN cannot obtain an accurate solution, while the results of PINN+Projection depends heavily on the choice of \(c\) in (4.6). When \(c\) is not equal to \(1/\alpha=1000\), the PINN+Projection method is not guaranteed to converge to the reference solution. This also confirms that the variational loss (3.7) brings great difficulties to neural network training of the KKT system (4.6), unless \(c=1/\alpha\), in which case the control function \(u\) is canceled out inside the projection operator.

\[\mathbf{P}_{U_{ad}(\boldsymbol{\mu})}\left(u(\boldsymbol{\mu})-c\mathrm{d}_{u }J(y(\boldsymbol{\mu}),u(\boldsymbol{\mu});\boldsymbol{\mu})\right)=\mathbf{P }_{U_{ad}(\boldsymbol{\mu})}\left(u(\boldsymbol{\mu})-c(\alpha u(\boldsymbol {\mu})+p(\boldsymbol{\mu}))\right)=\mathbf{P}_{U_{ad}(\boldsymbol{\mu})}\left( -\frac{1}{\alpha}p(\boldsymbol{\mu})\right).\]

However, for the next non-smooth test problem, \(u\) cannot be separated from the variational loss for any \(c\), which results in failure for the PINN+Projection method. To demonstrate that AONN can get all-at-once solutions, we first take a \(100\times 100\) grid of \(\mathcal{P}\) and choose several different realizations of \(\boldsymbol{\mu}\) to solve their corresponding OCP using the dolfin-adjoint solver. Then the parameters on the grid together with spatial coordinates are run through the trained neural networks obtained by Algorithm 1 to get the optimal solutions of OCP(\(\boldsymbol{\mu}\)) all at once. It is worth noting that using the dolfin-adjoint solver to compute the optimal solutions for all parameters on the \(100\times 100\) grid is computationally expensive since \(10000\) simulations are required. So we only take \(16\) representative points on the grid for the dolfin-adjoint solver (It still takes several hours). Nevertheless, all-at-once solutions can be computed effectively and efficiently through our AONN framework. Figure 8 displays three quantities with respect to \(\mu_{1},\mu_{2}\), where Figure 8(a) shows the objective functional \(J\), Figure 8(b) is the accessibility of the desired state and Figure 8(c) displays the \(L_{2}\)-norm of the optimal control \(u\). The red dots in Figure 8 show the results obtained by the dolfin-adjoint solver, where \(16\) simulations of OCP with \((\mu_{1},\mu_{2})\in\{0.05,0.1833,0.3167,0.45\}\times\{0.5,1.1667,1.8333,2.5\}\) are performed. From Figure 8, it is clear that AONN can obtain accurate solutions.

### Test 5: Optimal control for the semilinear elliptic equations with sparsity parametrization

In this test problem, we again consider a control problem for the semilinear elliptic equations as that in (5.1). However, this time we consider a sparse optimal control problem with sparsity parametrization. The sparse solution in optimal control is often achieved by \(L_{1}\)-control cost [5, 6, 4] and its application to the controller placement problems is well studied [46]. Specifically, we consider the following objective functional with \(L_{1}\)-control cost:

\[J(y,u)=\frac{1}{2}\left\|y-y_{d}\right\|_{L_{2}}^{2}+\frac{\alpha}{2}\|u\|_{L _{2}}^{2}+\beta\|u\|_{L_{1}},\]

where the coefficient \(\beta\) of the \(L_{1}\)-term controls the sparsity of the control function \(u\). With the increase of \(\beta\), the optimal control gradually becomes sparse and eventually reaches zero. In order to make continuous observation of this phenomenon, we need to solve the following parametric optimal control problem by setting \(\beta\) as a variable

Figure 5: Test 4: (a) The parametric computational domain \(\Omega(\boldsymbol{\mu})\). (b) \(N=40000\) training collocation points sampled in \(\Omega_{\mathcal{P}}\) (there are no points inside the frustum).
Figure 5: Test 4: the solution obtained by the doffin-adjoint solver for a fixed parameter \(\mathbf{\mu}=(0.3,2.5)\), the approximate solutions of \(u\) obtained by AONN, PINN, PINN+Projection (with different \(c=100,1000,10000\)), and the absolute errors of the AONN solution and the PINN+Projection solution with \(c=\frac{1}{\alpha}=1000\).

Figure 6: Test 4: the AONN solutions \(u(\mathbf{\mu})\) with several realizations of \(\mathbf{\mu}=(\mu_{1},\mu_{2})\).

Figure 7: Test 4: the solution obtained by the doffin-adjoint solver for a fixed parameter \(\mathbf{\mu}=(0.3,2.5)\), the approximate solutions of \(u\) obtained by AONN, PINN, PINN+Projection (with different \(c=100,1000,10000\)), and the absolute errors of the AONN solution and the PINN+Projection solution with \(c=\frac{1}{\alpha}=1000\).
parameter \(\boldsymbol{\mu}\),

\[\left\{\begin{aligned} &\min_{y(\boldsymbol{\mu}),u(\boldsymbol{\mu})}J(y (\boldsymbol{\mu}),u(\boldsymbol{\mu});\boldsymbol{\mu}):=\frac{1}{2}\left\|y( \boldsymbol{\mu})-y_{d}\right\|_{L_{2}(\Omega)}^{2}+\frac{\alpha}{2}\|u( \boldsymbol{\mu})\|_{L_{2}(\Omega)}^{2}+\boldsymbol{\mu}\|u(\boldsymbol{\mu}) \|_{L_{1}(\Omega)},\\ &\text{subject to }\begin{cases}-\Delta y(\boldsymbol{\mu})+y( \boldsymbol{\mu})^{3}=u(\boldsymbol{\mu})&\text{ in }\Omega,\\ &y(\boldsymbol{\mu})=0&\text{ on }\partial\Omega,\\ \end{cases}\\ \text{and}\quad u_{a}\leq u(\boldsymbol{\mu})\leq u_{b}&\text{ a.e. in }\Omega.\end{aligned}\right. \tag{13}\]

We fixed the other parameters

\[\Omega=B(0,1),\] \[\alpha=0.002,u_{a}=-12,u_{b}=12,\] \[y_{d}=4\sin\left(2\pi x_{1}\right)\sin\left(\pi x_{2}\right) \exp(x_{1}),\]

and the range of parameter is set to \(\boldsymbol{\mu}\in[0,\boldsymbol{\mu}_{max}]\). The upper bound \(\boldsymbol{\mu}_{max}=0.128\) ensures that for any \(\boldsymbol{\mu}\geq\boldsymbol{\mu}_{max}\) the optimal control \(u^{*}(\boldsymbol{\mu})\) is identically zero. We compute the generalized derivative

\[\mathrm{d}_{u}J(y(\boldsymbol{\mu}),u(\boldsymbol{\mu});\boldsymbol{\mu})= \alpha u(\boldsymbol{\mu})+p(\boldsymbol{\mu})+\boldsymbol{\mu}\text{ sign}(u( \boldsymbol{\mu})). \tag{14}\]

where \(p\) is the solution of the adjoint equation as defined in (11), and sign is an element-wise operator that extracts the sign of a function.

Since the optimal control function varies for different \(\boldsymbol{\mu}\), solving a series of sparse optimal control problems is straightforward in general. For example, Eduardo Cases [4] calculated the optimal solutions for \(\boldsymbol{\mu}=2^{i}\times 10^{-3},i=0,1,\ldots,8\). Here, we use AONN to compute all the optimal solutions for any \(\boldsymbol{\mu}\in[0,0.128]\) all at once. The length factor function for the Dirichlet boundary condition is chosen as \(\ell(\mathbf{x})=1-x_{0}^{2}-x_{1}^{2}\). The neural networks \(\hat{y}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{y_{I}} \right),\hat{p}_{I}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{p_{ I}}\right)\) and \(\hat{u}\left(\mathbf{x}(\boldsymbol{\mu});\boldsymbol{\theta}_{u}\right)\) are trained by AONN, which have the same configurations to those in the previous test, except for the input dimension being \(3\), resulting in \(3\times 3376=10128\) undecided parameters. To this end, we sample \(N=20000\) points in the spatio-parametric space \(\Omega_{\mathcal{P}}=B(0,1)\times[0,\boldsymbol{\mu}_{max}]\) by a uniform distribution. In order to capture the information at the boundary of \(\boldsymbol{\mu}\), \(2000\) of these \(20000\) points are sampled in \(B(0,1)\times\{0\}\) and \(B(0,1)\times\{\boldsymbol{\mu}_{max}\}\). We take \(500\) iteration steps and gradually increase the training epochs with \(n_{\text{aug}}=100\) after every \(100\) iteration. As a result, training epochs for the state function and the adjoint function increase from \(200\) to \(600\) during training. The step size \(c^{k}\) starts with \(c^{0}=10\) and decreases by a factor \(\gamma=0.985\) after every iteration.

The optimal control for some representative \(\boldsymbol{\mu}\in[0,\boldsymbol{\mu}_{max}]\) computed by AONN are displayed in Figure 10. The AONN results are consistent with the results presented in ref.[4], where the sparsity of optimal control increases as \(\boldsymbol{\mu}\) increases. As shown in Figure 10, the initial optimal control for \(\boldsymbol{\mu}=0\) has eight peaks and each peak disappears as \(\boldsymbol{\mu}\) increases. To determine where it is most efficient to put the control device, one might require some manual tuning of \(\boldsymbol{\mu}\) and thus need to solve OCP many times for different \(\boldsymbol{\mu}\). Determining these optimal locations is easy if we have obtained the parametric solutions \(u^{*}(x,y,\boldsymbol{\mu})\), which is exactly what AONN does. The coordinates of the eight peaks are obtained by evaluating the last vanishing positions of \(u^{*}(x,y,\boldsymbol{\mu})\) as \(\boldsymbol{\mu}\) increases at a uniform \(100^{3}\) grid on \([-1,1]\times[-1,1]\times[0,\boldsymbol{\mu}_{max}]\). Figure 10 shows the variation of control values at the eight peaks as a function with respect to \(\boldsymbol{\mu}\). We observe that the control function values at points \(P_{1},P_{2},P_{3}\) start with \(12\), and begin to decrease after \(\boldsymbol{\mu}\) reaches a certain value, finally drop to zero. Value at \(P_{4}\) starts to decrease from a number less than \(12\) until it reaches zero. The behavior of points \(P_{5}\sim P_{8}\) is completely symmetric.

To conclude, with these five numerical tests, we examine the efficiency of AONN and compare its performance with PINN, PINN+Projection, and the traditional solver. The numerical results indicate that the proposed AONN method is more advantageous than the PINN+Projection method and the PINN method in solving parametric optimal control problems. The PINN method cannot obtain accurate solutions for complex constrained problems, and the PINN+projection method improves the accuracy of the PINN method in general but has limitations on nonsmooth problems such as the sparse optimal control problems, while the AONN method is a general framework performing better on different types of parametric optimal control problems.

## 6 Conclusions

We have developed AONN, an adjoint-oriented neural network method, for computing all-at-once solutions to parametric optimal control problems. That is, the optimal control solutions for arbitrary parameters can be obtained by solving only once. The key idea of AONN is to employ three neural networks to approximate the control function, the adjoint function, and the state function in the optimality conditions, which
allows this method to integrate the idea of the direct-adjoint looping (DAL) approach in neural network approximation. In this way, three parametric surrogate models using neural networks provide all-at-once representations of optimal solutions, which avoids mesh generation for both spatial and parametric spaces and thus can be generalized to high-dimensional problems. With the integration of DAL, AONN also avoids the penalty-based loss function of the complex Karush-Kuhn-Tucker (KKT) system, thereby reducing the training difficulty of neural networks and improving the accuracy of solutions. Numerical experiments have shown that AONN can solve parametric optimal control problems all at once with high accuracy in several application scenarios, including control parameters, physical parameters, model parameters, and geometrical parameters.

Many questions remain open, e.g., choosing the step size and the scaling factor are heuristic, and solving some complex problems requires a higher computational cost. Future works could include the analysis of the convergence rate to better understand the properties of AONN, the introduction of adaptive sampling strategies

Figure 10: Test 5: the AONN solution \(u(\boldsymbol{\mu})\) of eight fixed peaks \(P_{1}\sim P_{8}\) as a function respect to \(\boldsymbol{\mu}\). The legend on the right is the coordinates of the eight points.
to further improve both robustness and effectiveness, and the generalization and application of AONN to more challenging problems such as shape or topology optimizations.

## Acknowledgments

This study was funded in part by National Natural Science Foundation of China (#12131002) and China Postdoctoral Science Foundation (2022M711730).
