I am trying to solve a simple optimal control problem using the Hamilton-Jacobi-Bellman equation, numerically in Python. This is proving to be rather difficult as I end up having to solve the following: $$ J_t - (J_x)^2 + x\cdot J_x = 0 $$ I believe this to be a non-linear first order PDE. Being the HJB, we are given boundary condition at terminal time. $$ \mbox{BC:} \quad J(x_f) = \dfrac{1}{4} x_f^2 $$ I have attempted this problem myself by simply putting in central Euler approximations for $J_x$ and using backward difference for $J_t$. I then solve backwards in time by making $J_{i}^{j-1}$ the subject of the formula. $$ J_{i}^{j-1} = J_{i}^{j}+\frac{k}{4h^2}\big[ J_{i+1}^{j}-J_{i-1}^{j} \big]^2 + \frac{k}{h}\big[ J_{i+1}^{j}-J_{i-1}^{j} \big]x_i $$ This only works for a few computations before the scheme becomes very unstable.
Would you please help guide me on how to solve this. I don't quite know whether my approach is correct. Please suggest resources or help me to develop a scheme that would work.
P.S. I've seen the HJB written like this $$ \rho J = \alpha J_x $$ Where $J_t$ has been replaced by $\rho J$. This has been done : http://www.princeton.edu/~moll/HACTproject/HACT_Additional_Codes.pdf
Your problem is that Hamilton-Jacobi type equations are not amenable to ordinary forward, backward or central differencing since the solutions can develop discontinuities in their derivatives, even when the initial (or terminal) states are smooth. Instead, you need to think about the characteristics and where the information is flowing (Godunov/Upwind scheme). Alternatively, you can use numerical diffusion to smooth the solutions (Lax-Friedrichs scheme).
Stan Osher and Chi-Wang Shu give a nice overview with several different suggestions for a numerical Hamiltonian, and provide a method for ensuring arbitrarily high order accuracy:
https://epubs.siam.org/doi/10.1137/0728049
EDIT: As an aside, Lax-Friedrichs scheme is almost always the easiest to implement, but if you are solving at first-order accuracy, the numerical diffusion can have effects that are visible even to the naked eye. Depending on application, non-diffuse schemes may be preferable.