#  Neural Network Approaches for Parameterized Optimal Control

Deepanshu Verma∗ ,  Nick Winovich† ,  Lars Ruthotto‡ and  Bart van Bloemen
Waanders† ∗Department of Mathematics, Emory University, Atlanta, GA
<EMAIL_ADDRESS>†Departments of Mathematics and Computer Science, Emory
University, Atlanta, GA<EMAIL_ADDRESS>‡Sandia National Laboratories,
Albuquerque, NM  {nwinovi<EMAIL_ADDRESS>
###### Abstract.

We consider numerical approaches for deterministic, finite-dimensional optimal
control problems whose dynamics depend on unknown or uncertain parameters. We
seek to amortize the solution over a set of relevant parameters in an offline
stage to enable rapid decision-making and be able to react to changes in the
parameter in the online stage. To tackle the curse of dimensionality arising
when the state and/or parameter are high-dimensional, we represent the policy
using neural networks. We compare two training paradigms: First, our model-
based approach leverages the dynamics and definition of the objective function
to learn the value function of the parameterized optimal control problem and
obtain the policy using a feedback form. Second, we use actor-critic
reinforcement learning to approximate the policy in a data-driven way. Using
an example involving a two-dimensional convection-diffusion equation, which
features high-dimensional state and parameter spaces, we investigate the
accuracy and efficiency of both training paradigms. While both paradigms lead
to a reasonable approximation of the policy, the model-based approach is more
accurate and considerably reduces the number of PDE solves.

###### Key words and phrases:

Hamilton-Jacobi-Bellman equation, Reinforcement Learning, high dimensional
optimal control, PDE constrained optimization, neural networks, policy
optimization, actor-critic

###### 2020 Mathematics Subject Classification:

35F21, 49M41, 68T07

##  1\. Introduction

We are interested in deterministic, finite-time optimal control problems in
which the dynamical system depends on parameters that are unknown or
uncertain. Such problems arise in decision-making for complex systems and are
prevalent across application domains including physics, engineering, finance,
robotics, economics, environmental sciences, and traffic flow [6, 34, 28, 2,
32, 37].

We aim to amortize the solution across multiple parameter uncertainties and
obtain a policy that maps the current state and parameter to approximate
optimal actions. While this introduces significant offline costs, the policy
is fast to evaluate online so that rapid decision-making is enabled. Moreover,
our capability supports changes in the parameters over time.

Since an effective policy must depend on the current state of the system and
the value of the parameter, amortizing control policies are affected by the
curse of dimensionality. When the sum of the dimensions of the state or
parameter exceeds approximately four, traditional solution approaches become
ineffective. Tackling the curse of dimensionality leads to two critical
challenges: First, one needs to determine policies using a high-dimensional
function approximator. Motivated by their universal approximation properties,
we consider neural networks. Second, identifying the relevant parts of the
state space is critical to avoid the infeasibility of globally sampling high-
dimensional spaces. Therefore, it is common to limit the search to parts of
the state space that are likely to be visited when following optimal policies
[20, 33].

Given a class of function approximators, we focus primarily on numerical
methods for learning effective policies and compare a model-based and a data-
driven approach. Our model-based approach approximates the value function of
the control problem from which the policy is obtained using a feedback form.
The feedback form requires a model of the dynamical system and knowledge of
the objective functional. This enables the incorporation of the physics of the
system and include known properties of the value function, via the Hamilton-
Jacobi-Bellman (HJB) equations, into the training problem. Our data-driven
approach uses the framework of actor-critic reinforcement learning to decouple
the networks for the policy (i.e., the actor) and the value function (i.e.,
the critic) and trains both networks purely from observations. Both approaches
are intimately related through the underlying control theory, especially
Pontryagin’s maximum principle theory and HJB equations.

To compare the accuracy and efficiency of both approaches, we implement a
model problem motivated by control of contaminant transport. Here, the state
is governed by the two-dimensional advection-diffusion partial differential
equation (PDE). The objective is to prevent a contaminant originating at a
source location from reaching a target region by controlling the location of a
sink. Uncertain parameters in this case consist of different scenarios for the
source location and the velocity field. Since both the state and the parameter
are infinite-dimensional, the problem is well-suited to create high-
dimensional problem instances through discretization. Both the state and the
parameters are infinite-dimensional and therefore motivate high-dimensional
problems after discretization.

Our main contributions and observations are summarized as follows:

  * •

We extend the optimal control approaches from [27] to parametric control
problems, which leads to a considerable increase in the dimensionality of the
problem.

  * •

We provide a direct comparison of the OC and RL approaches and including
resective sub-optimalities. For the parameterized OC approach, this experiment
increases the dimension from around 100 in [27] to about 1000 (see section 6).

  * •

We observe in our experiments that the OC approach achieves more accurate
policies with fewer simulations (20x less) of the dynamical system.

The remainder of the paper is organized as follows: Section 2 introduces the
parametric OC problem and provides a brief background on OC theory that
informs our approaches. Section 3 discusses related approaches for learning
amortized policies. Section 4 presents our neural network approach for
computing the value function. Section 5 revisits key RL concepts, laying the
groundwork for employing Proximal Policy Optimization (PPO) and Twin Delayed
Deep Deterministic Policy Gradients (TD3) algorithms in policy approximation.
Section 6 conducts a comprehensive numerical assessment of both the neural-HJB
and RL approaches, employing PDE-based dynamics. Finally, Section 6.4 provides
a thorough discussion of results, highlighting the strengths and limitations
of the proposed methodologies.

##  2\. Parameterized Optimal Control Problem

We consider a family of deterministic, finite horizon optimal control problems
of the form

(2.1a) $\displaystyle\min_{\boldsymbol{u}\in
U}J(t,\boldsymbol{x},\boldsymbol{u};\boldsymbol{y}):=\int_{t}^{T}L(s,\boldsymbol{z}(s),\boldsymbol{u}(s))\rm
ds+\,G(\boldsymbol{z}(T)),$ subject to a nonlinear dynamical constraint (2.1b)
$\displaystyle\rm
d_{s}\boldsymbol{z}(s)=f(s,\boldsymbol{z}(s);\boldsymbol{y})+g(s,\boldsymbol{z}(s);\boldsymbol{y})\;\boldsymbol{u}(s),\quad
s\in(t,T];\quad\boldsymbol{z}(t)=\boldsymbol{x},$

where, given a parameter $\boldsymbol{y}\in\eta$, the functions
$f:[0,T]\times\mathds{R}^{d}\to\mathds{R}^{d}$ and
$g:[0,T]\times\mathds{R}^{d}\to\mathds{R}^{d\times a}$ model the evolution of
the states $\boldsymbol{z}:[0,T]\to\mathds{R}^{d}$ in response to the controls
$\boldsymbol{u}:[0,T]\to U\subset\mathds{R}^{a}$. The initial state of the
system is denoted as $\boldsymbol{x}\in\mathds{R}^{d}$, and the initial time
is represented by $t\in[0,T]$. The function
$L:[0,T]\times\mathds{R}^{d}\times\mathds{R}^{a}\to\mathds{R}$ represents the
running cost, and $G:\mathds{R}^{d}\to\mathds{R}$ represents the terminal cost
in the control objective eq. 2.1a.

We assume that $f,L,G,U$ are sufficiently regular (see [39, Chapter 2] for a
list of assumptions)

We aim to determine an optimal control (policy) amortized over the parameters
$\boldsymbol{y}\in\eta$ that minimizes the overall cost. In particular, given
a specific parameter $\boldsymbol{y}$, we seek to find the value function

(2.2) $\Phi(t,\boldsymbol{x};\boldsymbol{y})=\inf_{\boldsymbol{u}\in
U}J(t,\boldsymbol{z},\boldsymbol{u};\boldsymbol{y})\quad\text{s.t.
\quad\lx@cref{creftype~refnum}{eq:sys_dyn}},$

which represents the minimum cost-to-go for any initial state
$(t,\boldsymbol{x})$. The control $\boldsymbol{u}^{*}$ for this minimum value
is referred to as optimal control, and the corresponding trajectory
$\boldsymbol{z}^{*}$ is called an optimal trajectory.

The problem eq. 2.1 is typically tackled using local solution methods. These
methodologies involve minimizing the cost functional by eliminating the
constraints, all while considering a fixed parameter $\boldsymbol{y}$.
However, these localized strategies necessitate reevaluating the solution each
time the $\boldsymbol{y}$ changes, rendering the previously computed solution
obsolete. This clearly highlights the need for devising a control policy that
is amortized over the parameters $\boldsymbol{y}\in\eta$.

Subsequently, we will see that the value function $\Phi$ contains complete
information about the optimal control, and we can straightforwardly deduce
$\boldsymbol{u}^{*}$ from $\Phi$. Notably, this connection is not explicitly
utilized by RL algorithms during policy approximation. This distinction
between the HJB approach and RL is significant, as leveraging the underlying
physics of the problem can lead to improved results as presented in section 6.

We now introduce the Hamiltonian,
$H:[0,T]\times\mathds{R}^{d}\times\mathds{R}^{d}\to\mathds{R}\cup\\{\infty\\}$
of the system eq. 2.1, which plays a pivotal role in both OC theory and our
numerical approach:

(2.3)
$H(s,\boldsymbol{z},\boldsymbol{p};\boldsymbol{y})=\sup_{\boldsymbol{u}\in
U}\mathcal{H}(s,\boldsymbol{z},\boldsymbol{p},\boldsymbol{u};\boldsymbol{y})$

Here $\boldsymbol{p}:[t_{0},T]\to\mathds{R}^{d}$ is referred to as the adjoint
or costates of the system, and

$\mathcal{H}(s,\boldsymbol{z},\boldsymbol{p},\boldsymbol{u};\boldsymbol{y})=\boldsymbol{p}\cdot\left(f(s,\boldsymbol{z};\boldsymbol{y})+g(s,\boldsymbol{z};\boldsymbol{y})\;\boldsymbol{u}\right)-L(s,\boldsymbol{z},\boldsymbol{u}).$

The Pontryagin maximum principle provides a set of first-order necessary
conditions and states that, see [9, Theorem I.6.3], the adjoint,
$\boldsymbol{p}$, satisfies the adjoint equation. Furthermore, when the value
function $\Phi$ is differentiable, the adjoint $\boldsymbol{p}$ and the
optimal control $\boldsymbol{u}^{*}$ can be derived from $\Phi$ using [9,
Theorem I.6.2]: for $s\in[t,T]$

(2.4)
$\boldsymbol{p}(s)=\nabla_{\boldsymbol{z}}\Phi\big{(}s,\boldsymbol{z}^{*}(s);\boldsymbol{y}\big{)},\quad\text{and}$
(2.5)
$\begin{split}\boldsymbol{u}^{*}(s,\boldsymbol{z}^{*}(s),\nabla_{\boldsymbol{z}}\Phi\big{(}s,\boldsymbol{z}^{*}(s));\boldsymbol{y})\in\operatorname*{arg\,max}_{\boldsymbol{u}\in
U}\mathcal{H}\big{(}s,\boldsymbol{z}^{*}(s),\nabla_{\boldsymbol{z}}\Phi\big{(}s,\boldsymbol{z}^{*}(s);\boldsymbol{y}\big{)},\boldsymbol{u}(s);\boldsymbol{y}).\end{split}$

The eq. 2.5 characterizes optimal control in feedback form, and states that it
can be readily computed at any space-time pair, for a given $\boldsymbol{y}$,
when the value function $\Phi$ and $\nabla_{\boldsymbol{z}}\Phi$ are
available. For a thorough exploration of the aforementioned relationships, we
refer the readers to [9, 20, 27].

The model-based approach exploits this fundamental relationship to approximate
a $\Phi$ amortized across diverse parameters $\boldsymbol{y}$. This ensures a
readily accessible approximate optimal control for each parameter setting. The
data-driven approach, on the other hand, is agnostic to this relationship and
hence suffers in terms of accuracy.

##  3\. Related Work

This section reviews some closely related approaches for control/policy
training. Existing approaches can roughly be grouped by the type of function
approximators used and whether they primarily use the model or data during
training. Specifically focusing on NNs as function approximators, we
distinguish between model-based and data-driven training algorithms.

Model-based approaches, like the one presented by Kang et al. [16] adopt a
supervised learning approach to learn a closed-form value function. The
physics-informed neural networks (PINNs) method [22, 25] integrates the cost
functional into the standard PINNs loss for solving OC problems. Other
techniques, such as deep-learning-based surrogate models [37, 23] and operator
learning methods [35, 13], focus on achieving fast OC solution inference
without intensive computations. In [8], an extended PINN approach utilizes
Karush–Kuhn–Tucker (KKT) conditions as the loss function. These approaches
primarily tackle OC problems with partial differential equations (PDEs) as
constraints.

A recent development is the adjoint-oriented neural network method proposed in
[38] for parametric OC problems. This method learns the states, adjoints, and
control of the system through three different neural networks using a direct-
adjoint looping method-type loss. While capable of handling parametric
problems simultaneously with random sampling instead of spatial domain
discretization, it faces challenges in the presence of changes in initial
states or disturbances when following the optimal solution. Moreover, training
separate networks introduces computational complexity, limiting the
applicability of this approach to less complex problems.

Reinforcement learning (RL) has emerged as a potent data-driven approach and
shown success in diverse applications such as robotics, autonomous driving,
games, and recommender systems [24, 30, 36, 31, 7, 14, 4]. Its non-intrusive
and gradient-free nature make RL suitable for problems lacking well-defined
models. However, RL often requires many samples, which are computationally
expensive in our setting. The process of hyperparameter tuning in RL
algorithms is laborious and computationally expensive, hindering their
applicability for parametric dynamics [26, 12, 15, 42, 1]. To overcome these
limitations, there is a growing interest in effectively integrating system
knowledge to overcome the constraints of data-driven methods for handling
complex dynamics. In [40, 41], the authors investigate solving OC problems
using actor-critic and policy gradient type methods within the realm of RL.
While successful for general OC problems, these methods have limited direct
applicability to parametric problems. The main challenge arises from the added
complexity introduced by parameters, in addition to the already expensive task
of evolving the dynamics over time.

Works closely related to our OC approach include [27, 19, 20], which employ
NNs to parameterize the value function and penalize the HJB equation satisfied
by the value function.

##  4\. Model-Based Approach

In this section, we present our model driven approach, inspired by [27], for
approximating the amortized value function of the parameterized OC problem eq.
2.1. The central concept revolves around utilizing an efficient function
approximator to model the amortized value function $\Phi$ in eq. 2.2 and
subsequently computing the control using the feedback form given by eq. 2.5.
The learning adopts an unsupervised approach, akin to RL. The loss function
consists of the sum of the expected cost over different parameters,
$\boldsymbol{y}\in\eta$, driving the trajectories and penalty terms that
enforce the HJB equations along the trajectories and at the final-time.

###  4.1. Learning Problem

Let $\Phi_{\boldsymbol{\theta}}$ denote the approximation of the value
function $\Phi$ with parameters $\boldsymbol{\theta}$. Ideally, we aim for a
choice of $\boldsymbol{\theta}$ where $\Phi_{\boldsymbol{\theta}}$ matches the
value function of the corresponding control problem for every given
$\boldsymbol{y}$. However, this problem suffers from the curse of
dimensionality so we resort to an approach which enforces this property in a
subset of space-time domain.

To learn the parameters in an unsupervised way (that is, no apriori data for
$\Phi$ and optimal control trajectories), we approximately solve the
minimization problem

(4.1)
$\displaystyle\min_{\boldsymbol{\theta}}\mathds{E}_{\boldsymbol{y}\sim\eta}\left\\{J(t,\boldsymbol{z},\boldsymbol{u};\boldsymbol{y})+P_{\rm
HJB,\,\boldsymbol{y}}(\boldsymbol{z})\right\\}$ s.t. $\displaystyle\rm
d_{s}\boldsymbol{z}=f(s,\boldsymbol{z}(s);\boldsymbol{y})+g(s,\boldsymbol{z}(s);\boldsymbol{y})\;\boldsymbol{u}(s),\quad
s\in(t,T];\quad\boldsymbol{z}(t)=\boldsymbol{x}.$

Here $P_{\rm HJB,\,\boldsymbol{y}}$ penalizes deviations from the HJB PDE
satisfied by the value function for each parameter $\boldsymbol{y}$. In the
learning process, the optimal control is computed through the feedback form
provided in eq. 2.5, hence, incorporating the model information.

Exploiting the fact that the value function satisfies the HJB PDE (see [9,
Theorems I.5.1, I.6.1]), we guide the approximation of $\Phi$ using the
penalty term $P_{\rm HJB,\,\boldsymbol{y}}$, defined as:

(4.2) $\begin{split}P_{\rm
HJB,\,\boldsymbol{y}}(\boldsymbol{z})=&\beta_{1}\int_{t}^{T}|H(s,\boldsymbol{z},\nabla\Phi_{\boldsymbol{\theta}}(s,\boldsymbol{z}(s);\boldsymbol{y});\boldsymbol{y})-\partial_{s}\Phi_{\boldsymbol{\theta}}(s,\boldsymbol{z}(s);\boldsymbol{y})|\rm
ds\\\
&+\beta_{2}|G(\boldsymbol{z}(T))-\Phi_{\boldsymbol{\theta}}(T,\boldsymbol{z}(T);\boldsymbol{y})|+\beta_{3}|\nabla_{\boldsymbol{z}}G(\boldsymbol{z}(T))-\nabla_{\boldsymbol{z}}\Phi_{\boldsymbol{\theta}}(T,\boldsymbol{z}(T);\boldsymbol{y})|,\end{split}$

where the relative influence of each term is controlled by the components of
$\beta=(\beta_{1},\beta_{2},\beta_{3})\in\mathds{R}^{3}_{+}$.

We further make the following assumption for our model driven approach.

######  Assumption 4.1.

There exists a closed form solution to eq. 2.5, which allow us to write the
control $\boldsymbol{u}$ as a function of $\nabla_{\boldsymbol{z}}\Phi$,
explicitly, i.e.

(4.3)
$\boldsymbol{u}^{*}(s,\boldsymbol{z}^{*}(s),\nabla_{\boldsymbol{z}}\Phi\big{(}s,\boldsymbol{z}^{*}(s);\boldsymbol{y});\boldsymbol{y})=\operatorname*{arg\,max}_{\boldsymbol{u}\in
U}\mathcal{H}\big{(}s,\boldsymbol{z}^{*}(s),\nabla_{\boldsymbol{z}}\Phi\big{(}s,\boldsymbol{z}^{*}(s);\boldsymbol{y}\big{)},\boldsymbol{u}(s);\boldsymbol{y}),$

A closed-form solution for the optimal control exists in a wide variety of OC
problems [21, 3, 18]. While we do not explicitly demonstrate it in this work,
it is worth noting that this assumption can be relaxed to include implicitly
defined functions as long as they can be efficiently obtained. This
flexibility allows for the modeling of more general convex running costs and
enhances the applicability of our approach to a wider range of problems.

The eqs. 4.3 and 4.1 provide a framework for obtaining the optimal control and
trajectory based on the value function $\Phi$, subject to certain smoothness
assumptions, by outlining the necessary conditions for optimality. Once a good
approximation of the value function $\Phi$ is computed, this framework can be
applied to any initial data and parameters, also allowing for adaptability to
perturbations in the system.

###  4.2. Function Value Approximation

In principle, one could employ any high-dimensional function approximator to
parameterize the value function. Yet, due to the universal approximation
properties inherent in NNs, we opt for using an NN to parameterize the value
function. Designing an effective neural network architecture is essential for
various learning tasks, and it remains an active area of research. In our
approach, we treat this as a modular component, providing flexibility. Our
framework can seamlessly integrate with any scalar-valued neural network that
accepts inputs in $\mathds{R}^{d+1}$ and possesses at least one continuous
derivatives concerning its first $d+1$ inputs, to allow computations of
$\nabla\Phi$.

In our experiments, we use a residual neural network given by

(4.4)
$\Phi_{\boldsymbol{\theta}}(\boldsymbol{h}_{0})=\boldsymbol{w}^{\top}{\rm
NN}(\boldsymbol{h}_{0};\boldsymbol{\theta}_{\rm NN}),$

with trainable weights $\boldsymbol{\theta}$ containing
$\boldsymbol{w}\,\,{\in}\,\,\mathds{R}^{m}$ and $\boldsymbol{\theta}_{\rm
NN}\,\,{\in}\,\,\mathds{R}^{p}$. Here the inputs
$\boldsymbol{h}_{0}=(t,\boldsymbol{z}(t);\boldsymbol{y})\in\mathds{R}^{d+1}$
correspond to time-space, and ${\rm
NN}(\boldsymbol{h}_{0};\boldsymbol{\theta}_{\rm
NN})\colon\mathds{R}^{d+1}\to\mathds{R}^{m}$ is a residual neural network
(ResNet) [11]

(4.5)
$\begin{split}\boldsymbol{h}_{1}&=\rm\sigma(\boldsymbol{K}_{0}\boldsymbol{h}_{0}+\boldsymbol{b}_{0})\\\
\boldsymbol{h}_{i+2}&=\boldsymbol{h}_{i+1}+\rm\sigma(\boldsymbol{K}_{i+1}\boldsymbol{h}_{i+1}+\boldsymbol{b}_{i+1}),\quad
0\leq i\leq M-2\\\ {\rm NN}(\boldsymbol{h}_{0};\boldsymbol{\theta}_{\rm
NN})&=\boldsymbol{h}_{M}+\rm\sigma(\boldsymbol{K}_{M}\boldsymbol{h}_{M}+\boldsymbol{b}_{M}),\end{split}$

with neural network weights $\boldsymbol{\theta}_{\rm
NN}{=}(\boldsymbol{K}_{0},\ldots,\boldsymbol{K}_{M},\boldsymbol{b}_{0},\ldots,\boldsymbol{b}_{M})$
where $\boldsymbol{b}_{i}\in\mathds{R}^{m}\;\forall i$,
$\boldsymbol{K}_{0}\in\mathds{R}^{m\times(d+1)}$, and
$\\{\boldsymbol{K}_{1},\ldots,\boldsymbol{K}_{M}\\}\in\mathds{R}^{m\times m}$
with $M$ being the depth of the network. We use the element-wise nonlinearity
$\sigma(x)=\log(\exp(x)+\exp(-x))$, which is the antiderivative of the
hyperbolic tangent, i.e., $\sigma^{\prime}(x)=\tanh(x)$. For the experiments,
we use $m=64$ nodes per layer and a network depth of $M=4$.

###  4.3. Numerical Implementation

We tackle the control problem eq. 4.1 using the discretize-then-optimize
approach. We begin by discretizing the constraints and then optimize. To do
so, we first sample parameters $\boldsymbol{y}\sim\eta$ and then use a time
discretization of $N+1$ equidistant time points $t=s_{0},\ldots,s_{N}$ with
step size $\Delta s=(T-t)/N$ to eliminate the constraints in eq. 2.1b. The
specific discretization details for each experiment vary and are discussed in
their respective sections. This yields a state trajectory starting at
$\boldsymbol{z}_{0}=\boldsymbol{x}$ via

(4.6)
$\boldsymbol{z}_{i+1}=\boldsymbol{z}_{i}+f(s_{i},\boldsymbol{z}_{i+1};\boldsymbol{y})\Delta
s+g(s_{i},\boldsymbol{z}_{i+1};\boldsymbol{y})\,\boldsymbol{u}_{i}\,\Delta
s,\quad i=0,\ldots,N-1,$

where $\boldsymbol{z}_{i}=\boldsymbol{z}(s_{i})$ and
$\boldsymbol{u}_{i}=\boldsymbol{u}^{*}(s_{i},\boldsymbol{z}_{i},\nabla_{z}\Phi_{\boldsymbol{\theta}}(s_{i},\boldsymbol{z}_{i};\boldsymbol{y}))$
is the optimal control obtained from the feedback form, that is, from eq. 4.3,
computed by equating $\nabla_{\boldsymbol{u}}\mathcal{H}=0$. Finally, we
approximate the objective functional via

(4.7) $J(s_{k},\boldsymbol{z},\boldsymbol{u};\boldsymbol{y})=\Delta
s\sum_{i=k}^{N}L(s_{i},\boldsymbol{z}_{i},\boldsymbol{u}_{i})+G(\boldsymbol{z}_{N}),$

and simultaneously, the penalty term $P_{\rm HJB,\boldsymbol{y}}$ in a similar
manner.

In principle, any stochastic approximation approach can be used to solve the
above optimization problem. Here, we use Adam [17] and sample a minibatch of
trajectories originating in i.i.d. samples from $\eta$ with an initial
learning rate of $0.075$. During training, the learning rate gradually
diminishes with an exponential decay rate of $0.975$ until it stabilizes at
$0.0025$. For each training iteration, we randomly select a batch of $20$
problem parameters from the distribution $\eta$. Following the HJB framework,
we evaluate controls for each problem and evolve the systems accordingly.
Subsequently, we compute the losses, utilizing them to update the network
weights, as illustrated in Algorithm 1.

Algorithm 1 Model-based training approach

1:Number of training problems $\rm P$, problem parameter distribution $\eta$,
initial network weights $\boldsymbol{\theta}$, Hamiltonian $H$, and loss
weighting factors $(\beta_{1},\beta_{2},\beta_{3})$

2:Initial Setup for Training

3:_Derive Feedback Form Expression_

4:$\operatorname{Feedback\\_Form}:\,(-\nabla_{\boldsymbol{z}}\Phi_{\boldsymbol{\theta}},\,f,\,g)\,\mapsto\,\boldsymbol{u}$

5:

6:_Assemble Randomized Problems_

7:Sample $\rm P$ parameters $\\{\boldsymbol{y}_{i}\\}_{i=1}^{\rm P}\sim\eta$

8:Store matrices for $f$ and $g$ corresponding to
$\\{\boldsymbol{y}_{i}\\}_{i=1}^{\rm P}$

9:

10:HJB Training Iteration

11:_Sample Batch of Problems_

12:Retrieve matrices for $\boldsymbol{y}_{p}$ with
$p\sim\operatorname{Uniform}(\\{1,\dots,{\rm P}\\})$

13:

14:_Assess HJB Control Performance_

15:Set $s_{0}=0$ , $\boldsymbol{z}_{0}=\boldsymbol{x}$ , $J=0$, and
$P_{\rm{HJB},\boldsymbol{y}_{p}}=0$

16:while $i\leq N$ do

17:  _Derive Control from Feedback Form_

18: Evaluate network
$\Phi_{\boldsymbol{\theta}}(s_{i},\boldsymbol{z}_{i},\boldsymbol{y}_{p})$ and
gradient
$\nabla_{\boldsymbol{z}}\Phi_{\boldsymbol{\theta}}(s_{i},\boldsymbol{z}_{i};\boldsymbol{y}_{p})$

19:
$\boldsymbol{u}_{i}=\operatorname{Feedback\\_Form}(-\nabla_{\boldsymbol{z}}\Phi_{\boldsymbol{\theta}}(s_{i},\boldsymbol{z}_{i};\boldsymbol{y}_{p}),\,f,\,g)$

20:

21:  _Evolve Systems in Time_

22:
$\boldsymbol{z}_{i+1}\leftarrow\operatorname{Time\\_Integrator}(\boldsymbol{z}_{i},\boldsymbol{u}_{i};\boldsymbol{y}_{p})$

23: $s_{i+1}\leftarrow s_{i}+\Delta s$

24:

25:  _Update Intermediate Losses_

26: $J\leftarrow J+L(s_{i},\boldsymbol{z}_{i},\boldsymbol{u}_{i})\Delta s$

27: $P_{\rm{HJB},\boldsymbol{y}_{p}}\leftarrow
P_{\rm{HJB},\boldsymbol{y}_{p}}+\beta_{1}|H(s_{i},\boldsymbol{z}_{i},\nabla_{\boldsymbol{z}}\Phi_{\boldsymbol{\theta}}(s_{i},\boldsymbol{z}_{i};\boldsymbol{y}_{p});\boldsymbol{y}_{p})\,-\,\partial_{s}\Phi_{\boldsymbol{\theta}}(s_{i},\boldsymbol{z}_{i},\boldsymbol{y}_{p})|\Delta
s$

28:end while

29:$J\leftarrow J+G(\boldsymbol{z}_{N})$

30:$P_{\rm{HJB},\boldsymbol{y}_{p}}\leftarrow
P_{\rm{HJB},\boldsymbol{y}_{p}}+\beta_{2}|G(\boldsymbol{z}_{N}-\Phi_{\boldsymbol{\theta}}(s_{N},\boldsymbol{z}_{N},\boldsymbol{y}_{p})|+\beta_{3}|\nabla_{\boldsymbol{z}}G(\boldsymbol{z}_{N})-\nabla_{\boldsymbol{z}}\Phi_{\boldsymbol{\theta}}(s_{N},\boldsymbol{z}_{N},\boldsymbol{y}_{p})|$

31:

32:_Update Networks Weights_

33:$\theta\leftarrow\operatorname{Optimizer}(J\,+\,P_{\rm{HJB},\boldsymbol{y}_{p}})$

##  5\. Data-Driven Approach

We briefly review the data-driven approach, in particular, the Reinforcement
Learning (RL) framework, with an emphasis on Actor-Critic algorithm used to
approximate the optimal policy in our work; our presentation follows [5, 33].

Reinforcement learning (RL) provides a general framework for approximating
policies through a series of interactions with an environment. More precisely,
an RL environment specifies all possible system states, a description of
admissible actions, along with a model governing the evolution of the system
space when an action is performed. In RL language, the controls are referred
to as the actions. The environment is paramount for RL algorithms, and its
proper definition is crucial.

Let ${\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}$ denotes the policy,
parameterized by parameters $\boldsymbol{\theta}_{\boldsymbol{a}}$, which is a
mapping from the state space to the action space. The policy can be either
deterministic, denoted as
$\boldsymbol{u}_{i}={\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}(\boldsymbol{z}_{i})$,
or stochastic, indicated as
$\boldsymbol{u}_{i}\sim{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}(\cdot|\boldsymbol{z}_{t})$.
Using eq. 4.6, a policy generates a sequence of states and actions,
${(\boldsymbol{z}_{i},\boldsymbol{u}_{i})}_{i=0}^{N}$, referred to as an
episode.

In order to evaluate the performance of a policy, it is necessary to create
functions step() and reset() governing the problem dynamics and the
interaction between the agent and the environment. Aiming to address the
optimization problem introduced in section 2, these functions operate as
follows

  * •

_reset()_ : This function initializes or resets the environment to its initial
state, $\boldsymbol{z}_{0}$, preparing it for the new episode.

  * •

_step()_ : Given an action $\boldsymbol{u}_{i}$ as input, determined by a
policy, it returns the following:

    * –

The next state $\boldsymbol{z}_{i+1}$ using the dynamics described in eq. 4.6.

    * –

A scalar feedback signal representing the objective function, known as the
reward,

$\displaystyle r_{i}=\begin{cases}\Delta
s\,L(s_{i},\boldsymbol{z}_{i},\boldsymbol{u}_{i}),&\quad i<N\\\ \Delta
s\,L(s_{N},\boldsymbol{z}_{N},\boldsymbol{u}_{N})+G(\boldsymbol{z}_{N}),&\quad
i=N\end{cases}.$

    * –

The termination status indicates the terminal time $T$ has been reached or the
current episode has terminated.

Specific environment details for the problems considered are provided in
section 6.

The goal is to identify an optimal policy in order to minimize the cumulative
reward, called return111Conventional RL literature assumes that the objective
function is being maximized; for the purposes of aligning with the optimal
control problem setup from previous sections, we instead aim to minimize the
loss (which can be interpreted as the negative reward in standard RL
literature)., given by the value function

(5.1)
$\Psi(\boldsymbol{\theta}_{\boldsymbol{a}})\,\,=\,\,\mathds{E}_{{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}},\boldsymbol{y}\sim\eta}\left\\{\,\sum_{i=0}^{N}r_{i}\right\\}\,\,=\,\,\mathds{E}_{{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}},\boldsymbol{y}\sim\eta}\left\\{\,J(s_{0},\boldsymbol{z},\boldsymbol{u};\boldsymbol{y})\right\\}.$

When the policy is deterministic, the expectation over
${\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}$ in eq. 5.1 is omitted, making
the return equivalent to the one in the neural-HJB training defined in eq.
4.7.

######  Remark 5.1.

When starting from the initial state $\boldsymbol{z}_{0}$ and following the
optimal policy ${\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}^{*}$, we recover
the original value function $\Phi$ in eq. 2.2.

In our work, we consider the policy to be stochastic as they provide
additional flexibility, robustness, and adaptability that can be advantageous
in complex and uncertain environments.

Now that we have established a loss function representative of the original
optimization problem, we must also review the expression for the gradients
with respect to $\boldsymbol{\theta}_{\boldsymbol{a}}$ for updating these
parameters. This expression is provided by the following foundational result
from the theory of RL.

######  Theorem 5.2.

[33, Policy Gradient Theorem, Chapter 13] The gradient of the cumulative
reward objective function eq. 5.1 can be expressed in terms of the policy
gradients as follows:

(5.2)
$\displaystyle\nabla_{\boldsymbol{\theta}_{\boldsymbol{a}}}\Psi(\boldsymbol{\theta}_{\boldsymbol{a}})\,\,=\,\,\mathbb{E}_{{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{u}},\,\boldsymbol{y}\sim\eta}\left\\{\,\sum_{i=0}^{N}\,\nabla_{\boldsymbol{\theta}_{\boldsymbol{a}}}\log{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}(\boldsymbol{u}_{i}|\boldsymbol{z}_{i}){J}(s_{i},\boldsymbol{z},\boldsymbol{u};\boldsymbol{y})\right\\}.$

While the direct use of Theorem 5.2 for policy optimization can work
effectively in some cases, the practical performance often suffers from high
variance in the gradient estimate, and thus produce slow learning. To reduce
variance effects and increase stability, many policy-based RL algorithms
subtract the cumulative reward by a baseline, $b(\boldsymbol{s})$:

(5.3)
$\nabla_{\boldsymbol{\theta}_{\boldsymbol{a}}}\Psi(\boldsymbol{\theta}_{\boldsymbol{a}})\,\,=\,\,\mathbb{E}_{{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}},\,\boldsymbol{y}\sim\eta}\left\\{\,\sum_{i=0}^{N}\,\left(J(s_{i},\boldsymbol{z}_{i},\boldsymbol{u}_{i};\boldsymbol{y})-b(\boldsymbol{z}_{i})\right)\nabla_{\boldsymbol{\theta}_{\boldsymbol{a}}}\log{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}(\boldsymbol{u}_{i}|\boldsymbol{z}_{i})\right\\}.$

Intuitively, making the cumulative reward smaller by subtracting it with a
baseline will make smaller gradients, and thus smaller and more stable
updates.

###  5.1. Actor-Critic Models

Selecting an effective baseline poses a challenge, but the state-value
function is a suitable choice, making both the value function and policy
learnable. This creates a two-component RL framework known as an actor-critic
model:

  * •

The actor ${{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}}$ determines actions
based on the current state (i.e. defines the policy).

  * •

The critic $\Psi_{\boldsymbol{\theta}_{\boldsymbol{c}}}$, parameterizing the
baseline with parameters ${\boldsymbol{\theta}_{\boldsymbol{c}}}$, assesses
the approximate value of the current state (i.e. provides value function
estimates).

and both are parameterized by neural networks.

Now that the value function is also made learnable, that means we must provide
an objective for the critic network to learn the parameters
$\boldsymbol{\theta}_{\boldsymbol{c}}$. Since the role of the critic is to
provide a baseline estimate for the return, a natural choice for this
objective is the expected value of the squared error:

(5.4) $P_{\rm
critic}(\boldsymbol{\theta}_{\boldsymbol{c}})\,\,=\,\,\sum_{i=0}^{N}\,\left|\Psi_{\boldsymbol{\theta}_{\boldsymbol{c}}}(\boldsymbol{z}_{i})\,-\,J(s_{i},\boldsymbol{z},\boldsymbol{u};\boldsymbol{y})\right|^{2}.$

In summary, for the actor-critic models, we are approximately solving the
minimization problem given by

$\min_{\boldsymbol{\theta}_{\boldsymbol{a}},\boldsymbol{\theta}_{\boldsymbol{c}}}\,\,\mathbb{E}_{{{\pi_{\boldsymbol{\theta}}}_{\boldsymbol{a}}},\,\boldsymbol{y}\sim\eta}\left\\{J(s_{0},\boldsymbol{z},\boldsymbol{u};\boldsymbol{y})+\beta_{5}\,P_{\rm
critic}(\boldsymbol{\theta}_{\boldsymbol{c}})\right\\},$

where $\beta_{5}\in\mathds{R}_{+}$ controls the relative contribution of
critic. Although there are several actor-critic algorithms available, in this
work we use two prominent, state-of-the-art algorithms: Proximal Policy
Optimization (PPO) and Twin Delayed Deep Deterministic Policy Gradients (TD3).
Both PPO and TD3 are advanced versions of the actor-critic framework discussed
above, each addressing critical aspects stability, mitigating overestimation
bias, or providing better exploration strategies. For in-depth explanations
and implementation specifics, we direct readers to [29, 10].

###  5.2. RL Network Architecture

Since the observation data for our experiments is spatially structured on a
two-dimensional grid, we employ convolutional network architectures to process
the system state at each time-step. Both the actor and critic networks share
the same network architecture with the exception of the final layer which is
adapted to the specific format of the actions and value outputs, respectively.
As illustrated in Figure 1, the networks consist of three convolutional layers
with $3\times 3$ kernels each followed by $2\times 2$ max-pooling layers; the
series of convolutional layers is then followed by two dense layers and a
final linear layer is used to produce the position/variance in the case of the
actor network and the value prediction in the case of the critic network. All
layers are equipped with hyperbolic tangent activation functions, and
orthogonal initialization is used for both the convolutional and dense
weights.

|
---|---
Figure 1. Network architectures for the actor (left) and critic (right)
components of the RL models. Both networks receive input arrays input
containing the values of $\boldsymbol{z}(s)$, $\boldsymbol{y}$, and $s$.
Convolutional and max-pooling layers (blue) process the data received from the
PDE environment to extract features. These features are then flattened and
passed to dense layers (grey) to form the position, variance, and value
predictions.

We train the networks using a pair of Adam optimizers, one used for the actor
network and the other used for the critic, and apply an exponential decay to
the learning rates which are both initially set to $10^{-4}$. The observation
states are also normalized using a exponential moving average estimate for the
mean and variance of the input data in each batch.

It is essential to highlight the difference in network architectures employed
for the OC and RL approaches. The OC approach necessitates an architecture
with at least one continuous derivative with respect to the network inputs. In
contrast, the RL approach only requires policies that are differentiable with
respect to the network weights. Because of this, RL is compatible with most
off-the-shelf architectures, while architecture choices for the OC approach
are slightly more restricted.

##  6\. Numerical Results

In this section we compare the model-based and data-driven training approaches
for a control problem involving the advection diffusion equation. Since the
ground-truth solutions for these problems are not available, we use
conventional, gradient-based solutions as baselines to approximate the optimal
solutions for fixed parameters in a test set. All models were trained using 10
cores of an Intel Xeon Platinum 8176 CPU at 2.10GHz.

In our initial experiments, we trained all models using the same network
architecture; however, we found the OC and RL approaches required distinct
design choices to attain the best results. When using a convolutional
architecture for the OC training procedure, the gradient signals appeared to
be too weak, and models failed to train from the start. For the RL models,
replacing convolutional layers with a residual network led to a noticeable
drop in performance that remained after several iterations of hyperparameter
tuning. To provide a fair assessment, we present numerical results using the
architecture that yields the best performance for each approach.

###  6.1. Advection-Diffusion Problem Formulation

For our numerical tests, we consider the solution $a(\cdot,\cdot)$ to an
advection-diffusion PDE defined on a square domain
$\Omega=[0,1]\times[0,1]\subset\mathbb{R}^{2}$ as part of the system state
$\boldsymbol{z}$. The control problem is based on a two-dimensional,
continuous control vector $\boldsymbol{u}(s)=[u_{1}(s),u_{2}(s)]^{\top}$,
aiming to prevent positive values of the solution $a(\cdot,\cdot)$ from
occurring in a target region $\Omega_{\rm
targ}=\\{(x_{1},x_{2})\in\Omega\,:\,x_{1}>0.75\\}$ at the final time $T=0.75$.
A motivating contaminant control example is illustrated in Figure 2 where $a$
represents contaminant concentration (in red) originating from a source
location corresponding to fixed $\boldsymbol{y}$. The parameters
$\boldsymbol{y}$ either contains information about the source location
(Experiment 1) or both the source location and phase of advecting velocity
(Experiment 2). The goal is to prevent this contamination from reaching the
target subdomain $\Omega_{\rm targ}$ which is shown to the right of the
vertical dashed line. To achieve this, an optimal sequence of locations and
magnitudes for the sink (in blue) must be found, both of which will depend
strongly on the parameter $\boldsymbol{y}$ under consideration.

For a given $\boldsymbol{y}\in\eta$, the objective function for the problem is
given explicitly by:

(6.1) $\frac{1}{2}\int_{0}^{T}|\boldsymbol{u}(s)|^{2}{\rm
d}s+\rho\int_{\Omega_{\rm
targ}}\left|\max\left\\{a(T,(x_{1},x_{2})),0\right\\}\right|\,{\rm
d}x_{1}\,{\rm d}x_{2}$

and part of the dynamics from eq. 2.1b are prescribed by an advection-
diffusion PDE of the form

(6.2) ${\partial_{s}a}+{\rm div}(\kappa\,\nabla
a)+\varepsilon_{\boldsymbol{y}}\,\nabla
a+\varphi_{\boldsymbol{y}}\,\,=\,\,u_{1}\,Q(\alpha)$

with initial condition $a(0,x_{1},x_{2})=0$ and homogeneous Neumann boundary
conditions.

The controls $\boldsymbol{u}$ influence system dynamics through the sink term
‘$Q$’. Here, $u_{1}$ determines the magnitude of the sink while $u_{2}$
controls the sink’s velocity via the location variable $\alpha$ by:

$\alpha(t)\,\,=\,\,\alpha(0)\,+\,\int_{0}^{t}u_{2}(s)\,ds\hskip
10.84006pt\mbox{with}\hskip 10.84006pt\alpha(0)=0.5.$

For the sink term $Q$, we consider

$Q(x_{1},x_{2},\alpha(t))\,\,=\,\,25\cdot\exp\left(|x_{1}-0.6|/0.025+|x_{2}-\alpha(t)|/0.15\right).$

Figure 2. Example evolution of advection diffusion system

Each parameter $\boldsymbol{y}\in\eta$ consists of three components
$\boldsymbol{y}=[\boldsymbol{y}_{x_{1}},\boldsymbol{y}_{x_{2}},\boldsymbol{y}_{v}]$.
The parameters $\boldsymbol{y}_{x_{1}}$ and $\boldsymbol{y}_{x_{2}}$ denote
the $x_{1}$ and $x_{2}$ coordinates of the location of the source term,
$\varphi_{\boldsymbol{y}}$, given by

(6.3)
$\varphi_{\boldsymbol{y}}(x_{1},x_{2})\,\,=\,\,(c/\sigma_{s})\,\exp\left(-(|x_{1}-\boldsymbol{y}_{x_{1}}|+|x_{2}-\boldsymbol{y}_{x_{2}}|)/\sigma_{s}\right)$

while $\boldsymbol{y}_{v}$ adjusts the phase of the velocity field,
$\varepsilon_{\boldsymbol{y}}$.

For our first family of problems, we consider a fixed velocity field, but
randomized source locations, in a horizontal setup with parameters

$\varepsilon_{\boldsymbol{y}}(x_{1},x_{2})=[25,0],\quad
c=5,\quad\sigma_{s}=0.01,\quad\text{and}\quad\kappa=0.008.$

We then consider a relatively challenging family of problems, sinusoidal
setup, which incorporates a more complex velocity field and extends problem
randomization to control the phase of the velocity as well:

(6.4)
$\varepsilon_{\boldsymbol{y}}(x_{1},x_{2})\,\,=\,\,\begin{bmatrix}(1+x_{1})\cdot\sqrt{0.9^{2}-(0.75\cdot\cos(1.1-x_{1})\cdot\sin(4\pi(x_{1}-\boldsymbol{y}_{v})))^{2}}\\\
\,-0.9\cdot\cos(1.1-x_{1})\cdot\sin(4\pi(x_{1}-\boldsymbol{y}_{v}))\end{bmatrix}$

For the sinusoidal setup, we also adjust the source magnitude to $c=0.5$ and
the source size to $\sigma_{s}=0.025$ to produce more natural flows for the
modified velocity field.

We sample the components of $\boldsymbol{y}$ from the uniform distributions

$\boldsymbol{y}_{x_{1}}\sim\operatorname{Unif}(0.1,0.25)\hskip
7.22743pt,\hskip
10.84006pt\boldsymbol{y}_{x_{2}}\sim\operatorname{Unif}(0.2,0.8)\hskip
7.22743pt,\hskip
10.84006pt\boldsymbol{y}_{v}\sim\operatorname{Unif}(-0.425,0.0)$

so that the source locations are varied in the horizontal and vertical
directions across the left side of the domain. For the sinusoidal setup, the
velocity field directs concentration near the source upward for values of
$\boldsymbol{y}_{v}\approx-0.425$ and directs downward for
$\boldsymbol{y}_{v}\approx 0.0$.

Finally, to balance the importance of the running cost and terminal cost terms
in the objective function, we choose a scaling factor of $\rho=40$ in eq. 6.1.
This is generally necessary since the scale of the action space does not
always match that of the state space, and one of the loss components can
dominate the other if left unweighted.

####  6.1.1. Discretization

Since the continuous formulation results in an infinite dimensional optimal
control problem, we cannot apply numerical methods to the problem described
above directly. Instead we consider an approximate formulation arising from
the spatial discretization of eq. 6.2 using finite elements.

Denoting the discretized, FEM representation of $a(s,x_{1},x_{2})$ by
$\boldsymbol{a}(s)$, the state vector for the problem is given by:

$\boldsymbol{z}(s)\,\,=\,\,\left[\boldsymbol{a}(s)\,,\,\alpha(s)\right]^{\top}.$

The dynamics for the $\boldsymbol{a}$ component of the state vector are
defined by the linear systems associated with the FEM formulation of eq. 6.2.
By construction, the state dynamics for $\alpha$ are simply $\rm
d_{s}\alpha(s)=u_{2}(s)$.

For the spatial discretization, we have selected linear, triangular elements
with nodes associated with a uniform $32\times 32$ grid. The FEM
representation for the concentration $\boldsymbol{a}$ is passed to the neural
networks by evaluating pointwise on the grid locations
$\\{(x_{1}^{i},x_{2}^{j})\\}_{i,j=1}^{32}$ with
$x_{1}^{i}=i\cdot{\Delta}x_{1}$, $x_{2}^{j}=j\cdot{\Delta}x_{2}$, and
${\Delta}x_{1}={\Delta}x_{2}=1/31$. We further discretize the PDE in time via
implicit Euler with a step count of $N=25$ and step size ${\Delta}t=0.02$ to
obtain the fully discretized control objective

(6.5)
$J(0,\boldsymbol{z},\boldsymbol{u};\boldsymbol{y}):=\frac{1}{2}\sum_{i=0}^{25}|\boldsymbol{u}_{i}|^{2}\Delta
t\,+\,\rho\sum_{i,j=1}^{32}1_{\Omega_{\rm
targ}}(x_{1}^{i},x_{2}^{j})\left|\max\left\\{\boldsymbol{a}(T),0\right\\}\right|\Delta
x_{1}\Delta x_{2}.$

####  6.1.2. Feedback Form for HJB

To apply the proposed model-based approach procedure, we first need to derive
the feedback form for the optimal control. This expression depends on the
specific PDE under consideration and provides a natural way of incorporating
system knowledge into the training procedure. As stated in eq. 2.5, the
optimal control is a maximizer for the Hamiltonian evaluated at
$\nabla_{\boldsymbol{z}}\Phi$. By taking the gradient of the Hamiltonian and
setting the resulting expression to zero, the feedback form for both problem
setups is given by:

$\boldsymbol{u}(s)=\left[-\boldsymbol{Q}^{T}\nabla_{\boldsymbol{a}}\Phi(s),\,-\nabla_{\alpha}\Phi(s,\boldsymbol{z}(s);\boldsymbol{y})\right]^{\top}$

where $\boldsymbol{Q}$ is the finite element matrix associated with the sink
term ‘$Q$’ in eq. 6.2.

####  6.1.3. Parallel Implementation of RL environments

The FEM calculations used to simulate the PDE environment present a
significant computational bottleneck in the training procedure. Fortunately,
both the PPO and TD3 algorithms can be trained using interactions with
multiple, distinct environments in parallel. This is possible since both
algorithms decouple the actor weights used during optimization from the
network weights used in the simulations222PPO decouples weights through an ‘
_old_ ’ policy, and TD3 uses a related technique using a ‘ _target_ ’ policy.
. In particular, since both algorithms update policy weights based on
interactions performed by decoupled policies, it is possible to run multiple
episodes in parallel before calculating losses and applying gradient updates.
To take advantage of this, we employ OpenMPI and MPI for Python to parallelize
the interactions of the actor network with several different environment
realizations simultaneously. The environments are initialized in parallel,
observation data is gathered to the root node and passed to the actor/critic
networks as a single batch, and the proposed actions are then broadcast back
to the respective environments. We repeat this process until all environments
have completed a full episode. After all episodes are complete, we compute the
associated returns, group transitions from all environments into minibatches,
and perform gradient updates.

###  6.2. Hyperparameter Selection

Determining hyperparameters, including network architectures and optimization
parameters, plays a crucial role in neural network methodologies. In our
observation, the selection of hyperparameters frequently depends on the
specific approach and problem at hand. Developing well-founded mathematical
principles for parameter selection in these models represents a significant
area for future research, but it exceeds the scope of this paper. Nonetheless,
we provide our insights into hyperparameter selection for the two approaches
outlined in this paper.

The proposed HJB approach requires tuning for only a small handful of
parameters: the networks’ width and depth, the learning rate and decay
schedule, the batch size, and three weights $(\beta_{1},\beta_{2},\beta_{3})$
which determine how much the HJB constraints contribute to the loss. The RL
models considered in this work have a much larger set of hyperparameters and
require a substantial number of system solves to identify the correct training
settings. This is particularly true of scaling and clipping parameters that
are used to prevent the RL training procedure from becoming unstable. When
these parameters are not set correctly, the models’ performance will often
deteriorate midway through training, and the procedure must be restarted.

###  6.3. Baseline via Gradient-Based Solutions

To understand how the proposed, amortized procedure compares with more
conventional approaches, we now review how gradient-based methods handle this
class of control problems. These methods apply to a single problem realization
(i.e., fixed value of $\boldsymbol{y}$) and rely on gradient information,
namely $\nabla_{\boldsymbol{u}}J$, to identify an optimal control. For this
work, we have implemented a gradient-based solver using FEM matrices extracted
from FEniCS and automatic differentiation provided by PyTorch to retrieve the
necessary gradients. We then use the limited memory Broyden-Fletcher-Goldfarb-
Shanno algorithm to search the control space for an optimal solution. We use
the solutions obtained from this gradient-based search as proxies for the true
optimal solutions (which are intractable for the problems considered in this
work).

Here, it is important to note that the control obtained from the gradient-
based method is specific to the given parameter. While the control may perform
well for nearby parameters, in general, changes in the parameter require
another solve. In many cases, the original control does not even provide an
effective initialization point for new problems (since the optimal control can
change considerably when system dynamics are modified). Because of this, the
initialization procedure is another important and non-trivial aspect of
gradient-based methods. Even though these methods are designed to solve a
single problem, unlike amortized approaches that target a number of problems
at once, the gradient-based search can fail to converge to an optimal control
without proper initialization.

###  6.4. Assessment of HJB and RL Performance

For the first set of experiments, we consider the horizontal velocity setup
with two controls corresponding to the magnitude and velocity of the sink
term. The RL and HJB models are tasked with positioning the sink in locations
that block the concentration from entering the target region while using as
little movement as possible. To evaluate each model’s performance, we assess
the accuracy of the models on a set of fixed validation problems throughout
training. For the validation set, we have selected parameter values
$\boldsymbol{y}_{x_{1}}\in\\{0.125,0.225\\}$ and
$\boldsymbol{y}_{x_{2}}\in\\{0.25,0.4,0.5,0.6,0.75\\}$. This corresponds to
$10$ realizations of the horizontal setup with source locations spread across
the left side of the domain. The results of this experiment are summarized in
Figure 3 where we observe that the HJB framework yields several, noticeable
improvements over both RL models.

The most clear, and practically significant, distinction between the HJB and
RL training procedures is the rapid drop in the HJB validation loss at the
very start of training. This is perhaps to be expected thanks to the gradient
information exposed by the HJB training formulation, while the RL models are
initially reliant on an uninformed, trial-and-error state of exploration. For
practical applications, this is a key feature of the HJB training procedure;
the model is able to identify effective control policies using far fewer
forward-model PDE solves than its RL counterparts. By relying on fewer forward
solves, the HJB approach has the potential to be applied to more
computationally expensive simulations where performing the tens of thousands
of queries required by RL models is simply impractical.

Figure 3. Horizontal problem setup: (left) Validation loss during training and
(right) number of PDE solves required for different target accuracies of
control objective.

To asses how well the model-based approach framework extends to more
complicated systems, we now turn to evaluating its performance on the
sinusoidal setup. This setup introduces more complex system dynamics which
incorporate non-uniform velocity fields for advection. We also add another
layer of randomization by varying the phase of the velocity field according to
an additional problem parameter $\boldsymbol{y}_{v}$. These variations in the
velocity field are designed to reflect uncertain ambient conditions (such as
weather) which may need to be taken into account when developing control
strategies for practical applications. This form of randomization influences
how systems evolve in time and will help us gauge how well the model-based
approach and RL approaches are able to amortize over problems with
substantially different dynamics.

For the sinusoidal setup, we define the validation set using parameter values
selected from $\boldsymbol{y}_{x_{1}}\in\\{0.125,0.225\\}$,
$\boldsymbol{y}_{x_{2}}\in\\{0.25,0.4,0.5,0.6,0.75\\}$, and
$\boldsymbol{y}_{v}\in\\{-0.35,-0.2125,-0.1\\}$. This yields $30$ problem
realizations with source locations spread across the left side of the domain
and velocity fields which direct concentration near these source terms upward,
horizontally, as well as downward.

The results for the sinusoidal setup are depicted in Figure 4, where we once
again observe a clear advantage in data efficiency for the HJB model. The HJB
model achieves an average control objective of $0.15$ using only 780 PDE
solves. However, both RL models fail to achieve an average control objective
below $0.25$ even with the full 15,000 solves.

In Figure 5, we also present the suboptimality of HJB and RL approaches
compared to the baseline, as discussed in Section 6.3. It can be observed that
the HJB model achieves much lower suboptimality than RL models with
significantly fewer PDE solves. This difference is particularly pronounced in
the sinusoidal setup, where the HJB approach can achieve a suboptimality of
$0.03$ in 4740 PDE solves, while RL approaches fail to go below $0.20$.

This faster convergence can be partially attributed to the training algorithms
used in RL. While both RL models succeed in preventing positive concentrations
from reaching the target region, they do so with unnecessary movement. The RL
policies often exhibit bias in initial movement in one direction (e.g.,
drifting downward at the start of a simulation) and then correct with a sharp
reversal in direction if needed. These inefficiencies result in much larger
running costs for the RL models and prevent them from reaching the control
objective levels attained by the HJB approach.

Figure 4. Sinusoidal problem setup: (left) Validation loss during training and
(right) number of PDE solves required for different target accuracies of
control objective. Figure 5. Suboptimality, relative to the baseline, on
validation problems for the horizontal (left column) and sinusoidal (right
column) problem setups.

##  7\. Discussion

We consider deterministic, finite-time optimal control problems with unknown
or uncertain parameters and aim to approximate optimal policies for all
parameters with neural networks. We compare to training paradises. On the one
hand, our model-based approach parameterizes the value function using a neural
network and derives the policy from the feedback form. By leveraging control
theory, this approach incorporates the system’s physics by penalizing the
Hamilton-Jacobi-Bellman equation satisfied by the value function during
training. On the other hand, the data-driven approach adopts state-of-the-art
actor-critic RL frameworks and trains separate networks for the policy (actor)
and the value function (critic). Both networks are trained exclusively from
observational data. Despite drawing inspiration from control theory,
particularly Dynamic Programming, both approaches exhibit distinct behaviors
in practical applications.

Although we strived to use almost identical learning problems in both
paradigms, we had to allow for some differences to account for their
assumptions. Most importantly, our model-based approach had to use a simpler
network architecture to ensure the differentiability with respect to the
inputs needed to apply the feedback form. While non-differentiable off-the-
shelf CNN architectures worked well for the data-driven approach, the training
did not yield adequate results using the simpler architecture.

We apply both approaches to a model problem inspired by contaminant
containment to assess accuracy and efficiency. In this scenario, the state of
the system is governed by a two-dimensional advection-diffusion partial
differential equation (PDE), and the objective is to prevent the contaminant
from reaching a target region by controlling the location and intensity of a
sink, with the parameter reflecting various scenarios for the source locations
and velocity fields.

While both approaches follow an unsupervised training regime, our numerical
results reveal that the model-based approach approximates a reasonable policy
with much lower sub-optimality, when compared to a gradient-based baseline.
Notably, it achieves this with far fewer PDE solves than RL (see Figure 5).
The superior performance of the model-based approach can be attributed to the
incorporation of physics and system derivatives guided by control theory. In
contrast, lacking awareness of this information, RL demands a larger number of
PDE solutions. During RL training, many models showed promising learning for
several thousand PDE solves but eventually diverged and were unable to
recover. This phenomenon never occurred in the model-based approach.

Despite RL’s data-prohibitive nature, its non-intrusive and gradient-free
characteristics make it widely applicable and suitable for experimenting with
new problems lacking well-defined models or facing challenges in efficiently
computing feedback forms and system derivatives. As part of future work, we
aim to extend our model-based approach to such scenarios and revisit the
comparison against RL.

Another advantage of the model-based approach lies in its simpler
hyperparameter setup, requiring fewer than half the number of hyperparameters
compared to RL approaches. Effective hyperparameters identified for one
problem setup could be seamlessly applied to other problems without
modification. Conversely, RL training procedures proved sensitive to
hyperparameter choices, often necessitating re-tuning between problems for
convergence. This again emphasizes the data efficiency of the model-based
approach.

As a part of future work, we plan to compare the methods for more complex
problems and investigate their scalability using increasingly finer
discretizations.

## Acknowledgments

LR’s and DV’s work was supported in part by NSF awards DMS 1751636 and DMS
2038118, AFOSR grant FA9550-20-1-0372, and US DOE Office of Advanced
Scientific Computing Research Field Work Proposal 20-023231. NW’s and BV’s
work was supported in part by Sandia National Laboratories which is a
multimission laboratory managed and operated by National Technology and
Engineering Solutions of Sandia LLC, a wholly owned subsidiary of Honeywell
International, Inc., for the U.S. Department of Energy’s National Nuclear
Security Administration under Contract DE-NA0003525, and US Department of
Energy, and Office of Advanced Scientific Computing Research, Field Work
Proposal 20-023231.

## References

  * [1] Susan Amin, Maziar Gomrokchi, Harsh Satija, Herke van Hoof, and Doina Precup.  A survey of exploration methods in reinforcement learning.  arXiv preprint arXiv:2109.00157, 2021.
  * [2] Francesco Ballarin, Elena Faggiano, Andrea Manzoni, Alfio Maria Quarteroni, Gianluigi Rozza, Sonia Ippolito, Carlo Antona, and Roberto Scrofani.  Numerical modeling of hemodynamics scenarios of patient-specific coronary artery bypass grafts.  Biomechanics and Modeling in Mechanobiology, 16:1373–1399, 2017\.
  * [3] Somil Bansal and Claire J. Tomlin.  DeepReach: A deep learning approach to high-dimensional reachability.  In IEEE International Conference on Robotics and Automation (ICRA), pages 1817–1824, 2021.
  * [4] Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Christopher Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub W. Pachocki, Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang.  Dota 2 with large scale deep reinforcement learning.  ArXiv, abs/1912.06680, 2019.
  * [5] Dimitri P. Bertsekas.  Reinforcement learning and optimal control.  Athena Scientific Optimization and Computation Series. Athena Scientific, Belmont, MA, [2019] ©2019.  Second printing with editorial revisions.
  * [6] Rodrigo C. Carlson, Ioannis Papamichail, Markos Papageorgiou, and Albert Messmer.  Optimal motorway traffic flow control involving variable speed limits and ramp metering.  Transportation Science, 44(2):238–253, 2010.
  * [7] Minmin Chen, Alex Beutel, Paul Covington, Sagar Jain, Francois Belletti, and Ed H. Chi.  Top-k off-policy correction for a reinforce recommender system.  Proceedings of the Twelfth ACM International Conference on Web Search and Data Mining, 2018.
  * [8] Nicola Demo, Maria Strazzullo, and Gianluigi Rozza.  An extended physics informed neural network for preliminary analysis of parametric optimal control problems.  arXiv:2110.13530, 2021.
  * [9] Wendell H. Fleming and H. Mete Soner.  Controlled Markov Processes and Viscosity Solutions, volume 25 of Stochastic Modelling and Applied Probability.  Springer, New York, second edition, 2006.
  * [10] Scott Fujimoto, Herke van Hoof, and David Meger.  Addressing function approximation error in actor-critic methods.  In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1587–1596. PMLR, 10–15 Jul 2018\.
  * [11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
  * [12] Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger.  Deep reinforcement learning that matters.  In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
  * [13] Rakhoon Hwang, Jae Yong Lee, Jin Young Shin, and Hyung Ju Hwang.  Solving PDE-constrained control problems using operator learning.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 4504–4512, 2022.
  * [14] Jemin Hwangbo, Joonho Lee, Alexey Dosovitskiy, Dario Bellicoso, Vassilios Tsounis, Vladlen Koltun, and Marco Hutter.  Learning agile and dynamic motor skills for legged robots.  Science Robotics, 4, 2019.
  * [15] Alex Irpan.  Deep reinforcement learning doesn’t work yet.  2018\.
  * [16] Wei Kang, Qi Gong, and Tenavi Nakamura-Zimmerer.  Algorithms of data development for deep learning and feedback design.  arXiv:1912.00492, 2019.
  * [17] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv:1412.6980, 2014.
  * [18] Karl Kunisch and Daniel Walter.  Semiglobal optimal feedback stabilization of autonomous systems via deep neural network approximation.  ESAIM: Control, Optimisation and Calculus of Variations, 27, 2021\.
  * [19] Karl Kunisch and Daniel Walter.  Optimal feedback control of dynamical systems via value-function approximation.  arXiv:2302.13122, 2023.
  * [20] Xingjian Li, Deepanshu Verma, and Lars Ruthotto.  A neural network approach for stochastic optimal control.  arXiv:2209.13104, 2022.
  * [21] Victor G Lopez, Frank L Lewis, Yan Wan, Edgar N Sanchez, and Lingling Fan.  Solutions for multiagent pursuit-evasion games on communication graphs: Finite-time capture and asymptotic behaviors.  IEEE Transactions on Automatic Control (TAC), 65(5):1911–1923, 2019\.
  * [22] Lu Lu, Raphael Pestourie, Wenjie Yao, Zhicheng Wang, Francesc Verdugo, and Steven G Johnson.  Physics-informed neural networks with hard constraints for inverse design.  SIAM Journal on Scientific Computing, 43(6):B1105–B1132, 2021.
  * [23] Kjetil O Lye, Siddhartha Mishra, Deep Ray, and Praveen Chandrashekar.  Iterative surrogate model optimization (ISMO): an active learning algorithm for PDE constrained optimization with deep neural networks.  Computer Methods in Applied Mechanics and Engineering, 374:113575, 2021.
  * [24] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller.  Playing atari with deep reinforcement learning.  arXiv, 1312.5602, 2013.
  * [25] Saviz Mowlavi and Saleh Nabi.  Optimal control of PDEs using physics-informed neural networks.  Journal of Computational Physics, page 111731, 2022.
  * [26] Evgenii Nikishin, Pavel Izmailov, Ben Athiwaratkun, Dmitrii Podoprikhin, Timur Garipov, Pavel Shvechikov, Dmitry Vetrov, and Andrew Gordon Wilson.  Improving stability in deep reinforcement learning with weight averaging.  In Uncertainty in artificial intelligence workshop on uncertainty in Deep learning, 2018.
  * [27] Derek Onken, Levon Nurbekyan, Xingjian Li, Samy Wu Fung, Stanley Osher, and Lars Ruthotto.  A neural network approach for high-dimensional optimal control.  arXiv:2104.03270, 2021.
  * [28] Gianluigi Rozza, Andrea Manzoni, and Federico Negri.  Reduction strategies for PDE-constrained optimization problems in haemodynamics.  In Proceedings of the 6th European Congress on Computational Methods in Applied Sciences and Engineering, number CONF, pages 1748–1769. Vienna Technical University, 2012.
  * [29] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  arXiv:1707.06347, 2017.
  * [30] David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, L. Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis.  Mastering the game of go with deep neural networks and tree search.  Nature, 529:484–489, 2016.
  * [31] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, L. Sifre, Dharshan Kumaran, Thore Graepel, Timothy P. Lillicrap, Karen Simonyan, and Demis Hassabis.  A general reinforcement learning algorithm that masters chess, shogi, and go through self-play.  Science, 362:1140 – 1144, 2018.
  * [32] Maria Strazzullo, Francesco Ballarin, Renzo Mosetti, and Gianluigi Rozza.  Model reduction for parametrized optimal control problems in environmental marine sciences and engineering.  SIAM Journal on Scientific Computing, 40(4):B1055–B1079, 2018.
  * [33] Richard S Sutton and Andrew G Barto.  Reinforcement learning: An introduction.  MIT press, 2018.
  * [34] Fredi Tröltzsch.  Optimal Control of Partial Differential Equations: Theory, Methods and Applications, volume 112 of Graduate Studies in Mathematics.  2010\.
  * [35] Sifan Wang, Mohamed Aziz Bhouri, and Paris Perdikaris.  Fast PDE-constrained optimization via self-supervised operator learning.  arXiv:2110.13297, 2021.
  * [36] Jiajun Wu, Erika Lu, Pushmeet Kohli, Bill Freeman, and Joshua B. Tenenbaum.  Learning to see physics via visual de-animation.  In NIPS, 2017.
  * [37] Mengfei Xu, Shufang Song, Xuxiang Sun, Wengang Chen, and Weiwei Zhang.  Machine learning for adjoint vector in aerodynamic shape optimization.  Acta Mechanica Sinica, pages 1–17, 2021.
  * [38] Peng-Heng Yin, Guangqiang Xiao, Keju Tang, and Chao Yang.  AONN: An adjoint-oriented neural network method for all-at-once solutions of parametric optimal control problems.  arXiv:2302.02076, 2023.
  * [39] Jiongmin Yong and Xun Yu Zhou.  Stochastic controls, volume 43 of Applications of Mathematics (New York).  Springer-Verlag, New York, 1999.  Hamiltonian systems and HJB equations.
  * [40] Mo Zhou, Jiequn Han, and Jianfeng Lu.  Actor-critic method for high dimensional static Hamilton–Jacobi–Bellman partial differential equations based on neural networks.  SIAM Journal on Scientific Computing, 43(6):A4043–A4066, jan 2021\.
  * [41] Mo Zhou and Jianfeng Lu.  A policy gradient framework for stochastic optimal control problems with global convergence guarantee.  arXiv:2302.05816, 2023.
  * [42] Hao Dong Zihan Ding.  Challenges of reinforcement learning.  In Shanghang Zhang Hao Dong, Zihan Ding, editor, Deep Reinforcement Learning: Fundamentals, Research, and Applications, chapter 7, pages 249–272. Springer Nature, 2020.  http://www.deepreinforcementlearningbook.org.