diff --git a/ext/MadNLPSolverExt/MadNLPSolverExt.jl b/ext/MadNLPSolverExt/MadNLPSolverExt.jl index 4aceef8..0774279 100644 --- a/ext/MadNLPSolverExt/MadNLPSolverExt.jl +++ b/ext/MadNLPSolverExt/MadNLPSolverExt.jl @@ -29,6 +29,8 @@ include("utils.jl") @test opts.max_iter == 3000 @test opts.print_level == 3 @test opts.hessian_approximation == "exact" + @test opts.intermediate_callback === nothing + @test opts.fixed_variable_treatment === nothing opts2 = DirectTrajOpt.MadNLPOptions(max_iter = 100, tol = 1e-6) @test opts2.max_iter == 100 @@ -36,6 +38,153 @@ include("utils.jl") @test opts isa Solvers.AbstractSolverOptions end +@testitem "MadNLP intermediate_callback (raw MadNLP callback) fires per iter" setup=[ + DTOTestHelpers, +] begin + import MadNLP + + mutable struct _IterCounter <: MadNLP.AbstractUserCallback + count::Base.RefValue{Int} + end + (cb::_IterCounter)(::MadNLP.AbstractMadNLPSolver, _) = (cb.count[] += 1; true) + + cb = _IterCounter(Ref(0)) + prob, _ = make_standard_prob() + solve!( + prob; + options = DirectTrajOpt.MadNLPOptions( + max_iter = 5, + intermediate_callback = cb, + fixed_variable_treatment = MadNLP.RelaxBound, + ), + verbose = false, + ) + @test cb.count[] > 0 +end + +@testitem "MadNLP intermediate_callback (AbstractIntermediateCallback) fires per iter" setup=[ + DTOTestHelpers, +] begin + import MadNLP + + mutable struct _AgnosticCounter <: DirectTrajOpt.AbstractIntermediateCallback + count::Base.RefValue{Int} + last_primal_len::Base.RefValue{Int} + end + function (cb::_AgnosticCounter)(primal::AbstractVector, iter::Integer) + cb.count[] += 1 + cb.last_primal_len[] = length(primal) + return true + end + + cb = _AgnosticCounter(Ref(0), Ref(0)) + prob, _ = make_standard_prob() + solve!( + prob; + options = DirectTrajOpt.MadNLPOptions( + max_iter = 5, + intermediate_callback = cb, + fixed_variable_treatment = MadNLP.RelaxBound, + ), + verbose = false, + ) + @test cb.count[] > 0 + # With RelaxBound, the primal vector matches the full NLP variable count. + @test cb.last_primal_len[] == + length(prob.trajectory.datavec) + prob.trajectory.global_dim +end + +@testitem "MadNLP intermediate_callback auto-couples RelaxBound" setup=[DTOTestHelpers] begin + import MadNLP + + mutable struct _AutoCoupleProbe <: DirectTrajOpt.AbstractIntermediateCallback + last_primal_len::Base.RefValue{Int} + end + function (cb::_AutoCoupleProbe)(primal::AbstractVector, _) + cb.last_primal_len[] = length(primal) + return true + end + + cb = _AutoCoupleProbe(Ref(0)) + prob, _ = make_standard_prob() + # Note: NOT passing fixed_variable_treatment. set_options! should auto-set it. + solve!( + prob; + options = DirectTrajOpt.MadNLPOptions(max_iter = 5, intermediate_callback = cb), + verbose = false, + ) + # If RelaxBound auto-coupled correctly, the primal includes fixed variables. + @test cb.last_primal_len[] == + length(prob.trajectory.datavec) + prob.trajectory.global_dim +end + +@testitem "MadNLP auto-couple respects MadNLP's conditional default" setup=[DTOTestHelpers] begin + import MadNLP + + mutable struct _PassthroughProbe <: DirectTrajOpt.AbstractIntermediateCallback + len::Base.RefValue{Int} + end + (cb::_PassthroughProbe)(primal, _) = (cb.len[] = length(primal); true) + + cb = _PassthroughProbe(Ref(0)) + prob, _ = make_standard_prob() + # With `kkt_system = SparseCondensedKKTSystem`, MadNLP's own conditional + # default for `fixed_variable_treatment` is already `RelaxBound`, so the + # auto-couple should not fire. Capture logs and assert our @info is absent. + logs, _ = Test.collect_test_logs() do + solve!( + prob; + options = DirectTrajOpt.MadNLPOptions( + max_iter = 5, + intermediate_callback = cb, + kkt_system = MadNLP.SparseCondensedKKTSystem, + ), + verbose = false, + ) + end + @test !any(l -> occursin("Setting fixed_variable_treatment", l.message), logs) + # MadNLP's untouched conditional default still yields the full primal. + @test cb.len[] == length(prob.trajectory.datavec) + prob.trajectory.global_dim +end + +@testitem "MadNLP intermediate_callback early termination via return false" setup=[ + DTOTestHelpers, +] begin + import MadNLP + + mutable struct _Stopper <: DirectTrajOpt.AbstractIntermediateCallback + max_iters::Int + count::Base.RefValue{Int} + end + function (cb::_Stopper)(_, _) + cb.count[] += 1 + return cb.count[] < cb.max_iters + end + + cb = _Stopper(3, Ref(0)) + prob, _ = make_standard_prob() + solve!( + prob; + options = DirectTrajOpt.MadNLPOptions(max_iter = 100, intermediate_callback = cb), + verbose = false, + ) + # Callback stopped the solve well before max_iter=100. + @test cb.count[] <= 5 +end + +@testitem "MadNLP intermediate_callback rejects invalid type" setup=[DTOTestHelpers] begin + prob, _ = make_standard_prob() + bogus_cb(args...) = true # bare Function — neither abstract nor MadNLP subtype + @test_throws ArgumentError solve!( + prob; + options = DirectTrajOpt.MadNLPOptions( + max_iter = 5, + intermediate_callback = bogus_cb, + ), + verbose = false, + ) +end + @testitem "MadNLP basic solve" setup=[DTOTestHelpers] begin prob, _ = make_standard_prob() traj_before = deepcopy(prob.trajectory.data) diff --git a/ext/MadNLPSolverExt/solver.jl b/ext/MadNLPSolverExt/solver.jl index ca60e53..573696a 100644 --- a/ext/MadNLPSolverExt/solver.jl +++ b/ext/MadNLPSolverExt/solver.jl @@ -200,9 +200,55 @@ end # ---------------------------------------------------------------------------- +""" + _MadNLPCallbackAdapter(inner) + +Wrap a `DirectTrajOpt.AbstractIntermediateCallback` so MadNLP can call it +with its native `(solver, mode)` signature. + +The adapter: +- **Filters mode to `UserCallbackRegular`.** MadNLP also invokes user + callbacks during feasibility restoration and robust mode; those phases + surface intermediate IPM state that's typically not meaningful to a + trajectory-level callback, so they're silently skipped (return `true`). + This makes `solver.cnt.k` monotonic from the callback's point of view. +- **Translates `solver.x` → `MadNLP.variable(solver.x)`** to strip the + slack tail and hand back just the NLP primal. +- **Forwards `solver.cnt.k`** as the iteration index. +""" +struct _MadNLPCallbackAdapter <: MadNLP.AbstractUserCallback + inner::DirectTrajOpt.AbstractIntermediateCallback +end + +function (a::_MadNLPCallbackAdapter)( + solver::MadNLP.AbstractMadNLPSolver, + mode::MadNLP.AbstractUserCallbackStatus, +) + mode isa MadNLP.UserCallbackRegular || return true + return a.inner(MadNLP.variable(solver.x), solver.cnt.k) +end + function DirectTrajOpt.set_options!(optimizer::AbstractOptimizer, options::MadNLPOptions) ignored_options = [:eval_hessian] + # Auto-couple: an AbstractIntermediateCallback needs the full primal vector, + # which requires fixed_variable_treatment = RelaxBound. We only override when + # MadNLP's own conditional default (`kkt_system <: SparseCondensedKKTSystem ? + # RelaxBound : MakeParameter`) would otherwise pick `MakeParameter` and break + # the callback. When the user has selected a kkt_system whose default is + # already `RelaxBound`, MadNLP's if-one-liner gets to do its job untouched. + # Raw MadNLP callbacks are presumed to manage this themselves. + if options.intermediate_callback isa DirectTrajOpt.AbstractIntermediateCallback && + options.fixed_variable_treatment === nothing + madnlp_default_is_relax_bound = + options.kkt_system isa Type && + options.kkt_system <: MadNLP.SparseCondensedKKTSystem + if !madnlp_default_is_relax_bound + @info "Setting fixed_variable_treatment = MadNLP.RelaxBound for AbstractIntermediateCallback (MadNLP's kkt_system default would otherwise eliminate fixed vars from solver.x)" + optimizer.options[:fixed_variable_treatment] = MadNLP.RelaxBound + end + end + for name in fieldnames(typeof(options)) value = getfield(options, name) if name in ignored_options @@ -221,6 +267,22 @@ function DirectTrajOpt.set_options!(optimizer::AbstractOptimizer, options::MadNL hessian_approximation = ((value == "compact_lbfgs") ? MadNLP.CompactLBFGS : hessian_approximation) optimizer.options[name] = hessian_approximation + elseif name == :intermediate_callback + if value isa DirectTrajOpt.AbstractIntermediateCallback + # Wrap solver-agnostic callbacks in the MadNLP-shaped adapter. + optimizer.options[name] = _MadNLPCallbackAdapter(value) + elseif value isa MadNLP.AbstractUserCallback + # Raw MadNLP callbacks pass through unwrapped. + optimizer.options[name] = value + else + throw( + ArgumentError( + "intermediate_callback must be a subtype of " * + "`DirectTrajOpt.AbstractIntermediateCallback` or " * + "`MadNLP.AbstractUserCallback`, got $(typeof(value))", + ), + ) + end else optimizer.options[name] = value end diff --git a/src/solvers/_solvers.jl b/src/solvers/_solvers.jl index e81bb2d..ee212f2 100644 --- a/src/solvers/_solvers.jl +++ b/src/solvers/_solvers.jl @@ -2,6 +2,7 @@ module Solvers export AbstractOptimizer export AbstractSolverOptions, DefaultSolverOptions, _DefaultSolverOptions +export AbstractIntermediateCallback export _solve export _solve_with_kwargs export solve! @@ -17,6 +18,46 @@ using TestItemRunner const AbstractOptimizer = MOI.AbstractOptimizer abstract type AbstractSolverOptions end +""" + AbstractIntermediateCallback + +Solver-agnostic per-iteration callback for trajectory optimization. + +Subtypes implement a callable with signature + + (cb::SubType)(primal::AbstractVector, iter::Integer) -> Bool + +where `primal` is the current full NLP primal vector and `iter` is the +iteration index from the solver's main optimization loop. Return `true` to +continue solving, `false` to stop early (the solver will report a +user-requested termination). + +Each solver extension wraps an `AbstractIntermediateCallback` instance in a +solver-specific adapter at solve time, so the same callback object works +with every backend (MadNLP, Ipopt, …). + +# Contract + +- **`primal` may alias the solver's internal vector.** Copy it (e.g. + `collect(primal)`) if you need to retain the data past the callback + invocation — its contents may shift on the next iteration. +- **`iter` is monotonic.** The callback is invoked only from the solver's + main IPM loop; auxiliary phases (e.g. MadNLP's feasibility restoration + or robust modes) do not fire it. + +# Required MadNLP setup + +When using MadNLP, the callback must receive the **full** primal vector +to reconstruct trajectories correctly. MadNLP's default +`fixed_variable_treatment = MakeParameter` eliminates variables with +`lb == ub` from the working primal, so any subtype that maps `primal` +back onto a `NamedTrajectory` needs `fixed_variable_treatment = +MadNLP.RelaxBound`. When an `AbstractIntermediateCallback` is installed +via `MadNLPOptions.intermediate_callback`, DTO sets this automatically +(with an `@info` log) unless the user has provided a value. +""" +abstract type AbstractIntermediateCallback end + struct DefaultSolverOptions <: AbstractSolverOptions end const _DefaultSolverOptions::Ref{Type{<:AbstractSolverOptions}} = Ref{Type{<:AbstractSolverOptions}}(DefaultSolverOptions) diff --git a/src/solvers/madnlp_solver/options.jl b/src/solvers/madnlp_solver/options.jl index 6d3a382..5f8a6d8 100644 --- a/src/solvers/madnlp_solver/options.jl +++ b/src/solvers/madnlp_solver/options.jl @@ -14,6 +14,34 @@ export MadNLPOptions kkt_system::Any = nothing # e.g. MadNLP.SparseUnreducedKKTSystem cudss_ordering::Any = nothing # e.g. MadNLPGPU.AMD_ORDERING + # Per-iteration user callback. Two accepted forms: + # + # 1. A subtype of `DirectTrajOpt.AbstractIntermediateCallback` (solver-agnostic). + # Signature: `(cb)(primal::AbstractVector, iter::Integer) -> Bool`. + # The MadNLP extension wraps it in an internal adapter at solve time. + # + # 2. A raw `MadNLP.AbstractUserCallback` subtype with native MadNLP signature + # `(cb)(solver::MadNLP.AbstractMadNLPSolver, mode) -> Bool` — passed through + # unwrapped for users who want full access to the IPM state. + # + # Return `false` to stop the solver (yields `USER_REQUESTED_STOP`). + intermediate_callback::Any = nothing + + # Controls how MadNLP handles variables with `lb == ub`. Mirrors MadNLP's + # own `fixed_variable_treatment::Type` field — must be a `Type` (typically + # `MadNLP.MakeParameter` or `MadNLP.RelaxBound`). Default (`nothing`) defers + # to MadNLP's kkt_system-aware conditional default: + # + # kkt_system <: SparseCondensedKKTSystem ? RelaxBound : MakeParameter + # + # When an `AbstractIntermediateCallback` is installed and this field is + # left at `nothing`, `set_options!` only overrides to `RelaxBound` if + # MadNLP's conditional default would otherwise be `MakeParameter` (which + # eliminates fixed boundary vars from `solver.x` and breaks trajectory + # reconstruction). The conditional default's `RelaxBound` branch is left + # untouched. + fixed_variable_treatment::Union{Type,Nothing} = nothing + # # Only supported by DirectTrajOpt._solve, as an optional kwarg override of `hessian_approximation`; # # `hessian_approximation = eval_hessian ? "exact" : "compact_lbfgs"` # eval_hessian::Bool = true