v0.2.0

Author: thowell

Project.toml

@@ -1,17 +1,15 @@
 name = "IterativeLQR"
 uuid = "605048dd-e178-462b-beb9-98a09398ef27"
 authors = ["thowell <thowell@stanford.edu>"]
-version = "0.1.3"
+version = "0.2.0"
 
 [deps]
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
-BenchmarkTools = "1.0"
 JLD2 = "0.4"
 Scratch = "1.0"
 Symbolics = "0.1.29 - 0.1.29"

README.md

@@ -31,8 +31,8 @@ using LinearAlgebra
 T = 11 
 
 # particle 
-nx = 2
-nu = 1 
+num_state = 2
+num_action = 1 
 
 function particle(x, u, w)
    A = [1.0 1.0; 0.0 1.0]
@@ -41,27 +41,27 @@ function particle(x, u, w)
 end
 
 # model
-dyn = Dynamics(particle, nx, nu)
+dyn = Dynamics(particle, num_state, num_action)
 model = [dyn for t = 1:T-1] 
 
 # initialization
 x1 = [0.0; 0.0] 
 xT = [1.0; 0.0]
-ū = [1.0e-1 * randn(nu) for t = 1:T-1] 
+ū = [1.0e-1 * randn(num_action) for t = 1:T-1] 
 x̄ = rollout(model, x1, ū)
 
 # 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)
-cT = Cost(oT, nx, 0)
+ct = Cost(ot, num_state, num_action)
+cT = Cost(oT, num_state, 0)
 obj = [[ct for t = 1:T-1]..., cT]
 
 # constraints
 goal(x, u, w) = x - xT
 
 cont = Constraint()
-conT = Constraint(goal, nx, 0)
+conT = Constraint(goal, num_state, 0)
 cons = [[cont for t = 1:T-1]..., conT] 
 
 # problem

examples/acrobot.jl

@@ -12,9 +12,9 @@ using Plots
 T = 101 
 
 # ## acrobot 
-nx = 4 
-nu = 1 
-nw = 0 
+num_state = 4 
+num_action = 1 
+num_parameter = 0 
 
 function acrobot(x, u, w)
     mass1 = 1.0  
@@ -89,28 +89,28 @@ function midpoint_explicit(x, u, w)
 end
 
 # ## model
-dyn = Dynamics(midpoint_explicit, nx, nu, nw)
+dyn = Dynamics(midpoint_explicit, num_state, num_action, num_parameter)
 model = [dyn for t = 1:T-1] 
 
 # ## initialization
 x1 = [0.0; 0.0; 0.0; 0.0] 
 xT = [0.0; π; 0.0; 0.0]
-ū = [1.0 * randn(nu) for t = 1:T-1] 
-w = [zeros(nw) for t = 1:T]
+ū = [1.0 * randn(num_action) for t = 1:T-1] 
+w = [zeros(num_parameter) for t = 1:T]
 x̄ = rollout(model, x1, ū, w)
 
 # ## objective 
 ot = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4]) + 0.1 * dot(u, u)
 oT = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4])
-ct = Cost(ot, nx, nu, nw)
-cT = Cost(oT, nx, 0, nw)
+ct = Cost(ot, num_state, num_action, num_parameter)
+cT = Cost(oT, num_state, 0, num_parameter)
 obj = [[ct for t = 1:T-1]..., cT]
 
 # ## constraints
 goal(x, u, w) = x - xT
 
 cont = Constraint()
-conT = Constraint(goal, nx, 0)
+conT = Constraint(goal, num_state, 0)
 cons = [[cont for t = 1:T-1]..., conT] 
 
 # ## problem

examples/car.jl

@@ -12,9 +12,9 @@ using Plots
 T = 51 
 
 # ## car 
-nx = 3
-nu = 2
-nw = 0 
+num_state = 3
+num_action = 2
+num_parameter = 0 
 
 function car(x, u, w)
     [u[1] * cos(x[3]); u[1] * sin(x[3]); u[2]]
@@ -26,7 +26,7 @@ function midpoint_explicit(x, u, w)
 end
 
 # ## model
-dyn = Dynamics(midpoint_explicit, nx, nu, nw)
+dyn = Dynamics(midpoint_explicit, num_state, num_action, num_parameter)
 model = [dyn for t = 1:T-1] 
 
 # ## initialization
@@ -35,19 +35,19 @@ xT = [1.0; 1.0; 0.0]
 
 # ## rollout
 ū = [1.0e-2 * [1.0; 0.1] for t = 1:T-1]
-w = [zeros(nw) for t = 1:T] 
+w = [zeros(num_parameter) for t = 1:T] 
 x̄ = rollout(model, x1, ū, w)
 
 # ## objective 
 ot = (x, u, w) -> 1.0 * dot(x - xT, x - xT) + 1.0e-2 * dot(u, u)
 oT = (x, u, w) -> 1000.0 * dot(x - xT, x - xT)
-ct = Cost(ot, nx, nu, nw)
-cT = Cost(oT, nx, 0, nw)
+ct = Cost(ot, num_state, num_action, num_parameter)
+cT = Cost(oT, num_state, 0, num_parameter)
 obj = [[ct for t = 1:T-1]..., cT]
 
 # ## constraints
-ul = -5.0 * ones(nu) 
-uu = 5.0 * ones(nu)
+ul = -5.0 * ones(num_action) 
+uu = 5.0 * ones(num_action)
 
 p_obs = [0.5; 0.5] 
 r_obs = 0.1
@@ -69,8 +69,8 @@ function terminal_con(x, u, w)
     ]
 end
 
-cont = Constraint(stage_con, nx, nu, indices_inequality=collect(1:5))
-conT = Constraint(terminal_con, nx, nu, indices_inequality=collect(3 .+ (1:1)))
+cont = Constraint(stage_con, num_state, num_action, indices_inequality=collect(1:5))
+conT = Constraint(terminal_con, num_state, num_action, indices_inequality=collect(3 .+ (1:1)))
 cons = [[cont for t = 1:T-1]..., conT] 
 
 # ## problem

examples/particle.jl

@@ -11,8 +11,8 @@ using LinearAlgebra
 T = 11 
 
 # ## acrobot 
-nx = 2
-nu = 1 
+num_state = 2
+num_action = 1 
 
 function particle(x, u, w)
    A = [1.0 1.0; 0.0 1.0]
@@ -21,27 +21,27 @@ function particle(x, u, w)
 end
 
 # ## model
-dyn = Dynamics(particle, nx, nu)
+dyn = Dynamics(particle, num_state, num_action)
 model = [dyn for t = 1:T-1] 
 
 # ## initialization
 x1 = [0.0; 0.0] 
 xT = [1.0; 0.0]
-ū = [1.0e-1 * randn(nu) for t = 1:T-1] 
+ū = [1.0e-1 * randn(num_action) for t = 1:T-1] 
 x̄ = rollout(model, x1, ū)
 
 # ## 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)
-cT = Cost(oT, nx, 0)
+ct = Cost(ot, num_state, num_action)
+cT = Cost(oT, num_state, 0)
 obj = [[ct for t = 1:T-1]..., cT]
 
 # ## constraints
 goal(x, u, w) = x - xT
 
 cont = Constraint()
-conT = Constraint(goal, nx, 0)
+conT = Constraint(goal, num_state, 0)
 cons = [[cont for t = 1:T-1]..., conT] 
 
 # ## problem

src/augmented_lagrangian.jl

@@ -14,16 +14,16 @@ function augmented_lagrangian(model::Model{T}, costs::Objective{T}, constraints:
     # horizon
     H = length(model) + 1
     # penalty
-    constraint_penalty = [ones(c.nc) for c in constraints]
-    constraint_penalty_matrix = [Diagonal(ones(c.nc)) for c in constraints]
+    constraint_penalty = [ones(c.num_constraint) for c in constraints]
+    constraint_penalty_matrix = [Diagonal(ones(c.num_constraint)) for c in constraints]
     # duals
-    constraint_dual = [zeros(c.nc) for c in constraints]
+    constraint_dual = [zeros(c.num_constraint) for c in constraints]
     # active set
-    active_set = [ones(Int, c.nc) for c in constraints]
+    active_set = [ones(Int, c.num_constraint) for c in constraints]
     # pre-allocated memory
-    constraint_tmp = [zeros(c.nc) for c in constraints]
-    constraint_jacobian_state_tmp = [zeros(c.nc, t < H ? model[t].nx : model[H-1].ny) for (t, c) in enumerate(constraints)]
-    constraint_jacobian_action_tmp = [zeros(c.nc, t < H ? model[t].nu : 0) for (t, c) in enumerate(constraints)]
+    constraint_tmp = [zeros(c.num_constraint) for c in constraints]
+    constraint_jacobian_state_tmp = [zeros(c.num_constraint, t < H ? model[t].num_state : model[H-1].num_next_state) for (t, c) in enumerate(constraints)]
+    constraint_jacobian_action_tmp = [zeros(c.num_constraint, t < H ? model[t].num_action : 0) for (t, c) in enumerate(constraints)]
     data = constraint_data(model, constraints)
     AugmentedLagrangianCosts(costs, 
         data, 
@@ -54,8 +54,8 @@ function cost(obj::AugmentedLagrangianCosts, states, actions, parameters)
 
     for t = 1:H
         J += λ[t]' * c[t]
-        nc = obj.constraint_data.constraints[t].nc 
-        for i = 1:nc 
+        num_constraint = obj.constraint_data.constraints[t].num_constraint 
+        for i = 1:num_constraint 
             if a[t][i] == 1
                 J += 0.5 * ρ[t][i] * c[t][i]^2.0
             end
@@ -98,8 +98,8 @@ function augmented_lagrangian_update!(obj::AugmentedLagrangianCosts;
     H = length(c)
 
     for t = 1:H
-        nc = cons[t].nc 
-        for i = 1:nc 
+        num_constraint = cons[t].num_constraint 
+        for i = 1:num_constraint 
             λ[t][i] += ρ[t][i] * c[t][i]
             if i in cons[t].indices_inequality
                 λ[t][i] = max(0.0, λ[t][i])

src/backward_pass.jl

@@ -36,8 +36,8 @@ function backward_pass!(policy::PolicyData, problem::ProblemData;
     Qux = policy.action_value.hessian_action_state
 
     # terminal value function
-    P[T] .= gxx[T]
-    p[T] .=  gx[T]
+    P[H] .= gxx[H]
+    p[H] .=  gx[H]
 
     for t = H-1:-1:1
         # Qx[t] .= gx[t] + fx[t]' * p[t+1]

src/constraints.jl

@@ -2,10 +2,10 @@ struct Constraint{T}
     val 
     jacobian_state 
     jacobian_action
-    nc::Int
-    nx::Int 
-    nu::Int 
-    nw::Int
+    num_constraint::Int
+    num_state::Int 
+    num_action::Int 
+    num_parameter::Int
     val_cache::Vector{T} 
     jacobian_state_cache::Matrix{T}
     jacobian_action_cache::Matrix{T}
@@ -14,12 +14,12 @@ end
 
 Constraints{T} = Vector{Constraint{T}} where T
 
-function Constraint(f::Function, nx::Int, nu::Int; 
+function Constraint(f::Function, num_state::Int, num_action::Int; 
     indices_inequality::Vector{Int}=collect(1:0), 
-    nw::Int=0)
+    num_parameter::Int=0)
 
     #TODO: option to load/save methods
-    @variables x[1:nx], u[1:nu], w[1:nw]
+    @variables x[1:num_state], u[1:num_action], w[1:num_parameter]
 
     val = f(x, u, w)
     jacobian_state = Symbolics.jacobian(val, x)
@@ -29,13 +29,13 @@ function Constraint(f::Function, nx::Int, nu::Int;
     jacobian_state_func = eval(Symbolics.build_function(jacobian_state, x, u, w)[2])
     jacobian_action_func = eval(Symbolics.build_function(jacobian_action, x, u, w)[2])
 
-    nc = length(val) 
+    num_constraint = length(val) 
 
     return Constraint(
         val_func, 
         jacobian_state_func, jacobian_action_func,
-        nc, nx, nu, nw,  
-        zeros(nc), zeros(nc, nx), zeros(nc, nu), 
+        num_constraint, num_state, num_action, num_parameter,  
+        zeros(num_constraint), zeros(num_constraint, num_state), zeros(num_constraint, num_action), 
         indices_inequality)
 end
 
@@ -50,7 +50,7 @@ end
 
 function constraints!(violations, constraints::Constraints{T}, states, actions, parameters) where T
     for (t, con) in enumerate(constraints)
-        con.nc == 0 && continue
+        con.num_constraint == 0 && continue
         con.val(con.val_cache, states[t], actions[t], parameters[t])
         @views violations[t] .= con.val_cache
         fill!(con.val_cache, 0.0) # TODO: confirm this is necessary 
@@ -60,7 +60,7 @@ end
 function jacobian!(jacobian_states, jacobian_actions, constraints::Constraints{T}, states, actions, parameters) where T
     H = length(constraints)
     for (t, con) in enumerate(constraints)
-        con.nc == 0 && continue
+        con.num_constraint == 0 && continue
         con.jacobian_state(con.jacobian_state_cache, states[t], actions[t], parameters[t])
         @views jacobian_states[t] .= con.jacobian_state_cache
         fill!(con.jacobian_state_cache, 0.0) # TODO: confirm this is necessary

src/costs.jl

@@ -14,9 +14,9 @@ struct Cost{T}
     hessian_action_state_cache::Matrix{T}
 end
 
-function Cost(f::Function, nx::Int, nu::Int; nw::Int=0)
+function Cost(f::Function, num_state::Int, num_action::Int; num_parameter::Int=0)
     #TODO: option to load/save methods
-    @variables x[1:nx], u[1:nu], w[1:nw]
+    @variables x[1:num_state], u[1:num_action], w[1:num_parameter]
 
     val = f(x, u, w)
     gradient_state = Symbolics.gradient(val, x)
@@ -36,8 +36,8 @@ function Cost(f::Function, nx::Int, nu::Int; nw::Int=0)
         gradient_state_func, gradient_action_func, 
         hessian_state_state_func, hessian_action_action_func, hessian_action_state_func,
         zeros(1), 
-        zeros(nx), zeros(nu), 
-        zeros(nx, nx), zeros(nu, nu), zeros(nu, nx))
+        zeros(num_state), zeros(num_action), 
+        zeros(num_state, num_state), zeros(num_action, num_action), zeros(num_action, num_state))
 end
 
 Objective{T} = Vector{Cost{T}} where T

src/data/constraints.jl

@@ -10,9 +10,9 @@ end
 
 function constraint_data(model::Model, cons::Constraints) 
     H = length(cons)
-    c = [zeros(cons[t].nc) for t = 1:H]
-    cx = [zeros(cons[t].nc, t < H ? model[t].nx : model[H-1].ny) for t = 1:H]
-    cu = [zeros(cons[t].nc, model[t].nu) for t = 1:H-1]
+    c = [zeros(cons[t].num_constraint) for t = 1:H]
+    cx = [zeros(cons[t].num_constraint, t < H ? model[t].num_state : model[H-1].num_next_state) for t = 1:H]
+    cu = [zeros(cons[t].num_constraint, model[t].num_action) for t = 1:H-1]
     ConstraintsData(cons, c, cx, cu)
 end
 
@@ -27,9 +27,9 @@ function constraint_violation(constraint_data::ConstraintsData;
     H = length(constraints)
     max_violation = 0.0
     for t = 1:H
-        nc = constraints[t].nc 
+        num_constraint = constraints[t].num_constraint 
         ineq = constraints[t].indices_inequality
-        for i = 1:nc 
+        for i = 1:num_constraint 
             c = constraint_data.violations[t][i]
             cti = (i in ineq) ? max(0.0, c) : abs(c)
             max_violation = max(max_violation, cti)

src/data/model.jl

@@ -10,10 +10,10 @@ struct ModelData{T,X,U,W}
 end
 
 function model_data(dynamics::Vector{Dynamics{T}}) where T
-	jacobian_state = [zeros(d.ny, d.nx) for d in dynamics]
-    jacobian_action = [zeros(d.ny, d.nu) for d in dynamics]
-	jacobian_parameter = [zeros(d.ny, d.nw) for d in dynamics]
-    ModelData(model, jacobian_state, jacobian_action, jacobian_parameter)
+	jacobian_state = [zeros(d.num_next_state, d.num_state) for d in dynamics]
+    jacobian_action = [zeros(d.num_next_state, d.num_action) for d in dynamics]
+	jacobian_parameter = [zeros(d.num_next_state, d.num_parameter) for d in dynamics]
+    ModelData(dynamics, jacobian_state, jacobian_action, jacobian_parameter)
 end
 
 function reset!(data::ModelData)