Merge pull request #10 from Robbybp/rp/atlas

Use VectorNonlinearOracle for Atlas balancing

Author: andrewrosemberg

examples/Atlas/Project.toml

@@ -1,9 +1,15 @@
 [deps]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HSL_jll = "017b0a0e-03f4-516a-9b91-836bbd1904dd"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+MathProgIncidence = "892fab00-3092-4bd0-9c46-66676a93f84e"
 MeshCat = "283c5d60-a78f-5afe-a0af-af636b173e11"
 MeshCatMechanisms = "6ad125db-dd91-5488-b820-c1df6aab299d"
+Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 RigidBodyDynamics = "366cf18f-59d5-5db9-a4de-86a9f6786172"
 Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

examples/Atlas/atlas_balancing.jl

@@ -2,7 +2,9 @@ using Revise
 using JLD2
 using SparseArrays
 using MathOptInterface.Nonlinear.SymbolicAD
+import MathOptInterface as MOI
 using JuMP, Ipopt, HSL_jll
+using LinearAlgebra
 
 # include(joinpath(@__DIR__, "../mpc_utils.jl"))
 include(joinpath(@__DIR__, "atlas_utils.jl"))
@@ -138,7 +140,8 @@ function atlas_dynamics(xu::T...) where {T<:Real}
     h = 0.01
     x = collect(xu[1:atlas.nx])
     u = collect(xu[atlas.nx+1:end])
-    return rk4(atlas, x, u, h)
+    #return rk4(atlas, x, u, h)
+    return explicit_euler(atlas, x, u, h)
 end
 
 # Lets define a jump model that solves the MPC problem
@@ -149,7 +152,8 @@ end
 function build_and_solve_mpc(atlas_obj::Atlas, x_ref::Vector{Float64}, X_start::Vector{Float64}, N::Int; optimizer=optimizer_with_attributes(Ipopt.Optimizer, 
         # "print_level" => 0,
         "hsllib" => HSL_jll.libhsl_path,
-        "linear_solver" => "MA27"
+        "linear_solver" => "MA27",
+        "max_iter" => 20,
     )
 )
     model = Model()
@@ -167,23 +171,66 @@ function build_and_solve_mpc(atlas_obj::Atlas, x_ref::Vector{Float64}, X_start::
 
     # Build the memoized function array for the multi-output dynamics
     # so we can call the i-th dimension of next state individually
-    dyn_ops = memoize_dynamics_and_jacobian(atlas_dynamics,
-                                            atlas_obj.nx + atlas_obj.nu,
-                                            atlas_obj.nx)
+    #dyn_ops = memoize_dynamics_and_jacobian(atlas_dynamics,
+    #                                        atlas_obj.nx + atlas_obj.nu,
+    #                                        atlas_obj.nx)
+
+    VNO_dim = 2*atlas_obj.nx + atlas_obj.nu
+    # The input variable to this function:
+    # [ x[t+1,:], x[t,:], u[t,:] ]
+    jacobian_structure = Tuple{Int,Int}[]
+    append!(jacobian_structure, map(i -> (i,i), 1:atlas_obj.nx))
+    for i in 1:atlas_obj.nx, j in 1:(atlas_obj.nx + atlas_obj.nu)
+        push!(jacobian_structure, (i, j + atlas_obj.nx))
+    end
+    hessian_lagrangian_structure = [
+        (i, j)
+        for i in atlas_obj.nx+1:VNO_dim
+        for j in atlas_obj.nx+1:VNO_dim
+    ]
+    VNO = MOI.VectorNonlinearOracle(;
+        dimension = VNO_dim,
+        l = zeros(atlas_obj.nx),
+        u = zeros(atlas_obj.nx),
+        eval_f = (ret, x) -> begin
+            ret[1:atlas_obj.nx] .= x[1:atlas_obj.nx] - atlas_dynamics(x[atlas_obj.nx + 1:VNO_dim]...)
+            return
+        end,
+        jacobian_structure,
+        eval_jacobian = (ret, x) -> begin
+            dyn_jac = ForwardDiff.jacobian(x -> atlas_dynamics(x...), x[atlas_obj.nx + 1:VNO_dim])
+            jnnz = length(jacobian_structure)
+            ret[1:atlas_obj.nx] .= ones(atlas_obj.nx)
+            nx = atlas_obj.nx; nu = atlas_obj.nu
+            ret[atlas_obj.nx+1:jnnz] .= - reshape(dyn_jac', nx * (nx + nu))
+            return
+        end,
+        hessian_lagrangian_structure,
+        eval_hessian_lagrangian = (ret, x, λ) -> begin
+            hess = ForwardDiff.hessian(x -> dot(λ, atlas_dynamics(x...)), x[atlas_obj.nx + 1:VNO_dim])
+            hnnz = length(hessian_lagrangian_structure)
+            ret[1:hnnz] .= - reshape(hess, hnnz)
+            return
+        end
+    )
+    for t in 1:(N-1)
+        vars = vcat(x[t+1,:], x[t, :], u[t, :])
+        JuMP.@constraint(model, vars in VNO)
+    end
 
     # Add constraints: x[t+1] = f(x[t], u[t]) for t=1 to N-1
-    for t in 1:(N-1), i in 1:atlas_obj.nx
-        scalar_fun, grad_fun! = dyn_ops[i]
-        op = add_nonlinear_operator(
-            model,
-            atlas_obj.nx + atlas_obj.nu,  # number of inputs
-            scalar_fun,                   # f_i
-            grad_fun!,                    # gradient callback
-            name = Symbol("dyn_$(t)_$(i)")
-        )
-        # Bind the operator to the constraint x[t+1, i] == f_i( [x[t,:]; u[t,:]]... )
-        @constraint(model, x[t+1,i] == op( [x[t, :]; u[t, :]]... ))
-    end
+    #for t in 1:(N-1), i in 1:atlas_obj.nx
+    #    scalar_fun, grad_fun! = dyn_ops[i]
+    #    op = add_nonlinear_operator(
+    #        model,
+    #        atlas_obj.nx + atlas_obj.nu,  # number of inputs
+    #        scalar_fun,                   # f_i
+    #        grad_fun!,                    # gradient callback
+    #        name = Symbol("dyn_$(t)_$(i)")
+    #    )
+    #    # Bind the operator to the constraint x[t+1, i] == f_i( [x[t,:]; u[t,:]]... )
+    #    @constraint(model, x[t+1,i] == op( [x[t, :]; u[t, :]]... ))
+    #end
 
     # Possibly set an initial condition constraint: x[1, :] = x_ref
     @constraint(model, [j in 1:atlas_obj.nx], x[1,j] == X_start[j])
@@ -193,15 +240,16 @@ function build_and_solve_mpc(atlas_obj::Atlas, x_ref::Vector{Float64}, X_start::
     return model, value.(x), value.(u)
 end
 
-N = 100
+N = 10
 X_start = deepcopy(x_ref)
 X_start[atlas.nq + 5] = 1.3
 model, X_solve, U_solve = build_and_solve_mpc(atlas, x_ref, X_start, N; optimizer=optimizer_with_attributes(Ipopt.Optimizer, 
         # "print_level" => 0,
         "linear_solver" => "ma97",
-        "hessian_approximation" => "limited-memory",
-        "max_iter" => 300,
+        #"hessian_approximation" => "limited-memory",
+        #"max_iter" => 20,
         "mu_target" => 1e-8,
+        "print_user_options" => "yes",
     )
 )
 termination_status(model)
@@ -269,92 +317,95 @@ error = sum( norm(X_solve[t,:] .- X[t]) for t in 1:N )
     # global X[k + 1] = rk4(atlas, X[k], clamp.(u_ref + U[k], -atlas.torque_limits, atlas.torque_limits), h)
 # end
 
-X = [value.(x[t,:]) for t=1:N]
-animate!(atlas, mvis, X, Δt=h);
-readline()
-
-
-###############
-include(joinpath(@__DIR__, "rigidbodyutils.jl"))
-include(joinpath(@__DIR__, "atlas_utils.jl"))
-
-model = Model()
-optimizer=optimizer_with_attributes(Ipopt.Optimizer, 
-    # "print_level" => 0,
-    "hsllib" => HSL_jll.libhsl_path,
-    "linear_solver" => "MA27"
-)
-set_optimizer(model, optimizer)
-
-@variable(model, x[1:N, 1:atlas.nx])
-@variable(model, -atlas.torque_limits[i] <= u[t=1:N-1,i=1:atlas.nu] <= atlas.torque_limits[i])
-
-dyn_result = dynamics(atlas, convert.(NonlinearExpr, x[1, :]), convert.(NonlinearExpr,u[1, :]))
-
-function atlas_dynamics(x, u)
-    h = 0.01
-    return rk4(atlas, x, u, h)
-end
-
-function build_and_solve_mpc(atlas_obj::Atlas, x_ref::Vector{Float64}, X_start::Vector{Float64}, N::Int; optimizer=optimizer_with_attributes(Ipopt.Optimizer, 
-        # "print_level" => 0,
-        "hsllib" => HSL_jll.libhsl_path,
-        "linear_solver" => "MA27"
-    )
-)
-    model = Model()
-    set_optimizer(model, optimizer)
-
-    # For demonstration, let's define:
-    #   x[t=1:N, 1:atlas.nx]
-    #   u[t=1:N-1, 1:atlas.nu]
-    @variable(model, x[1:N, 1:atlas_obj.nx])
-    @variable(model, -atlas_obj.torque_limits[i] <= u[t=1:N-1,i=1:atlas.nu] <= atlas_obj.torque_limits[i])
-    # ^ adapt the indexing if torque_limits is an array of length nu, etc.
-
-    # Objective: sum of squared error from x_ref
-    @objective(model, Min, sum( (x[t,j] - x_ref[j])^2 for t in 2:N for j in 1:atlas_obj.nx ) )
-
-    for t in 1:(N-1)
-        # Bindconstraint x[t+1, i] == f_i( [x[t,:]; u[t,:]]... )
-        @constraint(model, x[t+1, :] == atlas_dynamics( convert.(NonlinearExpr, x[t, :]), convert.(NonlinearExpr,u[t, :])))
-    end
-
-    # Possibly set an initial condition constraint: x[1, :] = x_ref
-    @constraint(model, [j in 1:atlas_obj.nx], x[1,j] == X_start[j])
-
-    # Solve
-    optimize!(model)
-    return model, value.(x), value.(u)
-end
-
-N = 2
-X_start = deepcopy(x_ref)
-X_start[atlas.nq + 5] = 1.3
-model, X_solve, U_solve = build_and_solve_mpc(atlas, x_ref, X_start, N; optimizer=optimizer_with_attributes(Ipopt.Optimizer, 
-        # "print_level" => 0,
-        "linear_solver" => "ma97",
-        "hessian_approximation" => "limited-memory",
-        "max_iter" => 300,
-        "mu_target" => 1e-8,
-    )
-)
-termination_status(model)
-
-# Now we can simulate the system with the controls U_solve
-X = [zeros(atlas.nx) for _ = 1:N];
-X[1] = deepcopy(x_ref);
-X[1][atlas.nq + 5] = 1.3; # Perturb i.c.
-
-for k = 1:N - 1
-    X[k + 1] = rk4(atlas, X[k], U_solve[k, :], h)
-end
-
-animate!(atlas, mvis, X, Δt=h);
-
-error = sum( norm(X_solve[t,:] .- X[t]) for t in 1:N )
-
-
-function foo!(q::AbstractVector{T}) where {T}
-    q .= zero(T)
-end
+#X = [value.(x[t,:]) for t=1:N]
+#X = [X_solve[t, :] for t=1:N]
+#
+#animate!(atlas, mvis, X, Δt=h);
+#readline()
+#
+#
+################
+#include(joinpath(@__DIR__, "rigidbodyutils.jl"))
+#include(joinpath(@__DIR__, "atlas_utils.jl"))
+#
+#model = Model()
+#optimizer=optimizer_with_attributes(Ipopt.Optimizer, 
+#    # "print_level" => 0,
+#    "hsllib" => HSL_jll.libhsl_path,
+#    "linear_solver" => "MA27"
+#)
+#set_optimizer(model, optimizer)
+#
+#@variable(model, x[1:N, 1:atlas.nx])
+#@variable(model, -atlas.torque_limits[i] <= u[t=1:N-1,i=1:atlas.nu] <= atlas.torque_limits[i])
+#
+#dyn_result = dynamics(atlas, convert.(NonlinearExpr, x[1, :]), convert.(NonlinearExpr,u[1, :]))
+#
+#function atlas_dynamics(x, u)
+#    h = 0.01
+#    return rk4(atlas, x, u, h)
+#end
+#
+##function build_and_solve_mpc(atlas_obj::Atlas, x_ref::Vector{Float64}, X_start::Vector{Float64}, N::Int; optimizer=optimizer_with_attributes(Ipopt.Optimizer, 
+##        # "print_level" => 0,
+##        "hsllib" => HSL_jll.libhsl_path,
+##        "linear_solver" => "MA27"
+##    )
+##)
+##    model = Model()
+##    set_optimizer(model, optimizer)
+##
+##    # For demonstration, let's define:
+##    #   x[t=1:N, 1:atlas.nx]
+##    #   u[t=1:N-1, 1:atlas.nu]
+##    @variable(model, x[1:N, 1:atlas_obj.nx])
+##    @variable(model, -atlas_obj.torque_limits[i] <= u[t=1:N-1,i=1:atlas.nu] <= atlas_obj.torque_limits[i])
+##    # ^ adapt the indexing if torque_limits is an array of length nu, etc.
+##
+##    # Objective: sum of squared error from x_ref
+##    @objective(model, Min, sum( (x[t,j] - x_ref[j])^2 for t in 2:N for j in 1:atlas_obj.nx ) )
+##
+##    for t in 1:(N-1)
+##        # Bindconstraint x[t+1, i] == f_i( [x[t,:]; u[t,:]]... )
+##        @constraint(model, x[t+1, :] == atlas_dynamics( convert.(NonlinearExpr, x[t, :]), convert.(NonlinearExpr,u[t, :])))
+##    end
+##
+##    # Possibly set an initial condition constraint: x[1, :] = x_ref
+##    @constraint(model, [j in 1:atlas_obj.nx], x[1,j] == X_start[j])
+##
+##    # Solve
+##    optimize!(model)
+##    return model, value.(x), value.(u)
+##end
+#
+#N = 2
+#X_start = deepcopy(x_ref)
+#X_start[atlas.nq + 5] = 1.3
+#model, X_solve, U_solve = build_and_solve_mpc(atlas, x_ref, X_start, N; optimizer=optimizer_with_attributes(Ipopt.Optimizer, 
+#        # "print_level" => 0,
+#        "linear_solver" => "ma97",
+#        "hessian_approximation" => "limited-memory",
+#        "max_iter" => 300,
+#        "mu_target" => 1e-8,
+#        "print_info_string" => "yes",
+#    )
+#)
+#termination_status(model)
+#
+## Now we can simulate the system with the controls U_solve
+#X = [zeros(atlas.nx) for _ = 1:N];
+#X[1] = deepcopy(x_ref);
+#X[1][atlas.nq + 5] = 1.3; # Perturb i.c.
+#
+#for k = 1:N - 1
+#    X[k + 1] = rk4(atlas, X[k], U_solve[k, :], h)
+#end
+#
+#animate!(atlas, mvis, X, Δt=h);
+#
+#error = sum( norm(X_solve[t,:] .- X[t]) for t in 1:N )
+#
+#
+#function foo!(q::AbstractVector{T}) where {T}
+#    q .= zero(T)
+#end

examples/Atlas/atlas_ref.jld2

@@ -96,7 +96,7 @@ function dynamics(model::Atlas, x::AbstractVector{T1}, u::AbstractVector{T2}) wh
     T = promote_type(T1, T2)
     state = model.statecache[T]
     dyn_result = model.dynrescache[T]
-    println("dyn_result.M: ", dyn_result.massmatrix)
+    #println("dyn_result.M: ", dyn_result.massmatrix)
 
     # Set the mechanism state
     copyto!(state, x)
@@ -107,6 +107,10 @@ function dynamics(model::Atlas, x::AbstractVector{T1}, u::AbstractVector{T2}) wh
     return [dyn_result.q̇; dyn_result.v̇]
 end
 
+function explicit_euler(model::Atlas, x, u, h)
+    return x + h * dynamics(model, x, u)
+end
+
 function rk4(model::Atlas, x, u, h)
     k1 = dynamics(model, x, u)
     k2 = dynamics(model, x + h/2*k1, u)
@@ -136,4 +140,4 @@ function animate!(model::Atlas, mvis::MechanismVisualizer, qs; Δt=0.001)
     MeshCat.setanimation!(mvis, anim)
 
     return anim
-end
+end