Learning HJB Viscosity Solutions with PINNs for Continuous-Time Reinforcement Learning

We thank the reviewer for carefully reading our paper and pointing out some typos. We are happy to learn that you appreciated the idea of the paper and our approach to the problem. Below we try to address some of the questions raised.

The PINN method does not perform as well as existing DTRL methods (or at least PPO) in settings that require function approximation. While I appreciate that the authors give some insight about why that may be the case (and ideas for future research towards bridging the gap), I wish there was more motivating results: even if the highlighted issues with the approach are solved, what benefits should we expect to see relative to DTRL algorithms? Particularly, 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 we eventually find a more efficient training scheme.

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.

While the schedules that are presented seem intuitively reasonable, it would have been nice to see more discussion and/or analysis about those in the paper. I understand that space is limited, but 5 pages are spent before the method is even presented. Is there a theoretically principled way to choose the schedule, or is there a principled way to quantitatively compare them?

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.

When referring to the HJ-DQN algorithm (Kim et. al, 2021), it says “this approach is limited to Lipschitz continuous control”. I am familiar with this paper so I know what this means, but I definitely think it is worth clarifying further, since “Lipschitz continuous control” can be interpreted in at least a few different ways. On this note, is this really a limitation? Controls that vary smoothly in time are often desirable.

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.

We thank the reviewer for carefully reading the paper!

The main theoretical foundation is Lemma 3.1, which is a well-known results in optimal control community. For clarity, the authors have to spend more than half of the main context introduce this Lemma.

Indeed, this is a well-known result from optimal control and the paper mention that the proof can be found in Fleming & Soner (2006). We will make it more explicit for the final version. As this result together with uniqueness of viscosity solutions is necessary for introducing viscosity into PINNs we consider it is very important to provide enough background on optimal control.

This left the only novelty in the algorithm design where the authors introduced three decay schemes and the PDE solver (the latter of which has also been widely studied by PDE people). The overall significance of this work is therefore limited.

Do you refer to PINNs when you mention PDE solver? PINNs that we use to solve PDEs was introduced relatively recently. PINNs has not been studied enough yet and training PINNs is not easy. Moreover, PINNs was mostly studied in the context of PDE equations related to physical systems, where classical solutions exist and there is no need to consider viscosity solutions in general. Thus, introducing viscosity into PINNs framework is important and novel at the same time. Let us remark that this work has another significance as it is bridging two very different communities: Scientific Machine Learning and Optimal control, which is not an easy task and indeed there is not much work done in between those fields yet.

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.

We thank the reviewer for carefully reading our paper. We especially thank for additional related works pointed out by the reviewer, we will add them in our related works part. About comparisons to level-set methods, as far as we are familiar with those methods, similarly to FD or FEM methods they rely on discretization of the domain $O$, which means also limited scalability of the method. See our argument in Section 4.1 in the paper. Moreover, they seem to be more suitable for finite-horizon problems, where there exists a terminal condition that $V(x, t)$ should satisfy, like $V(x, T) = \Psi (x)$ with $T$ being a terminal time, while we consider infinite-horizon problems.

Generally, there lacks theoretical study for numerical solutions to a PDE. To ensure the necessary convergence, the conditions of Lemma 3.1 have to be verified, most importantly converges to uniformly. What presented in the subsection of 𝜖-schedulers is not adequate for this purpose.

We apologize for the confusion related to $\epsilon$-scheduling and uniform convergence. First of all, let us clarify that we should distinguish between the uniform convergence of equations and the uniform convergence of solutions and $\epsilon$-scheduling helps only uniform convergence of equations. Furthermore, our intention was not to say that $\epsilon$-scheduling can guarantee uniform convergence. The goal of $\epsilon$-scheduling is rather to encourage the uniform convergence. While a non-adaptive scheduler does not imply the uniform convergence of Eq.7 to Eq.4, the adaptive scheduler is designed in such a way to ensure that $\delta(\epsilon, \theta)$ converges to zero. Indeed each next $\epsilon_{n+1}$ should be small enough to ensure that $\delta(\epsilon_{n+1}, \theta_n) < \delta(\epsilon_{n}, \theta_{n-1})$. $\delta(\epsilon, \theta)$ is computed using a large number of samples from the domain $O$, samples being distributed uniformly in $O$ (related to this, we have noticed a typo with the formula of $\delta(\epsilon, \theta)$, where instead of $N_S$ should be $N_D$). Thus, $\delta(\epsilon, \theta)$ can be used as a good proxy to understand if RHS converges to 0. We will add a better intuition behind $\delta(\epsilon, \theta)$ in the paper.

As for uniform convergence of solutions, it is very hard to analyse the convergence of neural networks in terms of uniform convergence. We encourage it by introducing the regularization loss, see Eq.10. Together with uniform sampling, it helps the solution to stay close to the previously found solution for all points of the domain $O$.

Thank you for carefully reading the paper and highly evaluating its merit!

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?

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

It may be better called optimal control rather than reinforcement learning

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.

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?

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.

Further, we would like to address some common concerns. First, a few reviewers found that it was not clear how $\epsilon$-schedulers compare between each other and are connected to uniform convergence mentioned in Lemma 3.1. First of all, we would like to clarify that in order to satisfy conditions of Lemma 3.1 two types of uniform convergence are necessary: uniform convergence of equations (from Eq.7 to Eq.4) and uniform convergence of solutions ($W^{\epsilon} \to W$). While $\epsilon$-schedulers have only secondary effect on uniform convergence of solutions, they have a direct impact on the uniform convergence of equations. The non-adaptive scheduler is a simple way that performs regular $\epsilon$ updates at a frequency $N_{u}$, and $\epsilon$ goes to 0 with the speed $k_{\epsilon}$. On the one hand, this approach is simple and allows to control how often $\epsilon$ is updated. On the other hand, it does not depend on the value of RHS of Eq. 7, and if $k_{\epsilon}$ and $N_u$ are not chosen carefully then uniform convergence of Eq.7 to Eq.4 is not guaranteed. The adaptive scheduler does the updates according to the value of $\delta(\epsilon, \theta)$ that is a good proxy for RHS of Eq. 7. It is easy to show that $\delta(\epsilon_{n+1}, \theta_n) < \delta(\epsilon_{n}, \theta_{n-1})$, thus $\delta(\epsilon, \theta) \to 0$ when $\epsilon \to 0$. Therefore, the adaptive scheduler has more control over the uniform convergence of the equations, which is especially useful for smaller $\epsilon$. The hybrid is the combination of both schedulers using the non-adaptive at the beginning of the training for larger $\epsilon$, profiting from more regular updates, and the adaptive for smaller $\epsilon$ to have more guarantees for uniform convergence. According to our experiments, hybrid is the most robust to the choice of hyperparameters.

Another common concern is the absence of high-dimensional problems in our experimentation. Let us note that we have chosen problems of similar complexity and dimensionality as the ones considered in Lutter et al. (2020). The latter work also uses PINNs for HJB, but without viscosity considerations. While applying the method to more high-dimensional problems is very welcome, we argue that research on PINNs has not advanced enough to solve high-dimensional problems except some special cases. There are several reasons behind it. Firstly, they are difficult to train as one should find appropriate mixing coefficients for PDE loss and boundary losses. Secondly, most PINNs methods rely on uniform sampling, which does not cover enough the difficult zones, the problem that we also observed for our cases. The adaptive sampling schemes start to appear, but still are mostly tested on the low dimensional use cases and not suited for viscosity as they may break the uniform convergence of solutions. Finally, the choice of neural network architecture can be improved to process better non-smooth areas (see more details in replies to Reviewers N25p and 2hDe).

Finally, let us address the main concern of the reviewer N25p. It states that without theoretical analysis of uniform convergence or experiments on high-dimensional use cases, our work is not significant enough. We agree that both suggested improvements by Reviewer N25p are important and should be done, but we argue that those contributions can be left for future work. Indeed, the analysis of uniform convergence of NNs (NNs represent solutions) will be interesting, but it is a notoriously difficult problem. Going to high-dimensional use-cases is not yet straightforward as it is mentioned in the previous paragraph. Even without those potential improvements, we believe our work has a significant impact as it connects the popular PINN-based PDE solver that is not constrained by the dimensionality of the problem, with some core theoretical results known in optimal control (viscosity solutions and its properties). Despite the existing limitations, we think our work is pushing forward the state-of-the-art on this difficult problem and set the ground for others to advance this research topic at the crossroad between scientific machine learning, optimal control and reinforcement learning.