def get_nd_dynamics(state, u, force, dim=2, linearize_theta=None):
"""Return state space dyanmics in n dimensions.
Keyword arguments:
state -- the state of the cube
u -- the input torque on the wheel
force -- the 8 forces acting on the wheel corners
dim -- the dimension of the state space (default 2)
linearize_theta -- this is used only because the qp solver needed it this way
"""
# half state length
hl = len(state) / 2
m_t = m_c + m_w # total mass
I_t = I_c + I_w # total inertia
# gravity
g = 9.81
# unpack the states
# x = state[0]
# y = state[1]
if linearize_theta is not None:
theta = linearize_theta
else:
theta = state[dim]
# alpha = state[hl-1]
# xdot = state[0+hl]
# ydot = state[1+hl]
theta_dot = state[dim + hl]
alpha_dot = state[-1]
# derivative vector
derivs = np.zeros_like(state)
derivs[0:hl] = state[hl:] # set velocities
# ballistic dynamics
derivs[0 +
hl] = (force[1] - force[2] + force[6] - force[5]) * cos(theta) - (
force[0] + force[3] - force[4] - force[7]) * sin(
theta) # forces along x
derivs[1 +
hl] = (force[1] - force[2] + force[6] - force[5]) * sin(theta) + (
force[0] + force[3] - force[4] -
force[7]) * cos(theta) - g * m_t # forces in y direction
# cube angle acceleration
derivs[dim + hl] = (-u[0] + F_w * alpha_dot - F_c * theta_dot) / I_c + (
-force[0] + force[1] - force[2] + force[3] - force[4] + force[5] -
force[6] + force[7]) * .5
# wheel acceleration
derivs[-1] = (u[0] * I_t + F_c * theta_dot * I_w -
F_w * alpha_dot * I_t) / (I_w * I_c)
return derivs
def cube_dynamics(state, u, force):
# Need to grab important parameters
M_c = 1.0 # self.M_c
M_w = 1.0 # self.M_w
M_t = M_c + M_w
I_c = 1.0 #self.I_c
I_w = 1.0 #self.I_w
I_t = I_c + I_w
# Distance from edge to center of cube
L_t = np.sqrt(2) #np.sqrt.(2*self.L)
# Assuming friction is 0 right now
F_c = 0.5
F_w = 0.5
g = 9.81 # self.g
# Relevant states are x,z,thetay, phi
x = state[0]
z = state[2]
thetay = state[4]
phi = state[6]
# Velocity States
xdot = state[7]
zdot = state[9]
thetaydot = state[11]
phidot = state[13]
# Setup the derivative of the state vector
derivs = np.zeros_like(state)
derivs[0:7] = state[7:]
# Ballistic Dynamics
derivs[7] = (force[1] - force[2] + force[6] - force[5]) * cos(thetay) - (
force[0] + force[3] - force[4] - force[7]) * sin(
thetay) # forces along x
derivs[9] = (force[1] - force[2] + force[6] - force[5]) * sin(thetay) + (
force[0] + force[3] - force[4] -
force[7]) * cos(thetay) - g # forces in y direction
# Back torque due to wheel
derivs[11] = (-u[0] + F_w * phidot - F_c * thetaydot) / I_c + (
-force[0] + force[1] - force[2] + force[3] - force[4] + force[5] -
force[6] + force[7]) * .5
# Wheel accel
derivs[13] = (u[0] * I_t + F_c * thetaydot * I_w -
F_w * phidot * I_t) / (I_w * I_c)
return derivs
def complimentarity_constraint(mp, state, force, dim=2):
theta = state[dim]
distances = get_corner_distances(state[:], dim)
s = sin(theta)
c = cos(theta)
z_0 = force[0] * c + force[1] * s
z_1 = -force[2] * s + force[3] * c
z_2 = -force[4] * c - force[5] * s
z_3 = force[6] * s - force[7] * c
xy_0 = -force[0] * s + force[1] * c
xy_1 = -force[2] * c - force[3] * s
xy_2 = force[4] * s - force[5] * c
xy_3 = force[6] * c + force[7] * s
mp.AddConstraint(xy_0 <= z_0 * mu)
mp.AddConstraint(xy_0 >= -z_0 * mu)
mp.AddConstraint(xy_1 <= z_1 * mu)
mp.AddConstraint(xy_1 >= -z_1 * mu)
mp.AddConstraint(xy_2 <= z_2 * mu)
mp.AddConstraint(xy_2 >= -z_2 * mu)
mp.AddConstraint(xy_3 <= z_3 * mu)
mp.AddConstraint(xy_3 >= -z_3 * mu)
val_0 = np.asarray([force[0], force[2], force[4], force[6]])
val_1 = np.asarray([force[1], force[3], force[5], force[7]])
mp.AddConstraint(val_0.dot(distances) <= complimentarity_constraint_thresh)
mp.AddConstraint(
val_0.dot(distances) >= -complimentarity_constraint_thresh)
mp.AddConstraint(val_1.dot(distances) <= complimentarity_constraint_thresh)
mp.AddConstraint(
val_1.dot(distances) >= -complimentarity_constraint_thresh)
def _inertia_matrix(self,
state: np.ndarray,
dtype: Any = float) -> np.ndarray:
"""Helper method to calculate the inertia matrix."""
M12 = self._mp * self._l * cos(state[1])
return np.array(
[[self._mc + self._mp, M12], [M12, self._mp * self._l**2]],
dtype=dtype)
def EvaluateDynamics(x, u):
# """
# Evaluate dynamics for system xdot = f(x, u).
# Inputs:
# x (Array[Continuous Variables]): State Variables Array
# x = [r, alpha, beta, Vx, Vy, Vz, m, phi, psi]
# u (Array[Continuous Variables]): Input (control) Variables Array
# u = [T, omega_phi, omega_psi]
# Returns:
# xdot (Array[Continuous Variables]): Dyanmics of state variables.
# """
g_m = 3.71 # Mars
I_sp = 302.39
r = x[0]
alpha = x[1]
beta = x[2]
Vx = x[3]
Vy = x[4]
Vz = x[5]
m = x[6]
phi = x[7]
psi = x[8]
rdot = Vx
alphadot = Vy / (r * drake_math.cos(beta))
betadot = Vz / r
T = u[0]
omega_phi = u[1]
omega_psi = u[2]
mdot = -T / (I_sp * g_m)
Vxdot = T * drake_math.sin(phi) / m - g_m + (Vy**2 + Vz**2) / r
Vydot = T * drake_math.cos(phi) * drake_math.cos(psi) / m - (
Vx * Vy) / r + (Vy * Vz * drake_math.tan(beta)) / r
Vzdot = T * drake_math.cos(phi) * drake_math.sin(psi) / m - (
Vx * Vz) / r - (Vy**2 * drake_math.tan(beta)) / r
xdot = np.array([
rdot, alphadot, betadot, Vxdot, Vydot, Vzdot, mdot, omega_phi,
omega_psi
])
return xdot
def nominal_dynamics(self, s0, s1, F, delta):
""" Return the residuals for the nominal dynamics of the vehicle.
Args:
s0 (1D array of ContinuousVariable): [description]
s1 (1D array of ContinuousVariable): [description]
F (1D array of ContinuousVariable): [description]
delta (ContinuousVariable): [description]
Returns:
[type]: [description]
"""
s = unpack_state_vector(s0)
snext = unpack_state_vector(s1)
F = unpack_force_vector(F)
# Trigonometric functions of variables.
cos_d = pmath.cos(delta)
sin_d = pmath.sin(delta)
cos_psi = pmath.cos(s["psi"])
sin_psi = pmath.sin(s["psi"])
# Net longitudinal and lateral forces.
Fx = F["r_long"] + F["f_long"] * cos_d - F["f_lat"] * sin_d
Fy = F["r_lat"] + F["f_long"] * sin_d + F["f_lat"] * cos_d
derivs = np.array([
s["psidot"] * s["ydot"] + (Fx / self.m), #xddot
-s["psidot"] * s["xdot"] + (Fy / self.m), #yddot
s["psidot"], #psidot
(1.0 / self.Iz) *
(self.lf * (F["f_long"] * sin_d + F["f_lat"] * cos_d) -
self.lr * F["r_lat"]), #psiddot
s["xdot"] * cos_psi - s["ydot"] * sin_psi, #Xdot
s["xdot"] * sin_psi + s["ydot"] * cos_psi #Ydot
])
# Apply Euler Integration to get Residuals
residuals = np.array([
snext["xdot"] - (s["xdot"] + self.dt * derivs[0]),
snext["ydot"] - (s["ydot"] + self.dt * derivs[1]),
snext["psi"] - (s["psi"] + self.dt * derivs[2]),
snext["psidot"] - (s["psidot"] + self.dt * derivs[3]),
snext["X"] - (s["X"] + self.dt * derivs[4]),
snext["Y"] - (s["Y"] + self.dt * derivs[5])
])
return residuals
def jacobian(self, t: float, state: np.ndarray, u: np.float) -> np.ndarray:
c_theta = cos(state[1])
s_theta = sin(state[1])
theta = state[1]
omega = state[3]
# d(xdd)/dx, d(xdd/dxd), d(thetadd)/dx, d(theta_dd)/d_xd are all zero
jac = np.zeros((4, 4))
jac[0][2] = 1.0 # d/dxdot (xdot)
jac[1][3] = 1.0 # d/dthetadot (theta_dot)
#d(xdd)/dtheta
term1 = 1.0 / (self._mc + self._mp * s_theta**2)
dterm1dtheta = -(2.0 * self._mp * s_theta * c_theta) * (term1**2)
term2 = (u + self._mp * s_theta * (self._l *
(omega**2) + self._g * c_theta))
dterm2dtheta = (self._mp * c_theta * self._l * (omega**2) -
self._g * s_theta)
jac[2][1] = term1 * dterm2dtheta + term2 * dterm1dtheta
#d(xdd)/dthetad
jac[2][3] = term1 * 2.0 * self._mp * s_theta * self._l * state[3]
#d(thetadd)/dtheta
term3 = 1.0 / self._l * term1
dterm3dtheta = 1.0 / self._l * dterm1dtheta
term4 = (-u * c_theta - self._mp * self._l *
(state[3]**2) * c_theta * s_theta -
(self._mc + self._mp) * self._g * s_theta)
dterm4dtheta = (u * s_theta -
self._mp * self._l * state[3] * cos(2 * state[1]) -
(self._mc + self._mp) * self._g * c_theta)
jac[3][1] = term3 * dterm4dtheta + term4 * dterm3dtheta
#d(thetadd)/dthetad
jac[3][3] = term3 * (-self._mp * self._l * state[3] * sin(2 * theta))
return jac
def get_corner_x_positions(state, dim=2, linearize_theta=None):
"""Return x position of the corners in the order [0,1,2,3] as a numpy array.
Keyword arguments:
state -- the state of the cube
dim -- the dimension of the state space (default 2)
"""
x = state[0]
if linearize_theta is not None:
theta = linearize_theta
else:
theta = state[dim]
offset = .5 * np.sqrt(2) * cos(np.pi / 4.0 + theta)
val = cos(theta)
pos_0 = x - offset
pos_1 = pos_0 + val
pos_2 = x + offset
pos_3 = pos_2 - val
return np.asarray([pos_0, pos_1, pos_2, pos_3])
def linearization(self, t: float, state: np.ndarray,
u: float) -> Tuple[np.ndarray, np.ndarray]:
"""
Calculate linearization about arbitrary state.
Arguments:
t: time (s)
state: state at time ``t``
u: force control signal in N
Returns:
linearized "A" and "B" matrices according to xdot = Ax + Bu
"""
B = np.zeros(4)
B[2] = 1.0 / (self._mc + self._mp * sin(state[1])**2)
B[3] = 1 / self._l * B[2] * (
1.0 /
(self._l *
(self._mc + self._mp * sin(state[1])**2))) * (-cos(state[1]))
return (self.jacobian(t, state, u), B)
def satellite_dynamics(self, state, u):
'''
state: x, y, vx, vy, h, w
'''
satellite_position = state[0:2]
derivs = np.zeros_like(state)
derivs[0:2] = state[2:4]
satellite_heading = state[4]
satellite_angvel = state[5]
m = self.mass; l = self.len; I = self.MOI; dep = self.dep;
# Underactuated system via dep array
derivs[2] = (u[0]*dep[0]+u[1]*dep[1]-u[2]*dep[2]-u[3]*dep[3])*sin(satellite_heading) / m # solves for dvx
derivs[3] = (-u[0]*dep[0]-u[1]*dep[1]+u[2]*dep[2]+u[3]*dep[3])*cos(satellite_heading) / m # solves for dvy
derivs[4] = satellite_angvel # solves for dh
derivs[5] = (-u[0]*dep[0]+u[1]*dep[1]-u[2]*dep[2]+u[3]*dep[3]) * l / (2.0*I) # solves for dw
return derivs
def test_scalar_math(self):
a = AD(1, [1., 0])
self._compare_scalar(a, a)
b = AD(2, [0, 1.])
# Arithmetic
self._compare_scalar(a + b, AD(3, [1, 1]))
self._compare_scalar(a + 1, AD(2, [1, 0]))
self._compare_scalar(1 + a, AD(2, [1, 0]))
self._compare_scalar(a - b, AD(-1, [1, -1]))
self._compare_scalar(a - 1, AD(0, [1, 0]))
self._compare_scalar(1 - a, AD(0, [-1, 0]))
self._compare_scalar(a * b, AD(2, [2, 1]))
self._compare_scalar(a * 2, AD(2, [2, 0]))
self._compare_scalar(2 * a, AD(2, [2, 0]))
self._compare_scalar(a / b, AD(1. / 2, [1. / 2, -1. / 4]))
self._compare_scalar(a / 2, AD(0.5, [0.5, 0]))
self._compare_scalar(2 / a, AD(2, [-2, 0]))
# Logical
af = a.value()
self._check_logical(lambda x, y: x == y, a, a, True)
self._check_logical(lambda x, y: x != y, a, a, False)
self._check_logical(lambda x, y: x < y, a, b, True)
self._check_logical(lambda x, y: x <= y, a, b, True)
self._check_logical(lambda x, y: x > y, a, b, False)
self._check_logical(lambda x, y: x >= y, a, b, False)
# Additional math
self._compare_scalar(a**2, AD(1, [2., 0]))
# Test autodiff overloads.
# See `math_overloads_test` for more comprehensive checks.
c = AD(0, [1., 0])
d = AD(1, [0, 1.])
self._compare_scalar(drake_math.sin(c), AD(0, [1, 0]))
self._compare_scalar(drake_math.cos(c), AD(1, [0, 0]))
self._compare_scalar(drake_math.tan(c), AD(0, [1, 0]))
self._compare_scalar(drake_math.asin(c), AD(0, [1, 0]))
self._compare_scalar(drake_math.acos(c), AD(np.pi / 2, [-1, 0]))
self._compare_scalar(drake_math.atan2(c, d), AD(0, [1, 0]))
self._compare_scalar(drake_math.sinh(c), AD(0, [1, 0]))
self._compare_scalar(drake_math.cosh(c), AD(1, [0, 0]))
self._compare_scalar(drake_math.tanh(c), AD(0, [1, 0]))
def constraint_evaluator1(z):
return np.array([(-(z[6] + z[7])*sin(z[2])/m)*dt + z[3] - z[8],
((z[6] + z[7])*cos(z[2]) - m*g/m)*dt + z[4] - z[9],
((z[6] - z[7])*r/I)*dt + z[5] - z[10]
])
def airplaneLongDynamics(self, state, u, uncertainty=False):
"""
Compute the longitudinal dynamics of the airplane, as in xdot = f(x,u). Can compute with/without uncertainty,
where uncertainty=True is used for simulation of system with modeling errors. Accepts either pydrake::symbolic
or floats for state, u, and formats output accordingly.
state: 6D state (x, z, V, gamma, theta, q)
u: 2D input (thrust [N], elevator [deg])
"""
x = state[0] # horizontal position
z = state[1] # vertical position
V = state[2] # airspeed
gamma = state[3] # flight path angle
theta = state[4] # pitch angle
q = state[5] # pitch rate
FT = u[0] # thrust force
de = u[1] # elevator deflection
cg = cos(gamma)
sg = sin(gamma)
alpha = theta - gamma # angle of attack
ca = cos(alpha)
sa = sin(alpha)
m = self.m
g = self.g
S = self.S
rho = self.rho # assume stays near sea level
cbar = self.cbar
I_y = self.I_y
if uncertainty:
m = 1.05 * m # make heavier
S = 0.95 * S # make wing area smaller
I_y = 1.05 * I_y # increase moment of inertia
# assume constant aero moment coefficients
CM_alpha = self.CM_alpha
CM_q = self.CM_q
CM_de = self.CM_de
# CL, CD based on flat plate theory
CL = 2 * ca * sa
CD = 2 * sa**2
CM = CM_alpha * alpha + CM_q * q + CM_de * de
# lift force, drag force, moment
L = (0.5 * rho * S * V**2) * CL
D = (0.5 * rho * S * V**2) * CD
M = (0.5 * rho * S * V**2) * cbar * CM
if u.dtype == np.dtype('O'):
derivs = np.zeros_like(state, dtype=object)
else:
derivs = np.zeros_like(state)
# (dynamics from Stevens, Lewis, Johnson: "Aircraft Control and Simulation" CH2.5)
# positions
derivs[0] = state[2] * cg
derivs[1] = state[2] * sg
# vdot
derivs[2] = (1.0 / m) * (-D + FT * ca - m * g * sg)
# gammadot
derivs[3] = (1.0 / (m * V)) * (L + FT * sa - m * g * cg)
# thetadot
derivs[4] = q
# qdot
derivs[5] = M / I_y
return derivs
def mblock_example(initial_state,
final_state,
min_time,
max_time,
max_torque,
dim=2):
print("Initial State: {}".format(initial_state))
print("Final State: {}".format(final_state))
# a few checks
assert (len(initial_state) == len(final_state))
assert (min_time <= max_time)
# some values that can be changed if desired
N = 50 # number knot points
dynamics_error_thresh = 0.01 # error thresh on xdot = f(x,u,f)
floor_offset = -.01 # used to allow a little penitration
final_state_error_thresh = 0.01 # final state error thresh
max_ground_force = 100
# impose contraint to stay on the ground
stay_on_ground = False
stay_on_ground_tolerance = 0.1 # tolerance for leaving the ground if contstraint is used
complimentarity_constraint_thresh = 0.0001
fix_corner_on_ground = True
corner_fix = [2, 0.5] # corner index, location
mu = 0.1 # friction force
# use SNOPT in Drake
mp = MathematicalProgram()
# state length
state_len = len(initial_state)
# total time used (assuming equal time steps)
time_used = mp.NewContinuousVariables(1, "time_used")
dt = time_used / (N + 1)
# input torque decision variables
u = mp.NewContinuousVariables(1, "u_%d" % 0) # only one input for the cube
u_over_time = u
for k in range(1, N):
u = mp.NewContinuousVariables(1, "u_%d" % k)
u_over_time = np.vstack((u_over_time, u))
total_u = u_over_time
# contact force decision variables
f = mp.NewContinuousVariables(8, "f_%d" % 0) # only one input for the cube
f_over_time = f
for k in range(1, N):
f = mp.NewContinuousVariables(8, "f_%d" % k)
f_over_time = np.vstack((f_over_time, f))
total_f = f_over_time
# state decision variables
x = mp.NewContinuousVariables(state_len,
"x_%d" % 0) # for both input thrusters
x_over_time = x
for k in range(1, N + 1):
x = mp.NewContinuousVariables(state_len, "x_%d" % k)
x_over_time = np.vstack((x_over_time, x))
total_x = x_over_time
# impose dynamic constraints
for n in range(N):
state_next = total_x[n + 1]
dynamic_state_next = total_x[n, :] + get_nd_dynamics(
total_x[n, :], total_u[n, :], total_f[n, :], dim) * dt
# make sure the actual and predicted align to follow dynamics
for j in range(state_len):
state_error = state_next[j] - dynamic_state_next[j]
mp.AddConstraint(state_error <= dynamics_error_thresh)
mp.AddConstraint(state_error >= -dynamics_error_thresh)
# can't penitrate the floor and can't leave the floor
for n in range(N):
distances = get_corner_distances(total_x[n, :], dim)
mp.AddConstraint(distances[0] >= floor_offset)
mp.AddConstraint(distances[1] >= floor_offset)
mp.AddConstraint(distances[2] >= floor_offset)
mp.AddConstraint(distances[3] >= floor_offset)
# don't leave the ground if specified
if stay_on_ground == True:
mp.AddConstraint(
distances[0] <= np.sqrt(2) + stay_on_ground_tolerance)
mp.AddConstraint(
distances[1] <= np.sqrt(2) + stay_on_ground_tolerance)
mp.AddConstraint(
distances[2] <= np.sqrt(2) + stay_on_ground_tolerance)
mp.AddConstraint(
distances[3] <= np.sqrt(2) + stay_on_ground_tolerance)
if fix_corner_on_ground == True:
x_pos = get_corner_x_positions(total_x[n, :], dim)
# make left corner on the ground
num, loc = corner_fix
mp.AddConstraint(distances[num] == 1.0)
mp.AddConstraint(x_pos[num] == loc)
# ground forces can't pull on the ground
for n in range(N):
force = total_f[n]
for j in range(8):
mp.AddConstraint(force[j] <= max_ground_force)
mp.AddConstraint(force[j] >= 0)
# add complimentary constraint
for n in range(N):
force = total_f[n]
state = total_x[n]
theta = state[dim]
distances = get_corner_distances(total_x[n + 1, :], dim)
s = sin(theta)
c = cos(theta)
z_0 = force[0] * c + force[1] * s
z_1 = -force[2] * s + force[3] * c
z_2 = -force[4] * c - force[5] * s
z_3 = force[6] * s - force[7] * c
xy_0 = -force[0] * s + force[1] * c
xy_1 = -force[2] * c - force[3] * s
xy_2 = force[4] * s - force[5] * c
xy_3 = force[6] * c + force[7] * s
mp.AddConstraint(xy_0 <= z_0 * mu)
mp.AddConstraint(xy_0 >= -z_0 * mu)
mp.AddConstraint(xy_1 <= z_1 * mu)
mp.AddConstraint(xy_1 >= -z_1 * mu)
mp.AddConstraint(xy_2 <= z_2 * mu)
mp.AddConstraint(xy_2 >= -z_2 * mu)
mp.AddConstraint(xy_3 <= z_3 * mu)
mp.AddConstraint(xy_3 >= -z_3 * mu)
# vector_0 = force[0] * force[1]
# vector_1 = force[2] * force[3]
# vector_2 = force[4] * force[5]
# vector_3 = force[6] * force[7]
# val = np.asarray([vector_0, vector_1, vector_2, vector_3])
# mp.AddConstraint(val.dot(distances) <= complimentarity_constraint_thresh)
# mp.AddConstraint(val.dot(distances) >= -complimentarity_constraint_thresh)
val_0 = np.asarray([force[0], force[2], force[4], force[6]])
val_1 = np.asarray([force[1], force[3], force[5], force[7]])
mp.AddConstraint(
val_0.dot(distances) <= complimentarity_constraint_thresh)
mp.AddConstraint(
val_0.dot(distances) >= -complimentarity_constraint_thresh)
mp.AddConstraint(
val_1.dot(distances) <= complimentarity_constraint_thresh)
mp.AddConstraint(
val_1.dot(distances) >= -complimentarity_constraint_thresh)
# initial state, no state error allowed
for i in range(state_len):
initial_state_error = x_over_time[0, i] - initial_state[i]
mp.AddConstraint(initial_state_error == 0.0)
# don't care about final wheel angle, so skip restriction in it
state_indices = [i for i in range(0, state_len)]
a = state_indices[0:state_len / 2 - 1] + state_indices[state_len / 2:]
for i in a:
# final
final_state_error = x_over_time[-1, i] - final_state[i]
mp.AddConstraint(final_state_error <= final_state_error_thresh)
mp.AddConstraint(final_state_error >= -final_state_error_thresh)
# add time constraint
mp.AddConstraint(time_used[0] >= min_time)
mp.AddConstraint(time_used[0] <= max_time)
# add torque constraints
for n in range(N):
mp.AddConstraint(u_over_time[n, 0] <= max_torque)
mp.AddConstraint(u_over_time[n, 0] >= -max_torque)
# try to keep the velocity of the wheel in the correct direction
# mp.AddLinearCost(x_over_time[:,-1].sum())
# mp.AddLinearCost(-x_over_time[:,1].sum())
# mp.AddLinearCost(-x_over_time[N//2,1])
print "Number of decision vars", mp.num_vars()
print(mp.Solve())
trajectory = mp.GetSolution(x_over_time)
input_trajectory = mp.GetSolution(u_over_time)
force_trajectory = mp.GetSolution(f_over_time)
t = mp.GetSolution(time_used)
time_array = np.arange(0.0, t, t / (N + 1))
return trajectory, input_trajectory, force_trajectory, time_array
def next_state_with_contact(state, u, dt, dim=2):
"""Return the next state after computing the force for the given timestep.
Note that this is not guaranteed to work realtime or even every time.
Keyword arguments:
state -- the state of the cube
u -- the input torque on the wheel
dt -- the timestep size
dim -- the dimension of the state space (default 2)
"""
mu = 0.1 # friction force
max_ground_force = 200
complimentarity_constraint_thresh = 0.01
# use SNOPT in Drake
mp = MathematicalProgram()
floor_offset = -0.01 # used to allow a little penitration
x = mp.NewContinuousVariables(len(state), "x_%d" % 0)
u_decision = mp.NewContinuousVariables(len(u), "u_%d" % 0)
f = mp.NewContinuousVariables(8, "f_%d" % 0)
# starting values
for i in range(len(x)):
mp.AddConstraint(x[i] == state[i])
for i in range(len(u)):
mp.AddConstraint(u_decision[i] == u[i])
dynamic_state_next = x[:] + get_nd_dynamics(x[:], u_decision[:], f[:],
dim) * dt
# can't penitrate the floor
distances = get_corner_distances(dynamic_state_next, dim)
mp.AddConstraint(distances[0] >= floor_offset)
mp.AddConstraint(distances[1] >= floor_offset)
mp.AddConstraint(distances[2] >= floor_offset)
mp.AddConstraint(distances[3] >= floor_offset)
# ground forces can't pull on the ground
for j in range(8):
# mp.AddConstraint(f[j] <= max_ground_force)
mp.AddConstraint(f[j] >= 0)
# add complimentary constraint
theta = state[dim]
distances = get_corner_distances(state, dim)
s = sin(theta)
c = cos(theta)
z_0 = f[0] * c + f[1] * s
z_1 = -f[2] * s + f[3] * c
z_2 = -f[4] * c - f[5] * s
z_3 = f[6] * s - f[7] * c
xy_0 = -f[0] * s + f[1] * c
xy_1 = -f[2] * c - f[3] * s
xy_2 = f[4] * s - f[5] * c
xy_3 = f[6] * c + f[7] * s
mp.AddConstraint(xy_0 <= z_0 * mu)
mp.AddConstraint(xy_0 >= -z_0 * mu)
mp.AddConstraint(xy_1 <= z_1 * mu)
mp.AddConstraint(xy_1 >= -z_1 * mu)
mp.AddConstraint(xy_2 <= z_2 * mu)
mp.AddConstraint(xy_2 >= -z_2 * mu)
mp.AddConstraint(xy_3 <= z_3 * mu)
mp.AddConstraint(xy_3 >= -z_3 * mu)
vector_0 = f[0] * f[1]
vector_1 = f[2] * f[3]
vector_2 = f[4] * f[5]
vector_3 = f[6] * f[7]
val = np.asarray([vector_0, vector_1, vector_2, vector_3])
mp.AddConstraint(val.dot(distances) <= complimentarity_constraint_thresh)
mp.AddConstraint(val.dot(distances) >= -complimentarity_constraint_thresh)
# print "Number of decision vars", mp.num_vars()
# print(mp.Solve())
mp.Solve()
f_comp = mp.GetSolution(f)
return state + get_nd_dynamics(state, u, f_comp, dim) * dt
def rpy_to_rotation_matrix_symbolic(rpy):
"""
Creates 3x3 rotation matrix from rpy
See http://danceswithcode.net/engineeringnotes/rotations_in_3d/rotations_in_3d_part1.html
"""
from pydrake.math import sin, cos
u = rpy[0]
v = rpy[1]
w = rpy[2]
R = np.zeros([3, 3], dtype=rpy.dtype)
# first row
R[0, 0] = cos(v) * cos(w)
R[0, 1] = sin(u) * sin(v) * cos(w) - cos(u) * sin(w)
R[0, 2] = sin(u) * sin(w) + cos(u) * sin(v) * cos(w)
# second row
R[1, 0] = cos(v) * sin(w)
R[1, 1] = cos(u) * cos(w) + sin(u) * sin(v) * sin(w)
R[1, 2] = cos(u) * sin(v) * sin(w) - sin(u) * cos(w)
# third row
R[2, 0] = -sin(v)
R[2, 1] = sin(u) * cos(v)
R[2, 2] = cos(u) * cos(v)
return R
def periodic_motion(dim=2):
# print("Initial State: {}".format(initial_state))
# print("Final State: {}".format(final_state))
# a few checks
min_time = 0.5
max_time = 15.0
max_torque = 1000.0
# some values that can be changed if desired
N = 50 # number knot points
dynamics_error_thresh = 0.01 # error thresh on xdot = f(x,u,f)
floor_offset = -.01 # used to allow a little penitration
final_state_error_thresh = 0.01 # final state error thresh
max_ground_force = 100
# impose contraint to stay on the ground
stay_on_ground = True
stay_on_ground_tolerance = 0.1 # tolerance for leaving the ground if contstraint is used
complimentarity_constraint_thresh = 0.0001
fix_corner_on_ground = True
corner_fix = [0, -0.5] # corner index, location
mu = 0.1 # friction force
# use SNOPT in Drake
mp = MathematicalProgram()
# state length
state_len = 8
# total time used (assuming equal time steps)
time_used = mp.NewContinuousVariables(1, "time_used")
dt = time_used / (N + 1)
# input torque decision variables
u = mp.NewContinuousVariables(1, "u_%d" % 0) # only one input for the cube
u_over_time = u
for k in range(1, N):
u = mp.NewContinuousVariables(1, "u_%d" % k)
u_over_time = np.vstack((u_over_time, u))
total_u = u_over_time
# contact force decision variables
f = mp.NewContinuousVariables(8, "f_%d" % 0) # only one input for the cube
f_over_time = f
for k in range(1, N):
f = mp.NewContinuousVariables(8, "f_%d" % k)
f_over_time = np.vstack((f_over_time, f))
total_f = f_over_time
# state decision variables
x = mp.NewContinuousVariables(state_len,
"x_%d" % 0) # for both input thrusters
x_over_time = x
for k in range(1, N + 1):
x = mp.NewContinuousVariables(state_len, "x_%d" % k)
x_over_time = np.vstack((x_over_time, x))
total_x = x_over_time
# impose dynamic constraints
for n in range(N):
state_next = total_x[n + 1]
dynamic_state_next = total_x[n, :] + get_nd_dynamics(
total_x[n, :], total_u[n, :], total_f[n, :], dim) * dt
# make sure the actual and predicted align to follow dynamics
for j in range(state_len):
state_error = state_next[j] - dynamic_state_next[j]
mp.AddConstraint(state_error <= dynamics_error_thresh)
mp.AddConstraint(state_error >= -dynamics_error_thresh)
# can't penitrate the floor and can't leave the floor
for n in range(N):
distances = get_corner_distances(total_x[n, :], dim)
mp.AddConstraint(distances[0] >= floor_offset)
mp.AddConstraint(distances[1] >= floor_offset)
mp.AddConstraint(distances[2] >= floor_offset)
mp.AddConstraint(distances[3] >= floor_offset)
# don't leave the ground if specified
if stay_on_ground == True:
mp.AddConstraint(
distances[0] <= np.sqrt(2) + stay_on_ground_tolerance)
mp.AddConstraint(
distances[1] <= np.sqrt(2) + stay_on_ground_tolerance)
mp.AddConstraint(
distances[2] <= np.sqrt(2) + stay_on_ground_tolerance)
mp.AddConstraint(
distances[3] <= np.sqrt(2) + stay_on_ground_tolerance)
if fix_corner_on_ground == True:
x_pos = get_corner_x_positions(total_x[n, :], dim)
# make left corner on the ground
mp.AddConstraint(distances[0] == 0.0)
num, loc = corner_fix
mp.AddConstraint(x_pos[num] == loc)
# ground forces can't pull on the ground
for n in range(N):
force = total_f[n]
for j in range(8):
mp.AddConstraint(force[j] <= max_ground_force)
mp.AddConstraint(force[j] >= 0)
# add complimentary constraint
for n in range(N):
force = total_f[n]
state = total_x[n]
theta = state[dim]
distances = get_corner_distances(total_x[n + 1, :], dim)
s = sin(theta)
c = cos(theta)
z_0 = force[0] * c + force[1] * s
z_1 = -force[2] * s + force[3] * c
z_2 = -force[4] * c - force[5] * s
z_3 = force[6] * s - force[7] * c
xy_0 = -force[0] * s + force[1] * c
xy_1 = -force[2] * c - force[3] * s
xy_2 = force[4] * s - force[5] * c
xy_3 = force[6] * c + force[7] * s
mp.AddConstraint(xy_0 <= z_0 * mu)
mp.AddConstraint(xy_0 >= -z_0 * mu)
mp.AddConstraint(xy_1 <= z_1 * mu)
mp.AddConstraint(xy_1 >= -z_1 * mu)
mp.AddConstraint(xy_2 <= z_2 * mu)
mp.AddConstraint(xy_2 >= -z_2 * mu)
mp.AddConstraint(xy_3 <= z_3 * mu)
mp.AddConstraint(xy_3 >= -z_3 * mu)
val_0 = np.asarray([force[0], force[2], force[4], force[6]])
val_1 = np.asarray([force[1], force[3], force[5], force[7]])
mp.AddConstraint(
val_0.dot(distances) <= complimentarity_constraint_thresh)
mp.AddConstraint(
val_0.dot(distances) >= -complimentarity_constraint_thresh)
mp.AddConstraint(
val_1.dot(distances) <= complimentarity_constraint_thresh)
mp.AddConstraint(
val_1.dot(distances) >= -complimentarity_constraint_thresh)
# initial
mp.AddConstraint(x_over_time[0, 0] == 0.0) # x
mp.AddConstraint(x_over_time[0, 1] == 0.0) # y
mp.AddConstraint(x_over_time[0, 2] == 0.0) # 0 angle
mp.AddConstraint(x_over_time[0, 3] == 0.0) # alpha angle
# final
mp.AddConstraint(x_over_time[-1, 0] == -1.0) # x
mp.AddConstraint(x_over_time[-1, 1] == 0.0) # y
mp.AddConstraint(x_over_time[-1, 2] == np.pi / 2.0) # y
# connection constraint, don't care about inner wheel angle
for i in [4, 5, 6, 7]:
mp.AddConstraint(x_over_time[0, i] == x_over_time[-1, i])
# add time constraint
mp.AddConstraint(time_used[0] >= min_time)
mp.AddConstraint(time_used[0] <= max_time)
# add torque constraints
for n in range(N):
mp.AddConstraint(u_over_time[n, 0] <= max_torque)
mp.AddConstraint(u_over_time[n, 0] >= -max_torque)
# minimize input
# mp.AddQuadraticCost(u_over_time[:,0].dot(u_over_time[:,0]))
# maximize velocity in correct direction (left)
mp.AddLinearCost(-x_over_time[:, 4].sum())
# minimize the time
# mp.AddLinearCost(time_used[0])
print "Number of decision vars", mp.num_vars()
print(mp.Solve())
trajectory = mp.GetSolution(x_over_time)
input_trajectory = mp.GetSolution(u_over_time)
force_trajectory = mp.GetSolution(f_over_time)
t = mp.GetSolution(time_used)
time_array = np.arange(0.0, t, t / (N + 1))
return trajectory, input_trajectory, force_trajectory, time_array
