Optimization of states and parameters with Multibody.jl and ModelingToolkit

Specific Domains Modelling & Simulations
,
1

Hi all,
I’m working with a large model created with Multibody.jl and would like to use it to run some optimizations. The optimization variables could be model parameters as well as model states (see a simple problem example below).
I am facing a difficulty to use the auto-differentiation in my problem.

Basing my code on the documentation found here : docs.sciml.ai/ModelingToolkit/v9.59/examples/remake/ , I can’t find a way to use an AutoForwardDiff method. Indeed, on one hand, the replace/remake function are not happy with the JuliaSimCompiler system (used for Multibody.jl), and on the other hand, I can’t find a way to replace state values (equivalent to ‘setu’) with these methods. The functions setp/setu don’t allow a type change and I want to use ForwardDiff.Dual because my model is fully differentiable. Is there a way I could use AutoForwardDiff in my use case?

Below is a MWE of my problem type. I built this one because the solution is trivial and I can debug easily my syntax.

using ModelingToolkit
using Multibody
using OrdinaryDiffEq
using Optimization, OptimizationMOI, Ipopt
using SymbolicIndexingInterface

opt_solver = OptimizationMOI.MOI.OptimizerWithAttributes(Ipopt.Optimizer,
    "hessian_approximation" => "limited-memory",
    "print_level" => 5,
    "max_iter" => 10)

t = Multibody.t
Dt = Differential(t)

@mtkmodel DummyExampleModel begin
    @structural_parameters begin
        r_trans = [0.0, 0.0, -1.0]
    end
    @variables begin
        w_dot(t)
    end
    @components begin
        world = Multibody.World(n=[0.0, 0.0, -1.0])
        motion = Multibody.Revolute(n=[0, 1, 0])
        body = Multibody.Body(state_priority=true)
        trans = Multibody.FixedTranslation(r=r_trans)
    end
    @equations begin
        w_dot ~ Dt(motion.w)
        connect(world.frame_b, motion.frame_a)
        connect(motion.frame_b, trans.frame_a)
        connect(body.frame_a, trans.frame_b)
    end
end

@named mdl = DummyExampleModel()
mdl = complete(mdl)
ssys = structural_simplify(multibody(mdl))

modified_params = [
    mdl.motion.phi => deg2rad(90.0)
]

ode_prob = ODEProblem(ssys, modified_params, (0.0, 0.001))

function loss(x, p)
    ode_prob_loc = p[1]
    setter_states_loc! = p[2]
    setter_states_loc!(ode_prob_loc, x)
    sol = solve(p[1], Tsit5(), saveat=0.01)
    objective = -sol[mdl.w_dot][1]^2
    return objective
end
function cons!(r, x, p)
    ode_prob_loc = p[1]
    setter_states_loc! = p[2]
    setter_states_loc!(ode_prob_loc, x)
    sol = solve(p[1], Tsit5(), saveat=0.01)
    r .= [sol[mdl.motion.w][1]]
end

optim_var = [mdl.motion.phi]

setter_states! = setu(ssys, optim_var) # Limitation : 'setu' is not compatible with ForwardDiff.Dual according to doc
optim_var_x0 = [deg2rad(170.0)]

opt_fun = OptimizationFunction(loss, Optimization.AutoFiniteDiff(); cons=cons!) # I'd like to use Optimization.AutoForwardDiff() but it's not compatible with 'setu'
opt_prob = OptimizationProblem(opt_fun, optim_var_x0, (ode_prob, setter_states!); lb=[deg2rad(0.0)], ub=[deg2rad(180.0)], lcons=[0.0], ucons=[0.0])

sol = solve(opt_prob, opt_solver)

# sol_timeseries = solve(ode_prob, Tsit5(), saveat=0.01)
# import GLMakie
# Multibody.render(mdl, sol_timeseries, display=true, up=[0, 0, 1], x=0, y=5, z=0, lookat=[0, 0, 0], loop=true, cache=false, framerate=10)

I am using the current LTS v1.10.8 with :

  [e9467ef8] GLMakie v0.10.18
  [b6b21f68] Ipopt v1.7.2
  [8391cb6b] JuliaSimCompiler v0.1.26
  [961ee093] ModelingToolkit v9.59.0
  [e1cad5d1] Multibody v0.3.2
  [7f7a1694] Optimization v4.1.1
  [fd9f6733] OptimizationMOI v0.5.2
  [36348300] OptimizationOptimJL v0.4.1
  [1dea7af3] OrdinaryDiffEq v6.91.0
  [53ae85a6] SciMLStructures v1.6.1
  [2efcf032] SymbolicIndexingInterface v0.3.38