# Solving HJB PDE using deep learning

**URL:** <https://discourse.julialang.org/t/solving-hjb-pde-using-deep-learning/60639>\
**Category:** Numerics\
**Tags:** economics, pde\
**Created:** [May 6, 2021, 10:35am UTC](https://discourse.julialang.org/t/solving-hjb-pde-using-deep-learning/60639 "2021-05-06T10:35:48Z")\
**Posts on this page:** 1\
**Showing post:** 14

<div class="post-metadata">

**Author:** ![Albert\_Zevelev](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albert_zevelev/32/11844_2.png) [@Albert\_Zevelev](https://discourse.julialang.org/u/Albert_Zevelev)\
**Post date:** [January 9, 2022, 6:40am UTC](https://discourse.julialang.org/t/solving-hjb-pde-using-deep-learning/60639/14 "2022-01-09T06:40:11Z")

</div>

1. @Goosee I was able to use (a simplified version of) your code to solve another non-linear HJB problem w/o any change in neural network tuning!  
I got it to work in minutes (my coding time, not run time)!  
The firm-investment problem is described [here](https://discourse.julialang.org/t/solving-the-4-quadrants-of-dynamic-optimization-problems-in-julia-help-wanted/73285).

```julia
using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using ForwardDiff
import ModelingToolkit: Interval, infimum, supremum

@parameters x # x=K
@variables u(..), s(..) # u(x)=V(K) & s(x) = I(K) policy
Dx = Differential(x) # Dx(u(x)) == V_K(K)
δ=0.10; ρ=0.05; z= ρ + δ +.10; Iss = (z-ρ-δ)/(ρ+δ); kss = Iss/δ;
eq = [
        Dx(u(x))-1 - s(x) ~ 0.0, 
        z*x - s(x) - 0.5*s(x)^2 + Dx(u(x))*(s(x) - δ*x) - ρ*u(x) ~ 0.0
        ]
elems = size(eq)[1]
bcs = [Dx(u(kss)) ~ (1 + δ*kss)]
domains = [x ∈ Interval(kss*0.1,kss*1.5)]
neurons = 20
chain = [FastChain(FastDense(1,neurons,tanh)
                    ,FastDense(neurons,neurons,softplus)
                    ,FastDense(neurons,1)) for _ in 1:elems]
dx0 = kss*(1.5-0.1)/400
initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))
discretization = PhysicsInformedNN(chain, GridTraining(dx0),init_params =initθ)

@named pde_system = PDESystem(eq,bcs,domains,[x],[u(x),s(x)])
prob = discretize(pde_system,discretization)
const losses = []
cb = function (p,l)
    push!(losses, l)
    if length(losses)%500==0
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    return false
end
maxit1 = 1200
maxit2 = 4000
res = GalacticOptim.solve(prob, ADAM(0.06); atol=1e-7, cb = cb, maxiters=maxit1)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob,BFGS(); cb = cb, maxiters=maxit2)
phi = discretization.phi

using Plots
kgrid = LinRange(kss*0.1,kss*1.5,1_000)
I = Int.(collect(size(res.minimizer)./elems))
Vfunc(x) = first(phi[1](x,res.minimizer[1:I[1]]))
Cfunc(x) = first(phi[2](x,res.minimizer[I[1]+1:end]))

# Plot V
plot(kgrid ,Vfunc.(kgrid),
    line =:dashdot,linewidth=3,label="Neural net",
    legend=:topleft,xlabel = "k",ylabel="V(k)", 
    fg_legend = :false, xtick=0:2:10, ytick=-20:1:-8)
#display(plot!(kgrid,vfunc,label="Finite differences",legend=:topleft))

plot(kgrid ,Cfunc.(kgrid),
    line =:dashdot,linewidth=3,label="Neural net",
    legend=:topleft,xlabel = "k",ylabel="I(k)", 
    fg_legend = :false, xtick=0:2:10, ytick=0.5:0.5:2)
plot!(kgrid ,x -> Iss, lab="I_SS", ylims=(0,1))
#

# Compute HJB error
dVk = ForwardDiff.derivative.(Vfunc, Array(kgrid))
II = dVk .- 1
HJBerror = z*kgrid .- II .- .5*II.^2 + dVk .*(II-δ*kgrid) - ρ*Vfunc.(kgrid)
plot(kgrid,HJBerror,linewidth=3,ylims=(-1.2e-4,1e-4),
    xlabel = "k",ylabel="HJB error",label="Neural net", fg_legend = :false)
########
display(plot(losses,legend = false, yaxis=:log))

```

The value function is affine (as it should be)

 ![image](https://global.discourse-cdn.com/julialang/original/3X/7/f/7f0e21eb2aaaf2af5c364b455c049b33ea9c0c21.png)  
The policy function is constant (as it should be)  
 ![image](https://global.discourse-cdn.com/julialang/original/3X/3/6/3603ff0c20669c392ecfb303cb0c91b51f61935c.png)  
The HJB error  
 ![image](https://global.discourse-cdn.com/julialang/original/3X/8/9/8953cc5b1c0dff160f29510f552af0606c9f757e.png)  
The loss  
 ![image](https://global.discourse-cdn.com/julialang/original/3X/8/4/84caa7473082ba6b585f14123bd10157e820b2ff.png)

1. @Goosee echoing @ChrisRackauckas I really hope you post issues for the problems not working…  
If you want to “scale up to more complex versions, say by adding a diffusion in productivity or more state variables.” I recommend you extend my simple firm investment problem (instead of the NCGM) so you don’t need to worry about things like diving by zero etc…
2. If you’d like please feel free to add examples to my [Github repo](https://github.com/azev77/4Quadrants)

---

_[View the full topic](https://discourse.julialang.org/t/solving-hjb-pde-using-deep-learning/60639)._