Pontryagin Neural Operator for Solving General-Sum Differential Games with Parametric State Constraints

(1) Convergence without supervision: The convergence issue in solving discontinuous HJI values via PINN has been investigated in Zhang et al. (2023a), where the authors proposed to hybridize PINN using supervised equilibria data generated by solving games offline via PMP with sampled initial states. The major limitation of the hybrid method is in its assumption that knowledge already exists regarding where the informative equilibria should be sampled in the vast state space, e.g., collision and near-collision cases in vehicle interactions. In addition, generating supervisory data via solving PMP encounters its own numerical issues due to the existence of multiple equilibria and singular arcs. To the best of the authors’ knowledge, this paper is among the first to address the PINN convergence issue without using any supervised data or domain knowledge of corner cases. To do so, we sample costate trajectories based on a sampling distribution defined over the state space, and evolve this distribution along the training according to the costate loss landscape. Considering costate losses as PMP-driven constraints, the sampling distribution essentially captures the dual variables of the PINN. We conduct numerical studies using a two-vehicle game at an uncontrolled intersection. Our experimental results reveal that, with the same data complexity, PNO surpasses the hybrid neural operator (hybrid PINN + DeepONet) in safety performance across a range of parametric collision zones. (2) Solving parametric HJI via value decomposition: This paper is also among the first to study the efficacy of physics-informed neural operators in the context of parametric HJI equations. We show that the DeepONet architecture essentially learns a decomposition of the parametric value and identifies the major value basis functions. Empirical results on the same two-vehicle game show that the safety performance of the learned neural operator across a parametric set of collision zones closely aligns with the ground truth. This investigation opens up new directions towards explainable data-driven models for policy learning: Can we translate basis value (or costate) functions into explainable sub-policies (e.g., temporal logic rules) that together comprise generalizable policies for parametric sets of games?

The existence of value for zero-sum differential games with Markovian payoffs and state constraints (and temporal logic specifications in general) have been derived (Bettiol et al., 2006). The results have been extended to general-sum games with non-temporal state constraints (Zhang et al., 2018). To facilitate value approximation using level set methods, an epigraphical technique has been introduced to smooth discontinuous values when state constraints are present (Gammoudi and Zidani, 2023).

HJI equations resulting from differential games are first-order nonlinear PDEs that suffer from CoD when solved by conventional level set methods (Osher et al., 2004; Mitchell and Templeton, 2005) or essentially nonoscillatory schemes (Osher and Shu, 1991). Monte Carlo methods, such as variants of PINN (Krishnapriyan et al., 2021), have recently shown promise at solving high-dimensional PDEs, including HJ equations, provided that the solution is smooth (E et al., 2021). PINN uses a trainable neural net to represent the solution space, and optimizes PDE-driven errors to approximate the solution. Such errors include the boundary residual (Han and Long, 2020; Han et al., 2018), the PDE residual (Jagtap et al., 2020; Bansal and Tomlin, 2021), and supervisory data on the characteristic curves of the PDE (Nakamura-Zimmerer et al., 2021). Convergence and generalization of PINN have been analyzed under the assumption that both the solution and the network are Lipschitz continuous (Han and Long, 2020; Shin et al., 2020; Ito et al., 2021). Recent studies have explored the effectiveness of PINN for solving PDEs with discontinuous solutions (Jagtap et al., 2020) and solutions with large Kolmogorov n-width (Mojgani et al., 2023). However, solving PDEs with discontinuous solutions and only terminal boundaries (such as HJI equations with state constraints) is still an open challenge without prior structural knowledge about the value landscape (Zhang et al., 2023b).

Neural operators emerged as a promising approach to universally approximating mappings between functions (Kovachki et al., 2023), with particularly successful applications to solving parametric PDEs (Wang et al., 2021). The architecture most related to this paper is DeepONet (Lu et al., 2021), which is composed of a branch net and a trunk net. The branch net extracts informative features of discretized input functions that define parameters of PDEs and the trunk net captures the basis functions that comprise parametric PDE solutions. DeepONet is one of the neural operator architectures that allow point-wise prediction, instead of predicting the entire function over its input domain (e.g., FNO (Li et al., 2020) and GNOT (Hao et al., 2023)). The extension of neural operators from supervised to PINN training has been studied for solving parametric physics equations defined on 2D and 3D state spaces (Wang et al., 2021). This paper examines the efficacy of the DeepONet architecture in physics-informed training of values with large Lipschitz constants defined on a 5D state-time space and a 2D parameter space. We compare two different neural operators empirically: the proposed Pontryagin neural operator and a hybrid one as a baseline.

Let the state dynamics of Player $i$ be $\dot{x}_{i}=f_{i}(x_{i},u_{i})$, where $x_{i}\in\mathcal{X}_{i}\subseteq\mathbb{R}^{d_{x}}$ is the system state and $u_{i}\in\mathcal{U}_{i}\subseteq\mathbb{R}^{d_{u}}$ is the control input. The joint state space is $\mathcal{X}:=\mathcal{X}_{1}\times\mathcal{X}_{2}$. The instantaneous loss of Player $i$ is $l_{i}(\cdot,\cdot):\mathcal{X}\times\mathcal{U}_{i}\rightarrow\mathbb{R}$ and the terminal loss $g_{i}(\cdot):\mathcal{X}_{i}\rightarrow\mathbb{R}$. The fixed time horizon of the game is $[0,T]$. We use $\textbf{a}_{i}=(a_{i},a_{-i})$ to concatenate elements $a_{i}$ from Player $i$ and $a_{-i}$ from the fellow player, and use $\textbf{a}=(a_{1},a_{2})$. Denote $\alpha_{i}\in\mathcal{A}:\mathcal{X}\times[0,T]\rightarrow\mathcal{U}_{i}$ as Player $i$’s policy. We use $x_{s}^{x_{i},t,\alpha_{i}}$ as the state of Player $i$ at time $s$ if he follows policy $\alpha_{i}$, dynamics $f_{i}$, and starts at $(x_{i},t)$. $\textbf{x}_{s}^{\textbf{x}_{i},t,\boldsymbol{\alpha}_{i}}:=\left(x_{s}^{x_{i},t,\alpha_{i}},x_{s}^{x_{-i},t,\alpha_{-i}}\right)$. Let $c_{i}(\cdot):\mathcal{X}\rightarrow\mathbb{R}$ be a state penalty derived from Player $i$’s state constraints, i.e., for any $\textbf{x}_{i}\in\mathcal{X}$, $c_{i}(\textbf{x}_{i})=0$ if $\textbf{x}_{i}$ satisfies Player $i$’s state constraints, and otherwise $c_{i}(\textbf{x}_{i})$ becomes a large positive number. In this paper, we consider $c_{i}(\cdot)$ to be differentiable but with a large Lipschitz constant. The value of Player $i$ is denoted by $\vartheta_{i}(\cdot,\cdot):\mathcal{X}\times[0,T]\rightarrow\mathbb{R}$. We omit arguments to $f_{i}$, $l_{i}$, $g_{i}$, $c_{i}$, $\vartheta_{i}$ when possible, and decorate them with the superscript $\theta\in\Theta\subseteq\mathbb{R}^{d_{\theta}}$ when the corresponding functions are type-specific, where $\Theta$ is the type space. E.g., $c^{\theta}_{i}(\cdot,\cdot)$ is the state penalty of Player $i$ of type $\theta$. Throughout the paper, we assume that $\mathcal{U}_{i}$ is compact and convex; $f_{i}$ and $c_{i}$ are Lipschitz continuous; $l_{i}$ and $g_{i}$ are Lipschitz continuous and bounded.

Let $\boldsymbol{\alpha}^{\dagger}$ be a pair of equilibrium policies of the payoffs

for $i\in\{1,2\}$, so that

The value for Player $i$ is $\vartheta_{i}(\textbf{x}_{i},t)=J_{i}(\textbf{x}_{i},t,\boldsymbol{\alpha}^{\dagger}_{i})$. The HJI equations that govern the values, denoted by $L$, and the boundary conditions, by $D$, are the following Starr and Ho (1969):

Therefore, Player $i$’s equilibrium policy can be derived as $\alpha_{i}^{\dagger}(\textbf{x}_{i},t)=\operatorname*{arg\,max}_{u\in\mathcal{U}_{i}}\{\nabla_{\textbf{x}_{i}}\vartheta_{i}^{T}\textbf{f}_{i}-(l_{i}+c_{i})\}$ if Eq. (3) can be solved for $(\vartheta_{1},\vartheta_{2})$ (Bressan, 2010). When needed, we denote by $\mathcal{L}^{\boldsymbol{\theta}}:=(L^{\boldsymbol{\theta}},D^{\boldsymbol{\theta}})$ the HJI of a game parameterized by ${\boldsymbol{\theta}}$.

The characteristic curves of $\mathcal{L}$ are open-loop equilibrium trajectories governed by PMP, which entails the following for initial state $(\bar{x}_{1},\bar{x}_{2})\in\mathcal{X}$ and $t\in[0,T]$:

Here $\lambda_{i}=\nabla_{x_{i}}\vartheta_{i}$ is the costate of Player $i$. Solving Eq. (4) for a given initial states in $\mathcal{X}\times[0,T]$ is a boundary-value problem (BVP). Provided that convergence can be achieved, the solution to this BVP approximates $\boldsymbol{\vartheta}$ through the resultant characteristic trajectory and the boundary conditions $D$. Zhang et al. (2023a) uses these characteristic trajectories as supervisory data, but assumes that informative initial states, i.e., those for which $c_{i}$ changes significantly along the trajectories, are known. PNO exploits the method of characteristics without requiring prior knowledge. See Sec. 4.

Following standard treatment in neural operators, we introduce an input function $a(\cdot,\cdot):\mathcal{X}\times\Theta^{2}\rightarrow(0,1)$ to encode parametric settings of state constraints: $a(\textbf{x},\boldsymbol{\theta})=1$ if x violates the constraints according to $\boldsymbol{\theta}$, or otherwise $a(\textbf{x},\boldsymbol{\theta})=0$. Let $X\in\mathbb{R}^{L\times d_{x}}$ be a lattice of $\mathcal{X}$, we denote by $a(X,\boldsymbol{\theta})\in\{0,1\}^{L}$ the batch Boolean outputs at all $L$ lattice nodes. Then $\hat{\vartheta}$ is defined as a linear combination of function bases:

where the branch network $\{b_{k}\}:\{0,1\}^{L}\rightarrow\mathbb{R}$ encodes the PDE parameters into function coefficients, and the trunk network $\{t_{k}\}:\mathbb{R}^{d_{x}}\times[0,T]\rightarrow\mathbb{R}$ encodes the input information into basis function values.

We introduce an additional costate network $\hat{\lambda}(\cdot,\cdot):\mathcal{X}\times[0,T]\rightarrow\mathcal{X}_{*}$, where $\mathcal{X}_{*}$ is the costate space. We can now evaluate the state, approximate costate, and approximate value trajectories starting from some sampled state and time $(\textbf{x},t)$ following their corresponding dynamics defined in Eq. (4), first by sequentially maximizing the Hamiltonian using forward approximate $\hat{\lambda}$ to derive the state trajectory $\hat{\textbf{x}}\in\mathcal{X}^{K}$, and then using the transversality conditions of the approximate costate and value to derive approximate costate trajectory $\tilde{\lambda}\in\mathcal{X}_{*}^{K}$ and approximate value trajectory $\tilde{\vartheta}\in\mathbb{R}^{K}$ backward in time, where $K=\frac{T-t}{\Delta t}$ and $\Delta t$ is a small time interval. Given a batch of $N$ initial states, we collect forward approximate state trajectories $\hat{\mathcal{X}}:=\{\hat{\textbf{x}}^{(n)}\}_{n=1}^{N}$, forward and backward approximate costate trajectories, $\hat{\Lambda}:=\{\hat{\lambda}^{(n)}\}_{n=1}^{N}$ and $\tilde{\Lambda}:=\{\tilde{\lambda}^{(n)}\}_{n=1}^{N}$ respectively, and backward approximate value trajectories $\tilde{\mathcal{V}}:=\{\tilde{\vartheta}^{(n)}\}_{n=1}^{N}$. Lastly, we denote the conventional PINN dataset for formulating the PDE and boundary losses as $\mathcal{D}:=\{(\textbf{x},t)^{(n)}\in\mathcal{X}\times[0,T]\}_{n=1}^{N_{L}}$ and $\mathcal{D}_{D}:=\{\textbf{x}^{(n)}\in\mathcal{X}\}_{n=1}^{N_{D}}$. The PNO training loss with respect to value and costate functions of both players $(\hat{\boldsymbol{\vartheta}},\hat{\boldsymbol{\lambda}})$ is then defined as:

where $\hat{\vartheta}_{i}^{(n)}$, $\tilde{\vartheta}_{i}^{(n)}$, $\hat{\lambda}_{i}^{(n)}$, $\tilde{\lambda}_{i}^{(n)}$ are abbreviations for $\hat{\vartheta}(\textbf{x}_{i}^{(n)},t^{(n)},\boldsymbol{\theta}_{i}^{(n)})$, $\tilde{\vartheta}_{i}(\textbf{x}_{i}^{(n)},t^{(n)},\boldsymbol{\theta}_{i}^{(n)})$, $\hat{\lambda}_{i}(\textbf{x}_{i}^{(n)},t^{(n)})$, $\tilde{\lambda}(\textbf{x}_{i}^{(n)},t^{(n)})$ respectively. $C_{i}>0$ for $i=1,..,5$ are hyperparameters that balance the loss terms. We summarize the training of PNO in Alg. 5. Note that for the Pontryagin-mode losses to be differentiable with respect to $(\hat{\boldsymbol{\vartheta}},\hat{\boldsymbol{\lambda}})$, we use costate net and ODE solver (RK45) to compute the forward and backward trajectories $(\hat{\mathcal{X}},\hat{\Lambda},\tilde{\Lambda},\tilde{\mathcal{V}})$. Line 9-12 in Alg. 5 outline the computation.

For standard PINN residual and boundary datasets $\mathcal{D}$ and $\mathcal{D}_{D}$, we use a curriculum learning scheme following Bansal and Tomlin (2021); Krishnapriyan et al. (2021), where states are sampled from an expanding time window starting from the terminal time, to observe temporal causality of value functions. For Pontryagin-mode trajectories $(\hat{\mathcal{X}},\hat{\Lambda},\tilde{\Lambda},\tilde{\mathcal{V}})$, our experimental results show that an importance sampling strategy is critical. We start by sampling $N$ initial states uniformly in $\mathcal{X}$. Subsequently, we follow the evolutionary sampling algorithm (Daw et al., 2022) to compute absolute values of the PDE residual $r(\textbf{x}_{i}^{(n)})=\|L(\hat{\vartheta}_{i}^{(n)},\nabla_{\textbf{x}_{i}^{(n)}}\hat{\vartheta}_{i}^{(n)},\textbf{x}_{i}^{(n)},t^{(n)})\|_{1}$ for $n=1,...,N$ using Eq. (3)¹¹1Our empirical studies showed that evolution based on costate losses achieves similar safety performance. Analytical understanding about the sampling-performance relation is yet to be established.. We iteratively sample and remove samples with residuals lower than the average. The resultant batch represents initial states for which the current $\hat{\boldsymbol{\vartheta}}$ cannot generalize well. The sampled initial states gradually accumulate in the region with informative collision and near-collision knowledge during the training. Convergence of the evolutionary-based PINN has been proved in Daw et al. (2022).

PNO does not require offline data collection through solving BVPs and thus avoids typical convergence issues encountered in solving differential games with multiple local equilibria and slow convergence dynamics. While a full probably approximately correct (PAC) learning proof of PNO is not yet available, our results suggest that PNO achieves more data-efficient learning of values and policies than an existing hybrid method that takes advantage of informative samples that reveal structures of the value or costate landscapes. One hypothesis for its success is that PNO decomposes the value landscape into (terminal) boundaries and value dynamics along trajectories (a Lagrangian perspective), i.e., the hard-to-learn value with large Lipschitz constant can potentially be learned more effectively when each of its components, namely, the continuous boundary, continuous trajectories, and discontinuous value dynamics, are easy to learn given the others. We note that the improvement in data-efficiency of PINN through the method of characteristics has been discussed in Mojgani et al. (2023) for 2D convection-diffusion and Burger’s equations, but not yet in the context of high-dimensional HJ equations for optimal control or differential games.

We adopt settings of the uncontrolled intersection case studied in Zhang et al. (2023a): Each vehicle is represented by states $x_{i}:=(d_{i},v_{i})$. Both vehicles follow dynamics: $\dot{d}_{i}=v_{i}$ and $\dot{v}_{i}=u_{i}$, where $u_{i}\in[-5,10]m/s^{2}$ is the control input. Let $R$, $L$, and $W$ be the road length, vehicle length, and vehicle width, respectively. The instantaneous loss and state constraint are defined as:

where

Here shape parameter $\gamma=5$, and $b=10^{4}$ is sufficiently large to prevent collisions. Similar to treating collision as hard constraints, this setting causes the PINN convergence issue. $\theta\in\Theta:=\{1,2,3,4,5\}$ denotes the aggressiveness level (safety preference) of a player. The terminal loss is defined to encourage players to pass the intersection and regain nominal speed:

where $\mu=10^{-6}$, $\bar{v}=18m/s$, and $T=3s$.

To demonstrate the advantage of PNO, we perform a comparison in favor of the baseline (the hybrid neural operator): For the baseline, we generate 1k ground truth trajectories from initial states uniformly sampled in an informed state set $\mathcal{X}_{GT}:=[15,20]m\times[18,25]m/s$ by solving Eq. (4). This set is identified to contain a significant amount of corner cases (e.g., $61\%$ of $\mathcal{X}_{GT}$ yield collisions when $\boldsymbol{\theta}=(5,5)$, see Tab. 1) and will also be used for test. Each trajectory consists of 31 data points with a time interval of $0.1s$ for each of the players, resulting in a total of 62k data points. For PNO, the costate net explores the initial state within a larger state set $\mathcal{X}_{HJ}:=[15,105]m\times[15,32]m/s$ and generates 1k closed-loop trajectories including 62k data points. $\mathcal{X}_{HJ}$ is much less informed than $\mathcal{X}_{GT}$ (less collisions). To improve training efficiency, we eliminate data points generated by the costate net that fall beyond the state bounds and normalize the input data to $[-1,1]$ for the trunk and costate net of both neural operators. Additionally, the trunk net uniformly samples 60k states in $\mathcal{X}_{HJ}:=[15,105]m\times[15,32]m/s$ for model refinement. Lastly, we train the model using four player-type configurations $(\theta_{1},\theta_{2})=\{(1,1),(1,5),(5,1),(5,5)\}$ and evaluate the model generalization performance within the type space $\Theta$. We reiterate that the total number of data points for the two methods are approximately the same.

The neural operators utilize fully-connected networks with 3 hidden layers, each comprising 64 neurons with tanh activation. The hybrid neural operator uses the Adam optimizer with a fixed learning rate of $2\times 10^{-5}$ for pre-training the network over 100k iterations using supervised data. It then combines the supervised data with states sampled from an expanding time window, starting from the terminal time, to refine the model for an additional 200k iterations. PNO uses the Adam optimizer with an adaptive learning rate, starting from $2\times 10^{-5}$, and is initially trained to satisfy the boundary conditions over 50k iterations. Subsequently, the network undergoes refinement through 3k gradient steps per epoch, encompassing a total of 10 epochs for every 10 training iterations. The costate net resamples the initial state to generate closed-loop trajectories every 10 epochs, and these data points are used to train the model. The overall network refinement process spans 300 training iterations. Both PNO and hybrid incorporate adaptive activation functions (Jagtap et al., 2020). In terms of wall-clock computational cost, the hybrid neural operator requires approximately 37 hours, including 5.5 hours for data acquisition and 31.5 hours for model training. PNO requires 27 hours for model training. For comparison, solving all 25 individual HJI PDEs in $\Theta^{2}$ with an acceptable meshgrid resolution using an efficient level set package (Bui et al., 2022) would require $>50$ hours. The resultant value approximation does not guarantee safety performance since value gradient approximation error is not controlled.

We use the neural operators to compute $\nabla_{\textbf{x}_{i}}\hat{\vartheta}_{i},\forall i=1,2$ as closed-loop control on test cases uniformly sampled in $\mathcal{X}_{GT}$ and report their safety performance in percentage of collisions, where collisions are defined by $\boldsymbol{\theta}$. For better transparency, the comparison uses two sets of ground truth trajectories: The first contains 600 trajectories for each parameter configuration in $\Theta^{2}$; the second uses the same setting but only contains trajectories without physical collisions (for initial states from which collision cannot be avoided according to $\boldsymbol{\theta}=(1,1)$). Tab. 1 and Tab. 2 summarize safety performance comparisons with and without inevitable collisions, generalizing to testing cases across $\Theta^{2}$. The results show that PNO outperforms the hybrid neural operator in most cases. More importantly, it achieves this without relying on domain knowledge (informative trajectories in $\mathcal{X}_{GT}$). Trajectories from test player types are visualized in Fig. 2. It should be noted that in some cases neural operators perform better than the ground truth. This is because a multi-start BVP solver solves the ground truth trajectories with heuristic initial guesses. We also note that while BVP solutions are open-loop, their corresponding values are consistent with the HJI equations for this case study (see comparison between BVP and HJI value contours at $t=0$ in Fig. 4).

Since branch and trunk nets learn a decomposition of the value function in the parameter space, Fig. 3 investigates the learned value structure through visualization of (1) the major basis value functions associated with the 6 leading means of absolute coefficients ($b_{k}$ in Eq. (5)) over $\Theta^{2}$ (Fig. 3a), (2) the ground truth values solved using level set for $\boldsymbol{\theta}=(1,1),~{}(1,5),~{}(5,1),~{}(5,5)$ (Fig. 3b), and (3) the sorted mean and standard deviation of absolute coefficients over $\Theta^{2}$ along the output dimension of branch net, where the dominant basis value functions correspond to larger $|b_{k}|$ (Fig. 3c). Further regularization on the branch net to promote basis sparsity and extension of existing generalization results to neural operators are worth investigating in future studies.

Due to the reported importance of activation choice in PINN, we conduct ablation studies to understand the influence of activation on safety performance. Tab. 3 summarizes the safety performance using different activation functions for comparisons w/ and w/o inevitable collisions. The comparison of closed-loop trajectories is visualized in Fig. 5. The results confirm that the choice of activation affects safety performance: tanh outperforms sin and relu.