Optimal Control

mtk_optimal_control.jl

using DifferentialEquations
using ModelingToolkit

# Time and time derivative
@parameters t
D = Differential(t)

# Forced Duffing oscillator
function Duffing(; name)
    @parameters α β
    @variables x[1:2](t) u(t)

    eqs = [
        D(x[1]) ~ x[2]
        D(x[2]) ~ -α*x[1] - β*x[1]^3 + u
    ]

    ODESystem(eqs, t, [x..., u], [α, β]; name)
end

# Nonlinear optimal control via Pontryagin Maximum Principle
function Pontryagin(; name, sys, L)
    # Ugly stuff
    f = [eq.rhs for eq in sys.eqs]
    x = sys.states[findall(startswith('x'), string.(sys.states))]
    u = sys.states[findall(startswith('u'), string.(sys.states))]

    # Lagrange multipliers / costate variables
    @variables λ[eachindex(x)](t)

    # Hamiltonian
    H = λ'*f + L(x, u, t)

    # Costate equations
    λ_eqs = map(λ, x) do λᵢ, xᵢ
        ∂H_∂xᵢ = expand_derivatives(Differential(xᵢ)(H))
        D(λᵢ) ~ -∂H_∂xᵢ
    end
    
    # Control optimality
    opt_eqs = map(u) do uᵢ
        ∂H_∂uᵢ = expand_derivatives(Differential(uᵢ)(H))
        0.0 ~ ∂H_∂uᵢ
    end

    ODESystem([opt_eqs..., λ_eqs..., sys.eqs...], t, [u..., λ..., x...], sys.ps; name)
end

# Boundary conditions -- free final state
function bc!(residual, vars, p, t)
    vars0, varsf = @views vars[begin], vars[end]
    λ⁺f, x⁺f, x⁺0 = @views varsf[1:2], varsf[3:4], vars0[3:4]
    x0 = last.(@view ic[3:4])

    # Initial condition of state variables
    residual[1:2] .= x⁺0 .- x0

    # Final condition of costate variables
    residual[3:4] .= λ⁺f .- ∂Φ_∂xf(x⁺f, xf, p[3])
end

# Lagrangian
L(x, u, t) = 0.5sum(u.^2)

# Final state cost and derivative (should do this symbolically too)
Φ(x⁺f, xf, γ) = γ * 0.5*sum((x⁺f .- xf).^2)
∂Φ_∂xf(x⁺f, xf, γ) = γ * (x⁺f .- xf)

# Make the systems
@named duff = Duffing()
@named opt_duff = Pontryagin(; sys=duff, L)
simplified_duff = structural_simplify(opt_duff)

u, λ₁, λ₂, x₁, x₂ = opt_duff.states
x = [x₁, x₂]
α, β = opt_duff.ps

# Initial conditions
ic = [
    λ₁ => 1.0
    λ₂ => 1.0
    x₁ => 0.0
    x₂ => 0.0
]

# Parameters
p = [
    α => 1.0
    β => 0.5
]

# Desired final state
xf = [5.0, 2.0]

# Final state vs control cost weight (higher favors final state accuracy)
γ = 1.0

# Regular problem
prob = ODEProblem(simplified_duff, ic, (0.0, 2π), p)

# Boundary value problem
bvp = TwoPointBVProblem(prob.f, bc!, prob.u0, (0.0, 2π), [prob.p..., γ])

# Solve the boundary value problem
@time bvp_sol = solve(bvp, MIRK4(), dt=0.05)