(JuMP) Best way to implement sequential convex (or quadratic) programming

Hi all,

I work with sequential convex programming (SCP) for trajectory optimization, which is where a nonlinear programming problem is approximated by a sequence of convex linearizations. The structure of the problem does not change between iterations, but the coefficients do. These problems (and convexified constraints) can have SOC constraints and in some cases SDP constraints.

Currently, I use JuMP.jl to set up the convex subproblems at each iteration and manually delete and (re)add the linearized constraints at each iteration. A simple implementation can be found here:

For some solvers, such as MOSEK, this works well, and has seemingly low computational overhead, but others, e.g. Gurobi, seem to have much higher overheads (multiples of the total solve time). When profiling it seems this seems to take a significant amount of time with many calls to GRBupdatemodel.

It would be useful to know what the intended workflow within JuMP would be for these types of problems. I currently see 4 possible approaches:

  1. Create a full new JuMP problem at each iteration
  2. Remove and add new constraints at each iteration
  3. Use JuMP parameters (possibly with ParametricOptInterface)
  4. set_normalized_coefficient to modify

I would be interested in knowing if there are any performance benefits to taking approaches 3 or 4 over approach 2. Approach 3 may have issues because of the limited types of constraints supported by ParametricOptInterface, and approach 4 may be difficult to implement as the normalized forms of constraints will need to be considered (adding complication).

2 Likes

It’s common for users to have these concerns.

If the structure of the problem wouldn’t change, then you only need to build the JuMP Model only once at the beginning. Later on you can modify some terms.

Building a JuMP Model, or re-add constraints both have allocation which is undesirable because later on you have GC overhead.

I think use modification APIs are the best.

One more point is that you need to wrap your scripts inside functions.

You can expect that the approach 4 will be the fastest, especially if you call it only once with the complete A matrix. As always, you need to benchmark it to know to what extend, and keep in mind that the results will be different if you switch the solver. I think that 4 is the way to go for SCP.

The modification approach is:

using JuMP
import Clarabel
import DifferentiationInterface
import ForwardDiff
import LinearAlgebra
import OrdinaryDiffEq

function keplerian_dynamics(u, p, t)
    r_2_3 = sqrt(u[1]^2 + u[2]^2 + u[3]^2)^3
    return [u[4:6]; -(u[1:3] ./ r_2_3) .+ p]
end

function keplerian_dynamics_arc(u, c, tspan)
    prob = OrdinaryDiffEq.ODEProblem(keplerian_dynamics, u, tspan, c)
    sol = OrdinaryDiffEq.solve(
        prob,
        OrdinaryDiffEq.Tsit5();
        abstol = 1e-8,
        reltol = 1e-8,
    )
    return sol.u[end]
end

function keplerian_dynamics_stm(u, c, tspan)
    alg = DifferentiationInterface.AutoForwardDiff()
    return DifferentiationInterface.jacobian(alg, [u; c]) do x
        return keplerian_dynamics_arc(x[1:6], x[7:9], tspan)
    end
end

function main()
    segments = 40
    times = range(; start = 0, stop = 2π, length = segments + 1)
    u0 = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]
    uf = [1.5, 0.0, 0.0, 0.0, 1.0, 0.0]
    u_ref = [zeros(3) for _ in 1:segments]
    x_ref = [copy(u0) for _ in 1:(segments+1)]
    for j in 1:segments
        x_ref[j+1] =
            keplerian_dynamics_arc(x_ref[j], u_ref[j], (times[j], times[j+1]))
    end
    model = Model(Clarabel.Optimizer)
    set_silent(model)
    @variables(model, begin
        x[i in 1:6, j in 1:(segments+1)]
        u[i in 1:3, j in 1:segments]
        0 <= u_norm[j in 1:segments] <= 0.25
        l1_state_err >= 0
    end)
    stm_tmp = ones(6, 9)
    @constraints(model, begin
        [i in 1:6], x[i,1] == x_ref[1][i]
        [j in 1:segments], [u_norm[j]; u[:,j]] in SecondOrderCone()
        [l1_state_err; x[:,end] .- uf] in MOI.NormOneCone(7)
        # Need to write this constraint so all variables on the left-hand side
        # and all constants on the right-hand side
        c_dyn[j in 1:segments],
            x[:,j+1] .- stm_tmp * [x[:,j]; u[:,j]] .== x_ref[j+1] - stm_tmp * [x_ref[j]; u_ref[j]]
    end)
    @objective(model, Min, sum(u_norm) + 1e3 * l1_state_err)
    for i in 1:50
        for j in 1:segments
            stm = keplerian_dynamics_stm(x_ref[j], u_ref[j], (times[j], times[j+1]))
            for k in 1:6
                set_normalized_coefficient.(c_dyn[j][k], x[:,j], -stm[k,1:6])
                set_normalized_coefficient.(c_dyn[j][k], u[:,j], -stm[k,7:9])
            end
            set_normalized_rhs.(c_dyn[j], x_ref[j+1] - stm * [x_ref[j]; u_ref[j]])
        end
        optimize!(model)
        assert_is_solved_and_feasible(model)
        for j in 1:segments
            u_ref[j] .= value.(u[:,j])
            x_ref[j+1] = keplerian_dynamics_arc(x_ref[j], u_ref[j], (times[j], times[j+1]))
        end
        target_error = LinearAlgebra.norm(x_ref[end] .- uf)
        println("Iteration $i: Target error = $target_error")
        if target_error < 1e-6
            break
        end
    end
end

julia> @time main()
Iteration 1: Target error = 1.658220397768194
Iteration 2: Target error = 1.5027216178530514
Iteration 3: Target error = 0.48153100568227425
Iteration 4: Target error = 0.09145557903943502
Iteration 5: Target error = 0.05688198318461101
Iteration 6: Target error = 0.001920405192872992
Iteration 7: Target error = 0.00014490965298112078
Iteration 8: Target error = 4.648175785956403e-6
Iteration 9: Target error = 7.987865111919145e-8
  0.106031 seconds (1.86 M allocations: 185.475 MiB, 16.88% gc time)

It’s not super obvious which is better, because I don’t think Clarabel supports in-place modification, so internally JuMP will be deleting and re-adding the constraints anyway.

2 Likes

Thank you. I have now tested a similar implementation of the set_normalized_coefficient / set_normalized_rhs in the full code and now the results appear to be better. But there is definitely a quite a bit more code required. Maybe this can be automatically generated for simple constraints.

In a quick test, Gurobi now appears to be equal to or slightly better than Mosek for the SCP problems, which is along what I would expect. There were no improvements to the Mosek (or Clarabel) run times. I’m assuming that performance benefits to this method are heavily linked with (1) if the solver supports incremental modification and (2) exactly how incremental modification is supported.

1 Like