Trajectory Optimization using Pseudospectral collocation

Hi everyone!

I am currently trying to solve a simple trajectory optimization problem using pseudospectral collocation or orthogonal collocation.

I have written a simple UAV simulation using JuMP.jl with a moving obstacle where the UAV has to reach the final state in minimum travel time. I discretized the UAV dynamics explicitly using the trapezoidal method.

using JuMP, Ipopt

# Definition of global constants which are used throughout the model.
const gₒ     = 9.81  # Acceleration due to gravity (m/s²)
const m = 0.625  # Mass of UAV (kg)
const I_x = 2.3e-3  # Moment of inertia along x-axis (kg·m²)
const I_y = 2.3e-3  # Moment of inertia along y-axis (kg·m²)
const I_z = 2.3e-3  # Moment of inertia along z-axis (kg·m²)

# Initial conditions defined as constants for starting the optimization
const x_s = 0  # Initial x position (m)
const y_s = 0  # Initial y position (m)
const z_s = 30.0  # Initial z position (m)
const v_x_s = 0.0  # Initial x velocity (m/s)
const v_y_s = 0.0  # Initial y velocity (m/s)
const v_z_s = 0.0  # Initial z velocity (m/s)
const ϕ_s = 0.0  # Initial roll angle (rad)
const θ_s = 0.0  # Initial pitch angle (rad)
const ψ_s = 0.0  # Initial yaw angle (rad)
const p_s = 0.0  # Initial roll rate (rad/s)
const q_s = 0.0  # Initial pitch rate (rad/s)
const r_s = 0.0  # Initial yaw rate (rad/s)

# Final conditions defined as constants for ending the optimization
const x_t = 50.0  # Final x position (m)
const y_t = 50.0  # Final y position (m)
const z_t = 0.0  # Final z position (m)
const v_x_t = 0.0  # Final x velocity (m/s)
const v_y_t = 0.0  # Final y velocity (m/s)
const v_z_t = 0.0  # Final z velocity (m/s)
const ϕ_t = 0.0  # Final roll angle (rad)
const θ_t = 0.0  # Final pitch angle (rad)
const ψ_t = 0.0  # Final yaw angle (rad)
const p_t = 0.0  # Final roll rate (rad/s)
const q_t = 0.0  # Final pitch rate (rad/s)
const r_t = 0.0  # Final yaw rate (rad/s)

const epsilon = 0.01  # Minimum time step to avoid too small values


# Number of mesh points used in the optimization, also known as knots
const nₒ = 400

# Model initialization with Ipopt optimizer, passing user options as attributes
model = Model(optimizer_with_attributes(Ipopt.Optimizer))

# Declaration of variables within the model, each defined over the range from 1 to n (number of knots)
@variables(model, begin
   0 ≤x[1:nₒ]  # x position (m)
   0 ≤y[1:nₒ]  # y position (m)
   0 ≤z[1:nₒ] # z position (m)
    -5 ≤ v_x[1:nₒ] ≤ 5  # x velocity (m/s)
    -5 ≤ v_y[1:nₒ] ≤ 5  # y velocity (m/s)
    -5 ≤ v_z[1:nₒ] ≤ 5  # z velocity (m/s)
    -π/2 ≤ ϕ[1:nₒ] ≤ π/2  # roll angle (rad)
    -π/2 ≤ θ[1:nₒ] ≤ π/2  # pitch angle (rad)
    -π/2 ≤ ψ[1:nₒ] ≤ π/2  # yaw angle (rad)
    -2 ≤ p[1:nₒ] ≤ 2  # roll rate (rad/s)
    -2 ≤ q[1:nₒ] ≤ 2  # pitch rate (rad/s)
    -2 ≤ r[1:nₒ] ≤ 2  # yaw rate (rad/s)
    -21 ≤ U₁[1:nₒ] ≤ 21  # control input 1 (N)
    -0.67 ≤ U₂[1:nₒ] ≤ 0.67  # control input 2 (N·m)
    -0.67 ≤ U₃[1:nₒ] ≤ 0.67  # control input 3 (N·m)
    -0.11 ≤ U₄[1:nₒ] ≤ 0.11  # control input 4 (N·m)
    Δt[1:nₒ] >= epsilon  # time step (s)
end)

# Fixing initial conditions to match specified start points
fix(x[1], x_s; force = true)
fix(y[1], y_s; force = true)
fix(z[1], z_s; force = true)
fix(v_x[1], v_x_s; force = true)
fix(v_y[1], v_y_s; force = true)
fix(v_z[1], v_z_s; force = true)
fix(ϕ[1], ϕ_s; force = true)
fix(θ[1], θ_s; force = true)
fix(ψ[1], ψ_s; force = true)
fix(p[1], p_s; force = true)
fix(q[1], q_s; force = true)
fix(r[1], r_s; force = true)

# Fixing final conditions to ensure the optimization targets are met
fix(x[nₒ], x_t; force = true)
fix(y[nₒ], y_t; force = true)
fix(z[nₒ], z_t; force = true)
fix(v_x[nₒ], v_x_t; force = true)
fix(v_y[nₒ], v_y_t; force = true)
fix(v_z[nₒ], v_z_t; force = true)
fix(ϕ[nₒ], ϕ_t; force = true)
fix(θ[nₒ], θ_t; force = true)
fix(ψ[nₒ], ψ_t; force = true)
fix(p[nₒ], p_t; force = true)
fix(q[nₒ], q_t; force = true)
fix(r[nₒ], r_t; force = true)

# Initial guess strategy using linear interpolation between start and end conditions
function set_initial_guess(model, nₒ)
    # Interpolation for state variables
    x_vals = LinRange(x_s, x_t, nₒ)
    y_vals = LinRange(y_s, y_t, nₒ)
    z_vals = LinRange(z_s, z_t, nₒ)
    v_x_vals = LinRange(v_x_s, v_x_t, nₒ)
    v_y_vals = LinRange(v_y_s, v_y_t, nₒ)
    v_z_vals = LinRange(v_z_s, v_z_t, nₒ)
    ϕ_vals = LinRange(ϕ_s, ϕ_t, nₒ)
    θ_vals = LinRange(θ_s, θ_t, nₒ)
    ψ_vals = LinRange(ψ_s, ψ_t, nₒ)
    p_vals = LinRange(p_s, p_t, nₒ)
    q_vals = LinRange(q_s, q_t, nₒ)
    r_vals = LinRange(r_s, r_t, nₒ)

    # Set initial guesses for state variables
    set_start_value.(model[:x], x_vals)
    set_start_value.(model[:y], y_vals)
    set_start_value.(model[:z], z_vals)
    set_start_value.(model[:v_x], v_x_vals)
    set_start_value.(model[:v_y], v_y_vals)
    set_start_value.(model[:v_z], v_z_vals)
    set_start_value.(model[:ϕ], ϕ_vals)
    set_start_value.(model[:θ], θ_vals)
    set_start_value.(model[:ψ], ψ_vals)
    set_start_value.(model[:p], p_vals)
    set_start_value.(model[:q], q_vals)
    set_start_value.(model[:r], r_vals)

    # Assuming control inputs approach a stable condition, possibly zero
    # Interpolating from a middle range to zero if not specifically stated
    U₁_vals = LinRange(0, 0, nₒ)  # Assuming controls stabilize around zero or any nominal value
    U₂_vals = LinRange(0, 0, nₒ)
    U₃_vals = LinRange(0, 0, nₒ)
    U₄_vals = LinRange(0, 0, nₒ)
    
    # Set initial guesses for control variables
    set_start_value.(model[:U₁], U₁_vals)
    set_start_value.(model[:U₂], U₂_vals)
    set_start_value.(model[:U₃], U₃_vals)
    set_start_value.(model[:U₄], U₄_vals)
end

# Call the function after model and variable declarations
set_initial_guess(model, nₒ)

# Constraints applying the chosen numerical integration method to the motion equations
for j in 2:nₒ
    i = j - 1  # index of previous knot
        @constraint(model, x[j] == x[i] + 0.5 * Δt[i] * (v_x[j] + v_x[i]))
        @constraint(model, y[j] == y[i] + 0.5 * Δt[i] * (v_y[j] + v_y[i]))
        @constraint(model, z[j] == z[i] + 0.5 * Δt[i] * (v_z[j] + v_z[i]))
        @constraint(model, v_x[j] == v_x[i] + 0.5 * Δt[i] * (((cos(ϕ[j]) * sin(θ[j]) * cos(ψ[j]) + sin(ϕ[j]) * sin(ψ[j])) * U₁[j] / m + (cos(ϕ[i]) * sin(θ[i]) * cos(ψ[i]) + sin(ϕ[i]) * sin(ψ[i])) * U₁[i] / m)))
        @constraint(model, v_y[j] == v_y[i] + 0.5 * Δt[i] * (((cos(ϕ[j]) * sin(θ[j]) * sin(ψ[j]) - sin(ϕ[j]) * cos(ψ[j])) * U₁[j] / m + (cos(ϕ[i]) * sin(θ[i]) * sin(ψ[i]) - sin(ϕ[i]) * cos(ψ[i])) * U₁[i] / m)))
        @constraint(model, v_z[j] == v_z[i] + 0.5 * Δt[i] * ((-gₒ + (cos(ϕ[j]) * cos(θ[j])) * U₁[j] / m + (cos(ϕ[i]) * cos(θ[i])) * U₁[i] / m)))
        @constraint(model, ϕ[j] == ϕ[i] + 0.5 * Δt[i] * (p[j] + p[i]))
        @constraint(model, θ[j] == θ[i] + 0.5 * Δt[i] * (q[j] + q[i]))
        @constraint(model, ψ[j] == ψ[i] + 0.5 * Δt[i] * (r[j] + r[i]))
        @constraint(model, p[j] == p[i] + 0.5 * Δt[i] * (((U₂[j] + q[j] * (I_y - I_z)) / I_x + (U₂[i] + q[i] * (I_y - I_z)) / I_x)))
        @constraint(model, q[j] == q[i] + 0.5 * Δt[i] * (((U₃[j] + r[j] * (I_z - I_x)) / I_y + (U₃[i] + r[i] * (I_z - I_x)) / I_y)))
        @constraint(model, r[j] == r[i] + 0.5 * Δt[i] * (((U₄[j] + p[j] * (I_x - I_y)) / I_z + (U₄[i] + p[i] * (I_x - I_y)) / I_z)))
end

# Safety distance constant
const safe_distance = 1

# Obstacle Movement Parameters
const obs_start_time = 0
const obs_move_to_landing_start = 1
const obs_landing_start = 2
const obs_move_to_final_start = 6
const obs_final_start = 10

# Obstacle Positions
const obs_initial_pos = (50, 50, 1)  # (x, y, z)
const obs_landing_pos = (50, 50, 0)
const obs_final_pos = (0, 0, 30)



# Function to calculate linear interpolation between two points
function interpolate_position(start_pos, end_pos, start_time, end_time, current_time)
    if current_time < start_time
        return start_pos
    elseif current_time > end_time
        return end_pos
    else
        frac = (current_time - start_time) / (end_time - start_time)
        return start_pos + (end_pos - start_pos) * frac
    end
end

# Function to define the obstacle's x, y, and z position at a given time
function obstacle_position(t)
    x, y, z = 0, 0, 0

    if t < obs_move_to_landing_start
        x, y, z = obs_initial_pos
    elseif t < obs_landing_start
        x = interpolate_position(obs_initial_pos[1], obs_landing_pos[1], obs_move_to_landing_start, obs_landing_start, t)
        y = interpolate_position(obs_initial_pos[2], obs_landing_pos[2], obs_move_to_landing_start, obs_landing_start, t)
        z = interpolate_position(obs_initial_pos[3], obs_landing_pos[3], obs_move_to_landing_start, obs_landing_start, t)
    elseif t < obs_move_to_final_start
        x, y, z = obs_landing_pos
    elseif t < obs_final_start
        x = interpolate_position(obs_landing_pos[1], obs_final_pos[1], obs_move_to_final_start, obs_final_start, t)
        y = interpolate_position(obs_landing_pos[2], obs_final_pos[2], obs_move_to_final_start, obs_final_start, t)
        z = interpolate_position(obs_landing_pos[3], obs_final_pos[3], obs_move_to_final_start, obs_final_start, t)
    else
        x, y, z = obs_final_pos
    end

    return (x, y, z)
end

# Collision avoidance constraints
for i in 1:nₒ
    obs_x, obs_y, obs_z = obstacle_position(i * 0.1)
    @constraint(model, (x[i] - obs_x)^2 + (y[i] - obs_y)^2 + (z[i] - obs_z)^2 >= safe_distance^2)
end

# Objective function to minimize the total time taken to reach the final position
#@objective(model, Min, sum(Δt[i] for i in 1:nₒ))
@objective(model, Min, sum(Δt))

# Running the optimization
optimize!(model)

if termination_status(model) == MOI.OPTIMAL
    println("Optimal solution found.")
elseif termination_status(model) == MOI.LOCALLY_SOLVED
    println("Locally optimal solution found.")
elseif termination_status(model) == MOI.INFEASIBLE
    println("Problem is infeasible.")
elseif termination_status(model) == MOI.DUAL_INFEASIBLE
    println("Problem is dual infeasible.")
elseif termination_status(model) == MOI.UNBOUNDED
    println("Problem is unbounded.")
else
    println("No optimal solution found: ", termination_status)
end
  • Is it possible to use pseudospectral collocation or orthogonal collocation with JuMP.jl to solve this problem?
  • If yes, can you please give a simple example?
  • Are there better packages that are more suitable for this kind of problem?

Thank you!

I didn’t look at your code, but yes it’s generally possible to use JuMP for collocation. Did something go wrong with your formulation? It would help to describe a specific issue you encountered.

There is an example of rocket control in the JuMP documentation. Perhaps hard to find since it’s buried under “nonlinear.”

InfiniteOpt.jl is a wrapper for JuMP that makes such problems easy to deal with. In my experience it’s great. There are several others I’m less familiar with, e.g. OptimalControl.jl.

Thanks for your reply. Nothing wrong with my current formulation. I showed the code because I was hoping that someone could help me reformulate the problem from using the explicit trapezoidal discretization to pseudospectral collocation.

I have read the documentation and the examples, but no mention of collocation method there.

Hi @DTdev, welcome to the forum.

There is nothing built-in to JuMP to help with pseudospectral collocation.

InfiniteOpt is a good choice: Hovercraft Path Planning · InfiniteOpt.jl

@pulsipher might have other suggestions.

1 Like

There is nothing impeding you from modeling your problem by using JuMP. From your previous replies and by scrolling through your code, I can see that you are interpolating the trajectory with linear functions, thus your current implementation is not pseudospectral.

  • First, pick your orthogonal basis functions. Let us say you chose a Fourier series. I would define one or more Fourier series function(s) where the Fourier coefficients are JuMP variables. As long as your control and/or trajectories are written as a function of this (these) Fourier function(s), you are in the right direction.

  • I would consider other quadrature rules for the integration part. There are many quadrature rules typically employed in the field of numerical optimal control. Look at Gaussian Quadrature rules. I like this Julia library: FastGaussQuadrature.jl.

  • Finally, I recommend that you start with a “global” approach. Do not discretize the whole interval and try to fit many orthogonal functions in between them. Try to use only one function to approximate the whole control and/or trajectory (global orthogonal collocation). There is still going to be a discretization of the whole interval, but this one will be related to the numerical integration aspect of your problem.

See if you can find the book by J.T. Betts on Numerical Optimal Control (" Practical Methods for Optimal Control Using Nonlinear Programming"). I hope this helps.

2 Likes

Thanks!

Thanks for your reply! I will try to do that.

I had converted the Casadi example of pseudospectral collocation to JuMP a while back but it does not have final time free. So you will have to adapt it a bit.

using JuMP
using Ipopt
using FastGaussQuadrature
using Polynomials
using GLMakie

# d: Number of points for interpolation. Degree of interpolating polynomial = d-1
d = 4
x, w = gausslegendre(d-1)

# Collocation points
tau_root = [0; (1.0 .+ x) ./ 2]

# Coefficients of the collocation equation
C = zeros(d, d)

# Coefficients of continuity equation
D = zeros(d)

# Coefficients of quadrature function
B = zeros(d)

# Construct polynomial basis
Lj::Polynomial = Polynomial([1])
for j in 1:d
    # Construct lagrange polynomials to get the polynomial basis at the collocation point
    global tau_root, C, D, B, Lj
    Lj = Polynomial([1])
    for r in 1:d
        if r != j
            Lj = Lj*Polynomial([-tau_root[r], 1]) / (tau_root[j] - tau_root[r])
        end
    end

    # Evaluate the polynomial at the end of interval
    D[j] = Lj(1.0)

    # Evaluate the time derivatives at the collocation point
    Ljd = derivative(Lj)
    for r = 1:d
        C[j, r] = Ljd(tau_root[r])
    end

    # Evaluate the integral of the polynomial to get the coefficients of the quadrature function
    LjI = integrate(Lj)
    B[j] = LjI(1.0)
end

# Time Horizon
T = 10
N = 1000
h = T/N
time = range(start = 0, step = h, length = N)

ns = 2
nu = 1

# Declare Model variables
dc = Model(Ipopt.Optimizer)
@variable(dc, x[1:ns, 1:d, 1:N], start = 0)
@variable(dc, -1 <= u[1:nu, 1:N] <= 1, start = 0)

# State constraints
@constraint(dc, x[1, 1:d, 1:N] .>= -0.25)

# Initial state
@constraint(dc, x[1:ns, 1, 1] .== [0, 1])

# System model
dyn(x, u) = [((1 - x[2]^2)*x[1] - x[2] + u[1]), x[1]]

# Objective function
L(x, u) = x[1]^2 + x[2]^2 + u[1]^2

pdot = zeros(ns)
J = 0
Xend = zeros(ns)

for i in 1:N-1 
    global pdot, J, Xend
    Xend = x[1:ns, 1, i]*D[1]
    for r = 2:d
        pdot = zeros(ns)
        for j in 1:d 
            pdot = pdot + C[j, r] * x[1:ns, j, i]
        end
        @constraint(dc, pdot .== h*dyn(x[1:ns, r, i], u[1:nu, i]))
        J = J + h*B[r]*L(x[1:ns, r, i], u[1:nu, i])
        Xend = Xend + x[1:ns, r, i]*D[r]
    end
    @constraint(dc, Xend .== x[1:ns, 1, i+1])
end

@objective(dc, Min, J)

optimize!(dc)

un = u[1:nu, :]
xn = x[1:ns,1,:]

f1, ax1, s1 = stairs(time, value.(un)[1, :], color = :blue, label = "u")
lines!(ax1, time, value.(xn)[1, :], color = :black, label = "x1")
lines!(ax1, time, value.(xn)[2, :], color = :red, label = "x2")
axislegend()

The GPOPS paper and the references therein explain the method. The edition of the Betts’ book which I had referred did not have pseudospectral methods. It only mentioned trapezoidal and hermite-simpson. I do not know if pseudospectral methods have any advantage with respect to trapezoidal and hermite-simpson.

2 Likes

Thanks for the code!

I’m still new to Trajectory Optimization.
Is using pseudospectral collocation really has no notable advantage over hermite-simpson? I would prefer accuracy over compute time.

1 Like

Regarding InfiniteOpt.jl, among other reformulation techniques, it automates the use of orthogonal collocation over finite elements. Currently, it only natively supports Lobatto quadrature, but adding support for other polynomials would be straightforward (I just haven’t had time for it yet). See Quick Start · InfiniteOpt.jl.

Often in control problems, control variables are held constant over collocation points. InfiniteOpt.jl supports this on the master branch, but it hasn’t been formally released quite yet. See Quick Start · InfiniteOpt.jl.

The key advantage to using InfiniteOpt over JuMP for problems like this is to avoid the modelling errors commonly made when implementing collocation methods by hand. Moreover, the discretization scheme can be rapidly changed without having to restructure the model.

3 Likes