Solving HJB PDE using deep learning

Specific Domains Numerics
,
14
  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.
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
image1184×776 36.9 KB

The policy function is constant (as it should be)
image
image1200×784 25 KB

The HJB error
image
image1174×790 62.3 KB

The loss
image
image1172×768 49 KB

  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
3 Likes