From d7d970a7290e6d47873a60ae5cc369768005524a Mon Sep 17 00:00:00 2001 From: bbhattacharyya1729 Date: Sun, 20 Apr 2025 22:54:59 -0500 Subject: [PATCH 01/17] Duration Constraint --- src/constraints/linear_constraints.jl | 20 +++++++++++++++++++- src/solvers/ipopt_solver/constraints.jl | 12 ++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/constraints/linear_constraints.jl b/src/constraints/linear_constraints.jl index 382513d..a08692c 100644 --- a/src/constraints/linear_constraints.jl +++ b/src/constraints/linear_constraints.jl @@ -5,7 +5,8 @@ export GlobalBoundsConstraint export AllEqualConstraint export TimeStepsAllEqualConstraint export L1SlackConstraint - +export TotalConstraint +export DurationConstraint ### ### EqualityConstraint ### @@ -245,3 +246,20 @@ function L1SlackConstraint( label ) end + +struct TotalConstraint <: AbstractLinearConstraint + indices::Vector{Int} + value::Float64 + label::String +end + + +function DurationConstraint( + traj::NamedTrajectory, + value::Float64; + label="duration constraint of $value" +) + @assert traj.timestep isa Symbol + indices = [index(k, traj.components[traj.timestep][1], traj.dim) for k ∈ 1:traj.T] + return TotalConstraint(indices, value ,label) +end \ No newline at end of file diff --git a/src/solvers/ipopt_solver/constraints.jl b/src/solvers/ipopt_solver/constraints.jl index 38c62e7..38f2737 100644 --- a/src/solvers/ipopt_solver/constraints.jl +++ b/src/solvers/ipopt_solver/constraints.jl @@ -82,4 +82,16 @@ function (con::L1SlackConstraint)( MOI.EqualTo(0.0) ) end +end + +function (con::TotalConstraint)( + opt::Ipopt.Optimizer, + vars::Vector{MOI.VariableIndex} +) +MOI.add_constraints( + opt, + MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, vars[idx]) for idx in con.indices], 0.0), + MOI.EqualTo(con.value) +) + end \ No newline at end of file From b9981a797b18d0d57d8834be94720c1b9bdbe52f Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Fri, 2 May 2025 22:37:28 -0500 Subject: [PATCH 02/17] free phases --- src/objectives/knot_point_objectives.jl | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/objectives/knot_point_objectives.jl b/src/objectives/knot_point_objectives.jl index d4bd9ef..8ff8651 100644 --- a/src/objectives/knot_point_objectives.jl +++ b/src/objectives/knot_point_objectives.jl @@ -41,6 +41,7 @@ function KnotPointObjective( names::AbstractVector{Symbol}, traj::NamedTrajectory, params::AbstractVector; + global_names::AbstractVector{Symbol}=Symbol[], times::AbstractVector{Int}=1:traj.T, Qs::AbstractVector{Float64}=ones(traj.T), ) @@ -49,7 +50,8 @@ function KnotPointObjective( Z_dim = traj.dim * traj.T + traj.global_dim x_comps = vcat([traj.components[name] for name in names]...) - x_slices = [slice(t, x_comps, traj.dim) for t in times] + g_comps = vcat([traj.global_components[name] for name in global_names]...) + x_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] function L(Z⃗::AbstractVector{<:Real}) loss = 0.0 @@ -111,14 +113,16 @@ function TerminalObjective( ℓ::Function, name::Symbol, traj::NamedTrajectory; - Q::Float64=1.0 + Q::Float64=1.0, + kwargs... ) return KnotPointObjective( ℓ, name, traj; Qs=[Q], - times=[traj.T] + times=[traj.T], + kwargs... ) end From acf0de1a2f3359f877e70662661eb035974318f3 Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Mon, 5 May 2025 12:28:59 -0500 Subject: [PATCH 03/17] intro global constraints / objectives, constraints to use global dim in sizes --- src/constraints/nonlinear_constraints.jl | 128 +++++++++++++++++++++-- src/dynamics.jl | 17 +-- src/objectives/global_objectives.jl | 0 src/objectives/knot_point_objectives.jl | 61 +++++++++++ 4 files changed, 191 insertions(+), 15 deletions(-) create mode 100644 src/objectives/global_objectives.jl diff --git a/src/constraints/nonlinear_constraints.jl b/src/constraints/nonlinear_constraints.jl index 49311ea..dd7562d 100644 --- a/src/constraints/nonlinear_constraints.jl +++ b/src/constraints/nonlinear_constraints.jl @@ -1,6 +1,11 @@ export NonlinearKnotPointConstraint +export GlobalNonlinearConstraint +# ----------------------------------------------------------------------------- # +# NonlinearKnotPointConstraint +# ----------------------------------------------------------------------------- # + struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint g!::Function ∂g!::Function @@ -54,6 +59,7 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint names::AbstractVector{Symbol}, traj::NamedTrajectory, params::AbstractVector; + global_names::AbstractVector{Symbol}=Symbol[], times::AbstractVector{Int}=1:traj.T, equality::Bool=true, jacobian_structure::Union{Nothing, SparseMatrixCSC}=nothing, @@ -61,9 +67,10 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint ) @assert length(params) == length(times) "params must have the same length as times" - z_dim = traj.dim + Z_dim = traj.dim * traj.T + traj.global_dim x_comps = vcat([traj.components[name] for name in names]...) - x_slices = [slice(t, x_comps, traj.dim) for t in times] + g_comps = vcat([traj.global_components[name] for name in global_names]...) + x_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] @assert g(traj[times[1]].data[x_comps], params[1]) isa AbstractVector{Float64} g_dim = length(g(traj[times[1]].data[x_comps], params[1])) @@ -93,25 +100,25 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint ForwardDiff.hessian!( μ∂²g[x_comps, x_comps], x -> μ[slice(i, g_dim)]' * g(x, params[i]), - Z⃗[slice(k, x_comps, z_dim)] + Z⃗[slice(k, x_comps, traj.dim)] ) end end if isnothing(jacobian_structure) - jacobian_structure = spzeros(g_dim, z_dim) + jacobian_structure = spzeros(g_dim, Z_dim) jacobian_structure[:, x_comps] .= 1.0 else - @assert size(jacobian_structure) == (g_dim, z_dim) + @assert size(jacobian_structure) == (g_dim, Z_dim) end ∂gs = [copy(jacobian_structure) for _ ∈ times] if isnothing(hessian_structure) - hessian_structure = spzeros(z_dim, z_dim) + hessian_structure = spzeros(Z_dim, Z_dim) hessian_structure[x_comps, x_comps] .= 1.0 else - @assert size(hessian_structure) == (z_dim, z_dim) + @assert size(hessian_structure) == (Z_dim, Z_dim) end μ∂²gs = [copy(hessian_structure) for _ ∈ times] @@ -158,7 +165,8 @@ function get_full_jacobian( NLC::NonlinearKnotPointConstraint, traj::NamedTrajectory ) - ∂g_full = spzeros(NLC.dim, traj.dim * traj.T) + Z_dim = traj.dim * traj.T + traj.global_dim + ∂g_full = spzeros(NLC.dim, Z_dim) for (i, (k, ∂gₖ)) ∈ enumerate(zip(NLC.times, NLC.∂gs)) ∂g_full[slice(i, NLC.g_dim), slice(k, traj.dim)] = ∂gₖ end @@ -169,13 +177,115 @@ function get_full_hessian( NLC::NonlinearKnotPointConstraint, traj::NamedTrajectory ) - μ∂²g_full = spzeros(traj.dim * traj.T, traj.dim * traj.T) + Z_dim = traj.dim * traj.T + traj.global_dim + μ∂²g_full = spzeros(Z_dim, Z_dim) for (k, μ∂²gₖ) ∈ zip(NLC.times, NLC.μ∂²gs) μ∂²g_full[slice(k, traj.dim), slice(k, traj.dim)] = μ∂²gₖ end return μ∂²g_full end +# ----------------------------------------------------------------------------- # +# GlobalNonlinearConstraint +# ----------------------------------------------------------------------------- # + +struct GlobalNonlinearConstraint <: AbstractNonlinearConstraint + g!::Function + ∂g!::Function + ∂gs::SparseMatrixCSC + μ∂²g!::Function + μ∂²gs::SparseMatrixCSC + equality::Bool + dim::Int + + function GlobalNonlinearConstraint( + g::Function, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + equality::Bool=true, + ) + global_comps = vcat([traj.global_components[name] for name in global_names]...) + local_comps = global_comps .- traj.dim * traj.T + + g_eval = g(vec(traj)[global_comps]) + @assert g_eval isa AbstractVector{Float64} + g_dim = length(g_eval) + + @views function g!(δ::AbstractVector, Z⃗::AbstractVector) + δ[:] = g(Z⃗[global_comps]) + return nothing + end + + @views function ∂g!(∂g::AbstractMatrix, Z⃗::AbstractVector) + ForwardDiff.jacobian!( + ∂g, + x -> g(x), + Z⃗[global_comps] + ) + end + + # global subspace + jacobian_structure = spzeros(g_dim, traj.global_dim) + jacobian_structure[:, local_comps] .= 1.0 + + @views function μ∂²g!( + μ∂²g::AbstractMatrix, + Z⃗::AbstractVector, + μ::AbstractVector + ) + ForwardDiff.hessian!( + μ∂²g[local_comps, local_comps], + x -> μ'g(x), + Z⃗[global_comps] + ) + end + + # global subspace + hessian_structure = spzeros(traj.global_dim, traj.global_dim) + hessian_structure[local_comps, local_comps] .= 1.0 + + return new( + g!, + ∂g!, + jacobian_structure, + μ∂²g!, + hessian_structure, + equality, + g_dim, + ) + end + + function GlobalNonlinearConstraint( + g::Function, + global_name::Symbol, + traj::NamedTrajectory; + kwargs... + ) + return GlobalNonlinearConstraint(g, [global_name], traj; kwargs...) + end +end + +function get_full_jacobian( + NLC::GlobalNonlinearConstraint, + traj::NamedTrajectory +) + Z_dim = traj.dim * traj.T + traj.global_dim + ∂g_full = spzeros(NLC.dim, Z_dim) + ∂g_full[:, traj.dim * traj.T + 1:Z_dim] = NLC.∂gs + return ∂g_full +end + +function get_full_hessian( + NLC::GlobalNonlinearConstraint, + traj::NamedTrajectory +) + Z_dim = traj.dim * traj.T + traj.global_dim + μ∂²g_full = spzeros(Z_dim, Z_dim) + g_comps = traj.dim * traj.T + 1:Z_dim + μ∂²g_full[g_comps, g_comps] = NLC.μ∂²gs + return μ∂²g_full +end + # ============================================================================= # @testitem "testing NonlinearConstraint" begin diff --git a/src/dynamics.jl b/src/dynamics.jl index bf4b101..e808eb8 100644 --- a/src/dynamics.jl +++ b/src/dynamics.jl @@ -51,6 +51,8 @@ function jacobian_structure( return ∂f end +# TODO: Where is get_full_jacobian? + function hessian_structure( integrators::Vector{<:AbstractIntegrator}, traj::NamedTrajectory @@ -62,8 +64,9 @@ function hessian_structure( return μ∂²f end -function get_full_hessian(μ∂²f::AbstractMatrix, traj::NamedTrajectory) - μ∂²F = spzeros(traj.dim * traj.T, traj.dim * traj.T) +function get_full_hessian(μ∂²f::AbstractMatrix, traj::NamedTrajectory) + Z_dim = traj.dim * traj.T + traj.global_dim + μ∂²F = spzeros(Z_dim, Z_dim) for k = 1:traj.T-1 μ∂²F[slice(k, 1:2traj.dim, traj.dim), slice(k, 1:2traj.dim, traj.dim)] .+= μ∂²f end @@ -213,16 +216,18 @@ function NullTrajectoryDynamics() ) end -function get_full_jacobian(D::TrajectoryDynamics, traj::NamedTrajectory) - ∂F = spzeros(D.dim * (traj.T - 1), traj.dim * traj.T) +function get_full_jacobian(D::TrajectoryDynamics, traj::NamedTrajectory) + Z_dim = traj.dim * traj.T + traj.global_dim + ∂F = spzeros(D.dim * (traj.T - 1), Z_dim) for k = 1:traj.T-1 ∂F[slice(k, D.dim), slice(k, 1:2traj.dim, traj.dim)] += D.∂fs[k] end return ∂F end -function get_full_hessian(D::TrajectoryDynamics, traj::NamedTrajectory) - μ∂²F = spzeros(traj.dim * traj.T, traj.dim * traj.T) +function get_full_hessian(D::TrajectoryDynamics, traj::NamedTrajectory) + Z_dim = traj.dim * traj.T + traj.global_dim + μ∂²F = spzeros(Z_dim, Z_dim) for k = 1:traj.T-1 μ∂²F[slice(k, 1:2traj.dim, traj.dim), slice(k, 1:2traj.dim, traj.dim)] .+= D.μ∂²fs[k] end diff --git a/src/objectives/global_objectives.jl b/src/objectives/global_objectives.jl new file mode 100644 index 0000000..e69de29 diff --git a/src/objectives/knot_point_objectives.jl b/src/objectives/knot_point_objectives.jl index 8ff8651..458e664 100644 --- a/src/objectives/knot_point_objectives.jl +++ b/src/objectives/knot_point_objectives.jl @@ -1,5 +1,10 @@ export KnotPointObjective export TerminalObjective +export GlobalObjective + +# ----------------------------------------------------------------------------- # +# KnotPointObjective +# ----------------------------------------------------------------------------- # """ KnotPointObjective( @@ -126,6 +131,62 @@ function TerminalObjective( ) end +# ----------------------------------------------------------------------------- # +# GlobalObjective +# ----------------------------------------------------------------------------- # + +""" + GlobalObjective( + ℓ::Function, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + kwargs... + ) + GlobalObjective( + ℓ::Function, + global_name::Symbol, + traj::NamedTrajectory; + kwargs... + ) + +Create an objective that only involves the global components. +""" +function GlobalObjective( + ℓ::Function, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + Q::Float64=1.0 +) + Z_dim = traj.dim * traj.T + traj.global_dim + g_comps = vcat([traj.global_components[name] for name in global_names]...) + + L(Z⃗::AbstractVector{<:Real}) = Q * ℓ(Z⃗[g_comps]) + + @views function ∇L(Z⃗::AbstractVector{<:Real}) + ∇ = zeros(Z_dim) + ∇[g_comps] = ForwardDiff.gradient(x -> Q * ℓ(x), Z⃗[g_comps]) + return ∇ + end + + function ∂²L_structure() + structure = spzeros(Z_dim, Z_dim) + structure[g_comps, g_comps] .= 1.0 + structure_pairs = collect(zip(findnz(structure)[1:2]...)) + return structure_pairs + end + + @views function ∂²L(Z⃗::AbstractVector{<:Real}) + ∂²ℓ = ForwardDiff.hessian(x -> Q * ℓ(x), Z⃗[g_comps]) + return ∂²ℓ[:] + end + + return Objective(L, ∇L, ∂²L, ∂²L_structure) +end + +function ℓ(ℓ::Function, global_name::Symbol, traj::NamedTrajectory; kwargs...) + return GlobalObjective(ℓ, [global_name], traj; kwargs...) +end + # ============================================================================ # @testitem "testing KnotPointObjective" begin From 55e0a8b7cf2c8e89492bacf0e890f086145bcf2a Mon Sep 17 00:00:00 2001 From: bbhattacharyya1729 Date: Tue, 6 May 2025 18:37:48 -0500 Subject: [PATCH 04/17] Added Pairwise --- src/dynamics.jl | 1 - src/objectives/knot_point_objectives.jl | 2 +- src/objectives/regularizers.jl | 31 ++++++++++++++++++++++++- src/problems.jl | 1 - src/solvers/ipopt_solver/solver.jl | 1 - 5 files changed, 31 insertions(+), 5 deletions(-) diff --git a/src/dynamics.jl b/src/dynamics.jl index bf4b101..91ab92e 100644 --- a/src/dynamics.jl +++ b/src/dynamics.jl @@ -154,7 +154,6 @@ struct TrajectoryDynamics end end - @views function μ∂²F!( μ∂²fs::Vector{SparseMatrixCSC{Float64, Int}}, Z⃗::AbstractVector, diff --git a/src/objectives/knot_point_objectives.jl b/src/objectives/knot_point_objectives.jl index d4bd9ef..917b1b0 100644 --- a/src/objectives/knot_point_objectives.jl +++ b/src/objectives/knot_point_objectives.jl @@ -50,7 +50,7 @@ function KnotPointObjective( Z_dim = traj.dim * traj.T + traj.global_dim x_comps = vcat([traj.components[name] for name in names]...) x_slices = [slice(t, x_comps, traj.dim) for t in times] - + function L(Z⃗::AbstractVector{<:Real}) loss = 0.0 for (i, x_slice) in enumerate(x_slices) diff --git a/src/objectives/regularizers.jl b/src/objectives/regularizers.jl index 27327d0..4f29cc5 100644 --- a/src/objectives/regularizers.jl +++ b/src/objectives/regularizers.jl @@ -1,5 +1,5 @@ export QuadraticRegularizer - +export PairwiseQuadraticRegularizer # ----------------------------------------------------------------------------- # # Quadratic Regularizer # ----------------------------------------------------------------------------- # @@ -122,3 +122,32 @@ end +##### +function PairwiseQuadraticRegularizer( + name1::Symbol, + name2::Symbol, + traj::NamedTrajectory, + R::Float64; +) + if traj.timestep isa Symbol + J = KnotPointObjective( + (v,p) -> 1/2 * sum(abs2, (v[1:length(v) ÷ 2]-v[length(v) ÷ 2+1:end] ) * p ) , + [name1,name2], + traj, + [t[1] for t in eachcol(traj[traj.timestep])]; + times = 1:traj.T, + Qs = R * ones(traj.T) + ) + + else + Δt = traj.timestep + J = KnotPointObjective( + v -> 1/2 * sum(abs2, (v[1:length(v) ÷ 2]-v[length(v) ÷ 2+1:end]) * Δt), + [name1,name2], + traj, + [nothing for _ in 1:traj.T]; + times = 1:traj.T, + Qs = R * ones(traj.T) + ) + end +end \ No newline at end of file diff --git a/src/problems.jl b/src/problems.jl index 144f7bb..d651a91 100644 --- a/src/problems.jl +++ b/src/problems.jl @@ -82,7 +82,6 @@ function get_trajectory_constraints(traj::NamedTrajectory) eq_con = EqualityConstraint(name, [traj.T], val, traj; label=label) push!(cons, eq_con) end - # add bounds constraints for (name, bound) ∈ pairs(traj.bounds) if name ∈ keys(traj.initial) && name ∈ keys(traj.final) diff --git a/src/solvers/ipopt_solver/solver.jl b/src/solvers/ipopt_solver/solver.jl index 1c725e3..cc1a181 100644 --- a/src/solvers/ipopt_solver/solver.jl +++ b/src/solvers/ipopt_solver/solver.jl @@ -255,4 +255,3 @@ end solve!(prob; max_iter=100) end - \ No newline at end of file From 66f694128d3f8e1e3361ca9368cfe6e32336ffb2 Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Wed, 7 May 2025 09:50:31 -0500 Subject: [PATCH 05/17] clean --- src/dynamics.jl | 2 -- src/objectives/global_objectives.jl | 0 src/problems.jl | 4 ---- 3 files changed, 6 deletions(-) delete mode 100644 src/objectives/global_objectives.jl diff --git a/src/dynamics.jl b/src/dynamics.jl index e808eb8..deab67b 100644 --- a/src/dynamics.jl +++ b/src/dynamics.jl @@ -51,8 +51,6 @@ function jacobian_structure( return ∂f end -# TODO: Where is get_full_jacobian? - function hessian_structure( integrators::Vector{<:AbstractIntegrator}, traj::NamedTrajectory diff --git a/src/objectives/global_objectives.jl b/src/objectives/global_objectives.jl deleted file mode 100644 index e69de29..0000000 diff --git a/src/problems.jl b/src/problems.jl index 144f7bb..65ac5b4 100644 --- a/src/problems.jl +++ b/src/problems.jl @@ -104,8 +104,4 @@ function get_trajectory_constraints(traj::NamedTrajectory) end - - - - end From 8d936fefd36635dd789710f694068c9104bc52c4 Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Fri, 9 May 2025 00:07:36 -0500 Subject: [PATCH 06/17] refactor to separate pure knot point from global mixed --- src/constraints/_constraints.jl | 1 + .../global_nonlinear_constraints.jl | 252 ++++++++++++++++++ src/constraints/nonlinear_constraints.jl | 188 +++---------- src/objectives/_objectives.jl | 1 + src/objectives/global_objectives.jl | 179 +++++++++++++ src/objectives/knot_point_objectives.jl | 79 ++---- 6 files changed, 483 insertions(+), 217 deletions(-) create mode 100644 src/constraints/global_nonlinear_constraints.jl create mode 100644 src/objectives/global_objectives.jl diff --git a/src/constraints/_constraints.jl b/src/constraints/_constraints.jl index 96a3b8a..73d75dd 100644 --- a/src/constraints/_constraints.jl +++ b/src/constraints/_constraints.jl @@ -21,5 +21,6 @@ abstract type AbstractNonlinearConstraint <: AbstractConstraint end include("linear_constraints.jl") include("nonlinear_constraints.jl") +include("global_nonlinear_constraints.jl") end diff --git a/src/constraints/global_nonlinear_constraints.jl b/src/constraints/global_nonlinear_constraints.jl new file mode 100644 index 0000000..7296ea8 --- /dev/null +++ b/src/constraints/global_nonlinear_constraints.jl @@ -0,0 +1,252 @@ +export GlobalNonlinearConstraint + + +# ----------------------------------------------------------------------------- # +# GlobalNonlinearConstraint +# ----------------------------------------------------------------------------- # + +struct GlobalNonlinearConstraint <: AbstractNonlinearConstraint + g!::Function + ∂g!::Function + ∂gs::SparseMatrixCSC + μ∂²g!::Function + μ∂²gs::SparseMatrixCSC + equality::Bool + dim::Int + + function GlobalNonlinearConstraint( + g::Function, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + equality::Bool=true, + ) + global_comps = vcat([traj.global_components[name] for name in global_names]...) + local_comps = global_comps .- traj.dim * traj.T + + g_eval = g(vec(traj)[global_comps]) + @assert g_eval isa AbstractVector{Float64} + g_dim = length(g_eval) + + @views function g!(δ::AbstractVector, Z⃗::AbstractVector) + δ[:] = g(Z⃗[global_comps]) + return nothing + end + + @views function ∂g!(∂g::AbstractMatrix, Z⃗::AbstractVector) + ForwardDiff.jacobian!( + ∂g, + x -> g(x), + Z⃗[global_comps] + ) + end + + # global subspace + jacobian_structure = spzeros(g_dim, traj.global_dim) + jacobian_structure[:, local_comps] .= 1.0 + + @views function μ∂²g!( + μ∂²g::AbstractMatrix, + Z⃗::AbstractVector, + μ::AbstractVector + ) + ForwardDiff.hessian!( + μ∂²g[local_comps, local_comps], + x -> μ'g(x), + Z⃗[global_comps] + ) + end + + # global subspace + hessian_structure = spzeros(traj.global_dim, traj.global_dim) + hessian_structure[local_comps, local_comps] .= 1.0 + + return new( + g!, + ∂g!, + jacobian_structure, + μ∂²g!, + hessian_structure, + equality, + g_dim, + ) + end + + function GlobalNonlinearConstraint( + g::Function, + global_name::Symbol, + traj::NamedTrajectory; + kwargs... + ) + return GlobalNonlinearConstraint(g, [global_name], traj; kwargs...) + end +end + +function get_full_jacobian( + NLC::GlobalNonlinearConstraint, + traj::NamedTrajectory +) + Z_dim = traj.dim * traj.T + traj.global_dim + ∂g_full = spzeros(NLC.dim, Z_dim) + ∂g_full[:, traj.dim * traj.T + 1:Z_dim] = NLC.∂gs + return ∂g_full +end + +function get_full_hessian( + NLC::GlobalNonlinearConstraint, + traj::NamedTrajectory +) + Z_dim = traj.dim * traj.T + traj.global_dim + μ∂²g_full = spzeros(Z_dim, Z_dim) + g_comps = traj.dim * traj.T + 1:Z_dim + μ∂²g_full[g_comps, g_comps] = NLC.μ∂²gs + return μ∂²g_full +end + +# --------------------------------------------------------------------------- # +# GlobalNonlinearKnotPointConstraint +# --------------------------------------------------------------------------- # + +# struct GlobalNonlinearKnotPointConstraint <: AbstractNonlinearConstraint +# g!::Function +# ∂g!::Function +# ∂gs::Vector{SparseMatrixCSC} +# μ∂²g!::Function +# μ∂²gs::Vector{SparseMatrixCSC} +# times::AbstractVector{Int} +# equality::Bool +# g_dim::Int +# dim::Int + +# """ +# TODO: Docstring +# """ +# function GlobalNonlinearKnotPointConstraint( +# g::Function, +# names::AbstractVector{Symbol}, +# traj::NamedTrajectory, +# params::AbstractVector; +# global_names::AbstractVector{Symbol}=Symbol[], +# times::AbstractVector{Int}=1:traj.T, +# equality::Bool=true, +# jacobian_structure::Union{Nothing, SparseMatrixCSC}=nothing, +# hessian_structure::Union{Nothing, SparseMatrixCSC}=nothing, +# ) +# @assert length(params) == length(times) "params must have the same length as times" + +# Z_dim = traj.dim * traj.T + traj.global_dim +# x_comps = vcat([traj.components[name] for name in names]...) +# g_comps = vcat([traj.global_components[name] for name in global_names]...) +# x_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] + +# @assert g(traj[times[1]].data[x_comps], params[1]) isa AbstractVector{Float64} +# g_dim = length(g(traj[times[1]].data[x_comps], params[1])) + +# @views function g!(δ::AbstractVector, Z⃗::AbstractVector) +# for (i, x_slice) ∈ enumerate(x_slices) +# δ[slice(i, g_dim)] = g(Z⃗[x_slice], params[i]) +# end +# end + +# @views function ∂g!(∂gs::Vector{<:AbstractMatrix}, Z⃗::AbstractVector) +# for (i, (x_slice, ∂g)) ∈ enumerate(zip(x_slices, ∂gs)) +# ForwardDiff.jacobian!( +# ∂g[:, x_comps], +# x -> g(x, params[i]), +# Z⃗[x_slice] +# ) +# end +# end + +# @views function μ∂²g!( +# μ∂²gs::Vector{<:AbstractMatrix}, +# Z⃗::AbstractVector, +# μ::AbstractVector +# ) +# for (i, (k, μ∂²g)) ∈ enumerate(zip(times, μ∂²gs)) +# ForwardDiff.hessian!( +# μ∂²g[x_comps, x_comps], +# x -> μ[slice(i, g_dim)]' * g(x, params[i]), +# Z⃗[slice(k, x_comps, traj.dim)] +# ) +# end +# end + +# if isnothing(jacobian_structure) +# jacobian_structure = spzeros(g_dim, Z_dim) +# jacobian_structure[:, x_comps] .= 1.0 +# else +# @assert size(jacobian_structure) == (g_dim, Z_dim) +# end + +# ∂gs = [copy(jacobian_structure) for _ ∈ times] + +# if isnothing(hessian_structure) +# hessian_structure = spzeros(Z_dim, Z_dim) +# hessian_structure[x_comps, x_comps] .= 1.0 +# else +# @assert size(hessian_structure) == (Z_dim, Z_dim) +# end + +# μ∂²gs = [copy(hessian_structure) for _ ∈ times] + +# return new( +# g!, +# ∂g!, +# ∂gs, +# μ∂²g!, +# μ∂²gs, +# times, +# equality, +# g_dim, +# g_dim * length(times) +# ) +# end + +# function GlobalNonlinearKnotPointConstraint( +# g::Function, +# names::AbstractVector{Symbol}, +# traj::NamedTrajectory; +# times::AbstractVector{Int}=1:traj.T, +# kwargs... +# ) +# params = [nothing for _ in times] +# g_param = (x, _) -> g(x) +# return GlobalNonlinearKnotPointConstraint( +# g_param, +# names, +# traj, +# params; +# times=times, +# kwargs... +# ) +# end + +# function GlobalNonlinearKnotPointConstraint(g::Function, name::Symbol, args...; kwargs...) +# return GlobalNonlinearKnotPointConstraint(g, [name], args...; kwargs...) +# end + +# end + +# function get_full_jacobian( +# NLC::GlobalNonlinearKnotPointConstraint, +# traj::NamedTrajectory +# ) +# Z_dim = traj.dim * traj.T + traj.global_dim +# ∂g_full = spzeros(NLC.dim, Z_dim) +# for (i, (k, ∂gₖ)) ∈ enumerate(zip(NLC.times, NLC.∂gs)) +# ∂g_full[slice(i, NLC.g_dim), slice(k, traj.dim)] = ∂gₖ +# end +# return ∂g_full +# end + +# function get_full_hessian( +# NLC::GlobalNonlinearKnotPointConstraint, +# traj::NamedTrajectory +# ) +# Z_dim = traj.dim * traj.T + traj.global_dim +# μ∂²g_full = spzeros(Z_dim, Z_dim) +# for (k, μ∂²gₖ) ∈ zip(NLC.times, NLC.μ∂²gs) +# μ∂²g_full[slice(k, traj.dim), slice(k, traj.dim)] = μ∂²gₖ +# end +# return μ∂²g_full +# end diff --git a/src/constraints/nonlinear_constraints.jl b/src/constraints/nonlinear_constraints.jl index dd7562d..9f88b7d 100644 --- a/src/constraints/nonlinear_constraints.jl +++ b/src/constraints/nonlinear_constraints.jl @@ -1,5 +1,4 @@ export NonlinearKnotPointConstraint -export GlobalNonlinearConstraint # ----------------------------------------------------------------------------- # @@ -12,45 +11,29 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint ∂gs::Vector{SparseMatrixCSC} μ∂²g!::Function μ∂²gs::Vector{SparseMatrixCSC} - times::AbstractVector{Int} equality::Bool + times::AbstractVector{Int} g_dim::Int dim::Int """ - NonlinearKnotPointConstraint( - g::Function, - names::AbstractVector{Symbol}, - traj::NamedTrajectory, - params::AbstractVector; - kwargs... - ) - NonlinearKnotPointConstraint( - g::Function, - names::AbstractVector{Symbol}, - traj::NamedTrajectory; - kwargs... - ) NonlinearKnotPointConstraint( g::Function, name::Symbol, - args...; + traj::NamedTrajectory; kwargs... ) - Create a NonlinearKnotPointConstraint object that represents a nonlinear constraint `g(x,p)` - on a trajectory variable `x` with parameters `p`. If the parameters argument is omitted, - `g(x)` is assumed to be a function of `x` only. + Create a NonlinearKnotPointConstraint object that represents a nonlinear constraint on a trajectory. # Arguments - - `g::Function`: Function that defines the constraint, g(x, p) or g(x). + - `g::Function`: Function that defines the constraint. If `equality=false`, the constraint is `g(x) ≤ 0`. - `name::Symbol`: Name of the variable to be constrained. - `traj::NamedTrajectory`: The trajectory on which the constraint is defined. - - `params::AbstractVector`: Parameters `p` for the constraint function `g`, for each time. # Keyword Arguments - - `times::AbstractVector{Int}=1:traj.T`: Time indices at which the constraint is enforced. - `equality::Bool=true`: If `true`, the constraint is `g(x) = 0`. Otherwise, the constraint is `g(x) ≤ 0`. + - `times::AbstractVector{Int}=1:traj.T`: Time indices at which the constraint is enforced. - `jacobian_structure::Union{Nothing, SparseMatrixCSC}=nothing`: Structure of the Jacobian matrix of the constraint. - `hessian_structure::Union{Nothing, SparseMatrixCSC}=nothing`: Structure of the Hessian matrix of the constraint. """ @@ -59,18 +42,15 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint names::AbstractVector{Symbol}, traj::NamedTrajectory, params::AbstractVector; - global_names::AbstractVector{Symbol}=Symbol[], - times::AbstractVector{Int}=1:traj.T, equality::Bool=true, + times::AbstractVector{Int}=1:traj.T, jacobian_structure::Union{Nothing, SparseMatrixCSC}=nothing, hessian_structure::Union{Nothing, SparseMatrixCSC}=nothing, ) @assert length(params) == length(times) "params must have the same length as times" - Z_dim = traj.dim * traj.T + traj.global_dim x_comps = vcat([traj.components[name] for name in names]...) - g_comps = vcat([traj.global_components[name] for name in global_names]...) - x_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] + x_slices = [slice(t, x_comps, traj.dim) for t in times] @assert g(traj[times[1]].data[x_comps], params[1]) isa AbstractVector{Float64} g_dim = length(g(traj[times[1]].data[x_comps], params[1])) @@ -83,6 +63,7 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint @views function ∂g!(∂gs::Vector{<:AbstractMatrix}, Z⃗::AbstractVector) for (i, (x_slice, ∂g)) ∈ enumerate(zip(x_slices, ∂gs)) + # Disjoint ForwardDiff.jacobian!( ∂g[:, x_comps], x -> g(x, params[i]), @@ -96,29 +77,30 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint Z⃗::AbstractVector, μ::AbstractVector ) - for (i, (k, μ∂²g)) ∈ enumerate(zip(times, μ∂²gs)) + for (i, (x_slice, μ∂²g)) ∈ enumerate(zip(x_slices, μ∂²gs)) + # Disjoint ForwardDiff.hessian!( μ∂²g[x_comps, x_comps], x -> μ[slice(i, g_dim)]' * g(x, params[i]), - Z⃗[slice(k, x_comps, traj.dim)] + Z⃗[x_slice] ) end end if isnothing(jacobian_structure) - jacobian_structure = spzeros(g_dim, Z_dim) + jacobian_structure = spzeros(g_dim, traj.dim) jacobian_structure[:, x_comps] .= 1.0 else - @assert size(jacobian_structure) == (g_dim, Z_dim) + @assert size(jacobian_structure) == (g_dim, traj.dim) end ∂gs = [copy(jacobian_structure) for _ ∈ times] if isnothing(hessian_structure) - hessian_structure = spzeros(Z_dim, Z_dim) + hessian_structure = spzeros(traj.dim, traj.dim) hessian_structure[x_comps, x_comps] .= 1.0 else - @assert size(hessian_structure) == (Z_dim, Z_dim) + @assert size(hessian_structure) == (traj.dim, traj.dim) end μ∂²gs = [copy(hessian_structure) for _ ∈ times] @@ -129,36 +111,35 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint ∂gs, μ∂²g!, μ∂²gs, - times, equality, + times, g_dim, g_dim * length(times) ) end +end - function NonlinearKnotPointConstraint( - g::Function, - names::AbstractVector{Symbol}, - traj::NamedTrajectory; - times::AbstractVector{Int}=1:traj.T, +function NonlinearKnotPointConstraint( + g::Function, + names::AbstractVector{Symbol}, + traj::NamedTrajectory; + times::AbstractVector{Int}=1:traj.T, + kwargs... +) + params = [nothing for _ in times] + g_param = (x, _) -> g(x) + return NonlinearKnotPointConstraint( + g_param, + names, + traj, + params; + times=times, kwargs... ) - params = [nothing for _ in times] - g_param = (x, _) -> g(x) - return NonlinearKnotPointConstraint( - g_param, - names, - traj, - params; - times=times, - kwargs... - ) - end - - function NonlinearKnotPointConstraint(g::Function, name::Symbol, args...; kwargs...) - return NonlinearKnotPointConstraint(g, [name], args...; kwargs...) - end +end +function NonlinearKnotPointConstraint(g::Function, name::Symbol, args...; kwargs...) + return NonlinearKnotPointConstraint(g, [name], args...; kwargs...) end function get_full_jacobian( @@ -185,106 +166,7 @@ function get_full_hessian( return μ∂²g_full end -# ----------------------------------------------------------------------------- # -# GlobalNonlinearConstraint -# ----------------------------------------------------------------------------- # - -struct GlobalNonlinearConstraint <: AbstractNonlinearConstraint - g!::Function - ∂g!::Function - ∂gs::SparseMatrixCSC - μ∂²g!::Function - μ∂²gs::SparseMatrixCSC - equality::Bool - dim::Int - - function GlobalNonlinearConstraint( - g::Function, - global_names::AbstractVector{Symbol}, - traj::NamedTrajectory; - equality::Bool=true, - ) - global_comps = vcat([traj.global_components[name] for name in global_names]...) - local_comps = global_comps .- traj.dim * traj.T - g_eval = g(vec(traj)[global_comps]) - @assert g_eval isa AbstractVector{Float64} - g_dim = length(g_eval) - - @views function g!(δ::AbstractVector, Z⃗::AbstractVector) - δ[:] = g(Z⃗[global_comps]) - return nothing - end - - @views function ∂g!(∂g::AbstractMatrix, Z⃗::AbstractVector) - ForwardDiff.jacobian!( - ∂g, - x -> g(x), - Z⃗[global_comps] - ) - end - - # global subspace - jacobian_structure = spzeros(g_dim, traj.global_dim) - jacobian_structure[:, local_comps] .= 1.0 - - @views function μ∂²g!( - μ∂²g::AbstractMatrix, - Z⃗::AbstractVector, - μ::AbstractVector - ) - ForwardDiff.hessian!( - μ∂²g[local_comps, local_comps], - x -> μ'g(x), - Z⃗[global_comps] - ) - end - - # global subspace - hessian_structure = spzeros(traj.global_dim, traj.global_dim) - hessian_structure[local_comps, local_comps] .= 1.0 - - return new( - g!, - ∂g!, - jacobian_structure, - μ∂²g!, - hessian_structure, - equality, - g_dim, - ) - end - - function GlobalNonlinearConstraint( - g::Function, - global_name::Symbol, - traj::NamedTrajectory; - kwargs... - ) - return GlobalNonlinearConstraint(g, [global_name], traj; kwargs...) - end -end - -function get_full_jacobian( - NLC::GlobalNonlinearConstraint, - traj::NamedTrajectory -) - Z_dim = traj.dim * traj.T + traj.global_dim - ∂g_full = spzeros(NLC.dim, Z_dim) - ∂g_full[:, traj.dim * traj.T + 1:Z_dim] = NLC.∂gs - return ∂g_full -end - -function get_full_hessian( - NLC::GlobalNonlinearConstraint, - traj::NamedTrajectory -) - Z_dim = traj.dim * traj.T + traj.global_dim - μ∂²g_full = spzeros(Z_dim, Z_dim) - g_comps = traj.dim * traj.T + 1:Z_dim - μ∂²g_full[g_comps, g_comps] = NLC.μ∂²gs - return μ∂²g_full -end # ============================================================================= # diff --git a/src/objectives/_objectives.jl b/src/objectives/_objectives.jl index 32eab20..6fc3000 100644 --- a/src/objectives/_objectives.jl +++ b/src/objectives/_objectives.jl @@ -81,6 +81,7 @@ end # ----------------------------------------------------------------------------- # include("knot_point_objectives.jl") +include("global_objectives.jl") include("minimum_time_objective.jl") include("regularizers.jl") diff --git a/src/objectives/global_objectives.jl b/src/objectives/global_objectives.jl new file mode 100644 index 0000000..38ec1a9 --- /dev/null +++ b/src/objectives/global_objectives.jl @@ -0,0 +1,179 @@ +export GlobalObjective +export GlobalKnotPointObjective + + +# ----------------------------------------------------------------------------- # +# GlobalObjective +# ----------------------------------------------------------------------------- # + +""" + GlobalObjective( + ℓ::Function, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + kwargs... + ) + GlobalObjective( + ℓ::Function, + global_name::Symbol, + traj::NamedTrajectory; + kwargs... + ) + +Create an objective that only involves the global components. +""" +function GlobalObjective( + ℓ::Function, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + Q::Float64=1.0 +) + Z_dim = traj.dim * traj.T + traj.global_dim + g_comps = vcat([traj.global_components[name] for name in global_names]...) + + L(Z⃗::AbstractVector{<:Real}) = Q * ℓ(Z⃗[g_comps]) + + @views function ∇L(Z⃗::AbstractVector{<:Real}) + ∇ = zeros(Z_dim) + ∇[g_comps] = ForwardDiff.gradient(x -> Q * ℓ(x), Z⃗[g_comps]) + return ∇ + end + + function ∂²L_structure() + structure = spzeros(Z_dim, Z_dim) + structure[g_comps, g_comps] .= 1.0 + structure_pairs = collect(zip(findnz(structure)[1:2]...)) + return structure_pairs + end + + @views function ∂²L(Z⃗::AbstractVector{<:Real}) + ∂²ℓ = ForwardDiff.hessian(x -> Q * ℓ(x), Z⃗[g_comps]) + return ∂²ℓ[:] + end + + return Objective(L, ∇L, ∂²L, ∂²L_structure) +end + +function ℓ(ℓ::Function, global_name::Symbol, traj::NamedTrajectory; kwargs...) + return GlobalObjective(ℓ, [global_name], traj; kwargs...) +end + +# ----------------------------------------------------------------------------- # +# Global KnotPointObjective +# ----------------------------------------------------------------------------- # + +function KnotPointObjective( + ℓ::Function, + names::AbstractVector{Symbol}, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory, + params::AbstractVector; + times::AbstractVector{Int}=1:traj.T, + Qs::AbstractVector{Float64}=ones(traj.T), +) + @assert length(Qs) == length(times) "Qs must have the same length as times" + @assert length(params) == length(times) "params must have the same length as times" + + Z_dim = traj.dim * traj.T + traj.global_dim + x_comps = vcat([traj.components[name] for name in names]...) + g_comps = vcat([traj.global_components[name] for name in global_names]...) + xg_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] + + function L(Z⃗::AbstractVector{<:Real}) + loss = 0.0 + for (i, x_slice) in enumerate(xg_slices) + x = Z⃗[x_slice] + loss += Qs[i] * ℓ(x, params[i]) + end + return loss + end + + @views function ∇L(Z⃗::AbstractVector{<:Real}) + ∇ = zeros(Z_dim) + for (i, x_slice) in enumerate(xg_slices) + # Add because global params + ∇[x_slice] .+= ForwardDiff.gradient(x -> Qs[i] * ℓ(x, params[i]), Z⃗[x_slice]) + end + return ∇ + end + + function ∂²L_structure() + structure = spzeros(Z_dim, Z_dim) + for x_slice in xg_slices + structure[x_slice, x_slice] .= 1.0 + end + structure_pairs = collect(zip(findnz(structure)[1:2]...)) + return structure_pairs + end + + function ∂²L_structure_mapping() + # Build a mapping from (i, j) -> idx + structure_dict = Dict{Tuple{Int, Int}, Int}() + for (idx, pair) in enumerate(∂²L_structure()) + structure_dict[pair] = idx + end + return structure_dict + end + + @views function ∂²L(Z⃗::AbstractVector{<:Real}) + structure_dict = ∂²L_structure_mapping() + ∂²L_values = zeros(length(structure_dict)) + for (i, x_slice) in enumerate(xg_slices) + ∂²ℓ = ForwardDiff.hessian(x -> Qs[i] * ℓ(x, params[i]), Z⃗[x_slice]) + + # TODO: Is there a more efficient way to do this? + for local_i in eachindex(x_slice) + global_i = x_slice[local_i] + for local_j in 1:local_i + global_j = x_slice[local_j] + + # Add to (i,j) + idx = structure_dict[(global_i, global_j)] + ∂²L_values[idx] += ∂²ℓ[local_i, local_j] + + # Add to (j,i) if off-diagonal + if local_i != local_j + idx_sym = structure_dict[(global_j, global_i)] + ∂²L_values[idx_sym] += ∂²ℓ[local_j, local_i] + end + end + end + end + return ∂²L_values + end + + return Objective(L, ∇L, ∂²L, ∂²L_structure) +end + +function KnotPointObjective( + ℓ::Function, + names::AbstractVector{Symbol}, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + times::AbstractVector{Int}=1:traj.T, + kwargs... +) + params = [nothing for _ in times] + ℓ_param = (x, _) -> ℓ(x) + return KnotPointObjective(ℓ_param, names, global_names, traj, params; times=times, kwargs...) +end + + +function TerminalObjective( + ℓ::Function, + name::Symbol, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + Q::Float64=1.0, + kwargs... +) + return KnotPointObjective( + ℓ, + name, + global_names, + traj; + Qs=[Q], + times=[traj.T], + kwargs... + ) +end \ No newline at end of file diff --git a/src/objectives/knot_point_objectives.jl b/src/objectives/knot_point_objectives.jl index 458e664..c650cd9 100644 --- a/src/objectives/knot_point_objectives.jl +++ b/src/objectives/knot_point_objectives.jl @@ -1,6 +1,6 @@ export KnotPointObjective export TerminalObjective -export GlobalObjective + # ----------------------------------------------------------------------------- # # KnotPointObjective @@ -46,7 +46,6 @@ function KnotPointObjective( names::AbstractVector{Symbol}, traj::NamedTrajectory, params::AbstractVector; - global_names::AbstractVector{Symbol}=Symbol[], times::AbstractVector{Int}=1:traj.T, Qs::AbstractVector{Float64}=ones(traj.T), ) @@ -55,8 +54,7 @@ function KnotPointObjective( Z_dim = traj.dim * traj.T + traj.global_dim x_comps = vcat([traj.components[name] for name in names]...) - g_comps = vcat([traj.global_components[name] for name in global_names]...) - x_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] + x_slices = [slice(t, x_comps, traj.dim) for t in times] function L(Z⃗::AbstractVector{<:Real}) loss = 0.0 @@ -70,8 +68,12 @@ function KnotPointObjective( @views function ∇L(Z⃗::AbstractVector{<:Real}) ∇ = zeros(Z_dim) for (i, x_slice) in enumerate(x_slices) - x = Z⃗[x_slice] - ∇[x_slice] = ForwardDiff.gradient(x -> Qs[i] * ℓ(x, params[i]), x) + # Disjoint + ForwardDiff.gradient!( + ∇[x_slice], + x -> Qs[i] * ℓ(x, params[i]), + Z⃗[x_slice] + ) end return ∇ end @@ -87,10 +89,14 @@ function KnotPointObjective( @views function ∂²L(Z⃗::AbstractVector{<:Real}) ∂²L_values = zeros(length(∂²L_structure())) + ∂²ℓ_length = length(x_comps)^2 for (i, x_slice) in enumerate(x_slices) - ∂²ℓ = ForwardDiff.hessian(x -> Qs[i] * ℓ(x, params[i]), Z⃗[x_slice]) - ∂²ℓ_length = length(∂²ℓ[:]) - ∂²L_values[(i - 1) * ∂²ℓ_length + 1:i * ∂²ℓ_length] = ∂²ℓ[:] + # Disjoint + ForwardDiff.hessian!( + ∂²L_values[(i - 1) * ∂²ℓ_length + 1:i * ∂²ℓ_length], + x -> Qs[i] * ℓ(x, params[i]), + Z⃗[x_slice] + ) end return ∂²L_values end @@ -131,61 +137,6 @@ function TerminalObjective( ) end -# ----------------------------------------------------------------------------- # -# GlobalObjective -# ----------------------------------------------------------------------------- # - -""" - GlobalObjective( - ℓ::Function, - global_names::AbstractVector{Symbol}, - traj::NamedTrajectory; - kwargs... - ) - GlobalObjective( - ℓ::Function, - global_name::Symbol, - traj::NamedTrajectory; - kwargs... - ) - -Create an objective that only involves the global components. -""" -function GlobalObjective( - ℓ::Function, - global_names::AbstractVector{Symbol}, - traj::NamedTrajectory; - Q::Float64=1.0 -) - Z_dim = traj.dim * traj.T + traj.global_dim - g_comps = vcat([traj.global_components[name] for name in global_names]...) - - L(Z⃗::AbstractVector{<:Real}) = Q * ℓ(Z⃗[g_comps]) - - @views function ∇L(Z⃗::AbstractVector{<:Real}) - ∇ = zeros(Z_dim) - ∇[g_comps] = ForwardDiff.gradient(x -> Q * ℓ(x), Z⃗[g_comps]) - return ∇ - end - - function ∂²L_structure() - structure = spzeros(Z_dim, Z_dim) - structure[g_comps, g_comps] .= 1.0 - structure_pairs = collect(zip(findnz(structure)[1:2]...)) - return structure_pairs - end - - @views function ∂²L(Z⃗::AbstractVector{<:Real}) - ∂²ℓ = ForwardDiff.hessian(x -> Q * ℓ(x), Z⃗[g_comps]) - return ∂²ℓ[:] - end - - return Objective(L, ∇L, ∂²L, ∂²L_structure) -end - -function ℓ(ℓ::Function, global_name::Symbol, traj::NamedTrajectory; kwargs...) - return GlobalObjective(ℓ, [global_name], traj; kwargs...) -end # ============================================================================ # From 5c009ebcdea194e9696c5bf1caaf0bc54dc590b4 Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Fri, 9 May 2025 19:27:41 -0500 Subject: [PATCH 07/17] check overlapping constraint for zero reset --- src/constraints/_constraints.jl | 2 +- .../global_nonlinear_constraints.jl | 252 ---------------- src/constraints/nonlinear_constraints.jl | 7 +- .../nonlinear_global_constraints.jl | 268 ++++++++++++++++++ src/objectives/global_objectives.jl | 69 ++--- 5 files changed, 305 insertions(+), 293 deletions(-) delete mode 100644 src/constraints/global_nonlinear_constraints.jl create mode 100644 src/constraints/nonlinear_global_constraints.jl diff --git a/src/constraints/_constraints.jl b/src/constraints/_constraints.jl index 73d75dd..6344c00 100644 --- a/src/constraints/_constraints.jl +++ b/src/constraints/_constraints.jl @@ -21,6 +21,6 @@ abstract type AbstractNonlinearConstraint <: AbstractConstraint end include("linear_constraints.jl") include("nonlinear_constraints.jl") -include("global_nonlinear_constraints.jl") +include("nonlinear_global_constraints.jl") end diff --git a/src/constraints/global_nonlinear_constraints.jl b/src/constraints/global_nonlinear_constraints.jl deleted file mode 100644 index 7296ea8..0000000 --- a/src/constraints/global_nonlinear_constraints.jl +++ /dev/null @@ -1,252 +0,0 @@ -export GlobalNonlinearConstraint - - -# ----------------------------------------------------------------------------- # -# GlobalNonlinearConstraint -# ----------------------------------------------------------------------------- # - -struct GlobalNonlinearConstraint <: AbstractNonlinearConstraint - g!::Function - ∂g!::Function - ∂gs::SparseMatrixCSC - μ∂²g!::Function - μ∂²gs::SparseMatrixCSC - equality::Bool - dim::Int - - function GlobalNonlinearConstraint( - g::Function, - global_names::AbstractVector{Symbol}, - traj::NamedTrajectory; - equality::Bool=true, - ) - global_comps = vcat([traj.global_components[name] for name in global_names]...) - local_comps = global_comps .- traj.dim * traj.T - - g_eval = g(vec(traj)[global_comps]) - @assert g_eval isa AbstractVector{Float64} - g_dim = length(g_eval) - - @views function g!(δ::AbstractVector, Z⃗::AbstractVector) - δ[:] = g(Z⃗[global_comps]) - return nothing - end - - @views function ∂g!(∂g::AbstractMatrix, Z⃗::AbstractVector) - ForwardDiff.jacobian!( - ∂g, - x -> g(x), - Z⃗[global_comps] - ) - end - - # global subspace - jacobian_structure = spzeros(g_dim, traj.global_dim) - jacobian_structure[:, local_comps] .= 1.0 - - @views function μ∂²g!( - μ∂²g::AbstractMatrix, - Z⃗::AbstractVector, - μ::AbstractVector - ) - ForwardDiff.hessian!( - μ∂²g[local_comps, local_comps], - x -> μ'g(x), - Z⃗[global_comps] - ) - end - - # global subspace - hessian_structure = spzeros(traj.global_dim, traj.global_dim) - hessian_structure[local_comps, local_comps] .= 1.0 - - return new( - g!, - ∂g!, - jacobian_structure, - μ∂²g!, - hessian_structure, - equality, - g_dim, - ) - end - - function GlobalNonlinearConstraint( - g::Function, - global_name::Symbol, - traj::NamedTrajectory; - kwargs... - ) - return GlobalNonlinearConstraint(g, [global_name], traj; kwargs...) - end -end - -function get_full_jacobian( - NLC::GlobalNonlinearConstraint, - traj::NamedTrajectory -) - Z_dim = traj.dim * traj.T + traj.global_dim - ∂g_full = spzeros(NLC.dim, Z_dim) - ∂g_full[:, traj.dim * traj.T + 1:Z_dim] = NLC.∂gs - return ∂g_full -end - -function get_full_hessian( - NLC::GlobalNonlinearConstraint, - traj::NamedTrajectory -) - Z_dim = traj.dim * traj.T + traj.global_dim - μ∂²g_full = spzeros(Z_dim, Z_dim) - g_comps = traj.dim * traj.T + 1:Z_dim - μ∂²g_full[g_comps, g_comps] = NLC.μ∂²gs - return μ∂²g_full -end - -# --------------------------------------------------------------------------- # -# GlobalNonlinearKnotPointConstraint -# --------------------------------------------------------------------------- # - -# struct GlobalNonlinearKnotPointConstraint <: AbstractNonlinearConstraint -# g!::Function -# ∂g!::Function -# ∂gs::Vector{SparseMatrixCSC} -# μ∂²g!::Function -# μ∂²gs::Vector{SparseMatrixCSC} -# times::AbstractVector{Int} -# equality::Bool -# g_dim::Int -# dim::Int - -# """ -# TODO: Docstring -# """ -# function GlobalNonlinearKnotPointConstraint( -# g::Function, -# names::AbstractVector{Symbol}, -# traj::NamedTrajectory, -# params::AbstractVector; -# global_names::AbstractVector{Symbol}=Symbol[], -# times::AbstractVector{Int}=1:traj.T, -# equality::Bool=true, -# jacobian_structure::Union{Nothing, SparseMatrixCSC}=nothing, -# hessian_structure::Union{Nothing, SparseMatrixCSC}=nothing, -# ) -# @assert length(params) == length(times) "params must have the same length as times" - -# Z_dim = traj.dim * traj.T + traj.global_dim -# x_comps = vcat([traj.components[name] for name in names]...) -# g_comps = vcat([traj.global_components[name] for name in global_names]...) -# x_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] - -# @assert g(traj[times[1]].data[x_comps], params[1]) isa AbstractVector{Float64} -# g_dim = length(g(traj[times[1]].data[x_comps], params[1])) - -# @views function g!(δ::AbstractVector, Z⃗::AbstractVector) -# for (i, x_slice) ∈ enumerate(x_slices) -# δ[slice(i, g_dim)] = g(Z⃗[x_slice], params[i]) -# end -# end - -# @views function ∂g!(∂gs::Vector{<:AbstractMatrix}, Z⃗::AbstractVector) -# for (i, (x_slice, ∂g)) ∈ enumerate(zip(x_slices, ∂gs)) -# ForwardDiff.jacobian!( -# ∂g[:, x_comps], -# x -> g(x, params[i]), -# Z⃗[x_slice] -# ) -# end -# end - -# @views function μ∂²g!( -# μ∂²gs::Vector{<:AbstractMatrix}, -# Z⃗::AbstractVector, -# μ::AbstractVector -# ) -# for (i, (k, μ∂²g)) ∈ enumerate(zip(times, μ∂²gs)) -# ForwardDiff.hessian!( -# μ∂²g[x_comps, x_comps], -# x -> μ[slice(i, g_dim)]' * g(x, params[i]), -# Z⃗[slice(k, x_comps, traj.dim)] -# ) -# end -# end - -# if isnothing(jacobian_structure) -# jacobian_structure = spzeros(g_dim, Z_dim) -# jacobian_structure[:, x_comps] .= 1.0 -# else -# @assert size(jacobian_structure) == (g_dim, Z_dim) -# end - -# ∂gs = [copy(jacobian_structure) for _ ∈ times] - -# if isnothing(hessian_structure) -# hessian_structure = spzeros(Z_dim, Z_dim) -# hessian_structure[x_comps, x_comps] .= 1.0 -# else -# @assert size(hessian_structure) == (Z_dim, Z_dim) -# end - -# μ∂²gs = [copy(hessian_structure) for _ ∈ times] - -# return new( -# g!, -# ∂g!, -# ∂gs, -# μ∂²g!, -# μ∂²gs, -# times, -# equality, -# g_dim, -# g_dim * length(times) -# ) -# end - -# function GlobalNonlinearKnotPointConstraint( -# g::Function, -# names::AbstractVector{Symbol}, -# traj::NamedTrajectory; -# times::AbstractVector{Int}=1:traj.T, -# kwargs... -# ) -# params = [nothing for _ in times] -# g_param = (x, _) -> g(x) -# return GlobalNonlinearKnotPointConstraint( -# g_param, -# names, -# traj, -# params; -# times=times, -# kwargs... -# ) -# end - -# function GlobalNonlinearKnotPointConstraint(g::Function, name::Symbol, args...; kwargs...) -# return GlobalNonlinearKnotPointConstraint(g, [name], args...; kwargs...) -# end - -# end - -# function get_full_jacobian( -# NLC::GlobalNonlinearKnotPointConstraint, -# traj::NamedTrajectory -# ) -# Z_dim = traj.dim * traj.T + traj.global_dim -# ∂g_full = spzeros(NLC.dim, Z_dim) -# for (i, (k, ∂gₖ)) ∈ enumerate(zip(NLC.times, NLC.∂gs)) -# ∂g_full[slice(i, NLC.g_dim), slice(k, traj.dim)] = ∂gₖ -# end -# return ∂g_full -# end - -# function get_full_hessian( -# NLC::GlobalNonlinearKnotPointConstraint, -# traj::NamedTrajectory -# ) -# Z_dim = traj.dim * traj.T + traj.global_dim -# μ∂²g_full = spzeros(Z_dim, Z_dim) -# for (k, μ∂²gₖ) ∈ zip(NLC.times, NLC.μ∂²gs) -# μ∂²g_full[slice(k, traj.dim), slice(k, traj.dim)] = μ∂²gₖ -# end -# return μ∂²g_full -# end diff --git a/src/constraints/nonlinear_constraints.jl b/src/constraints/nonlinear_constraints.jl index 9f88b7d..ae6bd96 100644 --- a/src/constraints/nonlinear_constraints.jl +++ b/src/constraints/nonlinear_constraints.jl @@ -52,8 +52,9 @@ struct NonlinearKnotPointConstraint <: AbstractNonlinearConstraint x_comps = vcat([traj.components[name] for name in names]...) x_slices = [slice(t, x_comps, traj.dim) for t in times] - @assert g(traj[times[1]].data[x_comps], params[1]) isa AbstractVector{Float64} - g_dim = length(g(traj[times[1]].data[x_comps], params[1])) + # inspect view of knot point data + @assert g(traj.datavec[x_slices[1]], params[1]) isa AbstractVector{Float64} + g_dim = length(g(traj.datavec[x_slices[1]], params[1])) @views function g!(δ::AbstractVector, Z⃗::AbstractVector) for (i, x_slice) ∈ enumerate(x_slices) @@ -149,6 +150,7 @@ function get_full_jacobian( Z_dim = traj.dim * traj.T + traj.global_dim ∂g_full = spzeros(NLC.dim, Z_dim) for (i, (k, ∂gₖ)) ∈ enumerate(zip(NLC.times, NLC.∂gs)) + # Disjoint ∂g_full[slice(i, NLC.g_dim), slice(k, traj.dim)] = ∂gₖ end return ∂g_full @@ -161,6 +163,7 @@ function get_full_hessian( Z_dim = traj.dim * traj.T + traj.global_dim μ∂²g_full = spzeros(Z_dim, Z_dim) for (k, μ∂²gₖ) ∈ zip(NLC.times, NLC.μ∂²gs) + # Disjoint μ∂²g_full[slice(k, traj.dim), slice(k, traj.dim)] = μ∂²gₖ end return μ∂²g_full diff --git a/src/constraints/nonlinear_global_constraints.jl b/src/constraints/nonlinear_global_constraints.jl new file mode 100644 index 0000000..aeb9584 --- /dev/null +++ b/src/constraints/nonlinear_global_constraints.jl @@ -0,0 +1,268 @@ +export NonlinearGlobalConstraint +export NonlinearGlobalKnotPointConstraint + + +# ----------------------------------------------------------------------------- # +# NonlinearGlobalConstraint +# ----------------------------------------------------------------------------- # + +struct NonlinearGlobalConstraint <: AbstractNonlinearConstraint + g!::Function + ∂g!::Function + ∂gs::SparseMatrixCSC + μ∂²g!::Function + μ∂²gs::SparseMatrixCSC + equality::Bool + dim::Int + + function NonlinearGlobalConstraint( + g::Function, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + equality::Bool=true, + ) + global_comps = vcat([traj.global_components[name] for name in global_names]...) + local_comps = global_comps .- traj.dim * traj.T + + g_eval = g(vec(traj)[global_comps]) + @assert g_eval isa AbstractVector{Float64} + g_dim = length(g_eval) + + @views function g!(δ::AbstractVector, Z⃗::AbstractVector) + δ[:] = g(Z⃗[global_comps]) + return nothing + end + + @views function ∂g!(∂g::AbstractMatrix, Z⃗::AbstractVector) + ForwardDiff.jacobian!( + ∂g, + x -> g(x), + Z⃗[global_comps] + ) + end + + # global subspace + jacobian_structure = spzeros(g_dim, traj.global_dim) + jacobian_structure[:, local_comps] .= 1.0 + + @views function μ∂²g!( + μ∂²g::AbstractMatrix, + Z⃗::AbstractVector, + μ::AbstractVector + ) + ForwardDiff.hessian!( + μ∂²g[local_comps, local_comps], + x -> μ'g(x), + Z⃗[global_comps] + ) + end + + # global subspace + hessian_structure = spzeros(traj.global_dim, traj.global_dim) + hessian_structure[local_comps, local_comps] .= 1.0 + + return new( + g!, + ∂g!, + jacobian_structure, + μ∂²g!, + hessian_structure, + equality, + g_dim, + ) + end + + function NonlinearGlobalConstraint( + g::Function, + global_name::Symbol, + traj::NamedTrajectory; + kwargs... + ) + return NonlinearGlobalConstraint(g, [global_name], traj; kwargs...) + end +end + +function get_full_jacobian( + NLC::NonlinearGlobalConstraint, + traj::NamedTrajectory +) + Z_dim = traj.dim * traj.T + traj.global_dim + ∂g_full = spzeros(NLC.dim, Z_dim) + ∂g_full[:, traj.dim * traj.T + 1:Z_dim] = NLC.∂gs + return ∂g_full +end + +function get_full_hessian( + NLC::NonlinearGlobalConstraint, + traj::NamedTrajectory +) + Z_dim = traj.dim * traj.T + traj.global_dim + μ∂²g_full = spzeros(Z_dim, Z_dim) + g_comps = traj.dim * traj.T + 1:Z_dim + μ∂²g_full[g_comps, g_comps] = NLC.μ∂²gs + return μ∂²g_full +end + +# --------------------------------------------------------------------------- # +# NonlinearGlobalKnotPointConstraint +# --------------------------------------------------------------------------- # + +struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint + g!::Function + ∂g!::Function + ∂gs::Vector{SparseMatrixCSC} + μ∂²g!::Function + μ∂²gs::Vector{SparseMatrixCSC} + times::AbstractVector{Int} + equality::Bool + g_dim::Int + dim::Int + + """ + TODO: Docstring + """ + function NonlinearGlobalKnotPointConstraint( + g::Function, + names::AbstractVector{Symbol}, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory, + params::AbstractVector; + times::AbstractVector{Int}=1:traj.T, + equality::Bool=true, + jacobian_structure::Union{Nothing, SparseMatrixCSC}=nothing, + hessian_structure::Union{Nothing, SparseMatrixCSC}=nothing, + ) + @assert length(params) == length(times) "params must have the same length as times" + + x_comps = vcat([traj.components[name] for name in names]...) + global_comps = vcat([traj.global_components[name] for name in global_names]...) + local_comps = global_comps .- traj.dim * traj.T + + # (Rebased) global data is appended to the knot point + xg_comps = vcat([x_comps, traj.dim .+ local_comps]...) + z_dim = traj.dim + traj.global_dim + + # Each slice indexes into Z⃗ + xg_slices = [vcat([slice(t, x_comps, traj.dim), global_comps]...) for t in times] + + Z⃗ = vec(traj) + @assert g(Z⃗[xg_slices[1]], params[1]) isa AbstractVector{Float64} + g_dim = length(g(Z⃗[xg_slices[1]], params[1])) + + @views function g!(δ::AbstractVector, Z⃗::AbstractVector) + for (i, xg_slice) ∈ enumerate(xg_slices) + δ[slice(i, g_dim)] = g(Z⃗[xg_slice], params[i]) + end + end + + @views function ∂g!(∂gs::Vector{<:AbstractMatrix}, Z⃗::AbstractVector) + for (i, (xg_slice, ∂g)) ∈ enumerate(zip(xg_slices, ∂gs)) + # Overlapping + ∂g[:, xg_comps] .+= ForwardDiff.jacobian( + xg -> g(xg, params[i]), + Z⃗[xg_slice] + ) + end + end + + @views function μ∂²g!( + μ∂²gs::Vector{<:AbstractMatrix}, + Z⃗::AbstractVector, + μ::AbstractVector + ) + for (i, (xg_slice, μ∂²g)) ∈ enumerate(zip(xg_slices, μ∂²gs)) + # Overlapping + μ∂²g[xg_comps, xg_comps] .+= ForwardDiff.hessian( + xg -> μ[slice(i, g_dim)]' * g(xg, params[i]), + Z⃗[xg_slice] + ) + end + end + + if isnothing(jacobian_structure) + jacobian_structure = spzeros(g_dim, z_dim) + jacobian_structure[:, xg_comps] .= 1.0 + else + @assert size(jacobian_structure) == (g_dim, z_dim) + end + + ∂gs = [copy(jacobian_structure) for _ ∈ times] + + if isnothing(hessian_structure) + hessian_structure = spzeros(z_dim, z_dim) + hessian_structure[xg_comps, xg_comps] .= 1.0 + else + @assert size(hessian_structure) == (z_dim, z_dim) + end + + μ∂²gs = [copy(hessian_structure) for _ ∈ times] + + return new( + g!, + ∂g!, + ∂gs, + μ∂²g!, + μ∂²gs, + times, + equality, + g_dim, + g_dim * length(times) + ) + end + + function NonlinearGlobalKnotPointConstraint( + g::Function, + names::AbstractVector{Symbol}, + global_names::AbstractVector{Symbol}, + traj::NamedTrajectory; + times::AbstractVector{Int}=1:traj.T, + kwargs... + ) + params = [nothing for _ in times] + g_param = (x, _) -> g(x) + return NonlinearGlobalKnotPointConstraint( + g_param, + names, + global_names, + traj, + params; + times=times, + kwargs... + ) + end + + function NonlinearGlobalKnotPointConstraint(g::Function, name::Symbol, args...; kwargs...) + return NonlinearGlobalKnotPointConstraint(g, [name], args...; kwargs...) + end + +end + +function get_full_jacobian( + NLC::NonlinearGlobalKnotPointConstraint, + traj::NamedTrajectory +) + Z_dim = traj.dim * traj.T + traj.global_dim + global_slice = traj.dim * traj.T .+ (1:traj.global_dim) + ∂g_full = spzeros(NLC.dim, Z_dim) + for (i, (k, ∂gₖ)) ∈ enumerate(zip(NLC.times, NLC.∂gs)) + # Overlapping + zg_slice = vcat([slice(k, traj.dim), global_slice]...) + ∂g_full[slice(i, NLC.g_dim), zg_slice] .+= ∂gₖ + end + return ∂g_full +end + +function get_full_hessian( + NLC::NonlinearGlobalKnotPointConstraint, + traj::NamedTrajectory +) + Z_dim = traj.dim * traj.T + traj.global_dim + global_slice = traj.dim * traj.T .+ (1:traj.global_dim) + μ∂²g_full = spzeros(Z_dim, Z_dim) + for (k, μ∂²gₖ) ∈ zip(NLC.times, NLC.μ∂²gs) + # Overlapping + zg_slice = vcat([slice(k, traj.dim), global_slice]...) + μ∂²g_full[zg_slice, zg_slice] .+= μ∂²gₖ + end + return μ∂²g_full +end diff --git a/src/objectives/global_objectives.jl b/src/objectives/global_objectives.jl index 38ec1a9..24a805b 100644 --- a/src/objectives/global_objectives.jl +++ b/src/objectives/global_objectives.jl @@ -62,7 +62,7 @@ end # Global KnotPointObjective # ----------------------------------------------------------------------------- # -function KnotPointObjective( +function GlobalKnotPointObjective( ℓ::Function, names::AbstractVector{Symbol}, global_names::AbstractVector{Symbol}, @@ -77,13 +77,13 @@ function KnotPointObjective( Z_dim = traj.dim * traj.T + traj.global_dim x_comps = vcat([traj.components[name] for name in names]...) g_comps = vcat([traj.global_components[name] for name in global_names]...) + xg_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] function L(Z⃗::AbstractVector{<:Real}) loss = 0.0 - for (i, x_slice) in enumerate(xg_slices) - x = Z⃗[x_slice] - loss += Qs[i] * ℓ(x, params[i]) + for (i, xg_slice) in enumerate(xg_slices) + loss += Qs[i] * ℓ(Z⃗[xg_slice], params[i]) end return loss end @@ -91,16 +91,16 @@ function KnotPointObjective( @views function ∇L(Z⃗::AbstractVector{<:Real}) ∇ = zeros(Z_dim) for (i, x_slice) in enumerate(xg_slices) - # Add because global params - ∇[x_slice] .+= ForwardDiff.gradient(x -> Qs[i] * ℓ(x, params[i]), Z⃗[x_slice]) + # Global parameters are shared + ∇[x_slice] .+= ForwardDiff.gradient(xg -> Qs[i] * ℓ(xg, params[i]), Z⃗[x_slice]) end return ∇ end function ∂²L_structure() structure = spzeros(Z_dim, Z_dim) - for x_slice in xg_slices - structure[x_slice, x_slice] .= 1.0 + for xg_slice in xg_slices + structure[xg_slice, xg_slice] .= 1.0 end structure_pairs = collect(zip(findnz(structure)[1:2]...)) return structure_pairs @@ -108,36 +108,29 @@ function KnotPointObjective( function ∂²L_structure_mapping() # Build a mapping from (i, j) -> idx - structure_dict = Dict{Tuple{Int, Int}, Int}() + structure_pairs = Dict{Tuple{Int, Int}, Int}() for (idx, pair) in enumerate(∂²L_structure()) - structure_dict[pair] = idx + structure_pairs[pair] = idx end - return structure_dict + + # Build a mapping from slice to structure + structure_map = [ + [structure_pairs[(i, j)] for j in xg_slice for i in xg_slice] + for xg_slice in xg_slices + ] + return structure_map end - @views function ∂²L(Z⃗::AbstractVector{<:Real}) - structure_dict = ∂²L_structure_mapping() - ∂²L_values = zeros(length(structure_dict)) - for (i, x_slice) in enumerate(xg_slices) - ∂²ℓ = ForwardDiff.hessian(x -> Qs[i] * ℓ(x, params[i]), Z⃗[x_slice]) + # precompute + ∂²L_slices = ∂²L_structure_mapping() + ∂²L_structure_length = length(∂²L_structure()) - # TODO: Is there a more efficient way to do this? - for local_i in eachindex(x_slice) - global_i = x_slice[local_i] - for local_j in 1:local_i - global_j = x_slice[local_j] - - # Add to (i,j) - idx = structure_dict[(global_i, global_j)] - ∂²L_values[idx] += ∂²ℓ[local_i, local_j] - - # Add to (j,i) if off-diagonal - if local_i != local_j - idx_sym = structure_dict[(global_j, global_i)] - ∂²L_values[idx_sym] += ∂²ℓ[local_j, local_i] - end - end - end + @views function ∂²L(Z⃗::AbstractVector{<:Real}) + ∂²L_values = zeros(∂²L_structure_length) + for (i, xg_slice) in enumerate(xg_slices) + ∂²ℓ = ForwardDiff.hessian(xg -> Qs[i] * ℓ(xg, params[i]), Z⃗[xg_slice]) + # Global parameters are shared + ∂²L_values[∂²L_slices[i]] .+= ∂²ℓ[:] end return ∂²L_values end @@ -145,7 +138,7 @@ function KnotPointObjective( return Objective(L, ∇L, ∂²L, ∂²L_structure) end -function KnotPointObjective( +function GlobalKnotPointObjective( ℓ::Function, names::AbstractVector{Symbol}, global_names::AbstractVector{Symbol}, @@ -155,10 +148,10 @@ function KnotPointObjective( ) params = [nothing for _ in times] ℓ_param = (x, _) -> ℓ(x) - return KnotPointObjective(ℓ_param, names, global_names, traj, params; times=times, kwargs...) + return GlobalKnotPointObjective(ℓ_param, names, global_names, traj, params; times=times, kwargs...) end - +# From KnotPointObjective function TerminalObjective( ℓ::Function, name::Symbol, @@ -167,9 +160,9 @@ function TerminalObjective( Q::Float64=1.0, kwargs... ) - return KnotPointObjective( + return GlobalKnotPointObjective( ℓ, - name, + [name], global_names, traj; Qs=[Q], From da8a5befc162d3ec73b9c5ef24e10c8b9ff2a42e Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Fri, 9 May 2025 20:03:04 -0500 Subject: [PATCH 08/17] add reset for global nonlinear constraint --- src/constraints/nonlinear_global_constraints.jl | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/constraints/nonlinear_global_constraints.jl b/src/constraints/nonlinear_global_constraints.jl index aeb9584..5b90a66 100644 --- a/src/constraints/nonlinear_global_constraints.jl +++ b/src/constraints/nonlinear_global_constraints.jl @@ -157,6 +157,13 @@ struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint @views function ∂g!(∂gs::Vector{<:AbstractMatrix}, Z⃗::AbstractVector) for (i, (xg_slice, ∂g)) ∈ enumerate(zip(xg_slices, ∂gs)) + # Reset + if i == 1 + ∂g[:, xg_comps] .= 0 + else + ∂g[:, x_comps] .= 0 + end + # Overlapping ∂g[:, xg_comps] .+= ForwardDiff.jacobian( xg -> g(xg, params[i]), @@ -171,6 +178,15 @@ struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint μ::AbstractVector ) for (i, (xg_slice, μ∂²g)) ∈ enumerate(zip(xg_slices, μ∂²gs)) + # Reset + if i == 1 + μ∂²g[xg_comps, xg_comps] .= 0 + else + ∂g[x_comps, x_comps] .= 0 + ∂g[x_comps, xg_comps] .= 0 + ∂g[xg_comps, xg_comps] .= 0 + end + # Overlapping μ∂²g[xg_comps, xg_comps] .+= ForwardDiff.hessian( xg -> μ[slice(i, g_dim)]' * g(xg, params[i]), From b2deb0fd9a36fee350662f54b6e2879c04a58ddc Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Mon, 19 May 2025 11:54:21 -0500 Subject: [PATCH 09/17] min time obj free time only --- src/objectives/minimum_time_objective.jl | 31 ++++++++---------------- 1 file changed, 10 insertions(+), 21 deletions(-) diff --git a/src/objectives/minimum_time_objective.jl b/src/objectives/minimum_time_objective.jl index d1f0221..0d1712c 100644 --- a/src/objectives/minimum_time_objective.jl +++ b/src/objectives/minimum_time_objective.jl @@ -9,30 +9,19 @@ A type of objective that counts the time taken to complete a task. `D` is a sca """ function MinimumTimeObjective( traj::NamedTrajectory; - D::Float64=1.0, - timesteps_all_equal::Bool=false, + D::Float64=1.0 ) @assert traj.timestep isa Symbol "Trajectory timestep must be a symbol (free time)" - - if !timesteps_all_equal - Δt_indices = [index(t, traj.components[traj.timestep][1], traj.dim) for t = 1:traj.T-1] - - L = Z⃗::AbstractVector{<:Real} -> D * sum(Z⃗[Δt_indices]) - - ∇L = (Z⃗::AbstractVector{<:Real}) -> begin - ∇ = zeros(eltype(Z⃗), length(Z⃗)) - ∇[Δt_indices] .= D - return ∇ - end - else - Δt_T_index = index(traj.T, traj.components[traj.timestep][1], traj.dim) - L = Z⃗::AbstractVector{<:Real} -> D * Z⃗[Δt_T_index] - ∇L = (Z⃗::AbstractVector{<:Real}) -> begin - ∇ = zeros(eltype(Z⃗), length(Z⃗)) - ∇[Δt_T_index] = D - return ∇ - end + Δt_index = traj.components[traj.timestep][1] + Δt_indices = [index(t, Δt_index, traj.dim) for t = 1:traj.T-1] + + L = Z⃗::AbstractVector{<:Real} -> D * sum(Z⃗[Δt_indices]) + + ∇L = (Z⃗::AbstractVector{<:Real}) -> begin + ∇ = zeros(eltype(Z⃗), length(Z⃗)) + ∇[Δt_indices] .= D + return ∇ end ∂²L = Z⃗::AbstractVector{<:Real} -> [] From 6f9db3403b294430c89edd37dcf150b11cd16685 Mon Sep 17 00:00:00 2001 From: bbhattacharyya1729 Date: Wed, 28 May 2025 12:26:23 -0500 Subject: [PATCH 10/17] added duration and symmetry consraints --- src/constraints/linear_constraints.jl | 39 +++++++++++++++++++++++++ src/solvers/ipopt_solver/constraints.jl | 21 +++++++++++++ 2 files changed, 60 insertions(+) diff --git a/src/constraints/linear_constraints.jl b/src/constraints/linear_constraints.jl index a08692c..3a94fda 100644 --- a/src/constraints/linear_constraints.jl +++ b/src/constraints/linear_constraints.jl @@ -7,6 +7,8 @@ export TimeStepsAllEqualConstraint export L1SlackConstraint export TotalConstraint export DurationConstraint +export SymmetryConstraint +export SymmetricControlConstraint ### ### EqualityConstraint ### @@ -262,4 +264,41 @@ function DurationConstraint( @assert traj.timestep isa Symbol indices = [index(k, traj.components[traj.timestep][1], traj.dim) for k ∈ 1:traj.T] return TotalConstraint(indices, value ,label) +end + + +struct SymmetryConstraint <: AbstractLinearConstraint + even_index_pairs::Vector{Tuple{Int64,Int64}} + odd_index_pairs::Vector{Tuple{Int64,Int64}} + label::String +end + +function SymmetricControlConstraint( + traj::NamedTrajectory, + name::Symbol, + idx::Vector{Int64}; + even = true, + label = "Symmetry Constraint on $name" +) + even_pairs = Vector{Tuple{Int64,Int64}}() + odd_pairs = Vector{Tuple{Int64,Int64}}() + + component_indicies = [slice(t, traj.components[name], traj.dim)[idx] for t ∈ 1:traj.T] + if(even) + even_pairs = vcat(even_pairs,reduce(vcat,[collect(zip(component_indicies[[idx,traj.T - idx+1]]...)) for idx in 1:traj.T ÷ 2])) + else + odd_pairs = vcat(odd_pairs,reduce(vcat,[collect(zip(component_indicies[[idx,traj.T - idx+1]]...)) for idx in 1:traj.T ÷ 2])) + end + + if traj.timestep isa Symbol + time_indices = [index(k, traj.components[traj.timestep][1], traj.dim) for k ∈ 1:traj.T] + even_pairs = vcat(even_pairs,[(time_indices[idx],time_indices[traj.T + 1 - idx]) for idx ∈ 1:traj.T÷2]) + end + + return SymmetryConstraint( + even_pairs, + odd_pairs, + label + ) + end \ No newline at end of file diff --git a/src/solvers/ipopt_solver/constraints.jl b/src/solvers/ipopt_solver/constraints.jl index 38f2737..7cd77fc 100644 --- a/src/solvers/ipopt_solver/constraints.jl +++ b/src/solvers/ipopt_solver/constraints.jl @@ -94,4 +94,25 @@ MOI.add_constraints( MOI.EqualTo(con.value) ) +end + +function (con::SymmetryConstraint)( + opt::Ipopt.Optimizer, + vars::Vector{MOI.VariableIndex} +) + + for (i1,i2) in con.even_index_pairs + MOI.add_constraints( + opt, + MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, vars[i1]) , MOI.ScalarAffineTerm(-1.0, vars[i2])], 0.0), + MOI.EqualTo(0.0) + ) + end + for (i1,i2) in con.odd_index_pairs + MOI.add_constraints( + opt, + MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, vars[i1]) , MOI.ScalarAffineTerm(1.0, vars[i2])], 0.0), + MOI.EqualTo(0.0) + ) + end end \ No newline at end of file From dcc5643a17afd0f66fe2345715385f1175a50605 Mon Sep 17 00:00:00 2001 From: BBhattacharyya1729 Date: Wed, 28 May 2025 12:28:45 -0500 Subject: [PATCH 11/17] Update regularizers.jl --- src/objectives/regularizers.jl | 33 --------------------------------- 1 file changed, 33 deletions(-) diff --git a/src/objectives/regularizers.jl b/src/objectives/regularizers.jl index 4f29cc5..89da865 100644 --- a/src/objectives/regularizers.jl +++ b/src/objectives/regularizers.jl @@ -1,5 +1,4 @@ export QuadraticRegularizer -export PairwiseQuadraticRegularizer # ----------------------------------------------------------------------------- # # Quadratic Regularizer # ----------------------------------------------------------------------------- # @@ -119,35 +118,3 @@ function QuadraticRegularizer( ) return QuadraticRegularizer(name, traj, R * ones(traj.dims[name]); kwargs...) end - - - -##### -function PairwiseQuadraticRegularizer( - name1::Symbol, - name2::Symbol, - traj::NamedTrajectory, - R::Float64; -) - if traj.timestep isa Symbol - J = KnotPointObjective( - (v,p) -> 1/2 * sum(abs2, (v[1:length(v) ÷ 2]-v[length(v) ÷ 2+1:end] ) * p ) , - [name1,name2], - traj, - [t[1] for t in eachcol(traj[traj.timestep])]; - times = 1:traj.T, - Qs = R * ones(traj.T) - ) - - else - Δt = traj.timestep - J = KnotPointObjective( - v -> 1/2 * sum(abs2, (v[1:length(v) ÷ 2]-v[length(v) ÷ 2+1:end]) * Δt), - [name1,name2], - traj, - [nothing for _ in 1:traj.T]; - times = 1:traj.T, - Qs = R * ones(traj.T) - ) - end -end \ No newline at end of file From d7781ab5301a077712d4f648cc03c14177caf24b Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Thu, 5 Jun 2025 09:22:05 -0500 Subject: [PATCH 12/17] update global data --- src/solvers/ipopt_solver/solver.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/solvers/ipopt_solver/solver.jl b/src/solvers/ipopt_solver/solver.jl index d782ec2..f2e02bb 100644 --- a/src/solvers/ipopt_solver/solver.jl +++ b/src/solvers/ipopt_solver/solver.jl @@ -196,6 +196,7 @@ function update_trajectory!( variables[1:n_vars] ) + update!(prob.trajectory, datavec) # get global variables after trajectory data global_keys = keys(prob.trajectory.global_data) @@ -210,8 +211,7 @@ function update_trajectory!( n_vars += n_global_vars end global_data = (; (global_keys .=> global_values)...) - - update!(prob.trajectory, datavec) + update!(prob.trajectory, global_data) # TODO: this results in a bug of shifted components when components are added after creating a trajectory, this affects constraints which store original componentes in probs # prob.trajectory = NamedTrajectory(datavec, global_data, prob.trajectory) From 5911dc1fd0e1f86aad40ca6776ade4e77e19eed5 Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Sat, 7 Jun 2025 18:02:14 +0200 Subject: [PATCH 13/17] refactor NT concrete --- src/constraints/nonlinear_constraints.jl | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/constraints/nonlinear_constraints.jl b/src/constraints/nonlinear_constraints.jl index 4e7749c..940cc4f 100644 --- a/src/constraints/nonlinear_constraints.jl +++ b/src/constraints/nonlinear_constraints.jl @@ -1,10 +1,5 @@ export NonlinearKnotPointConstraint - -# ----------------------------------------------------------------------------- # -# NonlinearKnotPointConstraint -# ----------------------------------------------------------------------------- # - # ----------------------------------------------------------------------------- # # NonlinearKnotPointConstraint # ----------------------------------------------------------------------------- # From 1ffee19142241af65de7f598e4eae2f590cd170e Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Fri, 13 Jun 2025 13:38:53 -0500 Subject: [PATCH 14/17] update for new NT --- src/constraints/linear_constraints.jl | 67 ++++++++++--------- .../nonlinear_global_constraints.jl | 26 +++---- .../time_dependent_bilinear_integrator.jl | 3 +- src/objectives/global_objectives.jl | 4 +- test/test_utils.jl | 5 ++ 5 files changed, 55 insertions(+), 50 deletions(-) diff --git a/src/constraints/linear_constraints.jl b/src/constraints/linear_constraints.jl index 242a218..8b1cef8 100644 --- a/src/constraints/linear_constraints.jl +++ b/src/constraints/linear_constraints.jl @@ -216,36 +216,37 @@ function TimeStepsAllEqualConstraint( return AllEqualConstraint(indices, bar_index, label) end -struct L1SlackConstraint <: AbstractLinearConstraint - x_indices::Vector{Int} - s1_indices::Vector{Int} - s2_indices::Vector{Int} - label::String -end - -function L1SlackConstraint( - name::Symbol, - traj::NamedTrajectory; - indices=1:traj.dims[name], - ts=(name ∈ keys(traj.initial) ? 2 : 1):(name ∈ keys(traj.final) ? traj.T-1 : traj.T), - label="L1 slack constraint on $name[$(indices)]" -) - @assert all(i ∈ 1:traj.dims[name] for i ∈ indices) - - s1_name = Symbol("s1_$name") - s2_name = Symbol("s2_$name") - - add_component!(traj, s1_name, rand(length(indices), traj.T)) - add_component!(traj, s2_name, rand(length(indices), traj.T)) - - x_indices = stack(slice(t, traj.components[name][indices], traj.dim) for t ∈ ts) - s1_indices = stack(slice(t, traj.components[s1_name], traj.dim) for t ∈ ts) - s2_indices = stack(slice(t, traj.components[s2_name], traj.dim) for t ∈ ts) - - return L1SlackConstraint( - x_indices, - s1_indices, - s2_indices, - label - ) -end +# TODO: Doesn't work with parametric trajectory +# struct L1SlackConstraint <: AbstractLinearConstraint +# x_indices::Vector{Int} +# s1_indices::Vector{Int} +# s2_indices::Vector{Int} +# label::String +# end + +# function L1SlackConstraint( +# name::Symbol, +# traj::NamedTrajectory; +# indices=1:traj.dims[name], +# ts=(name ∈ keys(traj.initial) ? 2 : 1):(name ∈ keys(traj.final) ? traj.T-1 : traj.T), +# label="L1 slack constraint on $name[$(indices)]" +# ) +# @assert all(i ∈ 1:traj.dims[name] for i ∈ indices) + +# s1_name = Symbol("s1_$name") +# s2_name = Symbol("s2_$name") + +# add_component!(traj, s1_name, rand(length(indices), traj.T)) +# add_component!(traj, s2_name, rand(length(indices), traj.T)) + +# x_indices = stack(slice(t, traj.components[name][indices], traj.dim) for t ∈ ts) +# s1_indices = stack(slice(t, traj.components[s1_name], traj.dim) for t ∈ ts) +# s2_indices = stack(slice(t, traj.components[s2_name], traj.dim) for t ∈ ts) + +# return L1SlackConstraint( +# x_indices, +# s1_indices, +# s2_indices, +# label +# ) +# end diff --git a/src/constraints/nonlinear_global_constraints.jl b/src/constraints/nonlinear_global_constraints.jl index 5b90a66..5aeb30a 100644 --- a/src/constraints/nonlinear_global_constraints.jl +++ b/src/constraints/nonlinear_global_constraints.jl @@ -22,14 +22,14 @@ struct NonlinearGlobalConstraint <: AbstractNonlinearConstraint equality::Bool=true, ) global_comps = vcat([traj.global_components[name] for name in global_names]...) - local_comps = global_comps .- traj.dim * traj.T + offset_global_comps = global_comps .+ traj.dim * traj.T - g_eval = g(vec(traj)[global_comps]) + g_eval = g(vec(traj)[offset_global_comps]) @assert g_eval isa AbstractVector{Float64} g_dim = length(g_eval) @views function g!(δ::AbstractVector, Z⃗::AbstractVector) - δ[:] = g(Z⃗[global_comps]) + δ[:] = g(Z⃗[offset_global_comps]) return nothing end @@ -37,13 +37,13 @@ struct NonlinearGlobalConstraint <: AbstractNonlinearConstraint ForwardDiff.jacobian!( ∂g, x -> g(x), - Z⃗[global_comps] + Z⃗[offset_global_comps] ) end # global subspace jacobian_structure = spzeros(g_dim, traj.global_dim) - jacobian_structure[:, local_comps] .= 1.0 + jacobian_structure[:, global_comps] .= 1.0 @views function μ∂²g!( μ∂²g::AbstractMatrix, @@ -51,15 +51,15 @@ struct NonlinearGlobalConstraint <: AbstractNonlinearConstraint μ::AbstractVector ) ForwardDiff.hessian!( - μ∂²g[local_comps, local_comps], + μ∂²g[global_comps, global_comps], x -> μ'g(x), - Z⃗[global_comps] + Z⃗[offset_global_comps] ) end # global subspace hessian_structure = spzeros(traj.global_dim, traj.global_dim) - hessian_structure[local_comps, local_comps] .= 1.0 + hessian_structure[global_comps, global_comps] .= 1.0 return new( g!, @@ -136,14 +136,14 @@ struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint x_comps = vcat([traj.components[name] for name in names]...) global_comps = vcat([traj.global_components[name] for name in global_names]...) - local_comps = global_comps .- traj.dim * traj.T + offset_global_comps = global_comps .+ traj.dim * traj.T - # (Rebased) global data is appended to the knot point - xg_comps = vcat([x_comps, traj.dim .+ local_comps]...) + # append global data to the knot point + xg_comps = vcat([x_comps, global_comps .+ traj.dim]...) z_dim = traj.dim + traj.global_dim - # Each slice indexes into Z⃗ - xg_slices = [vcat([slice(t, x_comps, traj.dim), global_comps]...) for t in times] + # append global data to the trajectory (each slice indexes into Z⃗) + xg_slices = [vcat([slice(t, x_comps, traj.dim), offset_global_comps]...) for t in times] Z⃗ = vec(traj) @assert g(Z⃗[xg_slices[1]], params[1]) isa AbstractVector{Float64} diff --git a/src/integrators/time_dependent_bilinear_integrator.jl b/src/integrators/time_dependent_bilinear_integrator.jl index f225de8..d394644 100644 --- a/src/integrators/time_dependent_bilinear_integrator.jl +++ b/src/integrators/time_dependent_bilinear_integrator.jl @@ -151,8 +151,7 @@ end @testitem "testing zoh TimeDependentBilinearIntegrator" begin include("../../test/test_utils.jl") - G, traj = bilinear_dynamics_and_trajectory() - add_component!(traj, :t, get_times(traj)) + G, traj = bilinear_dynamics_and_trajectory(time=true) # zero order hold B = TimeDependentBilinearIntegrator((a, t) -> G(a), traj, :x, :u, :t) diff --git a/src/objectives/global_objectives.jl b/src/objectives/global_objectives.jl index 24a805b..2dfecce 100644 --- a/src/objectives/global_objectives.jl +++ b/src/objectives/global_objectives.jl @@ -29,7 +29,7 @@ function GlobalObjective( Q::Float64=1.0 ) Z_dim = traj.dim * traj.T + traj.global_dim - g_comps = vcat([traj.global_components[name] for name in global_names]...) + g_comps = vcat([traj.dim * traj.T .+ traj.global_components[name] for name in global_names]...) L(Z⃗::AbstractVector{<:Real}) = Q * ℓ(Z⃗[g_comps]) @@ -76,7 +76,7 @@ function GlobalKnotPointObjective( Z_dim = traj.dim * traj.T + traj.global_dim x_comps = vcat([traj.components[name] for name in names]...) - g_comps = vcat([traj.global_components[name] for name in global_names]...) + g_comps = vcat([traj.dim * traj.T .+ traj.global_components[name] for name in global_names]...) xg_slices = [vcat([slice(t, x_comps, traj.dim), g_comps]...) for t in times] diff --git a/test/test_utils.jl b/test/test_utils.jl index 800ebc1..ae9d6ed 100644 --- a/test/test_utils.jl +++ b/test/test_utils.jl @@ -125,6 +125,7 @@ function bilinear_dynamics_and_trajectory(; Δt = 0.1, u_bound = 0.1, ω = 0.1, + time=false ) Gx = sparse(Float64[ 0 0 0 1; @@ -182,5 +183,9 @@ function bilinear_dynamics_and_trajectory(; ) ) + if time + traj = add_component(traj, :t, get_times(traj)) + end + return G, traj end From c0213eadc4a20c6b949b3afafa4cdc96e543ca8b Mon Sep 17 00:00:00 2001 From: BBhattacharyya1729 Date: Mon, 16 Jun 2025 16:06:12 -0500 Subject: [PATCH 15/17] added tests for symmetry and duration constraints --- src/constraints/linear_constraints.jl | 66 ++++++++++++++++++++++++++- 1 file changed, 65 insertions(+), 1 deletion(-) diff --git a/src/constraints/linear_constraints.jl b/src/constraints/linear_constraints.jl index ccb9965..7230ba9 100644 --- a/src/constraints/linear_constraints.jl +++ b/src/constraints/linear_constraints.jl @@ -305,4 +305,68 @@ function SymmetricControlConstraint( label ) -end \ No newline at end of file +end + +@testitem "testing symmetry constraint" begin + + include("../../../test/test_utils.jl") + + G, traj = bilinear_dynamics_and_trajectory() + + integrators = [ + BilinearIntegrator(G, traj, :x, :u), + DerivativeIntegrator(traj, :u, :du), + DerivativeIntegrator(traj, :du, :ddu) + ] + + J = TerminalObjective(x -> norm(x - traj.goal.x)^2, :x, traj) + J += QuadraticRegularizer(:u, traj, 1.0) + J += QuadraticRegularizer(:du, traj, 1.0) + J += MinimumTimeObjective(traj) + + g_u_norm = NonlinearKnotPointConstraint(u -> [norm(u) - 1.0], :u, traj; times=2:traj.T-1, equality=false) + + prob = DirectTrajOptProblem(traj, J, integrators;) + + + sym_constraint = SymmetricControlConstraint( + prob.trajectory, + :u, + [1]; + even = true + ); + push!(prob.constraints,sym_constraint); + + solve!(prob; max_iter=100) +end + +@testitem "testing duration constraint" begin + + include("../../../test/test_utils.jl") + + G, traj = bilinear_dynamics_and_trajectory() + + integrators = [ + BilinearIntegrator(G, traj, :x, :u), + DerivativeIntegrator(traj, :u, :du), + DerivativeIntegrator(traj, :du, :ddu) + ] + + J = TerminalObjective(x -> norm(x - traj.goal.x)^2, :x, traj) + J += QuadraticRegularizer(:u, traj, 1.0) + J += QuadraticRegularizer(:du, traj, 1.0) + J += MinimumTimeObjective(traj) + + g_u_norm = NonlinearKnotPointConstraint(u -> [norm(u) - 1.0], :u, traj; times=2:traj.T-1, equality=false) + + prob = DirectTrajOptProblem(traj, J, integrators;) + + + dur_constraint = DurationConstraint( + prob.trajectory, + 10.0; + ) + push!(prob.constraints,dur_constraint); + + solve!(prob; max_iter=100) +end From 9b6f1e75bfe4ab6f95223f448ba14f42517f0c5f Mon Sep 17 00:00:00 2001 From: BBhattacharyya1729 Date: Mon, 16 Jun 2025 16:06:57 -0500 Subject: [PATCH 16/17] Fixed duration/symmetry constraint tests --- src/constraints/linear_constraints.jl | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/constraints/linear_constraints.jl b/src/constraints/linear_constraints.jl index 7230ba9..4ed2f26 100644 --- a/src/constraints/linear_constraints.jl +++ b/src/constraints/linear_constraints.jl @@ -324,11 +324,8 @@ end J += QuadraticRegularizer(:du, traj, 1.0) J += MinimumTimeObjective(traj) - g_u_norm = NonlinearKnotPointConstraint(u -> [norm(u) - 1.0], :u, traj; times=2:traj.T-1, equality=false) - prob = DirectTrajOptProblem(traj, J, integrators;) - sym_constraint = SymmetricControlConstraint( prob.trajectory, :u, @@ -357,11 +354,8 @@ end J += QuadraticRegularizer(:du, traj, 1.0) J += MinimumTimeObjective(traj) - g_u_norm = NonlinearKnotPointConstraint(u -> [norm(u) - 1.0], :u, traj; times=2:traj.T-1, equality=false) - prob = DirectTrajOptProblem(traj, J, integrators;) - dur_constraint = DurationConstraint( prob.trajectory, 10.0; From 1b07173786427d475d2ec9c7129229d4b5adece6 Mon Sep 17 00:00:00 2001 From: Andy Goldschmidt Date: Mon, 16 Jun 2025 20:42:17 -0500 Subject: [PATCH 17/17] update for parametric NT, test global obj and constraints --- Project.toml | 4 +- src/constraints/linear_constraints.jl | 49 +++--- src/constraints/nonlinear_constraints.jl | 29 ++-- .../nonlinear_global_constraints.jl | 161 ++++++++++++++---- .../time_dependent_bilinear_integrator.jl | 4 +- src/objectives/global_objectives.jl | 73 +++++++- src/solvers/ipopt_solver/constraints.jl | 50 +++--- src/solvers/ipopt_solver/solver.jl | 60 ++----- test/test_utils.jl | 11 +- 9 files changed, 292 insertions(+), 149 deletions(-) diff --git a/Project.toml b/Project.toml index c9a59d0..befd45b 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "DirectTrajOpt" uuid = "c823fa1f-8872-4af5-b810-2b9b72bbbf56" authors = ["Aaron Trowbridge and contributors"] -version = "0.2.7" +version = "0.3.0" [deps] ExponentialAction = "e24c0720-ea99-47e8-929e-571b494574d3" @@ -31,7 +31,7 @@ JLD2 = "0.5" Libdl = "1.10, 1.11" LinearAlgebra = "1.10, 1.11" MathOptInterface = "1.40" -NamedTrajectories = "0.3" +NamedTrajectories = "0.4" OrdinaryDiffEq = "6.97" Reexport = "1.2" SciMLBase = "2.86" diff --git a/src/constraints/linear_constraints.jl b/src/constraints/linear_constraints.jl index 8b1cef8..8e3661c 100644 --- a/src/constraints/linear_constraints.jl +++ b/src/constraints/linear_constraints.jl @@ -64,8 +64,6 @@ function EqualityConstraint( return EqualityConstraint(name, ts, fill(val, traj.dims[name]), traj; label=label) end - - """ GlobalEqualityConstraint( name::Symbol, @@ -82,20 +80,36 @@ function GlobalEqualityConstraint( traj::NamedTrajectory; label="equality constraint on global variable $name" ) - @assert name ∈ keys(traj.global_data) + @assert name ∈ traj.global_names @assert length(val) == traj.global_dims[name] - indices = traj.global_components[name] + indices = traj.dim * traj.T .+ traj.global_components[name] return EqualityConstraint(indices, val, label) end +struct AllEqualConstraint <: AbstractLinearConstraint + indices::Vector{Int} + bar_index::Int + label::String +end +function TimeStepsAllEqualConstraint( + traj::NamedTrajectory; + label="timesteps all equal constraint" +) + @assert traj.timestep isa Symbol + indices = [index(k, traj.components[traj.timestep][1], traj.dim) for k ∈ 1:traj.T-1] + bar_index = index(traj.T, traj.components[traj.timestep][1], traj.dim) + return AllEqualConstraint(indices, bar_index, label) +end ### ### BoundsConstraint ### +# TODO: Refactor using `get_bounds_from_dims` from NamedTrajectories.jl + struct BoundsConstraint <: AbstractLinearConstraint indices::Vector{Int} bounds::Vector{Tuple{Float64, Float64}} @@ -158,11 +172,11 @@ function GlobalBoundsConstraint( traj::NamedTrajectory; label="bounds constraint on global variable $name" ) - @assert name ∈ keys(traj.global_data) + @assert name ∈ traj.global_names @assert length(bounds[1]) == length(bounds[2]) == traj.global_dims[name] @assert all(bounds[1] .<= bounds[2]) - indices = traj.global_components[name] + indices = traj.dim * traj.T .+ traj.global_components[name] bounds = collect(zip(bounds...)) @@ -175,7 +189,7 @@ function GlobalBoundsConstraint( traj::NamedTrajectory; label="bounds constraint on global variable $name" ) - @assert name ∈ keys(traj.global_data) + @assert name ∈ traj.global_names @assert length(bound) == traj.global_dims[name] @assert all(bound .>= 0) "bound must be non-negative when only one bound is provided" @@ -190,7 +204,7 @@ function GlobalBoundsConstraint( traj::NamedTrajectory; label="bounds constraint on global variable $name" ) - @assert name ∈ keys(traj.global_data) + @assert name ∈ traj.global_names @assert bound >= 0 "bound must be non-negative when only one bound is provided" bounds = (-fill(bound, traj.global_dims[name]), fill(bound, traj.global_dims[name])) @@ -199,22 +213,9 @@ function GlobalBoundsConstraint( end - -struct AllEqualConstraint <: AbstractLinearConstraint - indices::Vector{Int} - bar_index::Int - label::String -end - -function TimeStepsAllEqualConstraint( - traj::NamedTrajectory; - label="timesteps all equal constraint" -) - @assert traj.timestep isa Symbol - indices = [index(k, traj.components[traj.timestep][1], traj.dim) for k ∈ 1:traj.T-1] - bar_index = index(traj.T, traj.components[traj.timestep][1], traj.dim) - return AllEqualConstraint(indices, bar_index, label) -end +### +### L1SlackConstraint +### # TODO: Doesn't work with parametric trajectory # struct L1SlackConstraint <: AbstractLinearConstraint diff --git a/src/constraints/nonlinear_constraints.jl b/src/constraints/nonlinear_constraints.jl index 940cc4f..2de5550 100644 --- a/src/constraints/nonlinear_constraints.jl +++ b/src/constraints/nonlinear_constraints.jl @@ -26,7 +26,7 @@ struct NonlinearKnotPointConstraint{F1, F2, F3} <: AbstractNonlinearConstraint Create a NonlinearKnotPointConstraint object that represents a nonlinear constraint on a trajectory. # Arguments - - `g::Function`: Function that defines the constraint. If `equality=false`, the constraint is `g(x) ≤ 0`. + - `g::Function`: Function over knot point variable(s) that defines the constraint, `g(x)`. - `name::Symbol`: Name of the variable to be constrained. - `traj::NamedTrajectory`: The trajectory on which the constraint is defined. @@ -52,8 +52,9 @@ struct NonlinearKnotPointConstraint{F1, F2, F3} <: AbstractNonlinearConstraint x_slices = [slice(t, x_comps, traj.dim) for t in times] # inspect view of knot point data - @assert g(traj.datavec[x_slices[1]], params[1]) isa AbstractVector{Float64} - g_dim = length(g(traj.datavec[x_slices[1]], params[1])) + Z⃗ = vec(traj) + @assert g(Z⃗[x_slices[1]], params[1]) isa AbstractVector{Float64} + g_dim = length(g(Z⃗[x_slices[1]], params[1])) @views function g!(δ::AbstractVector, Z⃗::AbstractVector) for (i, x_slice) ∈ enumerate(x_slices) @@ -172,7 +173,7 @@ end # ============================================================================= # -@testitem "testing NonlinearConstraint" begin +@testitem "testing NonlinearKnotPointConstraint" begin using TrajectoryIndexingUtils @@ -182,31 +183,33 @@ end g(a) = [norm(a) - 1.0] + g_dim = 1 times = 1:traj.T NLC = NonlinearKnotPointConstraint(g, :u, traj; times=times, equality=false) + U_SLICE(k) = slice(k, traj.components[:u], traj.dim) - ĝ(Z⃗) = vcat([g(Z⃗[slice(k, traj.components[:u], traj.dim)]) for k ∈ times]...) + ĝ(Z⃗) = vcat([g(Z⃗[U_SLICE(k)]) for k ∈ times]...) - δ = zeros(length(times)) + δ = zeros(g_dim * traj.T) - NLC.g!(δ, traj.datavec) + NLC.g!(δ, vec(traj)) - @test δ ≈ ĝ(traj.datavec) + @test δ ≈ ĝ(vec(traj)) - NLC.∂g!(NLC.∂gs, traj.datavec) + NLC.∂g!(NLC.∂gs, vec(traj)) ∂g_full = Constraints.get_full_jacobian(NLC, traj) - ∂g_autodiff = ForwardDiff.jacobian(ĝ, traj.datavec) + ∂g_autodiff = ForwardDiff.jacobian(ĝ, vec(traj)) @test ∂g_full ≈ ∂g_autodiff - μ = randn(length(times)) + μ = randn(g_dim * traj.T) - NLC.μ∂²g!(NLC.μ∂²gs, traj.datavec, μ) + NLC.μ∂²g!(NLC.μ∂²gs, vec(traj), μ) - hessian_autodiff = ForwardDiff.hessian(Z -> μ'ĝ(Z), traj.datavec) + hessian_autodiff = ForwardDiff.hessian(Z -> μ'ĝ(Z), vec(traj)) μ∂²g_full = Constraints.get_full_hessian(NLC, traj) diff --git a/src/constraints/nonlinear_global_constraints.jl b/src/constraints/nonlinear_global_constraints.jl index 5aeb30a..cdd79a8 100644 --- a/src/constraints/nonlinear_global_constraints.jl +++ b/src/constraints/nonlinear_global_constraints.jl @@ -15,6 +15,17 @@ struct NonlinearGlobalConstraint <: AbstractNonlinearConstraint equality::Bool dim::Int + """ + Create a NonlinearGlobalConstraint object with global components. + + # Arguments + - `g::Function`: Function over knot point and global variable(s) that defines the constraint, `g(vcat(x, globals))`. + - `global_names::AbstractVector{Symbol}`: Name(s) of the variable(s) to be constrained. + - `traj::NamedTrajectory`: The trajectory on which the constraint is defined. + + # Keyword Arguments + - `equality::Bool=true`: If `true`, the constraint is `g(x) = 0`. Otherwise, the constraint is `g(x) ≤ 0`. + """ function NonlinearGlobalConstraint( g::Function, global_names::AbstractVector{Symbol}, @@ -22,9 +33,9 @@ struct NonlinearGlobalConstraint <: AbstractNonlinearConstraint equality::Bool=true, ) global_comps = vcat([traj.global_components[name] for name in global_names]...) - offset_global_comps = global_comps .+ traj.dim * traj.T + offset_global_comps = traj.dim * traj.T .+ global_comps - g_eval = g(vec(traj)[offset_global_comps]) + g_eval = g(traj.global_data[global_comps]) @assert g_eval isa AbstractVector{Float64} g_dim = length(g_eval) @@ -88,7 +99,7 @@ function get_full_jacobian( ) Z_dim = traj.dim * traj.T + traj.global_dim ∂g_full = spzeros(NLC.dim, Z_dim) - ∂g_full[:, traj.dim * traj.T + 1:Z_dim] = NLC.∂gs + ∂g_full[1:NLC.dim, traj.dim * traj.T + 1:Z_dim] = NLC.∂gs return ∂g_full end @@ -119,7 +130,9 @@ struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint dim::Int """ - TODO: Docstring + Create a NonlinearKnotPointConstraint object with global components. + + TODO: Consolidate with NonlinearKnotPointConstraint? """ function NonlinearGlobalKnotPointConstraint( g::Function, @@ -134,16 +147,17 @@ struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint ) @assert length(params) == length(times) "params must have the same length as times" + # collect the components and global components x_comps = vcat([traj.components[name] for name in names]...) global_comps = vcat([traj.global_components[name] for name in global_names]...) - offset_global_comps = global_comps .+ traj.dim * traj.T + offset_global_comps = traj.dim * traj.T .+ global_comps - # append global data to the knot point - xg_comps = vcat([x_comps, global_comps .+ traj.dim]...) - z_dim = traj.dim + traj.global_dim + # append global data to the trajectory (each slice indexes into Z⃗, creating xg) + xg_slices = [vcat(slice(t, x_comps, traj.dim), offset_global_comps) for t in times] - # append global data to the trajectory (each slice indexes into Z⃗) - xg_slices = [vcat([slice(t, x_comps, traj.dim), offset_global_comps]...) for t in times] + # append global data comps to each knot point (indexes into knot points) + xg_comps = vcat(x_comps, traj.dim .+ global_comps) + z_dim = traj.dim + traj.global_dim Z⃗ = vec(traj) @assert g(Z⃗[xg_slices[1]], params[1]) isa AbstractVector{Float64} @@ -157,16 +171,10 @@ struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint @views function ∂g!(∂gs::Vector{<:AbstractMatrix}, Z⃗::AbstractVector) for (i, (xg_slice, ∂g)) ∈ enumerate(zip(xg_slices, ∂gs)) - # Reset - if i == 1 - ∂g[:, xg_comps] .= 0 - else - ∂g[:, x_comps] .= 0 - end - - # Overlapping - ∂g[:, xg_comps] .+= ForwardDiff.jacobian( - xg -> g(xg, params[i]), + # Disjoint + ForwardDiff.jacobian!( + ∂g[:, xg_comps], + x -> g(x, params[i]), Z⃗[xg_slice] ) end @@ -178,17 +186,9 @@ struct NonlinearGlobalKnotPointConstraint <: AbstractNonlinearConstraint μ::AbstractVector ) for (i, (xg_slice, μ∂²g)) ∈ enumerate(zip(xg_slices, μ∂²gs)) - # Reset - if i == 1 - μ∂²g[xg_comps, xg_comps] .= 0 - else - ∂g[x_comps, x_comps] .= 0 - ∂g[x_comps, xg_comps] .= 0 - ∂g[xg_comps, xg_comps] .= 0 - end - - # Overlapping - μ∂²g[xg_comps, xg_comps] .+= ForwardDiff.hessian( + # Disjoint + ForwardDiff.hessian!( + μ∂²g[xg_comps, xg_comps], xg -> μ[slice(i, g_dim)]' * g(xg, params[i]), Z⃗[xg_slice] ) @@ -261,9 +261,9 @@ function get_full_jacobian( global_slice = traj.dim * traj.T .+ (1:traj.global_dim) ∂g_full = spzeros(NLC.dim, Z_dim) for (i, (k, ∂gₖ)) ∈ enumerate(zip(NLC.times, NLC.∂gs)) - # Overlapping - zg_slice = vcat([slice(k, traj.dim), global_slice]...) - ∂g_full[slice(i, NLC.g_dim), zg_slice] .+= ∂gₖ + # Disjoint + zg_slice = vcat(slice(k, traj.dim), global_slice) + ∂g_full[slice(i, NLC.g_dim), zg_slice] .= ∂gₖ end return ∂g_full end @@ -277,8 +277,99 @@ function get_full_hessian( μ∂²g_full = spzeros(Z_dim, Z_dim) for (k, μ∂²gₖ) ∈ zip(NLC.times, NLC.μ∂²gs) # Overlapping - zg_slice = vcat([slice(k, traj.dim), global_slice]...) + zg_slice = vcat(slice(k, traj.dim), global_slice) μ∂²g_full[zg_slice, zg_slice] .+= μ∂²gₖ end return μ∂²g_full end + +# ============================================================================ # + +@testitem "testing NonlinearGlobalConstraint" begin + include("../../test/test_utils.jl") + + _, traj = bilinear_dynamics_and_trajectory(add_global=true) + + g_fn(g) = [norm(g) - 1.0] + + g_dim = 1 + + NLC = NonlinearGlobalConstraint(g_fn, :g, traj; equality=false) + G_DIM = traj.dim * traj.T .+ traj.global_components[:g] + + ĝ(Z⃗) = g_fn(Z⃗[G_DIM]) + + δ = zeros(g_dim) + + NLC.g!(δ, vec(traj)) + + @test δ ≈ ĝ(vec(traj)) + + NLC.∂g!(NLC.∂gs, vec(traj)) + + ∂g_full = Constraints.get_full_jacobian(NLC, traj) + + ∂g_autodiff = ForwardDiff.jacobian(ĝ, vec(traj)) + + @test ∂g_full ≈ ∂g_autodiff + + μ = randn(g_dim) + + NLC.μ∂²g!(NLC.μ∂²gs, vec(traj), μ) + + hessian_autodiff = ForwardDiff.hessian(Z -> μ'ĝ(Z), vec(traj)) + + μ∂²g_full = Constraints.get_full_hessian(NLC, traj) + + @test μ∂²g_full ≈ hessian_autodiff +end + +@testitem "testing NonlinearGlobalKnotPointConstraint" begin + using TrajectoryIndexingUtils + + include("../../test/test_utils.jl") + + _, traj = bilinear_dynamics_and_trajectory(add_global=true) + + function g_fn(ug) + u, g = ug[1:traj.dims[:u]], ug[traj.dims[:u] + 1:end] + return [norm(u) - 1.0; norm(u) * norm(g) - 1.0] + end + + g_dim = 2 + times = 1:traj.T + + NLC = NonlinearGlobalKnotPointConstraint(g_fn, [:u], [:g], traj; times=times, equality=false) + U_DIM(k) = slice(k, traj.components[:u], traj.dim) + G_DIM = traj.dim * traj.T .+ traj.global_components[:g] + + ĝ(Z⃗) = vcat([g_fn(Z⃗[vcat(U_DIM(k), G_DIM)]) for k ∈ times]...) + + δ = zeros(g_dim * traj.T) + + NLC.g!(δ, vec(traj)) + + @test δ ≈ ĝ(vec(traj)) + + NLC.∂g!(NLC.∂gs, vec(traj)) + + ∂g_full = Constraints.get_full_jacobian(NLC, traj) + + ∂g_autodiff = ForwardDiff.jacobian(ĝ, vec(traj)) + + display(∂g_full) + println() + display(∂g_autodiff) + + @test ∂g_full ≈ ∂g_autodiff + + μ = randn(g_dim * traj.T) + + NLC.μ∂²g!(NLC.μ∂²gs, vec(traj), μ) + + hessian_autodiff = ForwardDiff.hessian(Z -> μ'ĝ(Z), vec(traj)) + + μ∂²g_full = Constraints.get_full_hessian(NLC, traj) + + @test μ∂²g_full ≈ hessian_autodiff +end \ No newline at end of file diff --git a/src/integrators/time_dependent_bilinear_integrator.jl b/src/integrators/time_dependent_bilinear_integrator.jl index d394644..cad7ad2 100644 --- a/src/integrators/time_dependent_bilinear_integrator.jl +++ b/src/integrators/time_dependent_bilinear_integrator.jl @@ -148,10 +148,12 @@ function hessian_structure(B::TimeDependentBilinearIntegrator) return μ∂²f end +# ============================================================================ # + @testitem "testing zoh TimeDependentBilinearIntegrator" begin include("../../test/test_utils.jl") - G, traj = bilinear_dynamics_and_trajectory(time=true) + G, traj = bilinear_dynamics_and_trajectory(add_time=true) # zero order hold B = TimeDependentBilinearIntegrator((a, t) -> G(a), traj, :x, :u, :t) diff --git a/src/objectives/global_objectives.jl b/src/objectives/global_objectives.jl index 2dfecce..8559a88 100644 --- a/src/objectives/global_objectives.jl +++ b/src/objectives/global_objectives.jl @@ -54,7 +54,7 @@ function GlobalObjective( return Objective(L, ∇L, ∂²L, ∂²L_structure) end -function ℓ(ℓ::Function, global_name::Symbol, traj::NamedTrajectory; kwargs...) +function GlobalObjective(ℓ::Function, global_name::Symbol, traj::NamedTrajectory; kwargs...) return GlobalObjective(ℓ, [global_name], traj; kwargs...) end @@ -169,4 +169,75 @@ function TerminalObjective( times=[traj.T], kwargs... ) +end + +# ============================================================================ # + +@testitem "testing GlobalObjective" begin + + using TrajectoryIndexingUtils + + include("../../test/test_utils.jl") + + _, traj = bilinear_dynamics_and_trajectory(add_global=true) + + L(g) = norm(g) + + Q = 2.0 + + OBJ = GlobalObjective(L, :g, traj, Q=Q) + G_COMP = traj.dim * traj.T .+ traj.global_components[:g] + L̂(Z⃗) = Q * L(Z⃗[G_COMP]) + + @test OBJ.L(vec(traj)) ≈ L̂(vec(traj)) + + ∂L_autodiff = ForwardDiff.gradient(L̂, vec(traj)) + @test OBJ.∇L(vec(traj)) ≈ ∂L_autodiff + + ∂²L_autodiff = ForwardDiff.hessian(L̂, vec(traj)) + + ∂²L_full = zeros(size(∂²L_autodiff)) + for (index, entry) in zip(OBJ.∂²L_structure(), OBJ.∂²L(vec(traj))) + i, j = index + ∂²L_full[i, j] = entry + end + + @test ∂²L_full ≈ ∂²L_autodiff +end + +@testitem "testing GlobalKnotPointObjective" begin + + using TrajectoryIndexingUtils + + include("../../test/test_utils.jl") + + _, traj = bilinear_dynamics_and_trajectory(add_global=true) + + function L(ug) + u, g = ug[1:traj.dims[:u]], ug[traj.dims[:u] .+ 1:end] + return norm(u) + norm(g) + end + + Qs = [1.0, 2.0] + times = [1, traj.T] + + OBJ = GlobalKnotPointObjective(L, [:u], [:g], traj, times=times, Qs=Qs) + G_COMP = traj.dim * traj.T .+ traj.global_components[:g] + U_COMP(k) = slice(k, traj.components[:u], traj.dim) + L̂(Z⃗) = sum(Q * L(Z⃗[vcat(U_COMP(k), G_COMP)]) for (Q, k) ∈ zip(Qs, times)) + + @test OBJ.L(vec(traj)) ≈ L̂(vec(traj)) + + ∂L_autodiff = ForwardDiff.gradient(L̂, vec(traj)) + @test OBJ.∇L(vec(traj)) ≈ ∂L_autodiff + + ∂²L_autodiff = ForwardDiff.hessian(L̂, vec(traj)) + + ∂²L_full = zeros(size(∂²L_autodiff)) + for (index, entry) in zip(OBJ.∂²L_structure(), OBJ.∂²L(vec(traj))) + i, j = index + ∂²L_full[i, j] = entry + end + + @test ∂²L_full ≈ ∂²L_autodiff end \ No newline at end of file diff --git a/src/solvers/ipopt_solver/constraints.jl b/src/solvers/ipopt_solver/constraints.jl index 4ecd2d7..f1107fc 100644 --- a/src/solvers/ipopt_solver/constraints.jl +++ b/src/solvers/ipopt_solver/constraints.jl @@ -59,28 +59,28 @@ function (con::AllEqualConstraint)( end end -function (con::L1SlackConstraint)( - opt::Ipopt.Optimizer, - vars::Vector{MOI.VariableIndex} -) - for (x, s1, s2) in zip(con.x_indices, con.s1_indices, con.s2_indices) - MOI.add_constraints( - opt, - vars[s1], - MOI.GreaterThan(0.0) - ) - MOI.add_constraints( - opt, - vars[s2], - MOI.GreaterThan(0.0) - ) - t1 = MOI.ScalarAffineTerm(1.0, vars[s1]) - t2 = MOI.ScalarAffineTerm(-1.0, vars[s2]) - t3 = MOI.ScalarAffineTerm(-1.0, vars[x]) - MOI.add_constraints( - opt, - MOI.ScalarAffineFunction([t1, t2, t3], 0.0), - MOI.EqualTo(0.0) - ) - end -end \ No newline at end of file +# function (con::L1SlackConstraint)( +# opt::Ipopt.Optimizer, +# vars::Vector{MOI.VariableIndex} +# ) +# for (x, s1, s2) in zip(con.x_indices, con.s1_indices, con.s2_indices) +# MOI.add_constraints( +# opt, +# vars[s1], +# MOI.GreaterThan(0.0) +# ) +# MOI.add_constraints( +# opt, +# vars[s2], +# MOI.GreaterThan(0.0) +# ) +# t1 = MOI.ScalarAffineTerm(1.0, vars[s1]) +# t2 = MOI.ScalarAffineTerm(-1.0, vars[s2]) +# t3 = MOI.ScalarAffineTerm(-1.0, vars[x]) +# MOI.add_constraints( +# opt, +# MOI.ScalarAffineFunction([t1, t2, t3], 0.0), +# MOI.EqualTo(0.0) +# ) +# end +# end \ No newline at end of file diff --git a/src/solvers/ipopt_solver/solver.jl b/src/solvers/ipopt_solver/solver.jl index f2e02bb..9b85bb1 100644 --- a/src/solvers/ipopt_solver/solver.jl +++ b/src/solvers/ipopt_solver/solver.jl @@ -152,33 +152,27 @@ function set_variables!( optimizer::Ipopt.Optimizer, traj::NamedTrajectory ) - # initialize n variables with trajectory data - n_traj_vars = traj.dim * traj.T - n_vars = n_traj_vars + traj.global_dim + data_dim = traj.dim * traj.T # add variables - variables = MOI.add_variables(optimizer, n_vars) + variables = MOI.add_variables(optimizer, data_dim + traj.global_dim) # set trajectory data MOI.set( optimizer, MOI.VariablePrimalStart(), - variables[1:n_traj_vars], + variables[1:data_dim], collect(traj.datavec) ) - # set global variables - running_vars = n_traj_vars - for global_vars_i ∈ values(traj.global_data) - n_global_vars = length(global_vars_i) - MOI.set( - optimizer, - MOI.VariablePrimalStart(), - variables[running_vars .+ (1:n_global_vars)], - global_vars_i - ) - running_vars += n_global_vars - end + # set global data + MOI.set( + optimizer, + MOI.VariablePrimalStart(), + variables[data_dim .+ (1:traj.global_dim)], + collect(traj.global_data) + ) + return variables end @@ -187,35 +181,11 @@ function update_trajectory!( optimizer::Ipopt.Optimizer, variables::Vector{MOI.VariableIndex} ) - - n_vars = prob.trajectory.dim * prob.trajectory.T - # get trajectory data - datavec = MOI.get( - optimizer, - MOI.VariablePrimal(), - variables[1:n_vars] + update!( + prob.trajectory, + MOI.get(optimizer, MOI.VariablePrimal(), variables), + type=:both ) - - update!(prob.trajectory, datavec) - - # get global variables after trajectory data - global_keys = keys(prob.trajectory.global_data) - global_values = [] - for global_var ∈ global_keys - n_global_vars = length(prob.trajectory.global_data[global_var]) - push!(global_values, MOI.get( - optimizer, - MOI.VariablePrimal(), - variables[n_vars .+ (1:n_global_vars)] - )) - n_vars += n_global_vars - end - global_data = (; (global_keys .=> global_values)...) - update!(prob.trajectory, global_data) - - # TODO: this results in a bug of shifted components when components are added after creating a trajectory, this affects constraints which store original componentes in probs - # prob.trajectory = NamedTrajectory(datavec, global_data, prob.trajectory) - return nothing end diff --git a/test/test_utils.jl b/test/test_utils.jl index ae9d6ed..2901b50 100644 --- a/test/test_utils.jl +++ b/test/test_utils.jl @@ -125,7 +125,8 @@ function bilinear_dynamics_and_trajectory(; Δt = 0.1, u_bound = 0.1, ω = 0.1, - time=false + add_time=false, + add_global=false ) Gx = sparse(Float64[ 0 0 0 1; @@ -183,8 +184,12 @@ function bilinear_dynamics_and_trajectory(; ) ) - if time - traj = add_component(traj, :t, get_times(traj)) + if add_time + traj = add_component(traj, :t, get_times(traj), type=:state) + end + + if add_global + traj = add_component(traj, :g, randn(N), type=:global) end return G, traj