A joint model of a 3D ball and a 2D catcher with free final time
In [1]:
# Plotting
%matplotlib inline
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.axes3d import Axes3D
# Numerics
import numpy as np
# CasADi
import casadi as ca
import casadi.tools as cat
In [2]:
# Simulation parameters
T_sim = 1.0 # simulation time
N_sim = 20 # time steps
h = T_sim/N_sim # time-step length
N_rk = 1 # number of RK4 steps (for each time-step)
# Model parameters
nx = 12; nu = 3
mu = 10.0/12; g0 = 9.81
a = 7.5; b = 2.5
wmin = -2 * ca.pi
wmax = 2 * ca.pi
thetamin = -ca.pi
thetamax = ca.pi
# Initial conditions
# Catcher
x_c0 = 7
y_c0 = 3
vx_c0 = -2
vy_c0 = 0
phi0 = 6./5 * ca.pi
# Ball
x_b0 = 0
y_b0 = 0
z_b0 = 0
vx_b0 = 5
vy_b0 = 6
vz_b0 = 10
# Create symbolic variables
states = cat.struct_symSX(["x_c","y_c",
"vx_c","vy_c","phi",
"x_b","y_b","z_b",
"vx_b","vy_b","vz_b",
"T"])
[x_c,y_c,vx_c,vy_c,phi,
x_b,y_b,z_b,vx_b,vy_b,vz_b,T] = states[...]
controls = cat.struct_symSX(["u","w","theta"])
[u,w,theta] = controls[...]
# ODE right-hand side
rhs = cat.struct_SX(states)
rhs["x_c"] = vx_c * T
rhs["y_c"] = vy_c * T
rhs["vx_c"] = (u * (ca.cos(theta) * ca.cos(phi) - \
ca.sin(theta) * ca.sin(phi)) - mu * vx_c) * T
rhs["vy_c"] = (u * (ca.cos(theta) * ca.sin(phi) + \
ca.sin(theta) * ca.cos(phi)) - mu * vy_c) * T
rhs["phi"] = w * T
rhs["x_b"] = vx_b * T
rhs["y_b"] = vy_b * T
rhs["z_b"] = vz_b * T
rhs["vx_b"] = 0
rhs["vy_b"] = 0
rhs["vz_b"] = -g0 * T
rhs["T"] = 0
# Degrees of freedom for the optimizer
V = cat.struct_symSX([
(
cat.entry("X",repeat=N_sim+1,struct=states),
cat.entry("U",repeat=N_sim,struct=controls)
)
])
# Continuous-time dynamics
f = ca.SXFunction([states, controls], [rhs])
f.init()
# Discrete-time dynamics
F = ca.simpleRK(f, N_rk)
In [3]:
# Objective function
J = 1e1 * (V["X",N_sim,"x_c"] - V["X",N_sim,"x_b"]) ** 2 + \
1e1 * (V["X",N_sim,"y_c"] - V["X",N_sim,"y_b"]) ** 2
# Regularize u
J += 1e-3 * ca.sum_square(ca.veccat(V["U",:,"u"])) * h * V["X",N_sim,"T"]
# Regularize theta
J += 1e-3 * ca.sum_square(ca.veccat(V["U",:,"theta"])) * h * V["X",N_sim,"T"]
# Build non-linear constraints and running costs
g = []
lbg = []
ubg = []
for i in range(N_sim):
# Multiple shooting
g.append(V["X",i+1] - \
F([V["X",i], V["U",i], h])[0])
lbg.append(ca.DMatrix.zeros(nx))
ubg.append(ca.DMatrix.zeros(nx))
# Control constraints
g.append(V["U",i,"u"] - a - b * ca.cos(V["U",i,"theta"]))
lbg.append(-ca.inf)
ubg.append(ca.DMatrix([0]))
# Encourage looking at the ball
d = ca.veccat((ca.cos(V["X",i,"phi"]), ca.sin(V["X",i,"phi"])))
r = ca.veccat((V["X",i,"x_b"]-V["X",i,"x_c"], V["X",i,"y_b"]-V["X",i,"y_c"]))
J -= 1e-2 * ca.mul(d.T,r)
g = ca.vertcat(g)
lbg = ca.veccat(lbg)
ubg = ca.veccat(ubg)
In [4]:
# Create NLP
nlp = ca.SXFunction(ca.nlpIn(x=V), ca.nlpOut(f=J,g=g))
nlp.init()
# Create solver
solver = ca.NlpSolver("ipopt",nlp)
solver.init()
# Set non-linear bounds
solver.setInput(lbg,"lbg")
solver.setInput(ubg,"ubg")
# Initialize simple bounds
lbx = V(-ca.inf)
ubx = V(ca.inf)
# ---------------------------------------------------------------------
# State limits
# ---------------------------------------------------------------------
# Initial condition of the catcher
lbx["X",0,"x_c"] = ubx["X",0,"x_c"] = x_c0
lbx["X",0,"y_c"] = ubx["X",0,"y_c"] = y_c0
lbx["X",0,"vx_c"] = ubx["X",0,"vx_c"] = vx_c0
lbx["X",0,"vy_c"] = ubx["X",0,"vy_c"] = vy_c0
lbx["X",0,"phi"] = ubx["X",0,"phi"] = phi0
# Initial condition of the ball
lbx["X",0,"x_b"] = ubx["X",0,"x_b"] = x_b0
lbx["X",0,"y_b"] = ubx["X",0,"y_b"] = y_b0
lbx["X",0,"z_b"] = ubx["X",0,"z_b"] = z_b0
lbx["X",0,"vx_b"] = ubx["X",0,"vx_b"] = vx_b0
lbx["X",0,"vy_b"] = ubx["X",0,"vy_b"] = vy_b0
lbx["X",0,"vz_b"] = ubx["X",0,"vz_b"] = vz_b0
# Require time to be strictly positive (>= 0.1)
lbx["X",:,"T"] = 0.1
# Simulation ends when the ball touches the ground
lbx["X",N_sim,"z_b"] = ubx["X",N_sim,"z_b"] = 0
# ---------------------------------------------------------------------
# Control limits
# ---------------------------------------------------------------------
# u >= 0 (the upper bound is set by non-linear constraints)
lbx["U",:,"u"] = 0
# wmin <= w <= wmax
lbx["U",:,"w"] = wmin
ubx["U",:,"w"] = wmax
# thetamin <= theta <= thetamax
lbx["U",:,"theta"] = thetamin
ubx["U",:,"theta"] = thetamax
# Set simple bounds
solver.setInput(lbx,"lbx")
solver.setInput(ubx,"ubx")
# Solve
with cat.nice_stdout():
solver.evaluate()
#solver.evaluate()
sol = V(solver.getOutput())
In [7]:
#%matplotlib inline
# Plot the trajectory in 3D
ts = np.linspace(0, float(sol["X",N_sim,"T"]), N_sim+1)
fig = plt.figure(figsize=(12,8))
ax = fig.add_subplot(111,projection='3d')
ax.set_zlim(0,6)
ax.scatter3D(sol["X",:,"x_b"],
sol["X",:,"y_b"],
sol["X",:,"z_b"],
c='b', lw='0.1', label='ball'
)
ax.scatter3D(sol["X",:,"x_c"],
sol["X",:,"y_c"],
ca.DMatrix.zeros(ca.veccat(sol["X",:,"x_c"]).size()),
c='r', lw='0.1', label='catcher'
)
ax.view_init(30, 20)
fig.show()
In [8]:
%matplotlib inline
# Plot the controls
fig, ax = plt.subplots(figsize=(10,4))
ax.grid(True)
ax.step(ts, ca.veccat((0,sol["U",:,"u"])), label='u')
ax.step(ts, ca.veccat((0,sol["U",:,"w"])), label='w')
ax.step(ts, ca.veccat((0,sol["U",:,"theta"])), label='theta')
ax.legend(loc=1);
In [9]:
# Plot the trajectory in 2D
fig, ax = plt.subplots(figsize=(10,10))
ax.grid(True)
ax.scatter(ca.vertcat(sol["X",:,"x_b"]),
ca.vertcat(sol["X",:,"y_b"]),
s=30, lw=0.02, c='b', label='ball'
)
ax.scatter(ca.vertcat(sol["X",:,"x_c"]),
ca.vertcat(sol["X",:,"y_c"]),
s=30, lw=0.02, c='c', label='catcher'
)
# Plot the arrows
Q = ax.quiver(ca.vertcat(sol["X",:,"x_c"]),
ca.vertcat(sol["X",:,"y_c"]),
ca.cos(ca.vertcat(sol["X",:,"phi"])),
ca.sin(ca.vertcat(sol["X",:,"phi"])),
units='xy', angles='xy',
scale=2.5, width=0.04,
headwidth=4, headlength=6, headaxislength=5,
color='m', lw='0.1'
)
qk = ax.quiverkey(Q, 0.04, 0.88, 1, 'gaze direction', labelpos='E',
fontproperties={'size': 12}
)
ax.legend(loc=2);
In [ ]:
In [ ]:
In [ ]:
In [ ]: