Optimal control (OC) problems are ubiquitous in pure and applied mathematics, physics, computer science, engineering, finance, and elsewhere [1], [2], [3]. Thus, theoretical and numerical analyses of OC problems have paramount importance across disciplines. We focus on developing numerical solution methods for general high-dimensional OC problems.
Two of the most common strategies to solve OC problems are Pontryagin's Maximum Principle (PMP) [4] and the Hamilton-Jacobi-Bellman (HJB) PDE [5] (Sec. II).
PMP is a local solution method because optimal controls correspond to fixed initial states. Necessary conditions render an ODE system for the state and adjoint variables and a maximum principle relating the adjoint variable with the optimal control. This approach is grid-free and thus suitable for highdimensional problems. However, the aforementioned ODE system is often challenging to solve due to its forwardbackward structure [6], [7]. Additionally, the OC problem's non-convexity renders the possibility of multiple non-optimal solutions, and additional considerations are necessary, such as the differentiability of the value function [3,Theorem 7.3.9]. Unfortunately, these conditions are virtually impossible to enforce or verify numerically. The effects of nonconvexity are especially pronounced in multi-agent collisionavoidance problems [8, I.A]. For OC problems that are convex, high-dimensional solvers can be devised via primaldual convex optimization methods [9], [10], [11], [12].
Because PMP is a local solution method, shocks or sudden changes in the system's initial conditions lead to a new optimization problem. PMP is unattractive for real-time applications because the control-search time is vital.
Alternatively, the HJB approach aims at solving the OC problem for all initial states at once. More precisely, the value function-also known as the optimal cost-to-go-of an OC problem is a solution of a suitable HJB PDE. After computing the value function, we can recover the optimal control at any state from the value function's gradient. Such controls are said to be in feedback form and are especially useful for real-time applications if calculation-times for the feedback form are short.
Although effective, HJB equations are challenging to solve numerically, especially when the state's dimension d ≥ 4. First, in a deterministic setup, HJB is a first-order non-linear equation and generally does not admit smooth solutions. Second, traditional numerical methods for HJB equations, such as ENO/WENO [13], rely on grids and therefore suffer from the curse of dimensionality [5].
We propose a machine learning framework to overcome the curse of dimensionality by approximating the value function with a neural network (NN). Our approach is a fusion of PMP and HJB. Combining the PMP and HJB approaches, we express the control cost in terms of the value function and search for an NN approximation that minimizes this cost on a cloud of initial states. Furthermore, we improve the NN training by adding HJB residual penalties (Fig. 1), similar to [14], [15], [16].
Our approach has several advantages. First, it applies to generic OC problems. Second, we find controls in a feedback form that is crucial for real-time applications. Training the NN on a cloud of initial states ensures the controls' optimality on a large portion of the state-space. Consequently, controls are robust to moderate disturbances or shocks to the system (Fig. 2). Third, our method is gridfree and suitable for high-dimensional problems. Finally, our approach incorporates machine learning techniques for solving large scale optimization problems. As a result, we are able to solve 150-dimensional OC problems (Sec. V-D).
Applications of deep learning techniques to OC problems appear in seminal works [17], [18], [19], where the authors apply backward stochastic differential equations techniques to solve high-dimensional stochastic OC problems. In [20], the authors extend these methods by introducing and analyzing different loss functions. These works consider stochastic problems with fixed initial states. We focus on finding solutions to deterministic problems that are robust to shocks.
In [21], the authors first generate optimal controls for a sample of initial states then train an NN to fit this data. The data generation phase is performed by a different algorithm [7]. A similar approach exists in [8] with a different data generation algorithm [22]. In contrast, our approach directly minimizes the cost function without a generation phase.
Our work is based on the same framework as [23], which approximates the feedback control with an NN then optimizes the control cost on a cloud of initial states and provides a theoretical analysis of OC solutions via NN approximations. We extend the framework to finite horizon problems with non-quadratic costs and parameterize the value function instead of the feedback function. This latter approach enables enforcing HJB conditions, which empirically improves numerical performance for solving high-dimensional meanfield games, mean-field control, and normalizing flows [14],
[15], [16]. We demonstrate similar advantages in the OC problems considered in this work.
We consider deterministic OC problems with a particular emphasis on centrally controlled multi-agent systems. For n agents in an q-dimensional space we obtain a d = n • qdimensional OC problem. Thus, even moderate n, q yield problems that are out of reach for traditional HJB solvers.
We demonstrate the effectiveness of our method by solving a 50-agent control problem in a 3-dimensional space with obstacle and collision avoidance (Fig. 3). The overall dimension of this problem is d=150. Additionally, we demonstrate our model's robustness to shocks and the effect of the penalizers on a 4-dimensional corridor problem.
Here, we briefly recall relevant OC theory. We refer to [1, Chapters I, II] for a detailed exposition. We are interested in deterministic fixed finite time-horizon problems. Consider the time-horizon [0, T ] and the system's dynamics given by
Here, z ∈ R d describes the state of the system, and
to the application of a control u : [t, T ] → U for initial time t ∈ [0, T ] and initial state x. We assume that f, L, G, U are sufficiently regular (see [1, Sec. I.3, I.8-9] for a list of assumptions). Next, assume that the control u yields cost
where L : [0, T ] × R d × U → R is the running cost or the Lagrangian, and G : R d → R is the terminal cost. OC problems seek the control that incurs the minimal cost; i.e.,
where Φ is called the value function. A solution u * of (3) is called an optimal control. Accordingly, the z * which corresponds to u * is called an optimal trajectory. Next, the Hamiltonian of the system is given by
where p is the adjoint state. The following is a standing assumption throughout the paper.
Standing Assumption 1: Assume that (4) admits a unique continuous closed-form solution u * (t, x, p). Under this assumption and denoting ∇ as the gradient with respect to the state variables only, one has that
The PMP [4] states that for a solution (z * , u * ) of (3) there exist p :
The HJB PDE, also known as the dynamic programming equation, corresponding to ( 3) is given by
The cornerstone of the HJB approach is that the value function Φ is the unique viscosity solution of (7). Moreover, p in ( 6) and Φ are related by
Thus, the optimal control u * is given in a feedback form
and z * evolves according to
We derive a machine learning formulation of (3) by leveraging the advantages of both the PMP and the HJB approaches. In particular, we directly optimize (2) subject to (1). Rather than solving for the controls, we first postulate the dependence of (z, u) on Φ according to ( 9), (10) and parameterize Φ by an NN. We also add penalizers that punish deviations from (7), similar to [14], [16].
Denote the control variable corresponding to initial position x ∈ R d at time t=0 by u x : [0, T ] → U . And denote the trajectory corresponding to initial data (0, x) by z x ; i.e.,
Furthermore, fix a ρ 0 ∈ P(R d ) and consider the problem (11) Note that {u *
x } solves (11) if and only if u *
x is an optimal control for ρ 0 almost everywhere x ∈ R d . Thus, solving (3) with initial points x in some domain Ω is equivalent to solving (11) for ρ 0 ∈ P(R d ) such that supp(ρ 0 )=Ω. Formulation (11) is employed in mean-field game and control systems [14], [15] and in stabilization problems [23].
Postulating ( 9), (10), we arrive at our main formulation
In each OC problem, we set up and solve (12), where f , L, and thus H vary with the problem. The L contains terms that reflect the features of the problem (Sec. III-D).
We approximate Φ(•) by an NN, Φ(•; θ) (Sec. IV), and turn (12) into a finite-dimensional optimization over the weights θ. For brevity, we often omit Φ's explicit dependence on θ.
We introduce three penalty terms c HJt,x , c HJfin,x , and c HJgrad,x derived from the HJB PDE (7) as follows:
The HJ t penalizer arises from the first equation in (7), whereas HJ fin and HJ grad are direct results of the finaltime condition in (7) and its gradient, respectively. Penalizers prove helpful in training NNs for solving problems similar to (12) [14], [15], [16], [24]. They improve the training convergence (Sec. III-C) without altering the solution of (12).
Adding c HJt,x , c HJfin,x , c HJgrad,x to (12) and rewriting the time-integral in terms of ODE constraints, we obtain
subject to
initialized with z x (0) = x and x (0) = c HJt,x (0) = 0, and
We note that reformulating the Lagrangian L x (s) uses (5).
The objective function thus contains the accumulated running cost x (T ), the HJB penalty along the trajectories c HJt,x (T ), the final-time HJB penalty c HJfin,x , and
We experimentally assess the effectiveness of the penalizers c HJt , c HJfin , c HJgrad . To this end, we define six models (various combinations of the three HJB penalizers and one with weight decay) and train each on the corridor problem (Sec. V-B). Using the HJB penalizers results in the quicker model convergence on a hold-out validation set (Fig. 1).
HJ t : We enforce the PDE (7) describing the time derivative of Φ along the trajectories. Including this penalizer improves regularity and reduces the necessary number of time steps when solving the dynamics [14], [15], [16], [25].
HJ fin : We enforce the final-time condition of the PDE (7). The inclusion of this penalizer helps the network achieve the target [14]. Experimentally, using HJ fin correlates with a slightly lower G value (Fig. 1).
HJ grad :
We enforce the tranversality condition ∇Φ(T, z(T ))=∇G(z(T )), ∀z, a consequence of the final-time HJB condition (7). Numerically, all conditions are enforced on a finite sample set. Therefore, higher-order regularization may help the generalization; i.e., achieving a better match of Φ(T, •) and G for samples not used during training (the hold-out validation set). We observe the latter experimentally; using HJ grad instead of HJ fin results in lower G (Fig. 1). [21] similarly enforces ∇Φ values.
For multi-agent problems, the Lagrangian consists of three terms: an energy term E : U → R penalizing how much the agents travel, an obstacle term Q : R d → R penalizing agents at certain spatial locations (i.e., a terrain function), and an interaction term W : R d → R penalizing the proximity among agents (i.e, collision avoidance).
The n agents have initial states x 1 , . . . , x n ∈ R q . For initial joint-state x = (x 1 , . . . , x n ) ∈ R d with d = q • n, we represent the joint-state z x ∈ R d at time t as
where z xi ∈ R q is the ith agent's state. The control follows
As a result, we define the Lagrangian as
where we omit the dependence on t for brevity. The scalars α 2 , α 3 calibrate the magnitude (relative to E) of the penalization of the obstacle and interactions, respectively.
We parameterize the value function as
The inputs s=(x, t) ∈ R d+1 correspond to space-time, N (s; θ N ) : R d+1 →R m is an NN, and θ contains trainable weights: In our experiments, for N , we use a simple two-layer residual neural network (ResNet) [26]
for
We use the element-wise nonlinearity σ(x)= log(exp(x) + exp(-x)), the antiderivative of hyperbolic tangent, i.e., σ (x)= tanh(x) [14], [16].
We solve the ODE-constrained optimization problem ( 14) using the discretize-then-optimize approach [27], [28], where we define a discretization of the ODE, then optimize on that discretization. The model's forward pass uses a Runge-Kutta 4 integrator with n t time steps to approximate the constraints (15). The objective function is then computed, and automatic differentiation [29] calculates the gradient of the objective function with respect to θ. We use ADAM [30], a gradient-based stochastic method with momentum, to update the parameters θ. We iterate this process a selected number of times. For the learning rate (step size) provided to ADAM, we follow a piece-wise constant decay schedule, e.g., we divide the rate by 10 every 800 iterations (Fig. 1).
The number of time steps n t is selected a priori as a model hyperparameter. Large n t leads to high computation and training time while reducing error; meanwhile, too small n t leads to overfitting to a refinement of the time discretization of the trajectories. To avoid overfitting, we use more time steps for the hold-out validation set. For the corridor problem (Sec. V-B), we use n t =20 for training and n t =50 for validation (Fig. 2a). For the swarm problem (Sec. V-D), we use n t =26 for training and n t =80 for validation (Fig. 3). Training on a single NVIDIA Quadro RTX 8000 GPU requires about ten minutes for the corridor problem (d=4) and less than one hour for the swarm problem (d=150).
Other hyperparameters include the width of the ResNet m and the multipliers β 1 , β 2 , β 3 . In contrast, some multipliers are inherent to the OC problem-e.g., α 1 , α 2 , α 3and are used for both the baseline and NN. Our Python implementation with all tuned hyperparameters is available at https://github.com/donken/NeuralOC.
We present a comparable baseline approach and solve two multi-agent OC problems of dimensions 4 and 150.
We provide a comparable local solution method that solves the OC problem for a fixed initial state z (0) = x 0 . Applying forward Euler to the state equation and a midpoint rule to the integrals, we obtain the discrete optimization problem
where h=T /n t . We use T =1 and n t =50 and solve (21) using standard nonlinear programming techniques.
We design a four-dimensional problem in which two agents attempt to reach fixed targets on the other side of two hills. Suppose the agents with radius r=0.5 start at
. Thus, the initial and target joint-states are x 0 =[-2, -2, 2, -2] and y=[2, 2, -2, 2] , respectively. We sample from ρ 0 , which is a Gaussian centered at x 0 with an identity covariance. These sampled initial positions form the training set X. We sample again to create the validation set.
The obstacles are defined by the spatio-temporal cost function Q i , which we define in this experiment as the sum of four Gaussians. The energy terms are given by
and the dynamics are given by f (t, x, u) = u; the controls are the velocities. We model interactions via
for agents with radius r. For terminal costs, we choose We select multipliers α 1 =100, α 2 =10, 000, and α 3 =300.
The baseline and the NN solve the problem with comparable trajectories (Fig. 2a) and accuracy (Table I). The NN achieves a marginally worse value but slightly better G value (Table I). We attribute the G improvement to the finaltime HJB penalizers (Fig. 1).
We observe that approximating the value function leads to a shock-robust model. Since the controls are in the feedback form ∂ s z = -∇ p H(s, z, ∇Φ(s, z)), the model can quickly calculate updated trajectories despite shocks to the system. As an example, we consider a shock ξ ∈ R d (implemented as a random shift) to the system at time t=0.1 when solving the corridor problem for t ∈ [0, T ] (Fig. 2).
Our method is designed to handle minor shocks that stay within the space of trajectories of the initial cloud about x 0 . Our model computes a trajectory to y for many initial points. Therefore, for point x ∈ X, the model provides dynamics f (t, z x (t), u x (t)) before the shock. After the shock, the state picks up and follows the trajectory of some other point x ∈ X (Fig. 2b). Thus, the total trajectory has two portions
before and after the shock, respectively. A minor shock can thus be categorized as switching from one trajectory to another. (Fig. 2b). The NN and baseline results of the control problem along t=[0.1, 1] are similar (Table I).
Interestingly, our model extends outside the training region (Fig. 2c). Although the vast majority of NNs cannot extrapolate, our NN still solves the control problem after a major shock, demonstrating some extrapolation capabilities. We note that the NN solves the original problem for x 0 to near optimality. However, after such a large shock, the NN solves the control problem, but sub-optimally. In our example, we compare the NN's solution (Fig. 2c) with the baseline's solution for t=[0.1, 1] (Fig. 2d). The NN learned a solution where agent 2 passes through the corridor before agent 1. After the major shock, the NN applies these dynamics (Fig. 2c), while the baseline finds a more optimal solution (Fig. 2d). The NN is roughly 100% suboptimal (Table I).
We attribute the shock robustness to the NN parameterization of the global value function. Experimentally, the shock robustness of our model (Fig. 2) does not noticeably differ from a model trained without penalization. Since the NN is trained prior and offline, it handles shocks in real-time. In contrast, methods that solve for a single trajectory-e.g., the baseline-must pause to recompute following a shock.
We demonstrate the high-dimensional capabilities of our model by solving a swarm trajectory planning problem in the spirit of [22]. The swarm problem contains 50 threedimensional agents that fly from initial to target positions while avoiding each other and obstacles. We construct Q i to model rectangular prism obstacles and use (22), (23), and (24) for energy, interaction, and terminal costs. The NN guides all agents around the obstacles (Fig. 3). Naturally, if the time discretization is too coarse (small n t ), the model may simulate collisions solely due to inaccurate integration. In validation, we see that the agents avoid the obstacles and each other by observing values for Q and W are exactly 0.
We formulate and demonstrate an efficient NN approach for solving high-dimensional OC problems. Our method aims at computing the optimal control in feedback form in the relevant subset of the space-time domain. It combines the high-dimensional scalability from PMP and the global nature from HJB approaches. Using a numerical example, we demonstrate that the obtained feedback form generalizes outside the training space, which allows the agents to react to unforeseen events such as shocks. Our future endeavors relate to further experimentation of our method on OC problems with more involved dynamics and blending our method (trained prior) with distributed approaches in deployment.
We thank Reza Karimi for assisting with figure creation.