From 7b4a2844e402dd3e0cba09d5f51bcb2d815e60bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Wed, 12 Nov 2025 14:35:22 +0200
Subject: [PATCH 01/15] update InfiniteOpt

---
 Project.toml             | 2 +-
 ext/MTKInfiniteOptExt.jl | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 1dd887a4e9..ab13e016a5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -123,7 +123,7 @@ FunctionWrappers = "1.1"
 FunctionWrappersWrappers = "0.1"
 Graphs = "1.5.2"
 ImplicitDiscreteSolve = "0.1.2, 1"
-InfiniteOpt = "0.5"
+InfiniteOpt = "0.6"
 InteractiveUtils = "1"
 JuliaFormatter = "1.0.47, 2"
 JumpProcesses = "9.19"
diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl
index e0f02c0436..8dacb3b42d 100644
--- a/ext/MTKInfiniteOptExt.jl
+++ b/ext/MTKInfiniteOptExt.jl
@@ -257,11 +257,11 @@ end
 
 # JuMP variables and Symbolics variables never compare equal. When tracing through dynamics, a function argument can be either a JuMP variable or A Symbolics variable, it can never be both.
 function Base.isequal(::SymbolicUtils.Symbolic,
-        ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr})
+        ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, JuMP.GenericNonlinearExpr})
     false
 end
 function Base.isequal(
-        ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, InfiniteOpt.AbstractInfOptExpr},
+        ::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, JuMP.GenericNonlinearExpr},
         ::SymbolicUtils.Symbolic)
     false
 end

From e535746b7e7d04a2e60c3dad69580908f9006425 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Wed, 12 Nov 2025 14:37:43 +0200
Subject: [PATCH 02/15] add tune_parameters for JuMPDynamicOptProblem

---
 ext/MTKInfiniteOptExt.jl                 | 32 ++++++++++++++----
 src/systems/optimal_control_interface.jl | 41 +++++++++++++++---------
 2 files changed, 51 insertions(+), 22 deletions(-)

diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl
index 8dacb3b42d..daef9399b8 100644
--- a/ext/MTKInfiniteOptExt.jl
+++ b/ext/MTKInfiniteOptExt.jl
@@ -2,6 +2,7 @@ module MTKInfiniteOptExt
 using ModelingToolkit
 using InfiniteOpt
 using DiffEqBase
+using SciMLStructures
 using LinearAlgebra
 using StaticArrays
 using UnPack
@@ -13,12 +14,13 @@ struct InfiniteOptModel
     model::InfiniteModel
     U::Vector{<:AbstractVariableRef}
     V::Vector{<:AbstractVariableRef}
+    P::Vector{<:AbstractVariableRef}
     tₛ::AbstractVariableRef
     is_free_final::Bool
 end
 
 struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
-       AbstractDynamicOptProblem{uType, tType, isinplace}
+       SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace}
     f::F
     u0::uType
     tspan::tType
@@ -33,7 +35,7 @@ struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
 end
 
 struct InfiniteOptDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
-       AbstractDynamicOptProblem{uType, tType, isinplace}
+       SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace}
     f::F
     u0::uType
     tspan::tType
@@ -57,6 +59,9 @@ end
 function MTK.generate_input_variable!(m::InfiniteModel, c0, nc, ts)
     @variable(m, V[i = 1:nc], Infinite(m[:t]), start=c0[i])
 end
+function MTK.generate_tunable_params!(m::InfiniteModel, p0, np)
+    @variable(m, P[i=1:np], start=p0[i])
+end
 
 function MTK.generate_timescale!(m::InfiniteModel, guess, is_free_t)
     @variable(m, tₛ ≥ 0, start = guess)
@@ -81,10 +86,11 @@ MTK.set_objective!(m::InfiniteOptModel, expr) = @objective(m.model, Min, expr)
 function MTK.JuMPDynamicOptProblem(sys::System, op, tspan;
         dt = nothing,
         steps = nothing,
+        tune_parameters = false,
         guesses = Dict(), kwargs...)
     prob,
     _ = MTK.process_DynamicOptProblem(JuMPDynamicOptProblem, InfiniteOptModel, sys,
-        op, tspan; dt, steps, guesses, kwargs...)
+        op, tspan; dt, steps, tune_parameters, guesses, kwargs...)
     prob
 end
 
@@ -125,13 +131,24 @@ function MTK.lowered_var(m::InfiniteOptModel, uv, i, t)
     t isa Union{Num, Symbolics.Symbolic} ? X[i] : X[i](t)
 end
 
+function f_wrapper(f, Uₙ, Vₙ, p, P, t)
+    if SciMLStructures.isscimlstructure(p)
+        _, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
+        p′ = repack(P)
+        f(Uₙ, Vₙ, p′, t)
+    else
+        f(Uₙ, Vₙ, P, t)
+    end
+end
+
 function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
     @unpack A, α, c = tableau
     @unpack wrapped_model, f, p = prob
-    @unpack tₛ, U, V, model = wrapped_model
+    @unpack tₛ, U, V, P, model = wrapped_model
     t = model[:t]
     tsteps = supports(t)
-    dt = tsteps[2] - tsteps[1]
+
+    dt = (tsteps[end] - tsteps[1]) / (length(tsteps) - 1)
 
     nᵤ = length(U)
     nᵥ = length(V)
@@ -142,7 +159,7 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
                 ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = zeros(nᵤ))
                 Uₙ = [U[i](τ) + ΔU[i] * dt for i in 1:nᵤ]
                 Vₙ = [V[i](τ) for i in 1:nᵥ]
-                Kₙ = tₛ * f(Uₙ, Vₙ, p, τ + h * dt)
+                Kₙ = tₛ * f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt)
                 push!(K, Kₙ)
             end
             ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)])
@@ -158,7 +175,7 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
             for (i, h) in enumerate(c)
                 ΔU = @view ΔUs[i, :]
                 Uₙ = U + ΔU * dt
-                @constraint(model, [j = 1:nᵤ], K[i, j]==(tₛ * f(Uₙ, V, p, τ + h * dt)[j]),
+                @constraint(model, [j = 1:nᵤ], K[i, j]==(tₛ * f_wrapper(f, Uₙ, V, p, P, τ + h * dt)[j]),
                     DomainRestrictions(t => τ), base_name="solve_K$i($τ)")
             end
             @constraint(model,
@@ -233,6 +250,7 @@ function MTK.get_U_values(m::InfiniteModel)
     U_vals = value.(m[:U])
     U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:nt]
 end
+MTK.get_P_values(m::InfiniteModel) = value(m[:P])
 MTK.get_t_values(m::InfiniteModel) = value(m[:tₛ]) * supports(m[:t])
 MTK.objective_value(m::InfiniteModel) = InfiniteOpt.objective_value(m)
 
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index 5a0ddbf8d5..79bed22767 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -1,6 +1,3 @@
-abstract type AbstractDynamicOptProblem{uType, tType, isinplace} <:
-              SciMLBase.AbstractODEProblem{uType, tType, isinplace} end
-
 abstract type AbstractCollocation end
 
 struct DynamicOptSolution
@@ -22,8 +19,8 @@ end
     JuMPDynamicOptProblem(sys::System, op, tspan; dt, steps, guesses, kwargs...)
 
 Convert an System 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 
+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 `steps`.
 
 To construct the problem, please load InfiniteOpt along with ModelingToolkit.
@@ -33,7 +30,7 @@ function JuMPDynamicOptProblem end
     InfiniteOptDynamicOptProblem(sys::System, op, tspan; dt)
 
 Convert an System representing an optimal control system into a InfiniteOpt model
-for solving using optimization. Must provide `dt` for determining the length 
+for solving using optimization. Must provide `dt` for determining the length
 of the interpolation arrays.
 
 Related to `JuMPDynamicOptProblem`, but directly adds the differential equations
@@ -46,8 +43,8 @@ function InfiniteOptDynamicOptProblem end
     CasADiDynamicOptProblem(sys::System, op, tspan; dt, steps, guesses, kwargs...)
 
 Convert an System representing an optimal control system into a CasADi 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 
+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 `steps`.
 
 To construct the problem, please load CasADi along with ModelingToolkit.
@@ -57,8 +54,8 @@ function CasADiDynamicOptProblem end
     PyomoDynamicOptProblem(sys::System, op, tspan; dt, steps)
 
 Convert an System representing an optimal control system into a Pyomo 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 
+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 `steps`.
 
 To construct the problem, please load Pyomo along with ModelingToolkit.
@@ -229,13 +226,15 @@ end
 ### MODEL CONSTRUCTION ###
 ##########################
 function process_DynamicOptProblem(
-        prob_type::Type{<:AbstractDynamicOptProblem}, model_type, sys::System, op, tspan;
+        prob_type::Type{<:SciMLBase.AbstractDynamicOptProblem}, model_type, sys::System, op, tspan;
         dt = nothing,
         steps = nothing,
+        tune_parameters = false,
         guesses = Dict(), kwargs...)
     warn_overdetermined(sys, op)
     ctrls = unbound_inputs(sys)
     states = unknowns(sys)
+    params = tune_parameters ? tunable_parameters(sys) : []
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     op = Dict([default_toterm(value(k)) => v for (k, v) in op])
@@ -253,14 +252,16 @@ function process_DynamicOptProblem(
     pmap = recursive_unwrap(AnyDict(pmap))
     evaluate_varmap!(pmap, keys(pmap))
     c0 = value.([pmap[c] for c in ctrls])
+    p0, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
 
     tsteps = LinRange(model_tspan[1], model_tspan[2], steps)
     model = generate_internal_model(model_type)
     generate_time_variable!(model, model_tspan, tsteps)
     U = generate_state_variable!(model, u0, length(states), tsteps)
     V = generate_input_variable!(model, c0, length(ctrls), tsteps)
+    P = generate_tunable_params!(model, p0, length(params))
     tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t)
-    fullmodel = model_type(model, U, V, tₛ, is_free_t)
+    fullmodel = model_type(model, U, V, P, tₛ, is_free_t)
 
     set_variable_bounds!(fullmodel, sys, pmap, tspan[2])
     add_cost_function!(fullmodel, sys, tspan, pmap)
@@ -274,6 +275,7 @@ function generate_time_variable! end
 function generate_internal_model end
 function generate_state_variable! end
 function generate_input_variable! end
+function generate_tunable_params! end
 function generate_timescale! end
 function add_initial_constraints! end
 function add_constraint! end
@@ -467,6 +469,7 @@ function prepare_and_optimize! end
 function get_t_values end
 function get_U_values end
 function get_V_values end
+function get_P_values end
 function successful_solve end
 
 """
@@ -474,17 +477,25 @@ function successful_solve end
 
 - kwargs are used for other options. For example, the `plugin_options` and `solver_options` will propagated to the Opti object in CasADi.
 """
-function DiffEqBase.solve(prob::AbstractDynamicOptProblem,
+function DiffEqBase.solve(prob::SciMLBase.AbstractDynamicOptProblem,
         solver::AbstractCollocation; verbose = false, kwargs...)
     solved_model = prepare_and_optimize!(prob, solver; verbose, kwargs...)
 
     ts = get_t_values(solved_model)
     Us = get_U_values(solved_model)
     Vs = get_V_values(solved_model)
+    Ps = get_P_values(solved_model)
     is_free_final(prob.wrapped_model) && (ts .+ prob.tspan[1])
 
-    ode_sol = DiffEqBase.build_solution(prob, solver, ts, Us)
-    input_sol = isnothing(Vs) ? nothing : DiffEqBase.build_solution(prob, solver, ts, Vs)
+    # update the parameters with the ones in the solved_model
+    if !isempty(Ps)
+        new_p = SciMLStructures.replace(SciMLStructures.Tunable(), prob.p, Ps)
+        new_prob = remake(prob, p=new_p)
+    else
+        new_prob = prob
+    end
+    ode_sol = SciMLBase.build_solution(new_prob, solver, ts, Us)
+    input_sol = isnothing(Vs) ? nothing : SciMLBase.build_solution(new_prob, solver, ts, Vs)
 
     if !successful_solve(solved_model)
         ode_sol = SciMLBase.solution_new_retcode(

From f4ce4f09527c21bd60e21291d38f1ed109b47e03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Wed, 12 Nov 2025 15:42:16 +0200
Subject: [PATCH 03/15] add test

---
 test/extensions/dynamic_optimization.jl | 30 +++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl
index eebdc0873c..cc8c822ef1 100644
--- a/test/extensions/dynamic_optimization.jl
+++ b/test/extensions/dynamic_optimization.jl
@@ -422,3 +422,33 @@ end
     psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(4)))
     @test psol.sol.u[end] ≈ [π, 0, 0, 0]
 end
+
+@testset "Parameter estimation - JuMP" begin
+    @parameters α = 1.5 β = 1.0 [tunable=false] γ = 3.0 δ = 1.0
+    @variables x(t) y(t)
+
+    eqs = [D(x) ~ α * x - β * x * y,
+        D(y) ~ -γ * y + δ * x * y]
+
+    @mtkcompile sys0 = System(eqs, t)
+    tspan = (0.0, 1.0)
+    u0map = [x => 4.0, y => 2.0]
+    parammap = [α => 1.8, β => 1.0, γ => 6.5, δ => 1.0]
+
+    oprob = ODEProblem(sys0, [u0map; parammap], tspan)
+    osol = solve(oprob, Tsit5())
+    ts = range(tspan..., length=51)
+    data = osol(ts, idxs=x).u
+
+    costs = [EvalAt(t)(x)-data[i] for (i, t) in enumerate(ts)]
+    consolidate(u, sub) = sum(abs2.(u))
+
+    @mtkcompile sys = System(eqs, t; costs, consolidate)
+
+    sys′ = subset_tunables(sys, [γ, α])
+    jprob = JuMPDynamicOptProblem(sys′, u0map, tspan; dt=1/50, tune_parameters=true)
+    jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5()))
+
+    @test jsol.sol.ps[γ] ≈ 6.5 rtol=1e-4
+    @test jsol.sol.ps[α] ≈ 1.8 rtol=1e-4
+end

From 0d29bae2f07b11e840c95eb29adc6997c6829061 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Wed, 12 Nov 2025 15:43:27 +0200
Subject: [PATCH 04/15] bump SciMLBase for AbstractDynamicOptProblem

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index ab13e016a5..7911ed181c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -152,7 +152,7 @@ RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
 SCCNonlinearSolve = "1.4.0"
-SciMLBase = "2.115.0"
+SciMLBase = "2.125.0"
 SciMLPublic = "1.0.0"
 SciMLStructures = "1.7"
 Serialization = "1"

From c72a4e231d24cd62a2235942fb4f75aa99cfa47a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Wed, 12 Nov 2025 17:01:19 +0200
Subject: [PATCH 05/15] refactor dynopt tests

---
 test/extensions/dynamic_optimization.jl | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl
index cc8c822ef1..1aa92c53c5 100644
--- a/test/extensions/dynamic_optimization.jl
+++ b/test/extensions/dynamic_optimization.jl
@@ -1,4 +1,5 @@
 using ModelingToolkit
+using ModelingToolkit: t_nounits as t, D_nounits as D
 import InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
@@ -133,8 +134,6 @@ end
 
 @testset "Linear systems" begin
     # Double integrator
-    t = M.t_nounits
-    D = M.D_nounits
     @variables x(..) v(..)
     @variables u(..) [bounds = (-1.0, 1.0), input = true]
     constr = [v(1.0) ~ 0.0]
@@ -239,8 +238,6 @@ end
 end
 
 @testset "Rocket launch" begin
-    t = M.t_nounits
-    D = M.D_nounits
 
     @parameters h_c m₀ h₀ g₀ D_c c Tₘ m_c
     @variables h(..) v(..) m(..) = m₀ [bounds = (m_c, 1)] T(..) [input = true, bounds = (0, Tₘ)]
@@ -306,8 +303,6 @@ end
 end
 
 @testset "Free final time problems" begin
-    t = M.t_nounits
-    D = M.D_nounits
 
     @variables x(..) u(..) [input = true, bounds = (0, 1)]
     @parameters tf

From e3e7e699325b23f1d19ef0920d8fd7434d09472e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Wed, 12 Nov 2025 17:01:43 +0200
Subject: [PATCH 06/15] use AbstractDynamicOptProblem from SciMLBase

---
 ext/MTKCasADiDynamicOptExt.jl | 2 +-
 ext/MTKPyomoDynamicOptExt.jl  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl
index addc478d98..54e581257d 100644
--- a/ext/MTKCasADiDynamicOptExt.jl
+++ b/ext/MTKCasADiDynamicOptExt.jl
@@ -35,7 +35,7 @@ mutable struct CasADiModel
 end
 
 struct CasADiDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
-       AbstractDynamicOptProblem{uType, tType, isinplace}
+       SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace}
     f::F
     u0::uType
     tspan::tType
diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl
index 5b4e9e7a1c..e737d310c6 100644
--- a/ext/MTKPyomoDynamicOptExt.jl
+++ b/ext/MTKPyomoDynamicOptExt.jl
@@ -39,7 +39,7 @@ struct PyomoDynamicOptModel
 end
 
 struct PyomoDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
-       AbstractDynamicOptProblem{uType, tType, isinplace}
+       SciMLBase.AbstractDynamicOptProblem{uType, tType, isinplace}
     f::F
     u0::uType
     tspan::tType

From 3f269a207932764cccebb31803bb76b4ddb1f8bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Wed, 12 Nov 2025 17:32:55 +0200
Subject: [PATCH 07/15] fix typo

---
 src/systems/optimal_control_interface.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index 79bed22767..4616ff6231 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -18,7 +18,7 @@ end
 """
     JuMPDynamicOptProblem(sys::System, op, tspan; dt, steps, guesses, kwargs...)
 
-Convert an System representing an optimal control system into a JuMP model
+Convert a System 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 `steps`.
@@ -29,7 +29,7 @@ function JuMPDynamicOptProblem end
 """
     InfiniteOptDynamicOptProblem(sys::System, op, tspan; dt)
 
-Convert an System representing an optimal control system into a InfiniteOpt model
+Convert a System representing an optimal control system into a InfiniteOpt model
 for solving using optimization. Must provide `dt` for determining the length
 of the interpolation arrays.
 
@@ -42,7 +42,7 @@ function InfiniteOptDynamicOptProblem end
 """
     CasADiDynamicOptProblem(sys::System, op, tspan; dt, steps, guesses, kwargs...)
 
-Convert an System representing an optimal control system into a CasADi model
+Convert a System representing an optimal control system into a CasADi 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 `steps`.
@@ -53,7 +53,7 @@ function CasADiDynamicOptProblem end
 """
     PyomoDynamicOptProblem(sys::System, op, tspan; dt, steps)
 
-Convert an System representing an optimal control system into a Pyomo model
+Convert a System representing an optimal control system into a Pyomo 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 `steps`.

From 0655fd23efcd23143ca245d7db8b40f156633c6e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Thu, 13 Nov 2025 18:19:43 +0200
Subject: [PATCH 08/15] fix parameter substitution

take the MTKParameters values into account and substitute
the tunable parameters with the appropriate symbolic expression
---
 src/systems/optimal_control_interface.jl | 44 ++++++++++++++++++++----
 1 file changed, 37 insertions(+), 7 deletions(-)

diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index 4616ff6231..07724ce3d1 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -234,7 +234,7 @@ function process_DynamicOptProblem(
     warn_overdetermined(sys, op)
     ctrls = unbound_inputs(sys)
     states = unknowns(sys)
-    params = tune_parameters ? tunable_parameters(sys) : []
+    tunable_params = tune_parameters ? tunable_parameters(sys) : []
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     op = Dict([default_toterm(value(k)) => v for (k, v) in op])
@@ -248,9 +248,16 @@ function process_DynamicOptProblem(
     model_tspan, steps, is_free_t = process_tspan(tspan, dt, steps)
     warn_overdetermined(sys, op)
 
-    pmap = filter(p -> (first(p) ∉ Set(unknowns(sys))), op)
-    pmap = recursive_unwrap(AnyDict(pmap))
-    evaluate_varmap!(pmap, keys(pmap))
+    # Build pmap for symbolic substitution in costs/constraints/bounds
+    all_parameters = default_toterm.(parameters(sys))
+    # Extract all parameter values from processed p (which has defaults filled in)
+    getter = SymbolicIndexingInterface.getp(sys, all_parameters)
+    pmap = AnyDict(all_parameters .=> getter(p))
+
+    # Filter out tunable parameters - they should remain symbolic
+    tunable_set = Set(default_toterm.(tunable_params))
+    pmap = filter(kvp -> first(kvp) ∉ tunable_set, pmap)
+
     c0 = value.([pmap[c] for c in ctrls])
     p0, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
 
@@ -259,10 +266,16 @@ function process_DynamicOptProblem(
     generate_time_variable!(model, model_tspan, tsteps)
     U = generate_state_variable!(model, u0, length(states), tsteps)
     V = generate_input_variable!(model, c0, length(ctrls), tsteps)
-    P = generate_tunable_params!(model, p0, length(params))
+    P = generate_tunable_params!(model, p0, length(tunable_params))
     tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t)
     fullmodel = model_type(model, U, V, P, tₛ, is_free_t)
 
+    # Add the symbolic representation of the tunable parameters to the map
+    # The order of the Tunable section is given by the tunable_parameters(sys)
+    # Some backends need symbolic accessors instead of raw variables
+    P_syms = [get_param_for_pmap(fullmodel, P, i) for i in eachindex(tunable_params)]
+    merge!(pmap, Dict(tunable_params .=> P_syms))
+
     set_variable_bounds!(fullmodel, sys, pmap, tspan[2])
     add_cost_function!(fullmodel, sys, tspan, pmap)
     add_user_constraints!(fullmodel, sys, tspan, pmap)
@@ -279,6 +292,22 @@ function generate_tunable_params! end
 function generate_timescale! end
 function add_initial_constraints! end
 function add_constraint! end
+# Default: return P[i] directly. Symbolic backends (like Pyomo) can override this.
+get_param_for_pmap(model, P, i) = P isa AbstractArray ? P[i] : P
+
+function f_wrapper(f, Uₙ, Vₙ, p, P, t)
+    if isempty(P)
+        # no tunable parameters
+        return f(Uₙ, Vₙ, p, t)
+    end
+    if SciMLStructures.isscimlstructure(p)
+        _, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
+        p′ = repack(P)
+        f(Uₙ, Vₙ, p′, t)
+    else
+        f(Uₙ, Vₙ, P, t)
+    end
+end
 
 function set_variable_bounds!(m, sys, pmap, tf)
     @unpack model, U, V, tₛ = m
@@ -394,6 +423,7 @@ function process_integral_bounds end
 function lowered_integral end
 function lowered_derivative end
 function lowered_var end
+function lowered_param end
 function fixed_t_map end
 
 function add_user_constraints!(model, sys, tspan, pmap)
@@ -445,8 +475,8 @@ function substitute_toterm(vars, exprs)
     exprs = map(c -> Symbolics.fast_substitute(c, toterm_map), exprs)
 end
 
-function substitute_params(pmap, exprs)
-    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
+function substitute_params(pmap::Dict, exprs)
+    exprs = map(c -> Symbolics.fixpoint_sub(c, pmap), exprs)
 end
 
 function check_constraint_vars(vars)

From 6d14af0bbfb750a3bc40d05659baa2744f91bdb2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Thu, 13 Nov 2025 18:25:40 +0200
Subject: [PATCH 09/15] add tune_parameters for the rest of the backends

Co-authored-by: Claude <noreply@anthropic.com>
---
 ext/MTKCasADiDynamicOptExt.jl | 29 ++++++++++++++++++++++-------
 ext/MTKInfiniteOptExt.jl      | 21 ++++++---------------
 ext/MTKPyomoDynamicOptExt.jl  | 33 ++++++++++++++++++++++++++++-----
 3 files changed, 56 insertions(+), 27 deletions(-)

diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl
index 54e581257d..cbaf3e36e4 100644
--- a/ext/MTKCasADiDynamicOptExt.jl
+++ b/ext/MTKCasADiDynamicOptExt.jl
@@ -21,16 +21,17 @@ function Base.getindex(m::MXLinearInterpolation, i...)
     length(i) == length(size(m.u)) ? m.u[i...] : m.u[i..., :]
 end
 
-mutable struct CasADiModel
+mutable struct CasADiModel{T}
     model::Opti
     U::MXLinearInterpolation
     V::MXLinearInterpolation
+    P::T
     tₛ::MX
     is_free_final::Bool
     solver_opti::Union{Nothing, Opti}
 
-    function CasADiModel(opti, U, V, tₛ, is_free_final, solver_opti = nothing)
-        new(opti, U, V, tₛ, is_free_final, solver_opti)
+    function CasADiModel(opti, U, V, P, tₛ, is_free_final, solver_opti = nothing)
+        new{typeof(P)}(opti, U, V, P, tₛ, is_free_final, solver_opti)
     end
 end
 
@@ -66,10 +67,11 @@ end
 function MTK.CasADiDynamicOptProblem(sys::System, op, tspan;
         dt = nothing,
         steps = nothing,
+        tune_parameters = false,
         guesses = Dict(), kwargs...)
     prob,
     _ = MTK.process_DynamicOptProblem(
-        CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, guesses, kwargs...)
+        CasADiDynamicOptProblem, CasADiModel, sys, op, tspan; dt, steps, tune_parameters, guesses, kwargs...)
     prob
 end
 
@@ -90,6 +92,14 @@ function MTK.generate_input_variable!(model::Opti, c0, nc, tsteps)
     MXLinearInterpolation(V, tsteps, tsteps[2] - tsteps[1])
 end
 
+function MTK.generate_tunable_params!(model::Opti, p0, np)
+    P = CasADi.variable!(model, np)
+    for i in 1:np
+        set_initial!(model, P[i], p0[i])
+    end
+    P
+end
+
 function MTK.generate_timescale!(model::Opti, guess, is_free_t)
     if is_free_t
         tₛ = variable!(model)
@@ -143,7 +153,7 @@ end
 function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau)
     @unpack A, α, c = tableau
     @unpack wrapped_model, f, p = prob
-    @unpack model, U, V, tₛ = wrapped_model
+    @unpack model, U, V, P, tₛ = wrapped_model
     solver_opti = copy(model)
 
     tsteps = U.t
@@ -160,7 +170,7 @@ function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau)
                 ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = MX(zeros(nᵤ)))
                 Uₙ = U.u[:, k] + ΔU * dt
                 Vₙ = V.u[:, k]
-                Kₙ = tₛ * f(Uₙ, Vₙ, p, τ + h * dt) # scale the time
+                Kₙ = tₛ * MTK.f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt) # scale the time
                 push!(K, Kₙ)
             end
             ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)])
@@ -176,7 +186,7 @@ function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau)
                 ΔU = ΔUs[i, :]'
                 Uₙ = U.u[:, k] + ΔU * dt
                 Vₙ = V.u[:, k]
-                subject_to!(solver_opti, Kᵢ[:, i] == tₛ * f(Uₙ, Vₙ, p, τ + h * dt))
+                subject_to!(solver_opti, Kᵢ[:, i] == tₛ * MTK.f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt))
             end
             ΔU_tot = dt * (Kᵢ * α)
             subject_to!(solver_opti, U.u[:, k] + ΔU_tot == U.u[:, k + 1])
@@ -228,6 +238,11 @@ function MTK.get_V_values(model::CasADiModel)
     end
 end
 
+function MTK.get_P_values(model::CasADiModel)
+    value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value
+    value_getter(model.solver_opti, model.P)
+end
+
 function MTK.get_t_values(model::CasADiModel)
     value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value
     ts = value_getter(model.solver_opti, model.tₛ) .* model.U.t
diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl
index daef9399b8..943afbf9da 100644
--- a/ext/MTKInfiniteOptExt.jl
+++ b/ext/MTKInfiniteOptExt.jl
@@ -97,10 +97,11 @@ end
 function MTK.InfiniteOptDynamicOptProblem(sys::System, op, tspan;
         dt = nothing,
         steps = nothing,
+        tune_parameters = false,
         guesses = Dict(), kwargs...)
     prob,
     pmap = MTK.process_DynamicOptProblem(InfiniteOptDynamicOptProblem, InfiniteOptModel,
-        sys, op, tspan; dt, steps, guesses, kwargs...)
+        sys, op, tspan; dt, steps, tune_parameters, guesses, kwargs...)
     MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan)
     prob
 end
@@ -131,16 +132,6 @@ function MTK.lowered_var(m::InfiniteOptModel, uv, i, t)
     t isa Union{Num, Symbolics.Symbolic} ? X[i] : X[i](t)
 end
 
-function f_wrapper(f, Uₙ, Vₙ, p, P, t)
-    if SciMLStructures.isscimlstructure(p)
-        _, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
-        p′ = repack(P)
-        f(Uₙ, Vₙ, p′, t)
-    else
-        f(Uₙ, Vₙ, P, t)
-    end
-end
-
 function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
     @unpack A, α, c = tableau
     @unpack wrapped_model, f, p = prob
@@ -159,7 +150,7 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
                 ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = zeros(nᵤ))
                 Uₙ = [U[i](τ) + ΔU[i] * dt for i in 1:nᵤ]
                 Vₙ = [V[i](τ) for i in 1:nᵥ]
-                Kₙ = tₛ * f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt)
+                Kₙ = tₛ * MTK.f_wrapper(f, Uₙ, Vₙ, p, P, τ + h * dt)
                 push!(K, Kₙ)
             end
             ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)])
@@ -175,12 +166,12 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
             for (i, h) in enumerate(c)
                 ΔU = @view ΔUs[i, :]
                 Uₙ = U + ΔU * dt
-                @constraint(model, [j = 1:nᵤ], K[i, j]==(tₛ * f_wrapper(f, Uₙ, V, p, P, τ + h * dt)[j]),
-                    DomainRestrictions(t => τ), base_name="solve_K$i($τ)")
+                @constraint(model, [j = 1:nᵤ], K[i, j]==(tₛ * MTK.f_wrapper(f, Uₙ, V, p, P, τ + h * dt)[j]),
+                    DomainRestriction(==(τ), t), base_name="solve_K$i($τ)")
             end
             @constraint(model,
                 [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tsteps[end])),
-                DomainRestrictions(t => τ), base_name="solve_U($τ)")
+                DomainRestriction(==(τ), t), base_name="solve_U($τ)")
         end
     end
 end
diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl
index e737d310c6..3230ef2eb0 100644
--- a/ext/MTKPyomoDynamicOptExt.jl
+++ b/ext/MTKPyomoDynamicOptExt.jl
@@ -16,12 +16,14 @@ const SPECIAL_FUNCTIONS_DICT = Dict([acos => Pyomo.py_acos,
     log => Pyomo.py_log,
     sin => Pyomo.py_sin,
     sqrt => Pyomo.py_sqrt,
-    exp => Pyomo.py_exp])
+    exp => Pyomo.py_exp,
+    abs2 => (x -> x^2)])
 
 struct PyomoDynamicOptModel
     model::ConcreteModel
     U::PyomoVar
     V::PyomoVar
+    P::PyomoVar
     tₛ::PyomoVar
     is_free_final::Bool
     solver_model::Union{Nothing, ConcreteModel}
@@ -30,10 +32,10 @@ struct PyomoDynamicOptModel
     t_sym::Union{Num, Symbolics.BasicSymbolic}
     dummy_sym::Union{Num, Symbolics.BasicSymbolic}
 
-    function PyomoDynamicOptModel(model, U, V, tₛ, is_free_final)
+    function PyomoDynamicOptModel(model, U, V, P, tₛ, is_free_final)
         @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) T_SYM DUMMY_SYM
         model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0)
-        new(model, U, V, tₛ, is_free_final, nothing,
+        new(model, U, V, P, tₛ, is_free_final, nothing,
             PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM)
     end
 end
@@ -60,11 +62,11 @@ end
 _getproperty(s, name::Val{fieldname}) where {fieldname} = getproperty(s, fieldname)
 
 function MTK.PyomoDynamicOptProblem(sys::System, op, tspan;
-        dt = nothing, steps = nothing,
+        dt = nothing, steps = nothing, tune_parameters = false,
         guesses = Dict(), kwargs...)
     prob,
     pmap = MTK.process_DynamicOptProblem(PyomoDynamicOptProblem, PyomoDynamicOptModel,
-        sys, op, tspan; dt, steps, guesses, kwargs...)
+        sys, op, tspan; dt, steps, tune_parameters, guesses, kwargs...)
     conc_model = prob.wrapped_model.model
     MTK.add_equational_constraints!(prob.wrapped_model, sys, pmap, tspan)
     prob
@@ -94,6 +96,13 @@ function MTK.generate_input_variable!(m::ConcreteModel, c0, nc, ts)
     PyomoVar(m.V)
 end
 
+function MTK.generate_tunable_params!(m::ConcreteModel, p0, np)
+    m.p_idxs = pyomo.RangeSet(1, np)
+    init_f = Pyomo.pyfunc((m, i) -> (p0[Pyomo.pyconvert(Int, i)]))
+    m.P = pyomo.Var(m.p_idxs, initialize = init_f)
+    PyomoVar(m.P)
+end
+
 function MTK.generate_timescale!(m::ConcreteModel, guess, is_free_t)
     m.tₛ = is_free_t ? pyomo.Var(initialize = guess, bounds = (0, Inf)) : Pyomo.Py(1)
     PyomoVar(m.tₛ)
@@ -169,6 +178,16 @@ function MTK.lowered_var(m::PyomoDynamicOptModel, uv, i, t)
     Symbolics.unwrap(var)
 end
 
+function MTK.lowered_param(m::PyomoDynamicOptModel, i)
+    P = Symbolics.value(pysym_getproperty(m.model_sym, :P))
+    Symbolics.unwrap(P[i])
+end
+
+# For Pyomo, return symbolic accessors instead of raw PyomoVar
+function MTK.get_param_for_pmap(m::PyomoDynamicOptModel, P, i)
+    MTK.lowered_param(m, i)
+end
+
 struct PyomoCollocation <: AbstractCollocation
     solver::Union{String, Symbol}
     derivative_method::Pyomo.DiscretizationMethod
@@ -208,6 +227,10 @@ function MTK.get_V_values(output::PyomoOutput)
     m = output.model
     [[Pyomo.pyconvert(Float64, pyomo.value(m.V[i, t])) for i in m.v_idxs] for t in m.t]
 end
+function MTK.get_P_values(output::PyomoOutput)
+    m = output.model
+    [Pyomo.pyconvert(Float64, pyomo.value(m.P[i])) for i in m.p_idxs]
+end
 function MTK.get_t_values(output::PyomoOutput)
     m = output.model
     Pyomo.pyconvert(Float64, pyomo.value(m.tₛ)) * [Pyomo.pyconvert(Float64, t) for t in m.t]

From 3d1766a94dae13de40d2180f6e81fafb7b5f89e0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Thu, 13 Nov 2025 18:25:52 +0200
Subject: [PATCH 10/15] update tests

---
 test/extensions/dynamic_optimization.jl | 105 ++++++++++++++++++++++--
 1 file changed, 96 insertions(+), 9 deletions(-)

diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl
index 1aa92c53c5..852a153163 100644
--- a/test/extensions/dynamic_optimization.jl
+++ b/test/extensions/dynamic_optimization.jl
@@ -8,6 +8,7 @@ using Ipopt
 using DataInterpolations
 using CasADi
 using Pyomo
+using Test
 
 import DiffEqBase: solve
 const M = ModelingToolkit
@@ -18,8 +19,6 @@ const ENABLE_CASADI = VERSION >= v"1.11"
     # Test solving without anything attached.
     @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
     @variables x(..) y(..)
-    t = M.t_nounits
-    D = M.D_nounits
 
     eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
         D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
@@ -418,7 +417,45 @@ end
     @test psol.sol.u[end] ≈ [π, 0, 0, 0]
 end
 
-@testset "Parameter estimation - JuMP" begin
+@testset "Parameter defaults usage" begin
+    # Test that parameter defaults are used when not explicitly provided in op
+    @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
+    @variables x(t) y(t)
+
+    eqs = [D(x) ~ α * x - β * x * y
+        D(y) ~ -γ * y + δ * x * y]
+
+    sys = mtkcompile(System(eqs, t, name=:sys))
+    tspan = (0.0, 1.0)
+    u0map = [x => 4.0, y => 2.0]
+
+    # Only provide initial conditions, rely on parameter defaults
+    jprob = JuMPDynamicOptProblem(sys, u0map, tspan, dt = 0.01)
+    jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRK4()))
+
+    # Compare with ODEProblem that also uses defaults
+    oprob = ODEProblem(sys, u0map, tspan)
+    osol = solve(oprob, SimpleRK4(), dt = 0.01)
+
+    @test jsol.sol.u ≈ osol.u
+
+    iprob = InfiniteOptDynamicOptProblem(sys, u0map, tspan, dt = 0.01)
+    isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3)))
+    @test isol.sol.u ≈ osol.u rtol=1e-4
+
+    if ENABLE_CASADI
+        cprob = CasADiDynamicOptProblem(sys, u0map, tspan, dt = 0.01)
+        csol = solve(cprob, CasADiCollocation("ipopt", constructRK4()))
+        @test csol.sol.u ≈ osol.u
+    end
+
+    pprob = PyomoDynamicOptProblem(sys, u0map, tspan, dt = 0.01)
+    psol = solve(pprob, PyomoCollocation("ipopt", BackwardEuler()))
+
+    @test psol.sol.u ≈ osol.u rtol=1e-2
+end
+
+@testset "Parameter estimation" begin
     @parameters α = 1.5 β = 1.0 [tunable=false] γ = 3.0 δ = 1.0
     @variables x(t) y(t)
 
@@ -428,22 +465,72 @@ end
     @mtkcompile sys0 = System(eqs, t)
     tspan = (0.0, 1.0)
     u0map = [x => 4.0, y => 2.0]
-    parammap = [α => 1.8, β => 1.0, γ => 6.5, δ => 1.0]
+    parammap = [α => 2.5, β => 1.0, γ => 3.0, δ => 1.8]
 
     oprob = ODEProblem(sys0, [u0map; parammap], tspan)
     osol = solve(oprob, Tsit5())
     ts = range(tspan..., length=51)
     data = osol(ts, idxs=x).u
 
-    costs = [EvalAt(t)(x)-data[i] for (i, t) in enumerate(ts)]
-    consolidate(u, sub) = sum(abs2.(u))
+    costs = [abs2(EvalAt(t)(x)-data[i]) for (i, t) in enumerate(ts)]
+    consolidate(u, sub) = sum(u)
 
     @mtkcompile sys = System(eqs, t; costs, consolidate)
 
-    sys′ = subset_tunables(sys, [γ, α])
+    sys′ = subset_tunables(sys, [δ, α])
     jprob = JuMPDynamicOptProblem(sys′, u0map, tspan; dt=1/50, tune_parameters=true)
     jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5()))
 
-    @test jsol.sol.ps[γ] ≈ 6.5 rtol=1e-4
-    @test jsol.sol.ps[α] ≈ 1.8 rtol=1e-4
+    @test jsol.sol.ps[δ] ≈ 1.8 rtol=1e-4
+    @test jsol.sol.ps[α] ≈ 2.5 rtol=1e-4
+
+    # test with different time stepping
+
+    jprob = JuMPDynamicOptProblem(sys′, u0map, tspan; dt=1/120, tune_parameters=true)
+    jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5()))
+
+    @test jsol.sol.ps[δ] ≈ 1.8 rtol=1e-4 broken=true
+    @test jsol.sol.ps[α] ≈ 2.5 rtol=1e-4 broken=true
+
+    iprob = InfiniteOptDynamicOptProblem(sys′, u0map, tspan, dt = 1/50, tune_parameters=true)
+    isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3)))
+
+    @test isol.sol.ps[δ] ≈ 1.8 rtol=1e-3
+    @test isol.sol.ps[α] ≈ 2.5 rtol=1e-3
+
+    # test with different time stepping
+
+    iprob = InfiniteOptDynamicOptProblem(sys′, u0map, tspan, dt = 1/120, tune_parameters=true)
+    isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3)))
+
+    @test isol.sol.ps[δ] ≈ 1.8 rtol=1e-3
+    @test isol.sol.ps[α] ≈ 2.5 rtol=1e-3
+
+    if ENABLE_CASADI
+        cprob = CasADiDynamicOptProblem(sys′, u0map, tspan; dt = 1/50, tune_parameters=true)
+        csol = solve(cprob, CasADiCollocation("ipopt", constructRK4()))
+        @test csol.sol.ps[δ] ≈ 1.8 rtol=1e-4
+        @test csol.sol.ps[α] ≈ 2.5 rtol=1e-4
+
+        # test with different time stepping
+
+        cprob = CasADiDynamicOptProblem(sys′, u0map, tspan; dt = 1/120, tune_parameters=true)
+        csol = solve(cprob, CasADiCollocation("ipopt", constructRK4()))
+        @test csol.sol.ps[δ] ≈ 1.8 rtol=1e-4 broken=false
+        @test csol.sol.ps[α] ≈ 2.5 rtol=1e-4 broken=true
+    end
+
+    pprob = PyomoDynamicOptProblem(sys′, u0map, tspan, dt = 1/50, tune_parameters=true)
+    psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(4)))
+
+    @test psol.sol.ps[δ] ≈ 1.8 rtol=1e-4
+    @test psol.sol.ps[α] ≈ 2.5 rtol=1e-4
+
+    # test with different time stepping
+
+    # pprob = PyomoDynamicOptProblem(sys′, u0map, tspan, dt = 1/120, tune_parameters=true)
+    # psol = solve(pprob, PyomoCollocation("ipopt", LagrangeLegendre(4)))
+
+    # @test psol.sol.ps[δ] ≈ 1.8 rtol=1e-4
+    # @test psol.sol.ps[α] ≈ 2.5 rtol=1e-4
 end

From 219ed56c365050300f05ba19859df45479db39aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Sat, 15 Nov 2025 03:37:34 +0200
Subject: [PATCH 11/15] fix for individual params needed in the model

---
 ext/MTKCasADiDynamicOptExt.jl            |  5 ++++-
 src/systems/optimal_control_interface.jl | 12 ++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl
index cbaf3e36e4..b49de5aa01 100644
--- a/ext/MTKCasADiDynamicOptExt.jl
+++ b/ext/MTKCasADiDynamicOptExt.jl
@@ -150,6 +150,9 @@ function MTK.lowered_integral(model::CasADiModel, expr, lo, hi)
     model.tₛ * total
 end
 
+MTK.needs_individual_tunables(::Opti) = true
+MTK.get_param_for_pmap(::Opti, P, i) = P[i]
+
 function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau)
     @unpack A, α, c = tableau
     @unpack wrapped_model, f, p = prob
@@ -240,7 +243,7 @@ end
 
 function MTK.get_P_values(model::CasADiModel)
     value_getter = MTK.successful_solve(model) ? CasADi.debug_value : CasADi.value
-    value_getter(model.solver_opti, model.P)
+    [value_getter(model.solver_opti, model.P[i]) for i in eachindex(model.P)]
 end
 
 function MTK.get_t_values(model::CasADiModel)
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index 07724ce3d1..64c0824e83 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -267,13 +267,15 @@ function process_DynamicOptProblem(
     U = generate_state_variable!(model, u0, length(states), tsteps)
     V = generate_input_variable!(model, c0, length(ctrls), tsteps)
     P = generate_tunable_params!(model, p0, length(tunable_params))
-    tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t)
-    fullmodel = model_type(model, U, V, P, tₛ, is_free_t)
-
     # Add the symbolic representation of the tunable parameters to the map
     # The order of the Tunable section is given by the tunable_parameters(sys)
     # Some backends need symbolic accessors instead of raw variables
-    P_syms = [get_param_for_pmap(fullmodel, P, i) for i in eachindex(tunable_params)]
+    P_syms = [get_param_for_pmap(model, P, i) for i in eachindex(tunable_params)]
+    P_backend = needs_individual_tunables(model) ? P_syms : P
+
+    tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t)
+    fullmodel = model_type(model, U, V, P_backend, tₛ, is_free_t)
+
     merge!(pmap, Dict(tunable_params .=> P_syms))
 
     set_variable_bounds!(fullmodel, sys, pmap, tspan[2])
@@ -294,6 +296,8 @@ function add_initial_constraints! end
 function add_constraint! end
 # Default: return P[i] directly. Symbolic backends (like Pyomo) can override this.
 get_param_for_pmap(model, P, i) = P isa AbstractArray ? P[i] : P
+# Some backends need symbolic accessors instead of raw variables (CasADi in particular)
+needs_individual_tunables(model) = false
 
 function f_wrapper(f, Uₙ, Vₙ, p, P, t)
     if isempty(P)

From 487d304aeaa975ce490fe31b0f1e58daec092f64 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Sat, 15 Nov 2025 03:38:18 +0200
Subject: [PATCH 12/15] switch to EvalAt in the Rocket launch and explicitely
 include all parameters in the system

---
 test/extensions/dynamic_optimization.jl | 36 ++++++++++++-------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl
index 852a153163..c6b654b1b9 100644
--- a/test/extensions/dynamic_optimization.jl
+++ b/test/extensions/dynamic_optimization.jl
@@ -238,49 +238,49 @@ end
 
 @testset "Rocket launch" begin
 
-    @parameters h_c m₀ h₀ g₀ D_c c Tₘ m_c
-    @variables h(..) v(..) m(..) = m₀ [bounds = (m_c, 1)] T(..) [input = true, bounds = (0, Tₘ)]
+    ps = @parameters h_c m₀ h₀ g₀ D_c c Tₘ m_c
+    vars = @variables h(t) v(t) m(t) = m₀ [bounds = (m_c, 1)] T(t) [input = true, bounds = (0, Tₘ)]
     drag(h, v) = D_c * v^2 * exp(-h_c * (h - h₀) / h₀)
     gravity(h) = g₀ * (h₀ / h)
 
-    eqs = [D(h(t)) ~ v(t),
-        D(v(t)) ~ (T(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)),
-        D(m(t)) ~ -T(t) / c]
+    eqs = [D(h) ~ v,
+        D(v) ~ (T - drag(h, v)) / m - gravity(h),
+        D(m) ~ -T / c]
 
     (ts, te) = (0.0, 0.2)
-    costs = [-h(te)]
-    cons = [T(te) ~ 0, m(te) ~ m_c]
-    @named rocket = System(eqs, t; costs, constraints = cons)
-    rocket = mtkcompile(rocket; inputs = [T(t)])
+    costs = [-EvalAt(te)(h)]
+    cons = [EvalAt(te)(T) ~ 0, EvalAt(te)(m) ~ m_c]
+    @named rocket = System(eqs, t, vars, ps; costs, constraints = cons)
+    rocket = mtkcompile(rocket; inputs = [T])
 
-    u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0]
+    u0map = [h => h₀, m => m₀, v => 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]
+        Tₘ => 3.5 * g₀ * m₀, T => 0.0, h₀ => 1, m_c => 0.6]
     jprob = JuMPDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false)
     jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructRadauIIA5()))
-    @test jsol.sol[h(t)][end] > 1.012
+    @test jsol.sol[h][end] > 1.012
 
     if ENABLE_CASADI
         cprob = CasADiDynamicOptProblem(
             rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false)
         csol = solve(cprob, CasADiCollocation("ipopt"))
-        @test csol.sol[h(t)][end] > 1.012
+        @test csol.sol[h][end] > 1.012
     end
 
     iprob = InfiniteOptDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001)
     isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer))
-    @test isol.sol[h(t)][end] > 1.012
+    @test isol.sol[h][end] > 1.012
 
     pprob = PyomoDynamicOptProblem(rocket, [u0map; pmap], (ts, te); dt = 0.001, cse = false)
     psol = solve(pprob, PyomoCollocation("ipopt", LagrangeRadau(4)))
-    @test psol.sol[h(t)][end] > 1.012
+    @test psol.sol[h][end] > 1.012
 
     # Test solution
     @parameters (T_interp::CubicSpline)(..)
-    eqs = [D(h(t)) ~ v(t),
-        D(v(t)) ~ (T_interp(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)),
-        D(m(t)) ~ -T_interp(t) / c]
+    eqs = [D(h) ~ v,
+        D(v) ~ (T_interp(t) - drag(h, v)) / m - gravity(h),
+        D(m) ~ -T_interp(t) / c]
     @mtkcompile rocket_ode = System(eqs, t)
     interpmap = Dict(T_interp => ctrl_to_spline(jsol.input_sol, CubicSpline))
     oprob = ODEProblem(rocket_ode, merge(Dict(u0map), Dict(pmap), interpmap), (ts, te))

From a57a464ca2b8b7d9980f877bbf937b566141540e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Mon, 17 Nov 2025 16:09:19 +0200
Subject: [PATCH 13/15] fix Pyomo extension

Co-authored-by: Claude <noreply@anthropic.com>
---
 ext/MTKPyomoDynamicOptExt.jl | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl
index 3230ef2eb0..adcd5f01a0 100644
--- a/ext/MTKPyomoDynamicOptExt.jl
+++ b/ext/MTKPyomoDynamicOptExt.jl
@@ -19,11 +19,11 @@ const SPECIAL_FUNCTIONS_DICT = Dict([acos => Pyomo.py_acos,
     exp => Pyomo.py_exp,
     abs2 => (x -> x^2)])
 
-struct PyomoDynamicOptModel
+struct PyomoDynamicOptModel{T}
     model::ConcreteModel
     U::PyomoVar
     V::PyomoVar
-    P::PyomoVar
+    P::T
     tₛ::PyomoVar
     is_free_final::Bool
     solver_model::Union{Nothing, ConcreteModel}
@@ -35,7 +35,7 @@ struct PyomoDynamicOptModel
     function PyomoDynamicOptModel(model, U, V, P, tₛ, is_free_final)
         @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) T_SYM DUMMY_SYM
         model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0)
-        new(model, U, V, P, tₛ, is_free_final, nothing,
+        new{typeof(P)}(model, U, V, P, tₛ, is_free_final, nothing,
             PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM)
     end
 end
@@ -178,15 +178,17 @@ function MTK.lowered_var(m::PyomoDynamicOptModel, uv, i, t)
     Symbolics.unwrap(var)
 end
 
-function MTK.lowered_param(m::PyomoDynamicOptModel, i)
-    P = Symbolics.value(pysym_getproperty(m.model_sym, :P))
-    Symbolics.unwrap(P[i])
+# For Pyomo, we need to create symbolic representations for pmap
+# This is needed because Pyomo uses symbolic building for constraints
+function MTK.get_param_for_pmap(m::ConcreteModel, P::PyomoVar, i)
+    # Create a symbolic variable that will be used in the pmap
+    # The actual PyomoVar will be accessed via the symbolic representation
+    @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel)
+    P_sym = Symbolics.value(pysym_getproperty(MODEL_SYM, :P))
+    Symbolics.unwrap(P_sym[i])
 end
 
-# For Pyomo, return symbolic accessors instead of raw PyomoVar
-function MTK.get_param_for_pmap(m::PyomoDynamicOptModel, P, i)
-    MTK.lowered_param(m, i)
-end
+MTK.needs_individual_tunables(m::ConcreteModel) = true
 
 struct PyomoCollocation <: AbstractCollocation
     solver::Union{String, Symbol}

From 3039293fd370edc4e58c9594e1b9e379a151a712 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Tue, 18 Nov 2025 05:48:15 +0200
Subject: [PATCH 14/15] Add error handling for JuMPDynamicOptProblem

When we have more collocation points than expected
the constraints are wrong
---
 ext/MTKCasADiDynamicOptExt.jl            | 10 +++--
 ext/MTKInfiniteOptExt.jl                 |  5 +++
 ext/MTKPyomoDynamicOptExt.jl             |  5 ++-
 src/systems/optimal_control_interface.jl | 53 +++++++++++++++++++++++-
 test/extensions/dynamic_optimization.jl  | 14 +++----
 5 files changed, 72 insertions(+), 15 deletions(-)

diff --git a/ext/MTKCasADiDynamicOptExt.jl b/ext/MTKCasADiDynamicOptExt.jl
index b49de5aa01..f61cec4274 100644
--- a/ext/MTKCasADiDynamicOptExt.jl
+++ b/ext/MTKCasADiDynamicOptExt.jl
@@ -28,10 +28,11 @@ mutable struct CasADiModel{T}
     P::T
     tₛ::MX
     is_free_final::Bool
+    tsteps::LinRange{Float64, Int}
     solver_opti::Union{Nothing, Opti}
 
-    function CasADiModel(opti, U, V, P, tₛ, is_free_final, solver_opti = nothing)
-        new{typeof(P)}(opti, U, V, P, tₛ, is_free_final, solver_opti)
+    function CasADiModel(opti, U, V, P, tₛ, is_free_final, tsteps, solver_opti = nothing)
+        new{typeof(P)}(opti, U, V, P, tₛ, is_free_final, tsteps, solver_opti)
     end
 end
 
@@ -160,7 +161,10 @@ function add_solve_constraints!(prob::CasADiDynamicOptProblem, tableau)
     solver_opti = copy(model)
 
     tsteps = U.t
-    dt = tsteps[2] - tsteps[1]
+    dt = (tsteps[end] - tsteps[1]) / (length(tsteps) - 1)
+
+    # CasADi uses linear interpolation for cost evaluations that are not on the collocation points
+    @assert tsteps == wrapped_model.tsteps "Collocation points mismatch."
 
     nᵤ = size(U.u, 1)
     nᵥ = size(V.u, 1)
diff --git a/ext/MTKInfiniteOptExt.jl b/ext/MTKInfiniteOptExt.jl
index 943afbf9da..bcfb7ce597 100644
--- a/ext/MTKInfiniteOptExt.jl
+++ b/ext/MTKInfiniteOptExt.jl
@@ -17,6 +17,7 @@ struct InfiniteOptModel
     P::Vector{<:AbstractVariableRef}
     tₛ::AbstractVariableRef
     is_free_final::Bool
+    tsteps::LinRange{Float64, Int}
 end
 
 struct JuMPDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
@@ -139,6 +140,10 @@ function add_solve_constraints!(prob::JuMPDynamicOptProblem, tableau)
     t = model[:t]
     tsteps = supports(t)
 
+    # InfiniteOpt can introduce additional collocation points
+    # Make sure that the collocation points are correct.
+    MTK.check_collocation_time_mismatch(prob.f.sys, wrapped_model.tsteps, tsteps)
+
     dt = (tsteps[end] - tsteps[1]) / (length(tsteps) - 1)
 
     nᵤ = length(U)
diff --git a/ext/MTKPyomoDynamicOptExt.jl b/ext/MTKPyomoDynamicOptExt.jl
index adcd5f01a0..a2a05d297a 100644
--- a/ext/MTKPyomoDynamicOptExt.jl
+++ b/ext/MTKPyomoDynamicOptExt.jl
@@ -26,16 +26,17 @@ struct PyomoDynamicOptModel{T}
     P::T
     tₛ::PyomoVar
     is_free_final::Bool
+    tsteps::LinRange{Float64, Int}
     solver_model::Union{Nothing, ConcreteModel}
     dU::PyomoVar
     model_sym::Union{Num, Symbolics.BasicSymbolic}
     t_sym::Union{Num, Symbolics.BasicSymbolic}
     dummy_sym::Union{Num, Symbolics.BasicSymbolic}
 
-    function PyomoDynamicOptModel(model, U, V, P, tₛ, is_free_final)
+    function PyomoDynamicOptModel(model, U, V, P, tₛ, is_free_final, tsteps)
         @variables MODEL_SYM::Symbolics.symstruct(ConcreteModel) T_SYM DUMMY_SYM
         model.dU = dae.DerivativeVar(U, wrt = model.t, initialize = 0)
-        new{typeof(P)}(model, U, V, P, tₛ, is_free_final, nothing,
+        new{typeof(P)}(model, U, V, P, tₛ, is_free_final, tsteps, nothing,
             PyomoVar(model.dU), MODEL_SYM, T_SYM, DUMMY_SYM)
     end
 end
diff --git a/src/systems/optimal_control_interface.jl b/src/systems/optimal_control_interface.jl
index 64c0824e83..ff29739b52 100644
--- a/src/systems/optimal_control_interface.jl
+++ b/src/systems/optimal_control_interface.jl
@@ -222,6 +222,56 @@ function process_tspan(tspan, dt, steps)
     end
 end
 
+function get_discrete_time_evaluations(expr)
+    vars = Symbolics.get_variables(expr)
+
+    # Group by base variable
+    result = Dict()
+
+    for v in vars
+        if iscall(v)
+            args = arguments(v)
+            if length(args) == 1 && args[1] isa Number
+                base_var = operation(v)
+                time_point = args[1]
+
+                if !haskey(result, base_var)
+                    result[base_var] = Float64[]
+                end
+                push!(result[base_var], time_point)
+            end
+        end
+    end
+
+    # Sort and unique the time points
+    for (var, times) in result
+        result[var] = sort!(unique!(times))
+    end
+
+    return result
+end
+
+function check_collocation_time_mismatch(sys, expected_tsteps, tsteps)
+    if collect(expected_tsteps)≠tsteps
+        eval_times = get_discrete_time_evaluations(cost(sys))
+        for (var, ts) in eval_times
+            tdiff = setdiff(ts, expected_tsteps)
+            @info tdiff
+            if !isempty(tdiff)
+                error("$var is evaluated inside the cost function at $(length(tdiff)) points " *
+                    "that are not in the $(length(expected_tsteps)) collocation points. " *
+                    "Cost evaluation points must align with the collocation grid. "*
+                    "Adjust the dt to match the time points used in the cost function.")
+            end
+        end
+        if length(expected_tsteps) ≠ length(tsteps)
+            error("Found extra $(abs(length(expected_tsteps) - length(tsteps))) collocation points.")
+        elseif maximum(abs.(expected_tsteps .- tsteps)) > 1e-10
+            error("The collocation points differ from the expected ones by more than 1e-10.")
+        end
+    end
+end
+
 ##########################
 ### MODEL CONSTRUCTION ###
 ##########################
@@ -274,7 +324,7 @@ function process_DynamicOptProblem(
     P_backend = needs_individual_tunables(model) ? P_syms : P
 
     tₛ = generate_timescale!(model, get(pmap, tspan[2], tspan[2]), is_free_t)
-    fullmodel = model_type(model, U, V, P_backend, tₛ, is_free_t)
+    fullmodel = model_type(model, U, V, P_backend, tₛ, is_free_t, tsteps)
 
     merge!(pmap, Dict(tunable_params .=> P_syms))
 
@@ -427,7 +477,6 @@ function process_integral_bounds end
 function lowered_integral end
 function lowered_derivative end
 function lowered_var end
-function lowered_param end
 function fixed_t_map end
 
 function add_user_constraints!(model, sys, tspan, pmap)
diff --git a/test/extensions/dynamic_optimization.jl b/test/extensions/dynamic_optimization.jl
index c6b654b1b9..f6b6d75382 100644
--- a/test/extensions/dynamic_optimization.jl
+++ b/test/extensions/dynamic_optimization.jl
@@ -487,10 +487,8 @@ end
     # test with different time stepping
 
     jprob = JuMPDynamicOptProblem(sys′, u0map, tspan; dt=1/120, tune_parameters=true)
-    jsol = solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5()))
-
-    @test jsol.sol.ps[δ] ≈ 1.8 rtol=1e-4 broken=true
-    @test jsol.sol.ps[α] ≈ 2.5 rtol=1e-4 broken=true
+    err_msg = "x is evaluated inside the cost function at 40 points that are not in the 121 collocation points."
+    @test_throws err_msg solve(jprob, JuMPCollocation(Ipopt.Optimizer, constructTsitouras5()))
 
     iprob = InfiniteOptDynamicOptProblem(sys′, u0map, tspan, dt = 1/50, tune_parameters=true)
     isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3)))
@@ -503,8 +501,8 @@ end
     iprob = InfiniteOptDynamicOptProblem(sys′, u0map, tspan, dt = 1/120, tune_parameters=true)
     isol = solve(iprob, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3)))
 
-    @test isol.sol.ps[δ] ≈ 1.8 rtol=1e-3
-    @test isol.sol.ps[α] ≈ 2.5 rtol=1e-3
+    @test isol.sol.ps[δ] ≈ 1.8 rtol=1e-4
+    @test isol.sol.ps[α] ≈ 2.5 rtol=1e-4
 
     if ENABLE_CASADI
         cprob = CasADiDynamicOptProblem(sys′, u0map, tspan; dt = 1/50, tune_parameters=true)
@@ -516,8 +514,8 @@ end
 
         cprob = CasADiDynamicOptProblem(sys′, u0map, tspan; dt = 1/120, tune_parameters=true)
         csol = solve(cprob, CasADiCollocation("ipopt", constructRK4()))
-        @test csol.sol.ps[δ] ≈ 1.8 rtol=1e-4 broken=false
-        @test csol.sol.ps[α] ≈ 2.5 rtol=1e-4 broken=true
+        @test csol.sol.ps[δ] ≈ 1.8 rtol=1e-4
+        @test csol.sol.ps[α] ≈ 2.5 rtol=1e-3
     end
 
     pprob = PyomoDynamicOptProblem(sys′, u0map, tspan, dt = 1/50, tune_parameters=true)

From df1728a4e6d946a415bcf4803dcf9946c7b137c5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Tue, 18 Nov 2025 06:28:48 +0200
Subject: [PATCH 15/15] add docs for parameter estimation with the dynamic
 optimization interface

Co-authored-by: Claude <noreply@anthropic.com>
---
 docs/src/tutorials/dynamic_optimization.md | 89 +++++++++++++++++++++-
 1 file changed, 88 insertions(+), 1 deletion(-)

diff --git a/docs/src/tutorials/dynamic_optimization.md b/docs/src/tutorials/dynamic_optimization.md
index ff1ffab3dc..5ba9598cba 100644
--- a/docs/src/tutorials/dynamic_optimization.md
+++ b/docs/src/tutorials/dynamic_optimization.md
@@ -101,7 +101,7 @@ The `tf` mapping in the parameter map is treated as an initial guess.
 Please note that, at the moment, free final time problems cannot support constraints defined at definite time values, like `x(3) ~ 2`.
 
 !!! warning
-    
+
     The Pyomo collocation methods (LagrangeRadau, LagrangeLegendre) currently are bugged for free final time problems. Strongly suggest using BackwardEuler() for such problems when using Pyomo as the backend.
 
 When declaring the problem in this case we need to provide the number of steps, since dt can't be known in advanced. Let's solve plot our final solution and the controller for the block, using InfiniteOpt as the backend:
@@ -126,3 +126,90 @@ axislegend(ax1)
 axislegend(ax2)
 fig
 ```
+
+### Parameter estimation
+
+The dynamic optimization framework can also be used for parameter estimation. In this approach, we treat unknown parameters as tunable variables and minimize the difference between model predictions and observed data.
+
+Let's demonstrate this with the Lotka-Volterra equations. First, we'll generate some synthetic data by solving the system with known parameter values:
+
+```@example dynamic_opt
+@parameters α = 1.5 β = 1.0 [tunable=false] γ = 3.0 δ = 1.0
+@variables x_pe(t) y_pe(t)
+
+eqs_pe = [D(x_pe) ~ α * x_pe - β * x_pe * y_pe,
+    D(y_pe) ~ -γ * y_pe + δ * x_pe * y_pe]
+
+@mtkcompile sys0_pe = System(eqs_pe, t)
+tspan_pe = (0.0, 1.0)
+u0map_pe = [x_pe => 4.0, y_pe => 2.0]
+
+# True parameter values (these are what we'll try to recover)
+parammap_pe = [α => 2.5, δ => 1.8]
+
+oprob_pe = ODEProblem(sys0_pe, [u0map_pe; parammap_pe], tspan_pe)
+osol_pe = solve(oprob_pe, Tsit5())
+
+# Generate synthetic data at 51 time points
+ts_pe = range(tspan_pe..., length=51)
+data_pe = osol_pe(ts_pe, idxs=x_pe).u
+```
+
+Now we'll set up the parameter estimation problem. We use `EvalAt` to evaluate the state at specific time points and construct a least-squares cost function:
+
+```@example dynamic_opt
+costs_pe = [abs2(EvalAt(ti)(x_pe) - data_pe[i]) for (i, ti) in enumerate(ts_pe)]
+
+@mtkcompile sys_pe = System(eqs_pe, t; costs = costs_pe)
+```
+
+By default the cost values are sumed up, if a different behaviour is desired, the `consolidate` keyword can be set in the `System` definition.
+
+Next, we select which parameters to tune using `subset_tunables`. Here we'll estimate `α` and `δ` while keeping `β` and `γ` fixed:
+
+```@example dynamic_opt
+sys_pe′ = subset_tunables(sys_pe, [α, δ])
+```
+
+Now we can solve the parameter estimation problem. Note the `tune_parameters=true` flag:
+
+```@example dynamic_opt
+iprob_pe = InfiniteOptDynamicOptProblem(sys_pe′, u0map_pe, tspan_pe; dt=1/50, tune_parameters=true)
+isol_pe = solve(iprob_pe, InfiniteOptCollocation(Ipopt.Optimizer, InfiniteOpt.OrthogonalCollocation(3)))
+
+println("Estimated α = ", isol_pe.sol.ps[α], " (true value: 2.5)")
+println("Estimated δ = ", isol_pe.sol.ps[δ], " (true value: 1.8)")
+```
+
+Let's visualize the fit:
+
+```@example dynamic_opt
+fig = Figure(resolution = (800, 400))
+ax = Axis(fig[1, 1], title = "Parameter Estimation Results", xlabel = "Time", ylabel = "Prey Population")
+scatter!(ax, ts_pe, data_pe, label = "Data", markersize = 8)
+lines!(ax, isol_pe.sol.t, isol_pe.sol[x_pe], label = "Fitted Model", linewidth = 2)
+axislegend(ax)
+fig
+```
+
+!!! note "Time Alignment for Cost Evaluation"
+    When using `EvalAt` for parameter estimation, different backends handle the case when evaluation times don't align with collocation points differently:
+
+    - **JuMP**: Will throw an error asking you to adjust `dt` if evaluation times don't match collocation points exactly.
+    - **CasADi**: Uses linear interpolation between collocation points for cost evaluations at intermediate times.
+    - **InfiniteOpt**: Automatically adds support points for the evaluation times, handling mismatched grids gracefully.
+
+    For example, InfiniteOpt can use a different `dt` than what the data spacing requires:
+
+    ```@example dynamic_opt
+    # With InfiniteOpt, dt doesn't need to match the data points:
+    iprob_pe2 = InfiniteOptDynamicOptProblem(sys_pe′, u0map_pe, tspan_pe,
+                                             dt = 1/120, tune_parameters=true)
+    isol_pe2 = solve(iprob_pe2, InfiniteOptCollocation(Ipopt.Optimizer,
+                                InfiniteOpt.OrthogonalCollocation(3)))
+
+    println("With dt=1/120: Estimated α = ", isol_pe2.sol.ps[α], " (true value: 2.5)")
+    println("With dt=1/120: Estimated δ = ", isol_pe2.sol.ps[δ], " (true value: 1.8)")
+    ```
+
+    This flexibility makes InfiniteOpt particularly convenient for parameter estimation when your data points don't naturally align with a uniform collocation grid.