fix: make free final time problems with constraints work

Author: vyudu

ext/MTKJuMPDynamicOptExt.jl

@@ -42,7 +42,7 @@ end
 Convert an ODESystem representing an optimal control system into a JuMP model
 for solving using optimization. Must provide either `dt`, the timestep between collocation 
 points (which, along with the timespan, determines the number of points), or directly 
-provide the number of points as `nsteps`.
+provide the number of points as `steps`.
 
 The optimization variables:
 - a vector-of-vectors U representing the unknowns as an interpolation array
@@ -61,7 +61,7 @@ function MTK.JuMPDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    pmap = MTK.todict(pmap)
+    pmap = Dict{Any, Any}(pmap)
     steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
@@ -87,7 +87,7 @@ function MTK.InfiniteOptDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    pmap = MTK.todict(pmap)
+    pmap = Dict{Any, Any}(pmap)
     steps, is_free_t = MTK.process_tspan(tspan, dt, steps)
     model = init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
 
@@ -103,21 +103,25 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
 
     if is_free_t
         (ts_sym, te_sym) = tspan
+        MTK.symbolic_type(ts_sym) !== MTK.NotSymbolic() && error("Free initial time problems are not currently supported.")
         @variable(model, tf, start=pmap[te_sym])
+        set_lower_bound(tf, ts_sym)
         hasbounds(te_sym) && begin
             lo, hi = getbounds(te_sym)
             set_lower_bound(tf, lo)
             set_upper_bound(tf, hi)
         end
-        pmap[ts_sym] = 0
-        pmap[te_sym] = 1
+        pmap[te_sym] = model[:tf]
         tspan = (0, 1)
     end
 
     @infinite_parameter(model, t in [tspan[1], tspan[2]], num_supports=steps)
     @variable(model, U[i = 1:length(states)], Infinite(t), start=u0[i])
     c0 = MTK.value.([pmap[c] for c in ctrls])
     @variable(model, V[i = 1:length(ctrls)], Infinite(t), start=c0[i])
+    for (i, ct) in enumerate(ctrls)
+        pmap[ct] = model[:V][i]
+    end
 
     set_jump_bounds!(model, sys, pmap)
     add_jump_cost_function!(model, sys, (tspan[1], tspan[2]), pmap; is_free_t)
@@ -131,6 +135,13 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
     return model
 end
 
+"""
+Modify the pmap by replacing controls with V[i](t), and tf with the model's final time variable for free final time problems.
+"""
+function modified_pmap(model, sys, pmap)
+    pmap = Dict{Any, Any}(pmap)
+end
+
 function set_jump_bounds!(model, sys, pmap)
     U = model[:U]
     for (i, u) in enumerate(unknowns(sys))
@@ -158,7 +169,7 @@ function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap; is_free
         @objective(model, Min, 0)
         return
     end
-    jcosts = substitute_jump_vars(model, sys, pmap, jcosts)
+    jcosts = substitute_jump_vars(model, sys, pmap, jcosts; is_free_t)
     tₛ = is_free_t ? model[:tf] : 1
 
     # Substitute integral
@@ -186,13 +197,14 @@ function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = fals
         for u in MTK.get_unknowns(conssys)
             x = MTK.operation(u)
             t = only(arguments(u))
-            MTK.symbolic_type(t) === NotSymbolic() &&
-                error("Provided specific time constraint in a free final time problem. This is not supported by the JuMP/InfiniteOpt collocation solvers. The offending variable is $u.")
+            if (MTK.symbolic_type(t) === MTK.NotSymbolic())
+                error("Provided specific time constraint in a free final time problem. This is not supported by the JuMP/InfiniteOpt collocation solvers. The offending variable is $u. Specific-time constraints can only be specified at the beginning or end of the timespan.")
+            end
         end
     end
 
     auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)])
-    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap)
+    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap, is_free_t)
 
     # Substitute to-term'd variables
     for (i, cons) in enumerate(jconstraints)
@@ -211,28 +223,33 @@ function add_initial_constraints!(model::InfiniteModel, u0, u0_idxs, ts)
     @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i])
 end
 
-function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict())
+function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict(), is_free_t = false)
     iv = MTK.get_iv(sys)
     sts = unknowns(sys)
     cts = MTK.unbound_inputs(sys)
     U = model[:U]
     V = model[:V]
+    x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
+    c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
+
     exprs = map(c -> Symbolics.fixpoint_sub(c, auxmap), exprs)
+    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
+    if is_free_t
+        tf = model[:tf]
+        free_t_map = Dict([[x(tf) => U[i](1) for (i, x) in enumerate(x_ops)]; [c(tf) => V[i](1) for (i, c) in enumerate(c_ops)]])
+        exprs = map(c -> Symbolics.fixpoint_sub(c, free_t_map), exprs)
+    end
 
     # for variables like x(t)
     whole_interval_map = Dict([[v => U[i] for (i, v) in enumerate(sts)];
                                [v => V[i] for (i, v) in enumerate(cts)]])
     exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs)
 
     # for variables like x(1.0)
-    x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
-    c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
     fixed_t_map = Dict([[x_ops[i] => U[i] for i in 1:length(U)];
                         [c_ops[i] => V[i] for i in 1:length(V)]])
 
     exprs = map(c -> Symbolics.fixpoint_sub(c, fixed_t_map), exprs)
-
-    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
     exprs
 end
 
@@ -246,8 +263,12 @@ function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap; is_free_
     diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)])
     tₛ = is_free_t ? model[:tf] : 1
 
+    @show diff_equations(sys)
+    @show pmap
     diff_eqs = substitute_jump_vars(model, sys, pmap, diff_equations(sys))
+    @show diff_eqs
     diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
+    @show diff_eqs
     @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs==tₛ * diff_eqs[i].rhs)
 
     # Algebraic equations

src/systems/optimal_control_interface.jl

@@ -173,7 +173,7 @@ function process_tspan(tspan, dt, steps)
     elseif symbolic_type(tspan[1]) === ScalarSymbolic() ||
            symbolic_type(tspan[2]) === ScalarSymbolic()
         isnothing(steps) &&
-            error("Free final time problems require specifying the number of steps, rather than dt.")
+            error("Free final time problems require specifying the number of steps using the keyword arg `steps`, rather than dt.")
         isnothing(dt) ||
             @warn "Specified dt for free final time problem. This will be ignored; dt will be determined by the number of timesteps."
 

test/downstream/jump_control.jl

@@ -178,22 +178,22 @@ end
 
     (ts, te) = (0.0, 0.2)
     costs = [-h(te)]
-    cons = [T(te) ~ 0]
+    cons = [T(te) ~ 0, m(te) ~ m_c]
     @named rocket = ODESystem(eqs, t; costs, constraints = cons)
     rocket, input_idxs = structural_simplify(rocket, ([T(t)], []))
 
     u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0]
     pmap = [
         g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5 * √(g₀ * h₀), D_c => 0.5 * 620 * m₀ / g₀,
         Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
-    jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
+    jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     jsol = solve(jprob, Ipopt.Optimizer, :RadauIIA5, silent = true)
     @test jsol.sol.u[end][1] > 1.012
 
     iprob = InfiniteOptDynamicOptProblem(
-        rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
+        rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     isol = solve(iprob, Ipopt.Optimizer,
-        derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
+                 derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true)
     @test isol.sol.u[end][1] > 1.012
 
     # Test solution
@@ -204,11 +204,16 @@ end
     @mtkbuild rocket_ode = ODESystem(eqs, t)
     interpmap = Dict(T_interp => ctrl_to_spline(jsol.input_sol, CubicSpline))
     oprob = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap))
-    osol = solve(oprob, RadauIIA5(); adaptive = false, dt = 0.005)
+    osol = solve(oprob, RadauIIA5(); adaptive = false, dt = 0.001)
     @test ≈(jsol.sol.u, osol.u, rtol = 0.02)
+
+    interpmap1 = Dict(T_interp => ctrl_to_spline(isol.input_sol, CubicSpline))
+    oprob1 = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap1))
+    osol1 = solve(oprob1, RadauIIA5(); adaptive = false, dt = 0.001)
+    @test ≈(isol.sol.u, osol1.u, rtol = 0.02)
 end
 
-@testset "Free final time problem" begin
+@testset "Free final time problems" begin
     t = M.t_nounits
     D = M.D_nounits
 
@@ -230,6 +235,27 @@ end
     iprob = InfiniteOptDynamicOptProblem(rocket, u0map, (0, tf), pmap; steps = 200)
     isol = solve(iprob, Ipopt.Optimizer)
     @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3)
+
+    @variables x(..) v(..)
+    @variables u(..) [bounds = (-1.0, 1.0), input = true]
+    @parameters tf
+    constr = [v(tf) ~ 0, x(tf) ~ 0]
+    cost = [tf] # Minimize time
+
+    @named block = ODESystem(
+        [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr)
+    block, input_idxs = structural_simplify(block, ([u(t)], []))
+    
+    u0map = [x(t) => 1.0, v(t) => 0.0]
+    tspan = (0.0, tf)
+    parammap = [u(t) => 0.0, tf => 1.0]
+    jprob = JuMPDynamicOptProblem(block, u0map, tspan, parammap; steps = 51)
+    jsol = solve(jprob, Ipopt.Optimizer, :Verner8, silent = true)
+    @test isapprox(jsol.sol.t[end], 2.0, atol = 1e-5)
+
+    iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; steps = 51)
+    isol = solve(iprob, Ipopt.Optimizer, silent = true)
+    @test isapprox(isol.sol.t[end], 2.0, atol = 1e-5)
 end
 
 @testset "Cart-pole problem" begin