general constraint run_ci

Author: thowell

README.md

@@ -9,8 +9,8 @@ minimize        gT(xT; wT) + sum(gt(xt, ut; wt))
 xt, ut
 subject to      xt+1 = ft(xt, ut; wt) , t = 1,...,T-1 
                 ct(xt, ut; wt) {<,=} 0, t = 1,...,T
-                xt_L <= xt <= xt_U, t = 1,...,T 
-                ut_L <= ut <= ut_U, t = 1,...,T-1.
+                xt_L < xt < xt_U, t = 1,...,T 
+                ut_L < ut < ut_U, t = 1,...,T-1.
 ```
 
 with direct trajectory optimization. 

examples/acrobot/acrobot.jl

@@ -4,9 +4,9 @@
 
 # ## Setup
 
-using DirectTrajectoryOptimization 
 using LinearAlgebra
 using Plots
+using DirectTrajectoryOptimization 
 
 # ## horizon 
 T = 101 
@@ -105,6 +105,34 @@ cons = [Constraint() for t = 1:T]
 # ## problem 
 p = ProblemData(obj, dyn, cons, bnds, options=Options())
 
+@variables z[1:p.nlp.num_var]
+
+function gen_con(z) 
+    z[1:2] - z[end-1:end]
+end
+
+gc = gen_con(z)
+jac = Symbolics.sparsejacobian(gc, z)
+gc_func = eval(Symbolics.build_function(gc, z)[2])
+gc_jac_func = eval(Symbolics.build_function(jac.nzval, z)[2])
+nc = length(gc) 
+nj = length(jac.nzval)
+sp_jac = [findnz(jac)[1:2]...]
+if false
+    @variables λ[1:nc]
+    lag_con = dot(λ, gc) 
+    hess = Symbolics.sparsehessian(lag_con, z)
+    hess_func = eval(Symbolics.build_function(hess.nzval, z, λ)[2])
+    sp_hess = [findnz(hess)[1:2]...]
+    nh = length(hess.nzval)
+else 
+    hess_func = Expr(:null) 
+    sp_hess = [Int[]]
+    nh = 0
+end
+
+
+
 # ## initialize
 x_interpolation = linear_interpolation(x1, xT, T)
 u_guess = [1.0 * randn(nu) for t = 1:T-1]

src/DirectTrajectoryOptimization.jl

@@ -11,6 +11,7 @@ const MOI = MathOptInterface
 include("objective.jl")
 include("constraints.jl")
 include("bounds.jl")
+include("general_constraint.jl")
 include("dynamics.jl")
 include("solver.jl")
 include("data.jl")
@@ -21,7 +22,7 @@ include("utils.jl")
 export Cost
 
 # constraints 
-export Bound, Bounds, Constraint, Constraints
+export Bound, Bounds, Constraint, Constraints, GeneralConstraint
 
 # dynamics 
 export Dynamics

src/constraints.jl

@@ -67,9 +67,11 @@ end
 
 function eval_hess_lag!(h, idx, cons::Constraints{T}, x, u, w, λ) where T
     for (t, con) in enumerate(cons)
-        con.hess(con.hess_cache, x[t], u[t], w[t], λ[t])
-        @views h[idx[t]] .+= con.hess_cache
-        fill!(con.hess_cache, 0.0) # TODO: confirm this is necessary
+        if !isempty(con.hess_cache)
+            con.hess(con.hess_cache, x[t], u[t], w[t], λ[t])
+            @views h[idx[t]] .+= con.hess_cache
+            fill!(con.hess_cache, 0.0) # TODO: confirm this is necessary
+        end
     end
 end
 
@@ -134,15 +136,3 @@ function hessian_indices(cons::Constraints{T}, key::Vector{Tuple{Int,Int}}, nx::
     end
     return idx
 end
-
-# struct ConstraintSet{T} 
-#     bounds::Bounds{T}
-#     stage::Constraints{T} 
-# end
-
-# ConstraintSet() = ConstraintSet([Bound()], [Constraint()]) 
-# ConstraintSet(stage::Constraints{T}) where T = ConstraintSet([Bound()], stage) 
-# ConstraintSet(bounds::Bounds{T}) where T = ConstraintSet(bounds, [Constraint()]) 
-
-# num_con(cons::ConstraintSet{T}) where T = sum([con.nc * length(con.idx) for con in cons.stage])
-# num_jac(cons::ConstraintSet{T}) where T = sum([con.nj * length(con.idx) for con in cons.stage])

src/data.jl

@@ -37,32 +37,42 @@ struct TrajectoryOptimizationIndices
     stage_con::Vector{Vector{Int}} 
     stage_jac::Vector{Vector{Int}} 
     stage_hess::Vector{Vector{Int}}
+    gen_con::Vector{Int}
+    gen_jac::Vector{Int}
+    gen_hess::Vector{Int}
     x::Vector{Vector{Int}}
     u::Vector{Vector{Int}}
     xu::Vector{Vector{Int}}
     xuy::Vector{Vector{Int}}
 end
 
-function indices(obj::Objective{T}, dyn::Vector{Dynamics{T}}, cons::Constraints{T}, key::Vector{Tuple{Int,Int}},
-        nx::Vector{Int}, nu::Vector{Int}) where T 
+function indices(obj::Objective{T}, dyn::Vector{Dynamics{T}}, cons::Constraints{T}, gc::GeneralConstraint{T},
+     key::Vector{Tuple{Int,Int}}, nx::Vector{Int}, nu::Vector{Int}, nz::Int) where T 
     # Jacobians
     dyn_con = constraint_indices(dyn, shift=0)
     dyn_jac = jacobian_indices(dyn, shift=0)
     stage_con = constraint_indices(cons, shift=num_con(dyn))
     stage_jac = jacobian_indices(cons, shift=num_jac(dyn)) 
+    gen_con = constraint_indices(gc, shift=(num_con(dyn) + num_con(cons)))
+    gen_jac = jacobian_indices(gc, shift=(num_jac(dyn) + num_jac(cons))) 
 
     # Hessian of Lagrangian 
     obj_hess = hessian_indices(obj, key, nx, nu)
     dyn_hess = hessian_indices(dyn, key, nx, nu)
     stage_hess = hessian_indices(cons, key, nx, nu)
+    gen_hess = hessian_indices(gc, key, nz)
 
     # indices
     x_idx = x_indices(dyn)
     u_idx = u_indices(dyn)
     xu_idx = xu_indices(dyn)
     xuy_idx = xuy_indices(dyn)
 
-    return TrajectoryOptimizationIndices(obj_hess, dyn_con, dyn_jac, dyn_hess, stage_con, stage_jac, stage_hess,
+    return TrajectoryOptimizationIndices(
+        obj_hess, 
+        dyn_con, dyn_jac, dyn_hess, 
+        stage_con, stage_jac, stage_hess,
+        gen_con, gen_jac, gen_hess,
         x_idx, u_idx, xu_idx, xuy_idx) 
 end
 
@@ -78,6 +88,9 @@ struct NLPData{T} <: MOI.AbstractNLPEvaluator
     sp_jac
     sp_hess_lag
     hess_lag::Bool 
+    gc::GeneralConstraint{T}
+    w::Vector{T}
+    λ_gen::Vector{T}
 end
 
 function primal_bounds(bnds::Bounds{T}, nz::Int, x_idx::Vector{Vector{Int}}, u_idx::Vector{Vector{Int}}) where T 
@@ -91,55 +104,79 @@ function primal_bounds(bnds::Bounds{T}, nz::Int, x_idx::Vector{Vector{Int}}, u_i
     return zl, zu
 end
 
-function constraint_bounds(cons::Constraints{T}, nc::Int, idx::TrajectoryOptimizationIndices) where T
+function constraint_bounds(cons::Constraints{T}, gc::GeneralConstraint{T}, 
+    nc_dyn::Int, nc_con::Int, idx::TrajectoryOptimizationIndices) where T
+    # total constraints
+    nc = nc_dyn + nc_con + gc.nc 
+    # bounds
     cl, cu = zeros(nc), zeros(nc) 
+    # stage
     for (t, con) in enumerate(cons) 
         cl[idx.stage_con[t][con.idx_ineq]] .= -Inf
     end
+    # general
+    cl[collect(nc_dyn + nc_con .+ gc.idx_ineq)] .= -Inf
     return cl, cu
 end 
 
-function NLPData(trajopt::TrajectoryOptimizationData; eval_hess=false) 
+function NLPData(trajopt::TrajectoryOptimizationData; 
+    eval_hess=false, 
+    general_constraint=GeneralConstraint()) 
+
     # number of variables
     nz = sum(trajopt.x_dim) + sum(trajopt.u_dim)
 
     # number of constraints
     nc_dyn = num_con(trajopt.dyn)
     nc_con = num_con(trajopt.cons)  
-    nc = nc_dyn + nc_con
+    nc_gen = num_con(general_constraint)
+    nc = nc_dyn + nc_con + nc_gen
 
     # number of nonzeros in constraint Jacobian
     nj_dyn = num_jac(trajopt.dyn)
     nj_con = num_jac(trajopt.cons)  
-    nj = nj_dyn + nj_con
+    nj_gen = num_jac(general_constraint)
+    nj = nj_dyn + nj_con + nj_gen
 
     # number of nonzeros in Hessian of Lagrangian
     nh = 0
 
     # constraint Jacobian sparsity
     sp_dyn = sparsity_jacobian(trajopt.dyn, trajopt.x_dim, trajopt.u_dim, row_shift=0)
     sp_con = sparsity_jacobian(trajopt.cons, trajopt.x_dim, trajopt.u_dim, row_shift=nc_dyn)
-    sp_jac = collect([sp_dyn..., sp_con...]) 
+    sp_gen = sparsity_jacobian(general_constraint, nz, row_shift=(nc_dyn + nc_con))
+    sp_jac = collect([sp_dyn..., sp_con..., sp_gen...]) 
 
     # Hessian of Lagrangian sparsity 
     sp_obj_hess = sparsity_hessian(trajopt.obj, trajopt.x_dim, trajopt.u_dim)
     sp_dyn_hess = sparsity_hessian(trajopt.dyn, trajopt.x_dim, trajopt.u_dim)
     sp_con_hess = sparsity_hessian(trajopt.cons, trajopt.x_dim, trajopt.u_dim)
-    sp_hess_lag = [sp_obj_hess..., sp_dyn_hess..., sp_con_hess...]
+    sp_gen_hess = sparsity_hessian(general_constraint, nz)
+    sp_hess_lag = [sp_obj_hess..., sp_dyn_hess..., sp_con_hess..., sp_gen_hess...]
     sp_hess_lag = !isempty(sp_hess_lag) ? sp_hess_lag : Tuple{Int,Int}[]
     sp_hess_key = sort(unique(sp_hess_lag))
-    # sp_hess_status = [!isempty(sp_obj_hess), !isempty(sp_dyn_hess), !isempty(sp_con_hess)]
 
     # indices 
-    idx = indices(trajopt.obj, trajopt.dyn, trajopt.cons, sp_hess_key, trajopt.x_dim, trajopt.u_dim)
+    idx = indices(trajopt.obj, trajopt.dyn, trajopt.cons, 
+        general_constraint, 
+        sp_hess_key, 
+        trajopt.x_dim, trajopt.u_dim, nz)
 
     # primal variable bounds
     zl, zu = primal_bounds(trajopt.bnds, nz, idx.x, idx.u) 
 
     # nonlinear constraint bounds
-    cl, cu = constraint_bounds(trajopt.cons, nc, idx) 
-
-    NLPData(trajopt, nz, nc, nj, nh, [zl, zu], [cl, cu], idx, sp_jac, sp_hess_key, eval_hess)
+    cl, cu = constraint_bounds(trajopt.cons, general_constraint, nc_dyn, nc_con, idx) 
+
+    NLPData(trajopt, 
+        nz, nc, nj, nh, 
+        [zl, zu], [cl, cu], 
+        idx, 
+        sp_jac, sp_hess_key,
+        eval_hess, 
+        general_constraint,
+        vcat(trajopt.w...),
+        zeros(general_constraint.nc))
 end
 
 struct SolverData 
@@ -186,14 +223,15 @@ function trajectory!(x::Vector{Vector{T}}, u::Vector{Vector{T}}, z::Vector{T},
     end
 end
 
-function duals!(λ_dyn::Vector{Vector{T}}, λ_stage::Vector{Vector{T}}, λ, 
-    idx_dyn::Vector{Vector{Int}}, idx_stage::Vector{Vector{Int}}) where T
+function duals!(λ_dyn::Vector{Vector{T}}, λ_stage::Vector{Vector{T}}, λ_gen::Vector{T}, λ, 
+    idx_dyn::Vector{Vector{Int}}, idx_stage::Vector{Vector{Int}}, idx_gen::Vector{Int}) where T
     for (t, idx) in enumerate(idx_dyn)
         λ_dyn[t] .= @views λ[idx]
     end 
     for (t, idx) in enumerate(idx_stage)
         λ_stage[t] .= @views λ[idx]
     end
+    λ_gen .= λ[idx_gen]
 end
 
 struct ProblemData{T} <: MOI.AbstractNLPEvaluator
@@ -202,11 +240,13 @@ struct ProblemData{T} <: MOI.AbstractNLPEvaluator
 end
 
 function ProblemData(obj::Objective{T}, dyn::Vector{Dynamics{T}}, cons::Constraints{T}, bnds::Bounds{T}; 
-    eval_hess=false, options=Options(),
+    eval_hess=false, 
+    general_constraint=GeneralConstraint(),
+    options=Options(),
     w=[[zeros(nw) for nw in dimensions(dyn)[3]]..., zeros(0)]) where T
 
     trajopt = TrajectoryOptimizationData(obj, dyn, cons, bnds, w=w)
-    nlp = NLPData(trajopt, eval_hess=eval_hess) 
+    nlp = NLPData(trajopt, general_constraint=general_constraint, eval_hess=eval_hess) 
     s_data = SolverData(nlp, options=options)
 
     ProblemData(nlp, s_data) 
@@ -236,49 +276,4 @@ end
 
 function solve!(p::ProblemData) 
     MOI.optimize!(p.s_data.solver) 
-end
-
-# struct Solver{T}
-#     p::Problem{T}
-#     nlp_bounds::Vector{MOI.NLPBoundsPair}
-#     block_data::MOI.NLPBlockData
-#     solver::Ipopt.Optimizer
-#     z::Vector{MOI.VariableIndex}
-# end
-
-# function Solver(trajopt::TrajectoryOptimizationProblem; eval_hess=false, options=Options()) 
-#     p = Problem(trajopt, eval_hess=eval_hess) 
-    
-#     nlp_bounds = MOI.NLPBoundsPair.(p.con_bnds...)
-#     block_data = MOI.NLPBlockData(nlp_bounds, NLPEvaluator(), true)
-    
-#     # instantiate NLP solver
-#     solver = Ipopt.Optimizer()
-
-#     # set NLP solver options
-#     for name in fieldnames(typeof(options))
-#         solver.options[String(name)] = getfield(options, name)
-#     end
-    
-#     z = MOI.add_variables(solver, nz)
-    
-#     for i = 1:p.num_var
-#         MOI.add_constraint(solver, z[i], MOI.LessThan(p.var_bnds[2][i]))
-#         MOI.add_constraint(solver, z[i], MOI.GreaterThan(p.var_bnds[1][i]))
-#     end
-    
-#     MOI.set(solver, MOI.NLPBlock(), block_data)
-#     MOI.set(solver, MOI.ObjectiveSense(), MOI.MIN_SENSE)
-    
-#     return Solver(p, nlp_bounds, block_data, solver, z) 
-# end
-
-# using MathOptInterface
-# const MOI = MathOptInterface
-
-# MOI.AbstractNLPEvaluator
-
-# nlp_bounds = MOI.NLPBoundsPair.([-ones(10), ones(10)]...)
-
-# struct MOIEval <: MOI.AbstractNLPEvaluator end
-# block_data = MOI.NLPBlockData(nlp_bounds, a(), true)
+end

src/dynamics.jl

@@ -85,12 +85,14 @@ function eval_jac!(j, idx, cons::Vector{Dynamics{T}}, x, u, w) where T
         fill!(con.jac_cache, 0.0) # TODO: confirm this is necessary
     end
 end
-
+isempty(ones(1))
 function eval_hess_lag!(h, idx, cons::Vector{Dynamics{T}}, x, u, w, λ) where T
     for (t, con) in enumerate(cons) 
-        con.hess(con.hess_cache, x[t+1], x[t], u[t], w[t], λ[t])
-        @views h[idx[t]] .+= con.hess_cache
-        fill!(con.hess_cache, 0.0) # TODO: confirm this is necessary
+        if !isempty(con.hess_cache)
+            con.hess(con.hess_cache, x[t+1], x[t], u[t], w[t], λ[t])
+            @views h[idx[t]] .+= con.hess_cache
+            fill!(con.hess_cache, 0.0) # TODO: confirm this is necessary
+        end
     end
 end