Paper: Learning HJB Viscosity Solutions with PINNs for Continuous-Time Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
The PINN method does not perform as well as existing DTRL methods (or at least PPO) in settings that require function approximation. While the authors give some insight about why that may be the case (and ideas for future research towards bridging the gap), there is a lack of motivating results: even if the highlighted issues with the approach are solved, it is unclear what benefits we should expect to see relative to DTRL algorithms. In particular, it would have been nice to see some results from the PINN method trained on much more data (i.e., to overcome the uniform sampling issue), to visualize the benefits we can expect if a more efficient training scheme is eventually found.

Author response:
We understand your point, but the problem is not in quantity of the data. Indeed, we tried to train our algorithm on larger datasets but after **beyond**  a certain dataset size there is little to no improvement on the result. The true limiting factor is the usage of uniform sampling in the domain. Indeed, "the hard areas", e.g. zones where the value function is particularly non-smooth or where distinguishing the best action is hard, require more samples to be learned properly than other zones. Similarly, DTRL algorithms learn more efficiently if the replay buffer has a considerable number of samples corresponding to high return trajectories.

Adaptive sampling approaches [RAR, R3] are emerging as a way to bypass this issue. However, those methods cannot be directly applied to our method. Such adaptive sampling methods may switch "focus" from one sub-area of $O$ to another throughout the entire training, which violates the uniform convergence of $V^{\epsilon} \to V$ and thus "breaks viscosity". It is not yet trivial how to do adaptive sampling while preserving viscosity, but it is what we want to explore in future work.

Reviewer comment (Theory):
While the presented $\epsilon$ schedules seem intuitively reasonable, there is limited discussion and/or analysis about them in the paper. Despite space limitations, there are 5 pages spent before the method is even presented. The review asks whether there is a theoretically principled way to choose the $\epsilon$ schedule or a principled way to quantitatively compare different schedules.

Author response:
We have provided in the paper a discussion related to schedulers and some complementary plots in Appendix A.5.2. In addition to that, as there is no viscosity metric that we could use to compare different schedulers, we cannot really quantitatively compare them on how well they find the viscosity solutions. The useful metric that we actually can apply is the value of right hand side of Eq. 7, i.e. $\text{Tr}(\nabla^2 W^{\epsilon}(x, \theta))$ or $\delta (\epsilon, \theta)$. However, for every scheduler, it is possible to choose the hyper-parameters that can lead to the convergence of $\text{Tr}(\nabla^2 W^{\epsilon}(x, \theta)) \to 0$ if $\epsilon \to 0$. Finally, for every scheduler we were able to find the set of hyperparameters that manages to achieve the good convergence, but as we observed if those hyperparameters are not chosen carefully then some schedules fail more than the others. According to our experiments, hybrid is the most robust to the choice of hyperparameters.

Reviewer comment (Theory):
When referring to the HJ-DQN algorithm (Kim et al., 2021), the paper states "this approach is limited to Lipschitz continuous control". While the reviewer understands what this means, they suggest that it is worth clarifying further since "Lipschitz continuous control" can be interpreted in different ways. Additionally, they question whether this is really a limitation, as controls that vary smoothly in time are often desirable.

Author response:
HJ-DQN is a good paper, but it is not exactly what we are looking at. Control that vary smoothly in time are desirable in critical applications such as robotic, self-driving cars. However, by constraining too much the set of admissible controls we may fail to find the optimal one. Moreover, there are systems for which the control varying smoothly in time is not feasible. Even though, we consider the control from the discrete space for now, we try to approach the problem of finding optimal control in a general way as possible and therefore without imposing too many constraints on the control space. Going to continuous control is important and interesting future work. At the same time, as we chose to work with more general settings in which $Q$ function degenerates, we cannot leverage on the works that are based on Q-learning and DQN-based algorithms, while HJ-DQN can be a way to adapt a lot of previous works on discrete-time to the continuous-time setting. 

We will change the text in our paper to better describe and compare to HJ-DQN.

Reviewer comment (Evaluation):
It should be made more clear why some HJB solutions are labeled as "bad" versus "good".

Author response:
We call viscosity solutions as "good" ones as they are the one that correspond to the optimal value function, while all other solutions are "bad". We will make it more explicit in the paper.

Reviewer comment (Writing):
The paper states that the work of Darbon et al. 2023 "work with min-plus algebra". Clarification is needed on what this means and what types of optimal control problems it is not suitable for.

Author response:
This paper assumes the linear dynamics and quadratic rewards settings, similar to LQR. Restricting their problems to this settings helps them to propose NN architectures with viscosity guarantees. As we have already mentioned, we would like to avoid introducing strong assumptions on the shape of dynamics function, we only assume the smoothness, which is important for the uniqueness of the viscosity solution. Let us also note that the dynamics of the real physical systems are often nonlinear and it is the case of pendulum task as well: $\frac{d^2 \theta}{dt^2} = - \frac{g}{l} \sin \theta$. We will make it more clear in the related works section of the paper.

Reviewer comment (Theory):
In equation 7, there is ambiguity about what $\nabla^2_x$ represents (is it the Laplacian?) and likewise for $\nabla^2_{xx}$. The significance of the RHS in equation 7 needs to be clarified, including whether other "perturbations" could be chosen instead, and how one can know if the RHS in equation 7 converges to 0 when $\epsilon \to 0$.

Author response:
Both $\nabla^2_x$ and $\nabla^2_{xx}$ denote the same thing ($\nabla^2_x$ being the typo, in most places we use $\nabla^2_{xx}$), which is the Hessian, thus $\text{Tr} \nabla^2_{xx} V(x)$ is a laplacian. We were considering before to use $\Delta V(x)$ as a laplacian, but this notation may be unfamiliar for some readers. We will make sure that there is no typo and less confusion about this notation in the paper. Thank you for noticing and pointing it out! 

>How can you know if the RHS in equation 7 converges to 0 when $\epsilon \to 0$?

In case of non-adaptive scheduler we cannot know it indeed, unless we choose very small $k_{\epsilon}$, but then it can create too huge gap between $V^{\epsilon_n}$ and $V^{\epsilon_{n-1}}$. The adaptive and hybrid schedulers update $\epsilon$ in accordance with $\delta(\epsilon, \theta)$ to make sure that it decreases. As $\delta(\epsilon, \theta)$ can be seen as mean squared value of RHS in equation 7, then it also decreases with each iteration.

Reviewer comment (Writing):
At the bottom of page 5, the phrase "dynamic programming approaches are able to find the solutions of the HJB equation that are intrinsically viscosity solutions" is unclear — clarification is needed on what "intrinsically viscosity solutions" means.

Author response:
We meant that some schemes are proven to converge to viscosity solutions for small enough discretization time step, even though the design of those schemes does not take viscosity into account. We will reformulate this sentence in the paper.

Reviewer comment (Evaluation):
It is unclear why PPO seemingly scales better than the PINNs method. Is this because PPO specializes to regions of the state space visited by nearly optimal policies, whereas the PINNs method trains on the whole space? Additionally, how do the methods based on the Pontryagin maximum principle mentioned in the related work compare?

Author response:
Yes, we believe that this is PPO's advantage over our approach.

>How do the methods based on the Pontryagin maximum principle that were mentioned in the related work compare?

We haven't compare to those methods. Paper based on the Pontryagin maximum principle supposes that the cost function is convex (so that the reward function is concave). Pontryagin maximum principle gives a necessary condition as boundary value problem (BVP) that the value function should verify. In general, BVP admits several solutions which can be sub-optimal. It is mentioned that optimality can be guaranteed under some convexity assumptions but in practice it is difficult to verify such conditions globally. In this work, they supposed the optimality of the solutions. Their sampling method is different to ours. Indeed, they use numerical methods to compute trajectory on short time interval by solving the BVP and then use them as a dataset to train a NN. As a result, the training procedure is comparable to supervised learning, while ours is comparable to unsupervised learning. A limitation of this method is that, since the optimality condition is a BVP, it is necessary to define the boundary conditions as equality (Dirichlet and Neumann boundary conditions). As mentioned in our work, it's not always easy to define such boundary conditions.

Reviewer comment (Theory):
If the authors want to enhance the theoretical contribution of this work, I would prefer to see how they would address the convergence in more detail. Lemma 3.1 is a general result, but in the PINN setting, what conditions on the NN can guarantee convergence and how to characterize approximation error or convergence speed, etc. These are challenging theoretical questions if the authors want to go multiple steps further than just citing Lemma 3.1 as the foundation of their method.

Author response:
> If the authors want to enhance the theoretical contribution of this work, I would prefer to see how they would address the convergence in more details. Lemma 3.1 is a general result, but in the PINN setting, what conditions on the NN can guarantee convergence and how to characterize approximation error or convergence speed, etc. These are challenging theoretical questions if the authors want to go multiple steps further than just citing Lemma 3.1 as the foundation of their method.


We agree that those are the important theoretical questions and they should be considered in the future work. However, performing this analysis is not straightforward, moreover analysing the convergence of neural networks is notoriously a challenging task, thus we leave it out of the scope of this paper for future work.

Reviewer comment (Experiments):
If the authors want to focus on the experimental aspect of their approach, I believe high dimensional control tasks may be the right playground for PINNs. The purpose of introducing NN in the PDE community is for solving large dimensional problems. The value of PINN for optimal control problems is likely to be similar. Otherwise, with known dynamics and easily observable rewards (as in this work), it is relatively easy to come up with efficient functional forms to approximate value functions without PINNs.

Author response:
> If the authors want to focus on the experimental aspect of their approach, I believe high dimensional control tasks may be the right playground for PINNs. The purpose of introducing NN in PDE community is for solving large dimensional problems. The value of PINN for optimal control problems are likely to be similar. 


PINNs is indeed a promising solver that was introduced to overcome the curse of dimensionality. Moreover, it has an advantage over the classical numerical solver as PINNs is not mesh-dependent and thus once trained it can directly output a solution at any point of the domain. However, it is not yet widely applied to high-dimensional problems, with an exception of a few simple cases. Typical application of PINNs is finding the solution of some PDE equation which depends on $x, y, z$ and $t$, i.e. at most 4-dimensional problem. 


There are a few reasons for that. PINNs are difficult to train, in particular due to a loss with different terms (mainly boundary condition term and PDE residual): controlling the respective weights of these terms in the loss is critical to convergence. Despite the  success for PINNS, this issue has limited the complexity of problems addressed through PINNS, in particular high-dimensional ones. In [3] the authors analyse this effect based on the training dynamics via the Neural Tangent Kernel (NTK) and propose a control scheme (validated in the paper with 1D PDEs). Other very recent works have addressed this issue through adaptive sampling of colocation points [6 and 7], reformulation of the loss to get a single term [5], NTK adaptive eigenvalue selection [4], taking a $l_{\infty}$ norm computed through an adversarial scheme to ensure convergence stability [1] or extension of SGD to dimensionality sampling [2]. If these approaches enable to improve PINNS training and address higher dimensional problems, we can stress that for the particular case of PINNs for HJB, authors (See [1,2] in particular) experiment on HJB with linear-quadratic-Gaussian (LQG) control guaranteeing a unique solution, and thus disregarding the difficulty of having a training scheme converging toward viscosity solutions. In this paper we take a more general stance and address this point with the $\epsilon$-schedulers. Experiments are indeed so far limited to low dimension experiments. Future work, as mentioned in the conclusion, will investigate how our algorithm could benefit from these recent approaches to improve PINNS training.

Progress could also be made on the side of the neural architecture, that is often simple for PINNS. We made attempts with different activation functions like sine functions [8], but without experiencing any significant gain.  

Overall, it is our goal to be able to target high-dimensional problems. However, the afore-mentioned existing limitations of PINNs should be addressed first, that may involve a sequence of significant improvements in the way PINNs are trained, thus it is left for the future work.

1. Is $L^2$ Physics Informed Loss Always Suitable for Training Physics Informed Neural Network? https://proceedings.neurips.cc/paper_files/paper/2022/hash/374050dc3f211267bd6bf0ea24eae184-Abstract-Conference.html
2. Tackling the Curse of Dimensionality with Physics-Informed Neural Networks - arXiv:2307.12306
3. WHEN AND WHY PINNS FAIL TO TRAIN: A NEURAL TANGENT KERNEL PERSPECTIVE. http://arxiv.org/abs/2007.14527
4. PINNACLE: PINN ADAPTIVE COLLOCATION AND EXPERIMENTAL POINTS SELECTION. https://openreview.net/pdf?id=GzNaCp6Vcg
5. About optimal loss function for training physics-informed neural networks under respecting causality. http://arXiv.org/abs/2304.02282
6. Mitigating Propagation Failures in Physics-informed Neural Networks using Retain-Resample-Release (R3) Sampling. http://arxiv.org/abs/2207.02338
7. DeepXDE: A deep learning library for solving differential equations. https://arxiv.org/abs/1907.04502
8. Implicit neural representations with periodic activation functions. https://arxiv.org/abs/2006.09661

Reviewer comment (Writing):
Although the title is "reinforcement learning," this method requires an analytic model of the system dynamics and the cost/reward function to apply PINN, and the states are sampled uniformly randomly in the state space. It may be better called optimal control rather than reinforcement learning, which usually assumes that the agent explore the environment by its own policy without explicit prior knowledge of the environment.

Author response:
We acknowledge that the problem considered in this paper indeed fits well into optimal control definition. At the same time Model-Based Continuous-Time Reinforcement Learning in its definition is very closely related to optimal control and thus if the model is learned our approach is also valid for CTRL. We indeed mention that in our introduction:

 "Thus, this work can be interesting for the optimal control community as we show how to use the neural networks to get the viscosity solutions and for the reinforcement learning community as our work can be used as a basis for Model Based Continuous Time Reinforcement Learning."

Even though, right now we consider the case of known dynamics, it is our goal to consider the case of unknown dynamics as well. Thus, we chose Continuous-Time Reinforcement Learning as the name.

Reviewer comment (Evaluation):
In Figure 1, why are the solutions have a dip at the origin, where the value function should be maximum. Is there any issue with dealing with the terminal cost/reward?

Author response:
It is an interesting question! One of the explanation (provided in the paper) is that the architecture that we use depends on smooth activation functions and has difficulty to approximate non-smooth areas. Another explanation, which concerns this zone in particular, is that only the actions that are deemed optimal by the neural network contribute to the loss computation, thus, if a chosen action is not optimal then it can introduce  error into the loss computation. This area is also the area where the best action is less distinguishable from suboptimal actions (the values in Eq.5 under argsup expression are very close to one another). A wrong choice can have a significant impact on which information is propagated during the backward propagation, e.g. different actions may have either $f(x,u) > 0$ or $f(x,u) < 0$, thus affecting the sign in front of $\nabla_x W(x)$ in $H(x, W(x), \nabla_x W(x))$. This can be observed on the image available at this link: https://anonymous.4open.science/r/review_iclr-D376/value_map_eps_%201.0seed_%201.pdf, where the action gap plot (the third one from the left) shows that areas with the lowest action gaps are the ones for which it is harder to identify the best action.

Reviewer comment (Presentation):
For people new to PINN, isn't it better to include the network architecture diagram with inputs, outputs, and how derivatives are combined for the objective function?

Author response:
Thank you for pointing it out, we will add a diagram to explain this into appendix!