Neural optimal control documentation

[ChrisRackauckas]

Here is a classical u(t) optimal control problem trained:

using DiffEqFlux, Flux, Optim, OrdinaryDiffEq, Plots
tspan = (0.0f0,1.0f0)
ann = FastChain(FastDense(1,64,sin), FastDense(64,64,sin), FastDense(64,1))
θ = initial_params(ann)
function dxdt_(dx,x,p,t)
    x1, x2 = x
    dx[1] = x[2]
    dx[2] = ann([t],p)[1]
end
x0 = [-4f0,0f0]
ts = Float32.(collect(0.0:0.01:1.0))
prob = ODEProblem(dxdt_,x0,tspan,θ)
concrete_solve(prob,Tsit5(),x0,θ,abstol=1e-8,reltol=1e-6)
function predict_adjoint(θ)
  Array(concrete_solve(prob,Tsit5(),x0,θ,saveat=ts))
end
function loss_adjoint(θ)
  x = predict_adjoint(θ)
  sum(abs2,4.0 .- x[1,:]) + 2sum(abs2,x[2,:]) + sum(abs2,[first(ann([t],θ)) for t in ts])
end
l = loss_adjoint(θ)
cb = function (θ,l)
  println(l)
  display(plot(solve(remake(prob,p=θ),Tsit5(),saveat=0.01),ylim=(-6,6)))
  return false
end
# Display the ODE with the current parameter values.
cb(θ,l)
loss1 = loss_adjoint(θ)
res = DiffEqFlux.sciml_train(loss_adjoint, θ, BFGS(), cb = cb)

should add it to the docs.

[AlCap23]

Just out of curiosity:
Is this a MWE which trains?

Works fine for me with

res = DiffEqFlux.sciml_train(loss_adjoint, θ, ADAM(5e-1),maxiters = 1000, cb = cb)

but does not really train with BFGS() as an optimizer. The optimizer seems to have issues to overcome 6.4e3.

I am using Julia 1.3 and the current releases of all packages.

[ChrisRackauckas]

It trains. BFGS is horrible to start with though, so I would suggest doing a few iterations of ADAM before BFGS.

[ChrisRackauckas]

I realized that the problem in the start isn't too interesting because the cost of using the control is so high that it pretty much never does anything... instead the following is much more interesting:

using DiffEqFlux, Flux, Optim, OrdinaryDiffEq, Plots, Statistics
tspan = (0.0f0,8.0f0)
ann = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
θ = initial_params(ann)
function dxdt_(dx,x,p,t)
    x1, x2 = x
    dx[1] = x[2]
    dx[2] = ann([t],p)[1]^3
end
x0 = [-4f0,0f0]
ts = Float32.(collect(0.0:0.01:tspan[2]))
prob = ODEProblem(dxdt_,x0,tspan,θ)
concrete_solve(prob,Vern9(),x0,θ,abstol=1e-10,reltol=1e-10)
function predict_adjoint(θ)
  Array(concrete_solve(prob,Vern9(),x0,θ,saveat=ts,sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true))))
end
function loss_adjoint(θ)
  x = predict_adjoint(θ)
  mean(abs2,4.0 .- x[1,:]) + 2mean(abs2,x[2,:]) + mean(abs2,[first(ann([t],θ)) for t in ts])/10
end
l = loss_adjoint(θ)
cb = function (θ,l)
  println(l)
  p = plot(solve(remake(prob,p=θ),Tsit5(),saveat=0.01),ylim=(-6,6),lw=3)
  plot!(p,ts,[first(ann([t],θ)) for t in ts],label="u(t)",lw=3)
  display(p)
  return false
end
# Display the ODE with the current parameter values.
cb(θ,l)
loss1 = loss_adjoint(θ)
res1 = DiffEqFlux.sciml_train(loss_adjoint, θ, ADAM(0.005), cb = cb,maxiters=100)
res2 = DiffEqFlux.sciml_train(loss_adjoint, res1.minimizer, BFGS(initial_stepnorm=0.01), cb = cb,maxiters=100)

function loss_adjoint(θ)
  x = predict_adjoint(θ)
  mean(abs2,4.0 .- x[1,:]) + 2mean(abs2,x[2,:]) + mean(abs2,[first(ann([t],θ)) for t in ts])
end

res3 = DiffEqFlux.sciml_train(loss_adjoint, res2.minimizer, BFGS(initial_stepnorm=0.01), cb = cb,maxiters=100)

l = loss_adjoint(res3.minimizer)
cb(res3.minimizer,l)
p = plot(solve(remake(prob,p=res3.minimizer),Tsit5(),saveat=0.01),ylim=(-6,6),lw=3)
plot!(p,ts,[first(ann([t],res3.minimizer)) for t in ts],label="u(t)",lw=3)
savefig("optimal_control.png")

We might want to pick a problem with a known analytical solution for the documentation.

[agerlach]

Am I reading this right that your dynamics are xdotdot = u(t)^3?

I would look at the examples here. In particular, the Van der Pol oscillator in Example 6.

[ChrisRackauckas]

Yup that's it.

[ChrisRackauckas]

I added this example to the documentation, but will leave this open because it would be nice to add a full exploration of that Van der Pol example.

[agerlach]

FYI, I am getting ReverseDiffVJP not defined and need to add using DiffEqSensitivity.

  [aae7a2af] DiffEqFlux v1.10.0
  [587475ba] Flux v0.10.4
  [429524aa] Optim v0.21.0
  [1dea7af3] OrdinaryDiffEq v5.37.0
  [91a5bcdd] Plots v1.2.4
  [10745b16] Statistics 

[ChrisRackauckas]

Oh, do using DiffEqSensitivity

[ChrisRackauckas]

Yeah it was missing from the tutorial. Added.