Indirect simple shooting

In this tutorial we present the indirect simple shooting method on a simple example.

Let us start by importing the necessary packages. We import the OptimalControl.jl package to define the optimal control problem. We import the Plots.jl package to plot the solution. The OrdinaryDiffEq.jl package is used to define the shooting function for the indirect method and the MINPACK.jl package permits to solve the shooting equation.

using BenchmarkTools    # to benchmark the methods
using OptimalControl    # to define the optimal control problem and its flow
using OrdinaryDiffEq    # to get the Flow function from OptimalControl
using MINPACK           # NLE solver: use to solve the shooting equation
using NonlinearSolve    # interface to NLE solvers
using Plots             # to plot the solution

Optimal control problem

Let us consider the following optimal control problem:

\[\left\{ \begin{array}{l} \min \displaystyle \frac{1}{2} \int_{t_0}^{t_f} u^2(t) \, \mathrm{d} t\\[1.0em] \dot{x}(t) = \displaystyle -x(t)+\alpha x^2(t)+u(t), \quad u(t) \in \R, \quad t \in [t_0, t_f] \text{ a.e.}, \\[0.5em] x(t_0) = x_0, \quad x(t_f) = x_f, \end{array} \right.%\]

with $t_0 = 0$, $t_f = 1$, $x_0 = -1$, $x_f = 0$, $\alpha=1.5$ and $\forall\, t \in [t_0, t_f]$, $x(t) \in \R$.

const t0 = 0
const tf = 1
const x0 = -1
const xf = 0
const α  = 1.5
ocp = @def begin

    t ∈ [t0, tf], time
    x ∈ R, state
    u ∈ R, control

    x(t0) == x0
    x(tf) == xf

    ẋ(t) == -x(t) + α * x(t)^2 + u(t)

    ∫( 0.5u(t)^2 ) → min

end

Boundary value problem

The pseudo-Hamiltonian of this problem is

\[ H(x,p,u) = p \, (-x+\alpha x^2+u) + p^0 u^2 /2,\]

where $p^0 = -1$ since we are in the normal case. From the Pontryagin Maximum Principle, the maximising control is given by

\[u(x, p) = p\]

since $\partial^2_{uu} H = p^0 = - 1 < 0$. Plugging this control in feedback form into the pseudo-Hamiltonian, and considering the limit conditions, we obtain the following two-points boundary value problem (BVP).

\[ \left\{ \begin{array}{l} \dot{x}(t) = \phantom{-} \nabla_p H[t] = -x(t) + \alpha x^2(t) + u(x(t), p(t)) = -x(t) + \alpha x^2(t) + p(t), \\[0.5em] \dot{p}(t) = - \nabla_x H[t] = (1 - 2 \alpha x(t))\, p(t), \\[0.5em] x(t_0) = x_0, \quad x(t_f) = x_f, \end{array} \right.\]

where $[t]~= (x(t),p(t),u(x(t), p(t)))$.

Shooting function

To achive our goal, let us first introduce the pseudo-Hamiltonian vector field

\[ \vec{H}(z,u) = \left( \nabla_p H(z,u), -\nabla_x H(z,u) \right), \quad z = (x,p),\]

and then denote by $\varphi_{t_0, x_0, p_0}(\cdot)$ the solution of the following Cauchy problem

\[\dot{z}(t) = \vec{H}(z(t), u(z(t))), \quad z(t_0) = (x_0, p_0).\]

Our goal becomes to solve

\[\pi( \varphi_{t_0, x_0, p_0}(t_f) ) = x_f,\]

where $\pi(x, p) = x$. To compute $\varphi$ with OptimalControl.jl package, we define the flow of the associated Hamiltonian vector field by:

u(x, p) = p
φ = Flow(ocp, u)

We define also the projection function on the state space.

π((x, p)) = x

Now, to solve the (BVP) we introduce the shooting function:

\[ \begin{array}{rlll} S \colon & \R & \longrightarrow & \R \\ & p_0 & \longmapsto & S(p_0) = \pi( \varphi_{t_0, x_0, p_0}(t_f) ) - x_f. \end{array}\]

S(p0) = π( φ(t0, x0, p0, tf) ) - xf    # shooting function

Resolution of the shooting equation

At the end, solving (BVP) is equivalent to solve $S(p_0) = 0$. This is what we call the indirect simple shooting method. We define an initial guess.

ξ = [0.1]    # initial guess

MINPACK.jl

We can use NonlinearSolve.jl package or, instead, MINPACK.jl to solve the shooting equation. To compute the Jacobian of the shooting function we use DifferentiationInterface.jl with ForwardDiff.jl backend.

using DifferentiationInterface
import ForwardDiff
backend = AutoForwardDiff()

Let us define the problem to solve.

nle!(s, ξ) = s[1] = S(ξ[1])                                 # auxiliary function
jnle!(js, ξ) = jacobian!(nle!, similar(ξ), js, backend, ξ)  # Jacobian of nle

We are now in position to solve the problem with the hybrj solver from MINPACK.jl through the fsolve function, providing the Jacobian. Let us solve the problem and retrieve the initial costate solution.

sol = fsolve(nle!, jnle!, ξ; show_trace=true)    # resolution of S(p0) = 0
p0_sol = sol.x[1]                                # costate solution
println("\ncostate:    p0 = ", p0_sol)
println("shoot: |S(p0)| = ", abs(S(p0_sol)), "\n")

Iter     f(x) inf-norm    Step 2-norm      Step time
------   --------------   --------------   --------------
     1     6.712249e-02     0.000000e+00         3.786963
     2     1.588980e-03     8.189120e-04         5.733771
     3     3.722710e-05     4.379408e-07         0.000424
     4     2.043880e-08     2.294999e-10         0.000366
     5     2.628775e-13     6.925509e-17         0.000387
     6     2.043561e-16     1.145743e-26         0.000374

costate:    p0 = 0.07202997482172421
shoot: |S(p0)| = 2.043561022960364e-16

NonlinearSolve.jl

Alternatively, we can use the NonlinearSolve.jl package to solve the shooting equation. The code is similar, but we use the solve function instead of fsolve. Let us define the problem.

nle!(s, ξ, λ) = s[1] = S(ξ[1])    # auxiliary function
prob = NonlinearProblem(nle!, ξ)  # NLE problem with initial guess

Now, let us solve the problem and retrieve the initial costate solution.

sol = solve(prob; show_trace=Val(true)) # resolution of S(p0) = 0
p0_sol = sol.u[1] # costate solution
println("\ncostate:    p0 = ", p0_sol)
println("shoot: |S(p0)| = ", abs(S(p0_sol)), "\n")

Algorithm: NewtonRaphson(
    descent = NewtonDescent(),
    autodiff = AutoForwardDiff(),
    vjp_autodiff = AutoReverseDiff(
        compile = false
    ),
    jvp_autodiff = AutoForwardDiff(),
    concrete_jac = Val{false}()
)

----    	-------------       	-----------
Iter    	f(u) inf-norm       	Step 2-norm
----    	-------------       	-----------
0       	6.71224852e-02      	0.00000000e+00
1       	1.58897952e-03      	2.86166386e-02
2       	8.72173248e-07      	6.46258280e-04
3       	2.62720451e-13      	3.55113679e-07
Final   	2.62720451e-13
----------------------

costate:    p0 = 0.0720299748216172
shoot: |S(p0)| = 2.6272045125485927e-13

Benchmarking

Let us benchmark the methods to solve the shooting equation.

@benchmark fsolve(nle!, jnle!, ξ; show_trace=false) # MINPACK

BenchmarkTools.Trial: 1615 samples with 1 evaluation per sample.
 Range (min … max):  2.196 ms … 95.631 ms  ┊ GC (min … max):  0.00% … 96.83%
 Time  (median):     2.329 ms              ┊ GC (median):     0.00%
 Time  (mean ± σ):   3.093 ms ±  4.268 ms  ┊ GC (mean ± σ):  13.15% ± 11.59%

  █▁   ▇                                                      
  ███▆████▇▁▁▁▅▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▅▅▄▃▅▆ █
  2.2 ms       Histogram: log(frequency) by time     13.1 ms <

 Memory estimate: 3.76 MiB, allocs estimate: 58469.

@benchmark solve(prob, SimpleNewtonRaphson(); show_trace=Val(false)) # NonlinearSolve

BenchmarkTools.Trial: 1297 samples with 1 evaluation per sample.
 Range (min … max):  2.615 ms … 123.324 ms  ┊ GC (min … max):  0.00% … 97.04%
 Time  (median):     3.147 ms               ┊ GC (median):     0.00%
 Time  (mean ± σ):   3.851 ms ±   5.309 ms  ┊ GC (mean ± σ):  13.54% ± 11.93%

  █▅   ▆▇▂▂                                                    
  ██▇▆▆████▄▁▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄▅▄▅▄▁▆▄▄ █
  2.62 ms      Histogram: log(frequency) by time      14.7 ms <

 Memory estimate: 4.49 MiB, allocs estimate: 66728.

According to the NonlinearSolve documentation, for small nonlinear systems, it could be faster to use the SimpleNewtonRaphson() descent algorithm.

@benchmark solve(prob, SimpleNewtonRaphson(); show_trace=Val(false)) # NonlinearSolve

BenchmarkTools.Trial: 1321 samples with 1 evaluation per sample.
 Range (min … max):  2.609 ms … 129.296 ms  ┊ GC (min … max):  0.00% … 91.71%
 Time  (median):     2.872 ms               ┊ GC (median):     0.00%
 Time  (mean ± σ):   3.782 ms ±   5.270 ms  ┊ GC (mean ± σ):  13.48% ± 12.08%

  █▅   ▂█▁▂                                                    
  ██▇▅▇████▄▄▁▁▁▄▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄▄▆▅▄▆▄ █
  2.61 ms      Histogram: log(frequency) by time      14.2 ms <

 Memory estimate: 4.49 MiB, allocs estimate: 66728.

Plot of the solution

The solution can be plot calling first the flow.

sol = φ((t0, tf), x0, p0_sol)
plot(sol)

In the indirect shooting method, the search for the optimal control is replaced by the computation of its associated extremal. This computation is equivalent to finding the initial costate (or covector) that solves the shooting function. Let us now plot the extremal trajectory in the phase space, along with the shooting function and its solution.

using Plots.PlotMeasures
function Plots.plot(S::Function, p0::Float64; Np0=20, kwargs...)

    # times for wavefronts
    times = range(t0, tf, length=3)

    # times for trajectories
    tspan = range(t0, tf, length=100)

    # interval of initial covector
    p0_min = -0.5
    p0_max = 2

    # covector solution
    p0_sol = p0

    # plot of the flow in phase space
    plt_flow = plot()
    p0s = range(p0_min, p0_max, length=Np0)
    for i ∈ eachindex(p0s)
        sol = φ((t0, tf), x0, p0s[i])
        x = state(sol).(tspan)
        p = costate(sol).(tspan)
        label = i==1 ? "extremals" : false
        plot!(plt_flow, x, p, color=:blue, label=label)
    end

    # plot of wavefronts in phase space
    p0s = range(p0_min, p0_max, length=200)
    xs  = zeros(length(p0s), length(times))
    ps  = zeros(length(p0s), length(times))
    for i ∈ eachindex(p0s)
        sol = φ((t0, tf), x0, p0s[i], saveat=times)
        xs[i, :] .= state(sol).(times)
        ps[i, :] .= costate(sol).(times)
    end
    for j ∈ eachindex(times)
        label = j==1 ? "flow at times" : false
        plot!(plt_flow, xs[:, j], ps[:, j], color=:green, linewidth=2, label=label)
    end

    #
    plot!(plt_flow, xlims=(-1.1, 1), ylims=(p0_min, p0_max))
    plot!(plt_flow, [0, 0], [p0_min, p0_max], color=:black, xlabel="x", ylabel="p", label="x=xf")

    # solution
    sol = φ((t0, tf), x0, p0_sol)
    x = state(sol).(tspan)
    p = costate(sol).(tspan)
    plot!(plt_flow, x, p, color=:red, linewidth=2, label="extremal solution")
    plot!(plt_flow, [x[end]], [p[end]], seriestype=:scatter, color=:green, label=false)

    # plot of the shooting function
    p0s = range(p0_min, p0_max, length=200)
    plt_shoot = plot(xlims=(p0_min, p0_max), ylims=(-2, 4), xlabel="p₀", ylabel="y")
    plot!(plt_shoot, p0s, S, linewidth=2, label="S(p₀)", color=:green)
    plot!(plt_shoot, [p0_min, p0_max], [0, 0], color=:black, label="y=0")
    plot!(plt_shoot, [p0_sol, p0_sol], [-2, 0], color=:black, label="p₀ solution", linestyle=:dash)
    plot!(plt_shoot, [p0_sol], [0], seriestype=:scatter, color=:green, label=false)

    # final plot
    plot(plt_flow, plt_shoot; layout=(1,2), leftmargin=15px, bottommargin=15px, kwargs...)

end

plot(S, p0_sol; size=(800, 450))