Skip to content

Instantly share code, notes, and snippets.

@alejandroastudillo
Last active November 15, 2019 13:59
Show Gist options
  • Save alejandroastudillo/3011dda8f5e94530583aaa7c0614dcdd to your computer and use it in GitHub Desktop.
Save alejandroastudillo/3011dda8f5e94530583aaa7c0614dcdd to your computer and use it in GitHub Desktop.
Problem-definition for NMPC toolchain for Serial Robots
import casadi as cs
import opengen as og

def rk4_explicit(fun,x,u,t_step):
    k1 = fun(x, u)
    k2 = fun(x + t_step*k1/2, u)
    k3 = fun(x + t_step*k2/2, u)
    k4 = fun(x + t_step*k3, u)
    return x + (t_step/6)*(k1 + 2*k2 + 2*k3 + k4);

# Define basic problem parameters
# ------------------------------------
ndof = 7
nv   = 2 # Number of path states
nw   = 1 # Number of path inputs

npos = 3
nrot = 3

nx   = 2*ndof + nv
nu   = ndof + nw
nl   = 2

pi = np.pi

q0 = [0, pi/6, 0, 4*pi/6, 0, -2*pi/6, -pi/2]

# Define OCP parameters
# ------------------------------------
# Sample time
dt          = 0.005
# Control horizon
N           = 16
# Tunnel radii
rho         = 1e-2
rho_rot     = 5e-2
# Quadratic cost weights
w_ds        = 1e1
w_pos       = 1e-1
w_rot       = 1e-1
w_u         = 4e-5
w_x         = 4e-5
w_l         = 1e-9
# Terminal cost weights
w_qN        = 1e-5
w_dqN       = 1e0
w_sN        = 1e-5
w_dsN       = 1e2
w_posN      = 1e2
w_rotN      = 1e2
# Linear cost weights
w_alphaL    = 1e1

# Define boundaries
# ------------------------------------
q_max = [np.inf, 126*pi/180, np.inf, 147*pi/180, np.inf, 117*pi/180, np.inf]
q_min = [-np.inf, -126*pi/180, -np.inf, -147*pi/180, -np.inf, -117*pi/180, -np.inf]
dq_max =  [0.8727, 0.8727, 0.8727, 0.8727, 0.8727, 0.8727, 0.8727]
dq_min =  [-0.8727, -0.8727, -0.8727, -0.8727, -0.8727, -0.8727, -0.8727]
s_max = [1]
s_min = [0]
ds_max = [np.inf]
ds_min = [0]

tau_max = [39, 39, 39, 39, 9, 9, 9]
tau_min = [-39, -39, -39, -39, -9, -9, -9]
dds_max = [np.inf]
dds_min = [-np.inf]

x_max = q_max + dq_max + s_max + ds_max # Concatenate q_max, dq_max, s_max, ds_max
x_min = q_min + dq_min + s_min + ds_min
u_max = tau_max + dds_max
u_min = tau_min + dds_min
l_max = [np.inf, np.inf]
l_min = [0, 0]


# Define decision variables and parameter x0
# ------------------------------------
dec_var = cs.MX.sym("dec_var", nx*(N+1) + nu*N + nl*(N-1),1)

x = cs.horzcat(dec_var[0:nx])
last_var_index = nx
u = cs.horzcat(dec_var[last_var_index:last_var_index+nu])
last_var_index += nu

for k in range(1, N+1):
    x = cs.horzcat(x, dec_var[last_var_index:nx+last_var_index])
    last_var_index += nx
    if k < N:
        if k == 1:
            l = dec_var[last_var_index:nl+last_var_index]
            last_var_index += nl
        else:
            l = cs.horzcat(l, dec_var[last_var_index:nl+last_var_index])
            last_var_index += nl
        u = cs.horzcat(u, dec_var[last_var_index:last_var_index+nu])
        last_var_index += nu

x0 = cs.MX.sym('x0', nx)


# Import robot functions
# ------------------------------------

expressions = cs.Function.load('robot_fncs/kin3_expressions_noG_simp.casadi').expand()
augstate = cs.SX.sym('augstate',nx)
notaugstate = augstate[0:ndof]
auginput = cs.SX.sym('auginput',nu)
ode_aug, ee_pos, ee_rot_n, ee_rot_s, ee_rot_a = expressions(augstate,auginput)
ode_aug     = cs.Function('ode_aug',[augstate, auginput],[ode_aug]).expand()
ee_pos      = cs.Function('ee_pos',[notaugstate],[ee_pos]).expand()
ee_rot_n    = cs.Function('ee_rot_n',[notaugstate],[ee_rot_n]).expand()
ee_rot_s    = cs.Function('ee_rot_s',[notaugstate],[ee_rot_s]).expand()
ee_rot_a    = cs.Function('ee_rot_a',[notaugstate],[ee_rot_a]).expand()
# Import path functions
# ------------------------------------
path_pos    = cs.Function.load('path_fncs/path_pos_simp.casadi').expand()
path_rot_n  = cs.Function.load('path_fncs/path_rot_n_simp.casadi').expand()
path_rot_s  = cs.Function.load('path_fncs/path_rot_s_simp.casadi').expand()
path_rot_a  = cs.Function.load('path_fncs/path_rot_a_simp.casadi').expand()
path_sdot   = cs.Function.load('path_fncs/path_sdot_smooth_01.casadi')

# Define errors
# ------------------------------------
xi = cs.SX.sym('xi',nx)
va = cs.SX.sym('va',3)
vb = cs.SX.sym('vb',3)
cross     = cs.Function('cross',[va, vb],[cs.vertcat(va[1]*vb[2] - va[2]*vb[1], va[2]*vb[0] - va[0]*vb[2],va[0]*vb[1] - va[1]*vb[0])]).expand()
pos_error = cs.Function('pos_error', [xi], [ee_pos(xi[0:ndof]) - path_pos(xi[nx-2])],['xi'],['pos_error']).expand()
rot_error = cs.Function('rot_error', [xi], [0.5*(cross(path_rot_n(xi[nx-2]), ee_rot_n(xi[0:ndof])) + cross(path_rot_s(xi[nx-2]), ee_rot_s(xi[0:ndof])) + cross(path_rot_a(xi[nx-2]), ee_rot_a(xi[0:ndof])))],['xi'],['rot_error']).expand()

# Define cost
# ------------------------------------
cost = 0
for k in range(0, N):
    u_k = u[:,k]
    x_k = x[:,k]

    cost += w_ds*(x_k[nx-1] - path_sdot(x_k[nx-2]))**2 \
            + w_pos*(cs.sumsqr(pos_error(x_k))) \
            + w_rot*(cs.sumsqr(rot_error(x_k))) \
            + w_u*(cs.sumsqr(u_k))

    if k == 0:
        cost += w_x*(cs.sumsqr(x_k))
    else:
        l_k = l[:,k-1]
        sum_l_k = 0
        for j in range(nl):
            sum_l_k += l_k[j]
        cost += w_l*(cs.sumsqr(l_k)) + w_x*(cs.sumsqr(x_k)) + w_alphaL*sum_l_k

x_N  = x[:,N]
q_N  = x_N[0:ndof]
dq_N = x_N[ndof:2*ndof]
s_N  = x_N[2*ndof]
ds_N = x_N[nx-1]

cost += w_qN*cs.sumsqr(q_N) + w_dqN*cs.sumsqr(dq_N) + w_sN*(1-s_N)**2 + w_dsN*cs.sumsqr(ds_N) + w_posN*cs.sumsqr(pos_error(x_N)) + w_rotN*cs.sumsqr(rot_error(x_N))


cost_fun = cs.Function('cost_fun',[dec_var],[cost]).expand()

cost = cost_fun(dec_var)
res_cost = cost_fun(initial_guess_primal)


# Define boundary constraints
# ------------------------------------
dec_var_min = x_min + u_min
for k in range(1, N):
    dec_var_min += x_min + l_min + u_min
dec_var_min += x_min


dec_var_max = x_max + u_max
for k in range(1, N):
    dec_var_max += x_max + l_max + u_max
dec_var_max += x_max

bounds = og.constraints.Rectangle(dec_var_min, dec_var_max)

# Define multiple shooting constraints
# ------------------------------------
ms_constraints = x[:,0] - x0

xs = cs.SX.sym("x",x.shape[0])
us = cs.SX.sym("u",u.shape[0])

rk4_explicit_casadi = cs.Function("rk4_explicit",[xs,us],[rk4_explicit(ode_aug,xs,us,dt)])

rk4_explicit_map = rk4_explicit_casadi.map(N, "openmp")
ms_constraints = cs.vec(cs.horzcat(x[:,0] - x0, x[:,1:]-rk4_explicit_map(x[:,:-1],u)))

set_y = og.constraints.Zero()
#set_y = og.constraints.BallInf(None, 1e-8)
ms_const_size = ms_constraints.numel()

# Define tunnel constraints
# ------------------------------------
tunnel_constraints_pos = cs.sumsqr(pos_error(x[:,1])) - l[0,0]
tunnel_constraints_rot = cs.sumsqr(rot_error(x[:,1])) - l[1,0]
for k in range(2, N):
    tunnel_constraints_pos = cs.vertcat(tunnel_constraints_pos, cs.sumsqr(pos_error(x[:,k])) - l[0,k-1])
    tunnel_constraints_rot = cs.vertcat(tunnel_constraints_rot, cs.sumsqr(rot_error(x[:,k])) - l[1,k-1])

tunnelpos_const_size = tunnel_constraints_pos.numel()
tunnelrot_const_size = tunnel_constraints_rot.numel()

set_y_pos = og.constraints.Ball2(None, rho**2)
set_y_rot = og.constraints.Ball2(None, rho_rot**2)

tunnel_fun = cs.Function('tunnel_fun',[dec_var],[tunnel_constraints_pos]).expand()
res_tunnel = tunnel_fun(initial_guess_primal) - rho**2
tunnel_constraints_pos = tunnel_fun(dec_var)

tunnelrot_fun = cs.Function('tunnelrot_fun',[dec_var],[tunnel_constraints_rot]).expand()
res_tunnel_rot = tunnelrot_fun(initial_guess_primal) - rho_rot**2
tunnel_constraints_rot = tunnelrot_fun(dec_var)

# Define Augmented Lagrangian constraints
AL_constraints = cs.vertcat(ms_constraints, tunnel_constraints_pos, tunnel_constraints_rot) # size MS_const = 272, size Tunnel_pos = 15, size Tunnel_rot = 15
segment_ids = [271, 286, 301]
dim = 302
set_AL = og.constraints.CartesianProduct(dim, segment_ids, [set_y, set_y_pos, set_y_rot])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment