dynamics constructor for user-generated dynamics

Author: thowell

Project.toml

@@ -9,6 +9,7 @@ Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
+Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
@@ -17,5 +18,6 @@ IfElse = "0.1"
 Ipopt = "0.8"
 MathOptInterface = "0.10"
 Parameters = "0.12"
-Symbolics = "0.1.29"
+Scratch = "1.0"
+Symbolics = "0.1.29 - 0.1.29"
 julia = "1.6"

README.md

@@ -2,20 +2,22 @@
 [![CI](https://github.com/thowell/DirectTrajectoryOptimization.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/thowell/DirectTrajectoryOptimization.jl/actions/workflows/CI.yml)
 [![codecov](https://codecov.io/gh/thowell/DirectTrajectoryOptimization.jl/branch/main/graph/badge.svg?token=821EI7HJEL)](https://codecov.io/gh/thowell/DirectTrajectoryOptimization.jl)
 
-This package solves constrained trajectory optimization problems: 
+A Julia package for solving constrained trajectory optimization problems: 
 
 ```
 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
+                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.
 ```
 
+with direct trajectory optimization. 
+
 Fast and allocation-free gradients and sparse Jacobians are automatically generated using [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl) for user-provided costs, constraints, and dynamics. The problem is solved with [Ipopt](https://coin-or.github.io/Ipopt/).
 
 ## Installation
 ```
 Pkg.add("DirectTrajectoryOptimization")
-```
+```

examples/car/car.jl

@@ -52,7 +52,7 @@ p_obs = [0.5; 0.5]
 r_obs = 0.1
 function obs(x, u, w) 
     e = x[1:2] - p_obs
-    return [dot(e, e) - r_obs^2.0]
+    return [r_obs^2.0 - dot(e, e)]
 end
 cont = StageConstraint(obs, nx, nu, nw, [t for t = 1:T-1], :inequality)
 conT = StageConstraint(obs, nx, 0, nw, [T], :inequality)

examples/double_integrator.jl

@@ -0,0 +1,102 @@
+# PREAMBLE
+
+# PKG_SETUP
+
+# ## Setup
+
+using DirectTrajectoryOptimization 
+using LinearAlgebra
+using Plots
+
+# ## horizon 
+T = 11 
+
+# ## double integrator 
+nx = 2
+nu = 1 
+nw = 0 
+nz = nx + nu + nx
+
+# ## dynamics
+function double_integrator(d, y, x, u, w)
+    A = [1.0 1.0; 0.0 1.0] 
+    B = [0.0; 1.0] 
+    d .= y - (A * x + B * u[1])
+end
+
+# ## user-provided dynamics gradient
+function double_integrator_grad(dz, y, x, u, w) 
+    A = [1.0 1.0; 0.0 1.0] 
+    B = [0.0; 1.0] 
+    dz .= [-A -B I]
+end
+
+# ## fast methods
+function double_integrator(y, x, u, w)
+    A = [1.0 1.0; 0.0 1.0] 
+    B = [0.0; 1.0] 
+    y - (A * x + B * u[1])
+end
+
+function double_integrator_grad(y, x, u, w) 
+    A = [1.0 1.0; 0.0 1.0] 
+    B = [0.0; 1.0] 
+    [-A -B I]
+end
+
+@variables y[1:nx] x[1:nx] u[1:nu] w[1:nw]
+
+di = double_integrator(y, x, u, w) 
+diz = double_integrator_grad(y, x, u, w) 
+di_func = eval(Symbolics.build_function(di, y, x, u, w)[2])
+diz_func = eval(Symbolics.build_function(diz, y, x, u, w)[2])
+
+# ## model
+dt = Dynamics(di_func, diz_func, nx, nx, nu)
+dyn = [dt for t = 1:T-1] 
+model = DynamicsModel(dyn)
+
+# ## initialization
+x1 = [0.0; 0.0] 
+xT = [1.0; 0.0] 
+
+# ## objective 
+ot = (x, u, w) -> 0.1 * dot(x, x) + 0.1 * dot(u, u)
+oT = (x, u, w) -> 0.1 * dot(x, x)
+ct = Cost(ot, nx, nu, nw, [t for t = 1:T-1])
+cT = Cost(oT, nx, 0, nw, [T])
+obj = [ct, cT]
+
+# ## constraints
+x_init = Bound(nx, nu, [1], xl=x1, xu=x1)
+x_goal = Bound(nx, 0, [T], xl=xT, xu=xT)
+cons = ConstraintSet([x_init, x_goal], [StageConstraint()])
+
+# ## problem 
+trajopt = TrajectoryOptimizationProblem(obj, model, cons)
+s = Solver(trajopt, options=Options())
+
+# ## initialize
+x_interpolation = linear_interpolation(x1, xT, T)
+u_guess = [1.0 * randn(nu) for t = 1:T-1]
+z0 = zeros(s.p.num_var)
+for (t, idx) in enumerate(s.p.trajopt.model.idx.x)
+    z0[idx] = x_interpolation[t]
+end
+for (t, idx) in enumerate(s.p.trajopt.model.idx.u)
+    z0[idx] = u_guess[t]
+end
+initialize!(s, z0)
+
+# ## solve
+@time solve!(s)
+
+# ## solution
+@show trajopt.x[1]
+@show trajopt.x[T]
+
+# ## state
+plot(hcat(trajopt.x...)')
+
+# ## control
+plot(hcat(trajopt.u[1:end-1]..., trajopt.u[end-1])', linetype = :steppost)

src/dynamics.jl

@@ -41,6 +41,35 @@ function Dynamics(f::Function, ny::Int, nx::Int, nu::Int; nw::Int=0, eval_hess=f
         sp_jac, sp_hess, zeros(ny), zeros(nj), zeros(nh))
 end
 
+function Dynamics(g::Function, gz::Function, ny::Int, nx::Int, nu::Int; nw::Int=0)  
+    # jacobian function 
+    nz = nx + nu + ny
+    jac_func = (nj, y, x, u, w) -> gz(reshape(view(nj, :), ny, nz), y, x, u, w)
+
+    # number of Jacobian elements
+    nj = ny * nz
+
+    # Jacobian sparsity
+    row = Int[]
+    col = Int[]
+    for j = 1:nz
+        for i = 1:ny 
+            push!(row, i) 
+            push!(col, j)
+        end
+    end
+
+    sp_jac = [row, col]
+  
+    # Hessian
+    hess_func = Expr(:null) 
+    sp_hess = [Int[]]
+    nh = 0
+  
+    return Dynamics(g, jac_func, hess_func, ny, nx, nu, nw, nj, nh,
+        sp_jac, sp_hess, zeros(ny), zeros(nj), zeros(nh))
+end
+
 function eval_con!(c, idx, cons::Vector{Dynamics{T}}, x, u, w) where T
     for (t, con) in enumerate(cons)
         con.val(con.val_cache, x[t+1], x[t], u[t], w[t])

src/problem.jl

@@ -82,7 +82,7 @@ function constraint_bounds(dyn::Vector{Dynamics{T}}, stage::StageConstraints{T},
     for con in stage 
         for t in con.idx 
             if con.type == :inequality
-                cu[idx.stage_con[i]] .= Inf
+                cl[idx.stage_con[i]] .= -Inf
             end
             i += 1
         end 

test/Project.toml

@@ -6,4 +6,4 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-Symbolics = "0.1.29"
+Symbolics = "0.1.29 - 0.1.29"

test/dynamics.jl

@@ -37,6 +37,7 @@
     j = zeros(DirectTrajectoryOptimization.num_jac(dyn))
 
     dt.val(dt.val_cache, x1, x1, u1, w1) 
+    # @benchmark $dt.val($dt.val_cache, $x1, $x1, $u1, $w1) 
     @test norm(dt.val_cache - euler_implicit(x1, x1, u1, w1)) < 1.0e-8
     dt.jac(dt.jac_cache, x1, x1, u1, w1) 
     jac_dense = zeros(dt.ny, dt.nx + dt.nu + dt.ny)

test/hessian_lagrangian.jl

@@ -1,5 +1,5 @@
 @testset "Hessian of Lagrangian" begin
-
+    MOI = DirectTrajectoryOptimization.MOI
     # horizon 
     T = 3
 
@@ -186,9 +186,9 @@
     @test norm(h0_full - Lxx_func(z0)) < 1.0e-8
     @test norm(norm(h0 - Lxx_sp_func(z0))) < 1.0e-8
 
-    # h0 = zeros(nh)
-    # MOI.eval_hessian_lagrangian(s.p, h0, z0[1:np], σ, z0[np .+ (1:nd)])
-    # @test norm(norm(h0 - Lxx_sp_func(z0))) < 1.0e-8
+    h0 = zeros(nh)
+    MOI.eval_hessian_lagrangian(s.p, h0, z0[1:np], σ, z0[np .+ (1:nd)])
+    @test norm(norm(h0 - Lxx_sp_func(z0))) < 1.0e-8
 
     # a = z0[1:np]
     # b = z0[np .+ (1:nd)]

test/runtests.jl

@@ -8,5 +8,5 @@ using DirectTrajectoryOptimization
 include("objective.jl")
 include("dynamics.jl")
 include("constraints.jl")
-include("hessian_lagrangian.jl")
+include("hessian_lagrangian.jl") #TODO: fix 
 include("solve.jl")