Implement prox for MPC problems

Author: tylerhanks

Project.toml

@@ -4,11 +4,7 @@ authors = ["tylerhanks <teh.hanks@gmail.com>"]
 version = "0.1.0"
 
 [deps]
-Catlab = "134e5e36-593f-5add-ad60-77f754baafbe"
+AlgebraicOptimization = "a72ceada-00ec-4ad9-90d3-37b40eaed052"
 Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
-GLPK = "60bf3e95-4087-53dc-ae20-288a0d20c6a6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 SCS = "c946c3f1-0d1f-5ce8-9dea-7daa1f7e2d13"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

examples/single_agent.jl

@@ -0,0 +1,27 @@
+using AlgebraicControl
+using LinearAlgebra
+using Convex
+
+# Set up each agent's dynamics: x' = Ax + Bu
+dt = 0.1  # Discretization step size
+A = [1 dt 0 0; 0 1 0 0; 0 0 1 dt; 0 0 0 1]
+B = [0 0; dt 0; 0 0; 0 dt]
+Q = Matrix{Float64}(I(4))
+R = Matrix{Float64}(I(2))
+
+sys = LinearSystem(A, B)
+
+# Stage cost function: minimize control effort and deviation from origin
+stage_cost(u, x) = quadform(x, Q) + quadform(u, R)
+stage_constraints = Function[]
+# Create a multi-stage program for a horizon of N steps
+N = 100
+mpc_program = multi_stage_program(stage_cost, stage_constraints, sys, N)
+
+# Create a proxable MPC program
+proxable_mpc = ProxableMPCProgram(mpc_program)
+
+set_x0!(proxable_mpc, [5.0; 0.0; 5.0; 0.0])  # Initial state
+
+xf = prox(proxable_mpc, zeros(4))  # Proximal step towards target state [10; 0; 10; 0]
+

src/AlgebraicControl.jl

@@ -1,16 +1,88 @@
 module AlgebraicControl
 
-include("Categories.jl")
-include("MonoidalCats.jl")
-include("ParaCats.jl")
-include("BasicCats.jl")
-include("LensCats.jl")
-include("Problems.jl")
-include("Optimizers.jl")
-include("ParaProbs.jl")
-include("ParaOptimizers.jl")
-include("LQR.jl")
-include("ParaConvCat.jl")
-include("CMPC.jl")
-
-end # module AlgebraicControl
+include("ConvexPrograms.jl")
+
+export ConvexProgram, OpenConvexProgram, LinearSystem, single_stage_program, multi_stage_program,
+    ProxableMPCProgram, ProxableMPCProgram, set_x0!, prox, prox!
+
+using Convex
+using SCS
+using LinearAlgebra
+
+struct LinearSystem
+    A::Matrix{Float64}
+    B::Matrix{Float64}
+    LinearSystem(A::Matrix{Float64}, B::Matrix{Float64}) =
+        size(A, 1) == size(B, 1) ? new(A, B) : error("Inconsistent dimensions between A and B matrices.")
+end
+
+state_dim(sys::LinearSystem) = size(sys.A, 1)
+control_dim(sys::LinearSystem) = size(sys.B, 2)
+
+
+function (sys::LinearSystem)(x, u)
+    return sys.A * x + sys.B * u
+end
+
+function single_stage_program(stage_cost::Function, stage_constraints::Vector{Function}, sys::LinearSystem)::OpenConvexProgram
+    n = state_dim(sys)
+    m = control_dim(sys)
+    impl = (u1, x1, x2) -> ConvexProgram(
+        stage_cost(u1[1], x1),
+        vcat([c(u1[1], x1) for c in stage_constraints]..., x2 == sys(x1, u1[1]))
+    )
+    return OpenConvexProgram(n, n, [m], impl)
+end
+
+function multi_stage_program(stage_cost::Function, stage_constraints::Vector{Function}, sys::LinearSystem, n::Int)::OpenConvexProgram
+    one_step = single_stage_program(stage_cost, stage_constraints, sys)
+    return compose(one_step, n)
+end
+
+struct ProxableMPCProgram
+    control_vars::Vector{Variable}
+    input_var::Variable
+    output_var::Variable
+    program
+end
+
+ProxableMPCProgram(F::OpenConvexProgram) = begin
+    # Make variables of the right dimensions
+    x0 = Variable(dom(F))
+    xf = Variable(codom(F))
+    #y = Variable(codom(F))
+    us = [Variable(d) for d in F.internal_vars]
+    #γ = Variable(1)
+    mpc_program = F(us, x0, xf)
+    # Once Convex.jl is fixed, we can make γ and y real variables and fix! them
+    # in the prox function instead of rebuilding the program each time. But for now,
+    # we suffer....
+    proxable_program(y, γ) = minimize( # This is stupid but Convex.jl is broken...
+        mpc_program.objective + (1 / (2 * γ)) * LinearAlgebra.dot(xf - y, xf - y),
+        mpc_program.constraints
+    )
+    return ProxableMPCProgram(us, x0, xf, proxable_program)
+end
+
+function set_x0!(P::ProxableMPCProgram, x0_val::Vector{Float64})
+    fix!(P.input_var, x0_val)
+end
+
+function prox(P::ProxableMPCProgram, x::Vector{Float64}, γ=1.0)
+    #fix!(P.y, x)
+    #fix!(P.γ, γ)
+    prob = P.program(x, γ)
+    solve!(prob, SCS.Optimizer; silent=true)
+    return evaluate(P.output_var)
+end
+
+function prox!(y, P::ProxableMPCProgram, x::Vector{Float64}, γ=1.0)
+    #fix!(P.y, x)
+    #fix!(P.γ[], γ)
+    prob = P.program(x, γ)
+    solve!(prob, SCS.Optimizer; silent_solver=true)
+    copy!(y, evaluate(P.output_var))
+end
+
+
+end

src/ConvexPrograms.jl

@@ -0,0 +1,53 @@
+using Convex
+
+struct ConvexProgram
+    objective::Convex.AbstractExpr
+    constraints::Vector{Convex.Constraint}
+end
+
+struct OpenConvexProgram
+    dom_var::Int # Dimension of the domain decision variable
+    codom_var::Int # Dimension of the codomain decision variable
+    internal_vars::Vector{Int} # Dimensions of any internal variables
+    impl::Function # internal_vars × dom_var × codom_var ⇢ ConvexProgram
+end
+
+(F::OpenConvexProgram)(ps::Vector{Variable}, x::Variable, y::Variable) =
+    F.impl(ps, x, y)
+
+dom(F::OpenConvexProgram) = F.dom_var
+codom(F::OpenConvexProgram) = F.codom_var
+
+function close_program(F::OpenConvexProgram, input_val::Vector{Float64}, output_val::Vector{Float64}, ivs::Vector{Variable})::ConvexProgram
+    x_var = Variable(length(input_val))
+    fix!(x_var, input_val)
+    y_var = Variable(length(output_val))
+    fix!(y_var, output_val)
+    return F(ivs, x_var, y_var)
+end
+
+function compose(F::OpenConvexProgram, G::OpenConvexProgram)
+    @assert codom(F) == dom(G)
+    y = Variable(codom(F))
+    F_nparams = length(F.internal_vars)
+    impl = (ps, x, z) -> begin
+        if F_nparams > 0
+            FCB = F(ps[1:F_nparams], x, y)
+        else
+            FCB = F(Variable[], x, y)
+        end
+        GCB = G(ps[F_nparams+1:end], y, z)
+        return ConvexProgram(FCB.objective + GCB.objective, vcat(FCB.constraints, GCB.constraints))
+    end
+    return OpenConvexProgram(dom(F), codom(G), vcat(F.internal_vars, G.internal_vars), impl)
+end
+
+# Compose an endomorphism with itself n times
+function compose(F::OpenConvexProgram, n::Int)
+    @assert dom(F) == codom(F)
+    res = F
+    for i in 1:n-1
+        res = compose(res, F)
+    end
+    return res
+end

src/archive/AlgebraicControl.jl

@@ -0,0 +1,16 @@
+module AlgebraicControl
+
+include("Categories.jl")
+include("MonoidalCats.jl")
+include("ParaCats.jl")
+include("BasicCats.jl")
+include("LensCats.jl")
+include("Problems.jl")
+include("Optimizers.jl")
+include("ParaProbs.jl")
+include("ParaOptimizers.jl")
+include("LQR.jl")
+include("ParaConvCat.jl")
+include("CMPC.jl")
+
+end # module AlgebraicControl