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 as the ODE solver backend by the Flow function from OptimalControl.jl 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    # ODE solver backend for Flow 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 achieve 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 the 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         2.478743
     2     1.588980e-03     8.189120e-04         4.190143
     3     3.722710e-05     4.379408e-07         0.000595
     4     2.043880e-08     2.294999e-10         0.000523
     5     2.628703e-13     6.925509e-17         0.000500
     6     2.307927e-16     1.145446e-26         0.000486

costate:    p0 = 0.07202997482172416
shoot: |S(p0)| = 2.307927285645836e-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.62870339e-13      	3.55113679e-07
Final   	2.62870339e-13
----------------------

costate:    p0 = 0.07202997482161713
shoot: |S(p0)| = 2.6287033906588936e-13

Benchmarking

Let us benchmark the methods to solve the shooting equation.

@benchmark fsolve(nle!, jnle!, ξ; tol=1e-8, show_trace=false) # MINPACK

BenchmarkTools.Trial: 1082 samples with 1 evaluation per sample.
 Range (min … max):  3.012 ms … 228.996 ms  ┊ GC (min … max):  0.00% … 98.01%
 Time  (median):     3.922 ms               ┊ GC (median):     0.00%
 Time  (mean ± σ):   4.617 ms ±  10.006 ms  ┊ GC (mean ± σ):  15.89% ±  8.12%

                           █▄▂                                 
  ▂▅▅▆▆▄▄▃▃▃▂▂▁▁▁▂▂▂▁▁▁▁▁▃████▆▄▃▃▃▄▃▄▄▄▃▂▂▂▂▁▂▂▁▁▁▁▁▁▂▃▃▄▃▃▂ ▃
  3.01 ms         Histogram: frequency by time        5.07 ms <

 Memory estimate: 4.16 MiB, allocs estimate: 71360.

@benchmark solve(prob; abstol=1e-8, reltol=1e-8, show_trace=Val(false)) # NonlinearSolve

BenchmarkTools.Trial: 712 samples with 1 evaluation per sample.
 Range (min … max):  4.402 ms … 227.526 ms  ┊ GC (min … max):  0.00% … 97.36%
 Time  (median):     5.912 ms               ┊ GC (median):     0.00%
 Time  (mean ± σ):   7.016 ms ±  13.129 ms  ┊ GC (mean ± σ):  16.03% ±  8.73%

                       ▂▃██▆▄▁                                 
  ▃▅▆▆▇▄▄▃▄▂▃▃▃▂▂▂▂▂▃▁▂███████▇▆▅▄▆▄▆▆▆▄▄▃▃▂▂▂▁▄▄▅▄▃▃▂▁▁▁▁▂▂▂ ▄
  4.4 ms          Histogram: frequency by time        7.93 ms <

 Memory estimate: 6.18 MiB, allocs estimate: 101906.

We can also use the SimpleNewtonRaphson() descent algorithm from NonlinearSolve as an alternative solver.

@benchmark solve(prob, SimpleNewtonRaphson(); abstol=1e-8, reltol=1e-8, show_trace=Val(false)) # NonlinearSolve

BenchmarkTools.Trial: 889 samples with 1 evaluation per sample.
 Range (min … max):  3.531 ms … 233.558 ms  ┊ GC (min … max):  0.00% … 97.60%
 Time  (median):     4.760 ms               ┊ GC (median):     0.00%
 Time  (mean ± σ):   5.619 ms ±  12.071 ms  ┊ GC (mean ± σ):  16.30% ±  7.86%

                           ▁▄▇█▆   ▁▁    ▁                     
  ▅▅▅▆▇▆▄▃▃▃▃▃▂▃▂▃▂▂▁▁▂▁▁▄▆█████▇▆▆██▇▄▄▄█▇▅▆▃▃▂▁▂▁▁▁▃▃▄▄▃▃▁▃ ▃
  3.53 ms         Histogram: frequency by time        6.04 ms <

 Memory estimate: 5.01 MiB, allocs estimate: 81458.

Plot of the solution

The solution can be plot calling first the flow.

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

We can verify that the solution satisfies the PMP by checking that the Hamiltonian is constant along the extremal. Let us define the Hamiltonian and compute its value along the solution.

H(x, p, u) = p * (-x + α * x^2 + u) - 0.5 * u^2  # pseudo-Hamiltonian with p^0 = -1

t = time_grid(sol)
x(t) = state(sol)(t)
p(t) = costate(sol)(t)
u(t) = control(sol)(t)

H_vals = [H(x(t), p(t), u(t)) for t in t]
println("Hamiltonian: H = ", H_vals[1], " (should be constant)")
println("Hamiltonian variation: max|H(t) - H(0)| = ", maximum(abs.(H_vals .- H_vals[1])))

Hamiltonian: H = 0.18266909569044426 (should be constant)
Hamiltonian variation: max|H(t) - H(0)| = 3.3298641621826164e-12

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 (flow at specific times)
    times = range(t0, tf, length=3)

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

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

    # covector solution (the one that solves the shooting equation)
    p0_sol = p0

    # LEFT PLOT: flow in phase space (x, p)
    # Plot extremals (trajectories) for different initial costates
    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 wavefronts: set of points reached by the flow at given times
    # Wavefronts help visualize how the flow evolves 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

    # Add boundary condition x = xf (vertical line)
    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")

    # Plot the extremal solution (in red) that satisfies the boundary condition
    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)

    # RIGHT PLOT: shooting function S(p₀)
    # The shooting function maps initial costate p₀ to the final state error
    # The zero of this function gives the correct initial costate
    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)

    # Combine both plots: left shows phase space, right shows shooting function
    plot(plt_flow, plt_shoot; layout=(1,2), leftmargin=15px, bottommargin=15px, kwargs...)

end

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