Bridging Physics-Informed Neural Networks with Reinforcement Learning:
  Hamilton-Jacobi-Bellman Proximal Policy Optimization (HJBPPO)
This paper introduces the Hamilton-Jacobi-Bellman Proximal Policy
Optimization (HJBPPO) algorithm into reinforcement learning. The
Hamilton-Jacobi-Bellman (HJB) equation is used in control theory to evaluate
the optimality of the value function. Our work combines the HJB equation with
reinforcement learning in continuous state and action spaces to improve the
training of the value network. We treat the value network as a Physics-Informed
Neural Network (PINN) to solve for the HJB equation by computing its
derivatives with respect to its inputs exactly. The Proximal Policy
Optimization (PPO)-Clipped algorithm is improvised with this implementation as
it uses a value network to compute the objective function for its policy
network. The HJBPPO algorithm shows an improved performance compared to PPO on
the MuJoCo environments.
Amartya Mukherjee, Jun Liu
Bridging Physics-Informed Neural Networks with Reinforcement Learning: Hamilton-Jacobi-Bellman Proximal Policy Optimization (HJBPPO)

###### Abstract

This paper introduces the Hamilton-Jacobi-Bellman Proximal Policy Optimization (HJBPPO) algorithm into reinforcement learning. The Hamilton-Jacobi-Bellman (HJB) equation is used in control theory to evaluate the optimality of the value function. Our work combines the HJB equation with reinforcement learning in continuous state and action spaces to improve the training of the value network. We treat the value network as a Physics-Informed Neural Network (PINN) to solve for the HJB equation by computing its derivatives with respect to its inputs exactly. The Proximal Policy Optimization (PPO)-Clipped algorithm is improvised with this implementation as it uses a value network to compute the objective function for its policy network. The HJBPPO algorithm shows an improved performance compared to PPO on the MuJoCo environments.

**keywords:** Continuous-Time Reinforcement Learning, Physics-Informed Neural Networks, Proximal Policy Optimization, Hamilton-Jacobi-Bellman Equation

## 1 Introduction

In recent years, there has been a growing interest in Reinforcement Learning (RL) for continuous control problems. RL has shown promising results in environments with unknown dynamics through a balance of exploration in the environment and exploitation of the learned policies. Since the advent of REINFORCE with Baseline, the value network in RL algorithms has shown to be useful towards finding optimal policies as a critic network [21]. This value network continues to be used in state-of-the-art RL algorithms today.

In discrete-time RL, the value function estimates returns from a given state as a sum of the returns over time steps. This value function is obtained from solving the Bellman Optimality Equation. On the other hand, in continuous-time RL, the value function estimates returns from a given state as an integral over time. This value function is obtained by solving a partial differential equation (PDE) known as the Hamilton-Jacobi-Bellman (HJB) equation [12]. Both equations are difficult to solve analytically and numerically, and therefore the RL agent must explore the environment and make successive estimations.

Currently existing algorithms in the RL literature such as Proximal Policy Optimization (PPO) aim to update the value function using the Bellman Optimality Equation so that it estimates the discrete-time returns for each state. However, we discovered that this value function, when trained on MuJoCo environments, does not show convergence towards the
optimal value function as described by the HJB equation (see Figure 4). This shows that information is lost when the value function is trained using discrete time steps rather than continuous time.

The introduction of physics-informed neural networks (PINNs) by [17] has led to significant advancements in scientific machine learning. PINNs leverage auto-differentiation to compute derivatives of neural networks with respect to their inputs and model parameters exactly. This enables the laws of physics (described by ODEs or PDEs) governing the dataset of interest to act as a regularization term for the neural network. As a result, PINNs outperform regular neural networks on such datasets by taking advantage of the underlying physics of the data.

To the best of our knowledge, this paper is the first to examine the intersection between PINNs and RL. In order to force the convergence of the value function in PPO towards the solution of the HJB equation, we utilize PINNs to encode this PDE and bridge the information gap between returns computed over discrete time and continuous time. This allows our algorithm to utilize auto-differentiation to eliminate the error associated with gradient computation and discretization of time. We propose the Hamilton-Jacobi-Bellman Proximal Policy Optimization (HJBPPO) algorithm, which demonstrates superior performance in terms of higher rewards, faster convergence, and greater stability compared to PPO on MuJoCo environments, making it a significant improvement.

## 2 Preliminaries

Consider a controlled dynamical system modeled by the following equation:

\[\dot{x}=f(x,u),\quad x(t_{0})=x_{0}, \tag{1}\]

where \(x(t)\) is the state and \(u(t)\) is the control input. In control theory, the optimal value function \(V^{*}(x)\) is useful towards finding a solution to control problems [8]:

\[V^{*}(x)=\sup_{u}\int_{t_{0}}^{\infty}\gamma^{t}R(x(\tau;t_{0},x_{0},u(\cdot)),u(\tau))d\tau, \tag{2}\]

where \(R(x,a)\) is the reward function and \(\gamma\) is the discount factor. The following theorem introduces a criteria for assessing the optimality of the value function [[11], [13]].

**Theorem 2.1**.: _A function \(V(x)\) is the optimal value function if and only if:_

1. \(V\in C^{1}(\mathbb{R}^{n})\) _and_ \(V\) _satisfies the Hamilton-Jacobi-Bellman (HJB) Equation_ \[V(x)\ln\gamma+\sup_{u\in U}\{R(x,u)+\nabla_{x}V^{T}(x)f(x,u)\}=0\] (3) _for all_ \(x\in\mathbb{R}^{n}\)_._
2. _For all_ \(x\in\mathbb{R}^{n}\)_, there exists a controller_ \(u^{*}(\cdot)\) _such that:_ \[V(x)\ln\gamma+R(x,u^{*}(x))+\nabla_{x}V^{T}(x)f(x,u^{*}(x))\] \[=V(x)\ln\gamma+\sup_{u\in U}\{R(x,\tilde{u})+\nabla_{x}V^{T}(x)f( x,\tilde{u})\}.\] (4)

Currently existing algorithms in RL do not focus on solving the HJB equation to maximize the total reward for each episode. For example, in PPO, the HJB equation does not seem to be satisfied when tested on MuJoCo environments.

To show this, we define the HJB loss at each episode as the following:

\[MSE_{f}\] \[=\frac{1}{T}\sum_{t=0}^{T-1}|V(x_{t})\ln\gamma+R(x_{t},a_{t})+ \nabla_{x}V^{T}(x_{t})f(x_{t},a_{t})|^{2}, \tag{5}\]

where \(T\) is the number of timesteps in the episode, \(x_{t}\) is the state of the environment at timestep \(t\), and \(a_{t}\) is the action taken at timestep \(t\). \(\nabla_{x}V^{T}(x_{t})\) is computed exactly using auto-differentiation. We approximate \(f(x_{t},a_{t})\) using finite differences:

\[MSE_{f}=\frac{1}{T}\sum_{t=0}^{T-1}|V(x_{t})\ln\gamma+R(x_{t},a _{t})\\ +\nabla_{x}V^{T}(x_{t})(\frac{x_{t+1}-x_{t}}{\Delta t})|^{2}, \tag{6}\]
where \(\Delta x\) is the time step size used in the environment. We have plotted the HJB loss for each environment using PPO in Figure 4. The mean HJB loss for each environment takes extremely high values and does not show convergence in 6 out of 10 of the environments, thus showing that the value function does not converge to the optimal value function as shown by the HJB equation.

As a comparison, we have plotted the graphs for the value network loss in Figure 5. The Bellman optimality loss shows convergence in 8 out of the 10 environments. This shows that information is lost when we solve the solve the Bellman optimality equation for a discrete-time value function compared to continuous-time value function. It also shows that convergence of the value function does not necessarily lead to convergence in the HJB loss.

To solve this problem as shown in Figure 4, we treat the value network as a PINN and use gradient-based methods to reduce the HJB loss.

## 3 Related Work

The use of HJB equations for continuous RL has sparked interest in recent years among the RL community as well as the control theory community, and has led to promising works. [10] introduced an alternate HJB equation for Q Networks and used it to derive a controller that is Lipschitz continuous in time. This algorithm has shown improved performance over Deep Deterministic Policy Gradient (DDPG) in three out of the four tested MuJoCo environments without the need for an actor network. [24] introduced a distributional HJB equation to train the FD-WGF Q-Learning algorithm. This models return distributions more accurately compared to Quantile Regression TD (QTD) for a particle-control task. Finite difference methods are used to solve this HJB equation numerically. We intend to build up from these works by adding a policy network and by incorporating PINNs to solve the HJB equation. Furthermore, the authors mentioned the use of auto-differentiation for increased accuracy of the distributional HJB equation as a potential area for future research in their conclusion. Our work relaxes the requirement for the controller to be Lipschitz, and it minimizes the computational error associated with finite difference methods.

The use of neural networks to solve the HJB equation has been an area of interest across multiple research projects. [7] uses a structured Recurrent Neural Network to solve for the HJB equation and achieve optimal control for the Dubins car problem. [22] uses the Pineda architecture [15] to estimate partial derivatives of the value function with respect to its inputs. They used iterative least squares method to solve for the HJB equation. This algorithm shows convergence in several control problems without the need for an initial stable policy. [20] develops the DGM algorithm to solve PDEs. They use auto-differentiation to compute first order derivatives and Monte-Carlo methods to estimate higher order derivatives. This algorithm was used to solve the HJB equation to achieve optimal control for a stochastic PDE and achieved an error of 0.1%. We intend to further advance from these works and use a PINN to solve for the HJB equation.

The use of a PINN to solve the HJB equation for the value network was done by [14] in an optimal feedback control problem setting. This mitigates the use of finite differences to compute derivatives of the value function. The paper achieves results similar to that of the true optimal control function in high dimensional problems. We intend to build up from this work by using PINNs in a RL setting where the dynamics are unknown and exploration is needed.

## 4 Hjbppo

To our knowledge, our work is the first to combine the HJB equation with a currently existing RL algorithm, PPO. It is also the first to use a PINN to solve the HJB equation in a RL setting.

The PPO-Clipped algorithm is improvised with this implementation because it uses a value network to compute advantages used to update its policy network [18]. PPO is an actor-critic method that limits the update of the policy network to a small trust region at every iteration. This ensures that the objective function of the policy network is a good ap
proximation of the true objective function and forces small updates to the value network as well. As a result, PPO shows state-of-the-art performance on deterministic RL environments by ensuring small and robust updates at every iteration.

A study by [9] presented a convergence analysis for two time scaled stochastic approximation with controlled noise. In actor-critic methods, the parameter updates in neural networks using optimizers such as stochastic gradient descent or Adam can be seen as numerical solutions to a stochastic ODE.

As such, this work has been utilized by [5] to introduce a convergence analysis for actor-critic methods. Furthermore, these results have been used to show the asymptotic convergence of PPO and RUD-DER [1]. The authors introduce model assumptions as well as loss function assumptions that need to be satisfied to ensure that parameters in PPO and RUD-DER may converge to a local minimum in a neighborhood near their initial values. The study shows that the parameters of the policy network and value network in PPO may converge to a local minimum in a neighborhood near their initial values.

Another theoretical study by [6] concludes that policy optimization methods including PPO shows a guaranteed convergence on LQR tasks through the use of gradient-based optimization by formulating it as a non-convex optimization problem. It also shows reliable performance on state-feedback control problems. The paper refers to advanced regularization techniques as a potential area for improving robustness. This further justifies the introduction of the HJB loss as a regularization term.

### PINNs for HJB equation

Our work combines the HJB equation with reinforcement learning in continuous state and action spaces to improve the training of the value network. On a stochastic system with infinite time horizon, the HJB equation is a second order elliptic equation [2]. A theoretical study by [19] shows that PINNs converge uniformly to the solution of second order linear elliptic equations, thus justifying the use of PINNs to solve the HJB equation.

We treat the value network as a PINN to solve for the HJB equation by computing its derivatives respect to its inputs exactly.

Note that the term

\[\sup_{\hat{u}\in U}\{R(x,\hat{u})+\nabla_{x}V^{T}(x)f(x,\hat{u})\}\]

in the HJB equation cannot be determined without exploration of the agent in its environment. From Theorem 2.1, the optimal policy \(\pi^{*}(a|x)\) and the optimal controller \(u^{*}(x)=\operatorname*{argmax}_{a}\pi^{*}(a|x)\) satisfies equation (4). The optimal policy is modeled by the policy network \(\pi_{\theta}\) parameterized by \(\theta\) and the optimal controller can be approximated using \(u(x)=\operatorname*{argmax}_{a}\pi_{\theta}(a|x)\). As a result, we can use the following approximation:

\[V(x)\ln\gamma+R(x,u(x))+\nabla_{x}V^{T}(x)f(x,u(x))\] \[\approx V(x)\ln\gamma+\sup_{\hat{u}\in U}\{R(x,\hat{u})+\nabla_{x }V^{T}(x)f(x,\hat{u})\}. \tag{7}\]

This, as a result, justifies the use of equation 6 as the HJB loss used to update the value function at each episode.

The loss function is computed as

\[J(\rho)=0.5MSE_{u}+\lambda_{HJB}MSE_{f},\]

where \(MSE_{f}\) is defined in equation (6) and \(MSE_{u}\) is the standard loss function for the value network used in PPO, and is used to improve the discrete time estimate of returns for the value function:

\[MSE_{u}=\frac{1}{T}\sum_{t=0}^{T-1}|V(x_{t})-(R(x_{t},a_{t})+\gamma V(x_{t+1} ))|^{2}, \tag{8}\]

where \(\{x_{t}\}_{t=1}^{T}\) is a batch of states explored in a single episode, and \(\{a_{t}=\operatorname*{argmax}_{a}\pi_{\theta}(a|x_{t})\}_{t=1}^{T}\) is a batch of actions executed at time step \(t\) following the policy \(\pi_{\theta}\). The hyperparameter \(\lambda_{HJB}\) is determined based on the magnitude of the HJB loss curves compared to the Bellman optimality loss curves so that both loss functions are given similar weight.

### Algorithm

The HJBPPO algorithm is provided in Algorithm 1. The policy update and the minimization of
for the value network is identical to PPO. In order to satisfy \(V(x)\in C^{1}(\mathbb{R}^{n})\) as stated in Theorem 2.1, we use the infinitely differentiable _tanh_ activation function for the value network.

Lines 3-7 in the algorithm are identical to the PPO algorithm. The advantage term \(A_{t}\) is computed as: \(A_{t}=\sum_{n=t}^{T-1}(\gamma\lambda)^{n-t}\delta_{n}\), where \(\delta_{t}=R_{t}+\gamma V_{\varphi}(s_{t+1})-V_{\varphi}(s_{t})\), \(\gamma\) is the discount factor, and \(\lambda\) is the generalized advantage estimation (GAE) parameter. \(\alpha_{1}\) and \(\alpha_{2}\) are learning rates for the policy network optimizer and value network optimizer respectively. \(\varepsilon\) is the clipping parameter.

Lines 8-9 in the algorithm are our modifications to the PPO algorithm. We treat the value network as a PINN and add a \(MSE_{f}\) term into its loss function. This way, the HJB equation is used as a regularization term for the value network.

We will compare the performance of HJBPPO to PPO on the MuJoCo environments for rewards as well as the HJB loss.

## 5 Results

### Training

The HJBPPO algorithm was implemented by modifying the code for PPO in the Stable Baselines 3 library by [16]. To ensure the reproducibility of our results, we have posted our code at (Github repository redacted, code provided in supplementary material) and we have provided our hyperparameters in Tables 1 and 2 in Appendix A.

The code was run on the Beluga cluster in Compute Canada. The cluster provided the MuJoCo envi

Figure 1: Comparison of learning curves for PPO (Red, Dashed) compared to HJBPPO (Blue, Smooth)
ronments for training. Training each algorithm over 1 million time steps took seven hours, and training over 10 million time steps took three days. The multiprocessing library from python was used to train each algorithm over multiple environments at the same time.

### Reward Curves

The reward graphs have been plotted in Figure 1, comparing HJBPPO to PPO on all the MuJoCo environments over a million time steps or ten million time steps. The line shows the total reward, averaged over 50 consecutive episodes, and the shaded area indicated the standard deviation of the total reward over 50 consecutive episodes. HJBPPO shows a significant improvement in Antv4 and HalfCheetah-v4. It shows faster convergence and stability in Reacher-v4, Swimmer-v4, and InvertedDoublePendulum-v4. And it shows a slight improvement in HumanoidStandup-v4, Hopper-v4, and Walker2d-v4. For the two remaining environments (Humanoid-v4 and InvertedPendulum-v4), it shows equal performance to PPO.

As a result, the graphs show that incorporating the continuous-time HJB equation into the PPO al

Figure 3: Bellman optimality loss curves for HJBPPO on MuJoCo environments

Figure 2: HJB loss curves for HJBPPO on MuJoCo environments
gorithm to train the value function leads to an improved learning curve for the agent. This is because HJBPPO uses a PINN to exploit the physics in the environment. It uses finite differences to approximate the underlying governing equation \(f(x,u)\) of the environment and uses auto-differentiation to solve the HJB equation to achieve optimal control.

### HJB Loss Curves

The HJB loss for each environment has been plotted in Figure 2. HJBPPO shows a significant decrease in the HJB loss compared to PPO. And the HJB loss shows convergence in 8 out of the 10 environments, including Ant-v4 and HalfCheetah-v4 where it performs significantly better than PPO in terms of rewards, thus showing that the value function converges to the optimal value function as shown in the HJB equation.

The HJB loss does not converge for Humanoid-v4 and HumanoidStandup-v4, even though the HJB loss for these environments is significantly lower than that as shown in Figure 4. In both of these environments, the reward curve for HJBPPO shows similar performance to PPO. This shows that HJBPPO does not show significantly improved performance

Figure 4: HJB loss curves for PPO on MuJoCo environments

Figure 5: Bellman optimality loss curves for PPO on MuJoCo environments
compared to PPO in general if the HJB loss does not show convergence.

### Bellman Optimality Loss Curves

The Bellman optimality loss for each environment has been plotted in Figure 3. The value network shows convergence in every environment. This shows that convergence of the value function in the continuous-time HJB equation also improves its convergence in the discrete-time Bellman optimality equation, while the converse may not necessarily be true.

Figure 6: Comparison of HJB loss curves for PPO (Red, Dashed) compared to HJBPPO (Blue, Smooth)

Figure 7: Comparison of Bellman optimality loss curves for PPO (Red, Dashed) compared to HJBPPO (Blue, Smooth)
### Comparison with PPO

For comparison, we have posted the HJB loss and Bellman optimality loss curved for PPO in Figures 4 and 5 below. A notable difference is that the HJB loss in HJBBPO takes significantly lower valued compared to PPO. This is because the HJB loss is actively being minimized during the training of HJBBPO. To make the difference clearer, we plotted the loss curves on the same graphs in Figure 6. It is clear that the HJB loss for HJBBPO takes extremely small values in comparison with PPO.

As shown in Figure 4, the HJB loss in PPO shows convergence in only 4 out of the 10 environments; InvertedPendulum-v4, InvertedDoublePendulum-v4, Reacher-v4, and Swimmer-v4. For the remaining environments, the HJB loss shows an overall increasing trend.

For the environments where the HJB loss converges for PPO, it also shows convergence for HJBBPPO as shown in Figure 2. While HJBBPO does not show convergence in HJB loss for HumanoidStandup-v4, the loss curve is an improvement compared to PPO, where the loss curve shows a significant increasing trend.

Figure 5 shows convergence in Bellman optimality loss for 8 out of the 10 MuJoCo environments using PPO. The convergence in the Bellman optimality loss is achieved by PPO by training the value function to solve for the Bellman optimality equation. However, despite the choice of this loss function, in Ant-v4 and Walker2d-v4, the Bellman optimality loss does not show convergence, and instead shows an increasing trend for large time steps.

This issue is solved in HJBBPO as shown in Figure 3. The Bellman optimality loss shows an overall decreasing trend in all environments including Ant-v4 and Walker2d-v4. To make the difference clearer, we plotted the Bellman optimality loss curves on the same graphs in Figure 7. The Bellman optimality loss curves for HJBBPO show equal performance in general compared to PPO with better convergence in Ant-v4, HumanoidStandup-v4, InvertedDoublePendulum-v4, Hopper-v4, and Walker2d-v4.

In summary, HJBBPPO shows an improved performance compared to PPO. It shows improvement in the rewards curves, the HJB loss curves, and the Bellman optimality loss curves. This is due to the fact that HJBBPO incorporates an HJB loss regularization term and uses works from optimal control to improve the learning of the value function, and thus, improve the convergence of the algorithm.

## 6 Conclusion

In this paper, we have introduced the HJBBPPO algorithm that improvises the PPO algorithm to solve the HJB equation. This paper is the first of its kind to combine PINNs with RL. We treat the value function as a PINN to solve the HJB equation in an RL setting. The HJBBPO algorithm shows an overall improvement in performance compared to PPO due to its ability to exploit the physics of the environment as well as optimal control to improve the learning curve of the agent. This paper also shows that convergence of the value function in the continuous-time HJB equation also improves its convergence in the discrete-time Bellman optimality equation.

## 7 Future Research

Despite showing an overall improvement in the reward curves, the HJBBPPO leaves room for improved RL algorithms using PINNs.

A limitation of the HJBBPO algorithm as shown in figure 2 is that the HJB loss does not always show convergence in the environments albeit showing a significant improvement compared to PPO. A potential area of further research could involve new optimization methods for PINNs that show improved convergence of the HJB loss.

The loss function \(MSE_{f}\) using in training the value network was derived as a result of the approximation used in equation 7. So this does not guarantee convergence of the policy network towards the optimal policy such that \(u(x)=\sup_{\hat{u}\in U}\{R(x,\hat{u})+\nabla_{x}V^{T}(x)f(x,\hat{u})\}\) where the controller \(u(x)\) is derived from the policy \(\pi_{\theta}(a|x)\). [5] proves the convergence of the policy network parameters in PPO to a local op
timum but it does not guarantee global convergence. Thus, a potential area of further research could involve alternate choices of HJB loss functions for the value network that relaxes this approximation.

In this paper, we have explored and compared two deterministic RL algorithms - HJBPPO and PPO. It will be interesting to see how this algorithm can be extended to a stochastic setting. In the SAC paper, [4] introduces an entropy-regularized stochastic policy that is less likely to overfit or stick to a local optima. As a consequence of the approximation used in equation (7), the HJBPPO algorithm also poses a risk that the HJB loss of the value function could lead it to overfit to a suboptimal policy. This risk could be lessened by introducing an alternate HJB equation that facilitates exploration and incorporates entropy maximization. As a result, improvising the SAC algorithm with this HJB equation using PINNs is a potential area for further exploration.

In the MuJoCo environments, the HJBPPO algorithm showed an improvement compared to PPO. But this is due to the fact that \(f(x,u)\) could be estimated through finite differences, thus allowing for the physics of the environment to be exploited. The environments give all the details of the state needed to choose an action. One limitation of HJBPPO is that it may not perform well in partially observable environments because the estimate of \(f(x,u)\) may be inaccurate. Deep Transformer Q Network (DTQN) was introduced by [3] and achieves state-of-the-art results in many partially observable environments. A potential area for further research may be the introduction of an alternate HJB equation that facilitates partial observability. The DTQN algorithm may be improvised by incorporating this HJB equation using PINNs.

Additionally, finite difference approximations become less accurate in environments with high dimensions [20]. This makes the HJB loss less reliable in environments such as Humanoid-v4 and HumanoidStandup-v4 where the state is a 376-dimensional vector. The finite difference approximation does not compute \(f(x,u)\) exactly because the environment uses semi-implicit Euler integration steps rather than Euler's method [23]. Thus, a potential area for future research could be combining HJBPPO with model-based RL so that \(f(x,u)\) can be estimated with a smaller error.

## Acknowledgments

The authors would like to thank Pascal Poupart, Ashish Gaurav, and Yanting Miao from the Department of Computer Science, University of Waterloo, for providing us with feedback for this paper.
