OCP <-> MTK

OCP-to-MTK/README.md

@@ -0,0 +1,30 @@
+# OCP to MTK
+
+Transforming optimal control problems from OptimalControl.jl (OCP) to ModelingToolkit.jl (MTK)
+
+## Transformation
+The transformation from a problem in OCP to a problem in MTK can be done with the file `convert_ocp_to_mtk.jl`.
+The folder `convert_tests/` contains tests for the conversion done with this file.
+
+## Examples
+In the folder `examples/`, there are examples of problems in OCP and their equivalents in MTK:
+
+- A problem with free final time
+- A problem with integral cost
+- A problem with a path constraint
+
+## Difficulties in transforming from OCP to MTK
+
+- Integral cost 
+    - Introducing a variable to write the problem with a Mayer formulation
+- Constraints
+    - Bounds
+        - For example $u\in[-10,10]$:  
+        `u(..), [input = true, bounds = (-10, 10)]`
+    - Other constraints
+        - For example the path constraint $x2(t)+u(t)\leq4$:  
+        `Symbolics.Inequality(x2(t) + u(t), 4, ≤)`
+- Dynamic optimization solvers
+    - [Here](https://docs.sciml.ai/ModelingToolkit/dev/API/dynamic_opt/#dynamic_opt_api) are the 4 solvers available for optimal control problems in MTK, but many parts of the documentation are still missing
+- Differences between `@variables` and `@parameters`
+    - `@parameters` (and other macro: `@independent_variables`, `@constants` and `@brownians`) works like `@variables` but allows MTK to attach additional metadata to it

OCP-to-MTK/convert_ocp_to_mtk.jl

@@ -0,0 +1,262 @@
+using OptimalControl
+
+using ModelingToolkit
+
+using Symbolics
+
+
+
+# 1. Récupération des paramètres du problème
+
+function translate_universal_ocp_to_mtk(ocp)
+
+    state_names = OptimalControl.state(ocp).components
+
+    control_names = OptimalControl.control(ocp).components
+
+    nx, nu = length(state_names), length(control_names)
+
+
+
+    ModelingToolkit.@independent_variables t
+
+    D = Differential(t)
+
+
+
+    x_vec = Num[only(ModelingToolkit.@variables $(Symbol(state_names[i]))(t)) for i in 1:nx]
+
+    u_vec = Num[only(ModelingToolkit.@variables $(Symbol(control_names[i]))(t) [input=true]) for i in 1:nu]
+
+    v_vec = Num[]
+
+
+
+    f_dyn = OptimalControl.dynamics(ocp)
+
+    dx_vec = Vector{Num}(undef, nx)
+
+    f_dyn(dx_vec, t, x_vec, u_vec, v_vec)
+
+
+
+    eqs = [D(x_vec[i]) ~ dx_vec[i] for i in 1:nx]
+
+
+
+    constraints_eqs = []
+
+    cm = OptimalControl.constraints(ocp)
+
+
+
+    if length(cm.state_box) == 4
+
+        lbs, indices, ubs, labels = cm.state_box
+
+        for i in 1:length(lbs)
+
+            idx = indices[i]
+
+            if lbs[i] > -Inf push!(constraints_eqs, ModelingToolkit.:≲(lbs[i], x_vec[idx])) end
+
+            if ubs[i] < Inf  push!(constraints_eqs, ModelingToolkit.:≲(x_vec[idx], ubs[i])) end
+
+        end
+
+    end
+
+    if length(cm.control_box) == 4
+
+        lbs, indices, ubs, labels = cm.control_box
+
+        for i in 1:length(lbs)
+
+            idx = indices[i]
+
+            if lbs[i] > -Inf push!(constraints_eqs, ModelingToolkit.:≲(lbs[i], u_vec[idx])) end
+
+            if ubs[i] < Inf  push!(constraints_eqs, ModelingToolkit.:≲(u_vec[idx], ubs[i])) end
+
+        end
+
+    end
+
+
+
+    return t, x_vec, u_vec, eqs, constraints_eqs
+
+end
+
+
+
+# 2. L'utilitaire d'appel
+
+function make_callables(sym_array)
+
+    return [(args...) -> Num(Symbolics.operation(Symbolics.unwrap(sym))(args...)) for sym in sym_array]
+
+end
+
+
+
+# 3. L'extracteur automatique de bornes
+
+function auto_extract_boundaries(ocp, x_sym, u_sym, v_sym, ts, te)
+
+    x_call = make_callables(x_sym)
+
+    x0 = [x(ts) for x in x_call]
+
+    xf = [x(te) for x in x_call]
+
+
+
+    boundary_cons = []
+
+    cm = OptimalControl.constraints(ocp)
+
+
+
+    # 1. Contraintes initiales
+
+    if hasproperty(cm, :initial) && length(cm.initial) == 4
+
+        lbs, f_init, ubs, labels = cm.initial
+
+        res = Vector{Num}(undef, length(lbs))
+
+        f_init(res, x0, v_sym)
+
+        for i in 1:length(lbs)
+
+            if lbs[i] == ubs[i] push!(boundary_cons, res[i] ~ lbs[i])
+
+            else
+
+                if lbs[i] > -Inf push!(boundary_cons, ModelingToolkit.:≲(lbs[i], res[i])) end
+
+                if ubs[i] < Inf  push!(boundary_cons, ModelingToolkit.:≲(res[i], ubs[i])) end
+
+            end
+
+        end
+
+    end
+
+
+
+    # 2. Contraintes finales
+
+    if hasproperty(cm, :final) && length(cm.final) == 4
+
+        lbs, f_final, ubs, labels = cm.final
+
+        res = Vector{Num}(undef, length(lbs))
+
+        f_final(res, xf, v_sym)
+
+        for i in 1:length(lbs)
+
+            if lbs[i] == ubs[i] push!(boundary_cons, res[i] ~ lbs[i])
+
+            else
+
+                if lbs[i] > -Inf push!(boundary_cons, ModelingToolkit.:≲(lbs[i], res[i])) end
+
+                if ubs[i] < Inf  push!(boundary_cons, ModelingToolkit.:≲(res[i], ubs[i])) end
+
+            end
+
+        end
+
+    end
+
+
+
+    # 3. Contraintes frontières globales
+
+    if hasproperty(cm, :boundary) && length(cm.boundary) == 4
+
+        lbs, f_bound, ubs, labels = cm.boundary
+
+        res = Vector{Num}(undef, length(lbs))
+
+        f_bound(res, x0, xf, v_sym)
+
+        for i in 1:length(lbs)
+
+            if lbs[i] == ubs[i] push!(boundary_cons, res[i] ~ lbs[i])
+
+            else
+
+                if lbs[i] > -Inf push!(boundary_cons, ModelingToolkit.:≲(lbs[i], res[i])) end
+
+                if ubs[i] < Inf  push!(boundary_cons, ModelingToolkit.:≲(res[i], ubs[i])) end
+
+            end
+        end
+    end
+
+
+
+    costs = Num[]
+
+    obj = ocp.objective
+
+    if hasproperty(obj, :mayer) && obj.mayer !== nothing
+        cost_expr = obj.mayer(x0, xf, v_sym)
+        push!(costs, obj.criterion == :max ? -cost_expr : cost_expr)
+
+    end
+
+
+
+    return costs, boundary_cons
+
+end
+
+
+
+# 4. L'assembleur de système
+
+function assemble_mtk_system(t_sym, u_sym, sys_eqs, sys_cons, costs, boundary_cons; sys_name=:opt_system)
+
+    toutes_les_contraintes = vcat(sys_cons, boundary_cons)
+
+    @named sys = System(sys_eqs, t_sym; costs=costs, constraints=toutes_les_contraintes)
+
+    return mtkcompile(sys, inputs = u_sym)
+
+end
+
+
+
+# 5. Le convertisseur
+
+function convert_ocp_to_mtk(ocp, ts, te; sys_name=:opt_system, extra_cons=[], init_state=[], init_control=[])
+
+    println("Conversion OCP -> MTK...")
+
+    t_sym, x_sym, u_sym, sys_eqs, sys_cons = translate_universal_ocp_to_mtk(ocp)
+
+    v_sym = Num[]
+
+    costs, boundary_cons = auto_extract_boundaries(ocp, x_sym, u_sym, v_sym, ts, te)
+
+    append!(boundary_cons, extra_cons)
+
+    # print(boundary_cons)
+    compiled_sys = assemble_mtk_system(t_sym, u_sym, sys_eqs, sys_cons, costs, boundary_cons, sys_name=sys_name)
+
+
+    u0map = []
+    for i in 1:length(x_sym) push!(u0map, x_sym[i] => (length(init_state) >= i ? Float64(init_state[i]) : 0.0)) end
+    for i in 1:length(u_sym) push!(u0map, u_sym[i] => (length(init_control) >= i ? Float64(init_control[i]) : 0.0)) end
+
+
+
+    println("Conversion terminée")
+    return compiled_sys, u0map
+
+end

OCP-to-MTK/convert_tests/test.jl

@@ -0,0 +1,62 @@
+using InfiniteOpt
+using Ipopt
+using OrdinaryDiffEq
+using DiffEqDevTools
+using CairoMakie
+
+include("../convert_ocp_to_mtk.jl")
+
+g0 = 1.0; m0 = 1.0; h0 = 1.0; hc = 500.0; mc = 0.6
+c  = 0.5 * sqrt(g0 * h0)
+Dc = 0.5 * 620 * m0 / g0
+Tm = 3.5 * g0 * m0
+t0 = 0.0; tf = 0.2
+
+drag(h, v) = Dc * v^2 * exp(-hc * (h - h0) / h0)
+grav(h)    = g0 * (h0 / h)^2
+
+ocp_test = @def begin
+    t ∈ [t0, tf], time  
+    x = (h, v, m) ∈ R^3, state
+    T ∈ R, control
+
+    0 ≤ T(t) ≤ Tm
+    mc ≤ m(t) ≤ m0
+
+    ∂(h)(t) == v(t)
+    ∂(v)(t) == (T(t) - drag(h(t), v(t))) / m(t) - grav(h(t))
+    ∂(m)(t) == -T(t) / c
+
+    h(t0) == 1.0
+    v(t0) == 0.0
+    m(t0) == 1.0
+    m(tf) == 0.6
+    h(tf) → max
+end
+
+rocket, u0map = convert_ocp_to_mtk(
+    ocp_test, t0, tf; 
+    sys_name = :rocket,
+    init_state = [1.0, 0.0, 1.0], 
+    init_control = [0.0]
+)
+
+jprob = JuMPDynamicOptProblem(rocket, u0map, (t0, tf); dt = 0.001)
+jsol  = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5()))
+
+println("Affichage des résultats...")
+
+fig = Figure(size = (900, 450))
+ax1 = Axis(fig[1, 1], title = "Trajectoire (h, v, m)", xlabel = "Temps (s)")
+ax2 = Axis(fig[1, 2], title = "Poussée", xlabel = "Temps (s)")
+
+for u in unknowns(rocket)
+    lines!(ax1, jsol.sol.t, jsol.sol[u], label = string(u), linewidth=2)
+end
+
+stairs!(ax2, jsol.sol.t, jsol.input_sol[1, :], label = "Poussée T(t)", color=:red, linewidth=2, step=:post)
+
+axislegend(ax1, position=:lt)
+axislegend(ax2)
+display(fig)
+