def foot_in_stepping_stone(prog, terrain, n_steps, decision_variables):
# unpack only decision variables needed in this function
position_left, position_right, stone_left, stone_right = decision_variables[:4]
# big-M vector
M = get_big_M(terrain)
# modify here
for t in range(n_steps + 1):
for i in range(len(terrain.stepping_stones)):
A = terrain.stepping_stones[i].A
b = terrain.stepping_stones[i].b
prog.AddLinearConstraint(le(np.dot(A, position_left[t]), b + (1 - stone_left[t, i])*M))
prog.AddLinearConstraint(le(np.dot(A, position_right[t]), b + (1 - stone_right[t, i])*M))
def add_arena_limits_nl(self, prog, state, N):
arena_lims = np.array([
self.params.arena_limits_x / 2.0, self.params.arena_limits_y / 2.0
])
for k in range(N + 1):
prog.AddLinearConstraint(le(state[k][0:2], arena_lims))
prog.AddLinearConstraint(ge(state[k][0:2], -arena_lims))
def add_input_limits_nl(self, prog, cmd, N):
"""Add saturation limits to cmd"""
for k in range(N):
prog.AddLinearConstraint(
le(cmd[k], self.params.input_limit * np.ones(2)))
prog.AddLinearConstraint(
ge(cmd[k], -self.params.input_limit * np.ones(2)))
def relative_position_limits(prog, n_steps, step_span, decision_variables):
# unpack only decision variables needed in this function
position_left, position_right = decision_variables[:2]
# modify here
for t in range(n_steps + 1):
prog.AddLinearConstraint(le(position_left[t], position_right[t] + step_span/2))
prog.AddLinearConstraint(ge(position_left[t], position_right[t] - step_span/2))
def step_sequence(prog, n_steps, step_span, decision_variables):
# unpack only decision variables needed in this function
position_left, position_right = decision_variables[:2]
first_left = decision_variables[-1]
# variable equal to one if first step is with right foot
first_right = 1 - first_left
# note that the step_span coincides with the maximum distance
# (both horizontal and vertical) between the position of
# a foot at step t and at step t + 1
step_limit = np.ones(2) * step_span
# sequence for the robot steps implied by the binaries first_left and first_right
# (could be written more compactly, but would be harder to read)
for t in range(n_steps):
# lengths of the steps
step_left = position_left[t + 1] - position_left[t]
step_right = position_right[t + 1] - position_right[t]
# for all even steps
if t % 2 == 0:
limit_left = step_limit * first_left # left foot can move iff first_left
limit_right = step_limit * first_right # right foot can move iff first_right
# for all odd steps
else:
limit_left = step_limit * first_right # left foot can move iff first_right
limit_right = step_limit * first_left # right foot can move iff first_left
# constraints on left-foot relative position
prog.AddLinearConstraint(le(step_left, limit_left))
prog.AddLinearConstraint(ge(step_left, - limit_left))
# constraints on right-foot relative position
prog.AddLinearConstraint(le(step_right, limit_right))
prog.AddLinearConstraint(ge(step_right, - limit_right))
def create_qp1(self, plant_context, V, q_des, v_des, vd_des):
# Determine contact points
contact_positions_per_frame = {}
active_contacts_per_frame = {} # Note this should be in frame space
for frame, contacts in self.contacts_per_frame.items():
contact_positions = self.plant.CalcPointsPositions(
plant_context, self.plant.GetFrameByName(frame), contacts,
self.plant.world_frame())
active_contacts_per_frame[
frame] = contacts[:,
np.where(contact_positions[2, :] <= 0.0)[0]]
N_c = sum([
active_contacts.shape[1]
for active_contacts in active_contacts_per_frame.values()
]) # num contact points
if N_c == 0:
print("Not in contact!")
return None
''' Eq(7) '''
H = self.plant.CalcMassMatrixViaInverseDynamics(plant_context)
# Note that CalcGravityGeneralizedForces assumes the form Mv̇ + C(q, v)v = tau_g(q) + tau_app
# while Eq(7) assumes gravity is accounted in C (on the left hand side)
C_7 = self.plant.CalcBiasTerm(
plant_context) - self.plant.CalcGravityGeneralizedForces(
plant_context)
B_7 = self.B_7
# TODO: Double check
Phi_foots = []
for frame, active_contacts in active_contacts_per_frame.items():
if active_contacts.size:
Phi_foots.append(
self.plant.CalcJacobianTranslationalVelocity(
plant_context, JacobianWrtVariable.kV,
self.plant.GetFrameByName(frame), active_contacts,
self.plant.world_frame(), self.plant.world_frame()))
Phi = np.vstack(Phi_foots)
''' Eq(8) '''
v_idx_act = self.v_idx_act
H_f = H[0:v_idx_act, :]
H_a = H[v_idx_act:, :]
C_f = C_7[0:v_idx_act]
C_a = C_7[v_idx_act:]
B_a = self.B_a
Phi_f_T = Phi.T[0:v_idx_act:, :]
Phi_a_T = Phi.T[v_idx_act:, :]
''' Eq(9) '''
# Assume flat ground for now
n = np.array([[0], [0], [1.0]])
d = np.array([[1.0, -1.0, 0.0, 0.0], [0.0, 0.0, 1.0, -1.0],
[0.0, 0.0, 0.0, 0.0]])
v = np.zeros((N_d, N_c, N_f))
for i in range(N_d):
for j in range(N_c):
v[i, j] = (n + mu * d)[:, i]
''' Quadratic Program I '''
prog = MathematicalProgram()
qdd = prog.NewContinuousVariables(
self.plant.num_velocities(),
name="qdd") # To ignore 6 DOF floating base
self.qdd = qdd
beta = prog.NewContinuousVariables(N_d, N_c, name="beta")
self.beta = beta
lambd = prog.NewContinuousVariables(N_f * N_c, name="lambda")
self.lambd = lambd
# Jacobians ignoring the 6DOF floating base
J_foots = []
for frame, active_contacts in active_contacts_per_frame.items():
if active_contacts.size:
num_active_contacts = active_contacts.shape[1]
J_foot = np.zeros((N_f * num_active_contacts, Atlas.TOTAL_DOF))
# TODO: Can this be simplified?
for i in range(num_active_contacts):
J_foot[N_f * i:N_f *
(i +
1), :] = self.plant.CalcJacobianSpatialVelocity(
plant_context, JacobianWrtVariable.kV,
self.plant.GetFrameByName(frame),
active_contacts[:,
i], self.plant.world_frame(),
self.plant.world_frame())[3:]
J_foots.append(J_foot)
J = np.vstack(J_foots)
assert (J.shape == (N_c * N_f, Atlas.TOTAL_DOF))
eta = prog.NewContinuousVariables(J.shape[0], name="eta")
self.eta = eta
x = prog.NewContinuousVariables(
self.x_size, name="x") # x_com, y_com, x_com_d, y_com_d
self.x = x
u = prog.NewContinuousVariables(self.u_size,
name="u") # x_com_dd, y_com_dd
self.u = u
''' Eq(10) '''
w = 0.01
epsilon = 1.0e-8
K_p = 10.0
K_d = 4.0
frame_weights = np.ones((Atlas.TOTAL_DOF))
q = self.plant.GetPositions(plant_context)
qd = self.plant.GetVelocities(plant_context)
# Convert q, q_nom to generalized velocities form
q_err = self.plant.MapQDotToVelocity(plant_context, q_des - q)
# print(f"Pelvis error: {q_err[0:3]}")
## FIXME: Not sure if it's a good idea to ignore the x, y, z position of pelvis
# ignored_pose_indices = {3, 4, 5} # Ignore x position, y position
ignored_pose_indices = {} # Ignore x position, y position
relevant_pose_indices = list(
set(range(Atlas.TOTAL_DOF)) - set(ignored_pose_indices))
self.relevant_pose_indices = relevant_pose_indices
qdd_ref = K_p * q_err + K_d * (v_des - qd) + vd_des # Eq(27) of [1]
qdd_err = qdd_ref - qdd
qdd_err = qdd_err * frame_weights
qdd_err = qdd_err[relevant_pose_indices]
prog.AddCost(
V(x, u) + w * ((qdd_err).dot(qdd_err)) +
epsilon * np.sum(np.square(beta)) + eta.dot(eta))
''' Eq(11) - 0.003s '''
eq11_lhs = H_f.dot(qdd) + C_f
eq11_rhs = Phi_f_T.dot(lambd)
prog.AddLinearConstraint(eq(eq11_lhs, eq11_rhs))
''' Eq(12) - 0.005s '''
alpha = 0.1
# TODO: Double check
Jd_qd_foots = []
for frame, active_contacts in active_contacts_per_frame.items():
if active_contacts.size:
Jd_qd_foot = self.plant.CalcBiasTranslationalAcceleration(
plant_context, JacobianWrtVariable.kV,
self.plant.GetFrameByName(frame), active_contacts,
self.plant.world_frame(), self.plant.world_frame())
Jd_qd_foots.append(Jd_qd_foot.flatten())
Jd_qd = np.concatenate(Jd_qd_foots)
assert (Jd_qd.shape == (N_c * 3, ))
eq12_lhs = J.dot(qdd) + Jd_qd
eq12_rhs = -alpha * J.dot(qd) + eta
prog.AddLinearConstraint(eq(eq12_lhs, eq12_rhs))
''' Eq(13) - 0.015s '''
def tau(qdd, lambd):
return self.B_a_inv.dot(H_a.dot(qdd) + C_a - Phi_a_T.dot(lambd))
self.tau = tau
eq13_lhs = self.tau(qdd, lambd)
prog.AddLinearConstraint(ge(eq13_lhs, -self.sorted_max_efforts))
prog.AddLinearConstraint(le(eq13_lhs, self.sorted_max_efforts))
''' Eq(14) '''
for j in range(N_c):
beta_v = beta[:, j].dot(v[:, j])
prog.AddLinearConstraint(eq(lambd[N_f * j:N_f * j + 3], beta_v))
''' Eq(15) '''
for b in beta.flat:
prog.AddLinearConstraint(b >= 0.0)
''' Eq(16) '''
prog.AddLinearConstraint(ge(eta, eta_min))
prog.AddLinearConstraint(le(eta, eta_max))
''' Enforce x as com '''
com = self.plant.CalcCenterOfMassPosition(plant_context)
com_d = self.plant.CalcJacobianCenterOfMassTranslationalVelocity(
plant_context, JacobianWrtVariable.kV, self.plant.world_frame(),
self.plant.world_frame()).dot(qd)
prog.AddLinearConstraint(x[0] == com[0])
prog.AddLinearConstraint(x[1] == com[1])
prog.AddLinearConstraint(x[2] == com_d[0])
prog.AddLinearConstraint(x[3] == com_d[1])
''' Enforce u as com_dd '''
com_dd = (
self.plant.CalcBiasCenterOfMassTranslationalAcceleration(
plant_context, JacobianWrtVariable.kV,
self.plant.world_frame(), self.plant.world_frame()) +
self.plant.CalcJacobianCenterOfMassTranslationalVelocity(
plant_context, JacobianWrtVariable.kV,
self.plant.world_frame(), self.plant.world_frame()).dot(qdd))
prog.AddLinearConstraint(u[0] == com_dd[0])
prog.AddLinearConstraint(u[1] == com_dd[1])
''' Respect joint limits '''
for name, limit in Atlas.JOINT_LIMITS.items():
# Get the corresponding joint value
joint_pos = self.plant.GetJointByName(name).get_angle(
plant_context)
# Get the corresponding actuator index
act_idx = getActuatorIndex(self.plant, name)
# Use the actuator index to find the corresponding generalized coordinate index
# q_idx = np.where(B_7[:,act_idx] == 1)[0][0]
q_idx = getJointIndexInGeneralizedVelocities(self.plant, name)
if joint_pos >= limit.upper:
# print(f"Joint {name} max reached")
prog.AddLinearConstraint(qdd[q_idx] <= 0.0)
elif joint_pos <= limit.lower:
# print(f"Joint {name} min reached")
prog.AddLinearConstraint(qdd[q_idx] >= 0.0)
return prog
def solveOptimization(state_init,
t_impact,
impact_combination,
T,
u_guess=None,
x_guess=None,
h_guess=None):
prog = MathematicalProgram()
h = prog.NewContinuousVariables(T, name='h')
u = prog.NewContinuousVariables(rows=T + 1,
cols=2 * n_quadrotors,
name='u')
x = prog.NewContinuousVariables(rows=T + 1,
cols=6 * n_quadrotors + 4 * n_balls,
name='x')
dv = prog.decision_variables()
prog.AddBoundingBoxConstraint([h_min] * T, [h_max] * T, h)
for i in range(n_quadrotors):
sys = Quadrotor2D()
context = sys.CreateDefaultContext()
dir_coll_constr = DirectCollocationConstraint(sys, context)
ind_x = 6 * i
ind_u = 2 * i
for t in range(T):
impact_indices = impact_combination[np.argmax(
np.abs(t - t_impact) <= 1)]
quad_ind, ball_ind = impact_indices[0], impact_indices[1]
if quad_ind == i and np.any(
t == t_impact
): # Don't add Direct collocation constraint at impact
continue
elif quad_ind == i and (np.any(t == t_impact - 1)
or np.any(t == t_impact + 1)):
prog.AddConstraint(
eq(
x[t + 1,
ind_x:ind_x + 3], x[t, ind_x:ind_x + 3] + h[t] *
x[t + 1, ind_x + 3:ind_x + 6])) # Backward euler
prog.AddConstraint(
eq(x[t + 1, ind_x + 3:ind_x + 6], x[t,
ind_x + 3:ind_x + 6])
) # Zero-acceleration assumption during this time step. Should maybe replace with something less naive
else:
AddDirectCollocationConstraint(
dir_coll_constr, np.array([[h[t]]]),
x[t, ind_x:ind_x + 6].reshape(-1, 1),
x[t + 1, ind_x:ind_x + 6].reshape(-1, 1),
u[t, ind_u:ind_u + 2].reshape(-1, 1),
u[t + 1, ind_u:ind_u + 2].reshape(-1, 1), prog)
for i in range(n_balls):
sys = Ball2D()
context = sys.CreateDefaultContext()
dir_coll_constr = DirectCollocationConstraint(sys, context)
ind_x = 6 * n_quadrotors + 4 * i
for t in range(T):
impact_indices = impact_combination[np.argmax(
np.abs(t - t_impact) <= 1)]
quad_ind, ball_ind = impact_indices[0], impact_indices[1]
if ball_ind == i and np.any(
t == t_impact
): # Don't add Direct collocation constraint at impact
continue
elif ball_ind == i and (np.any(t == t_impact - 1)
or np.any(t == t_impact + 1)):
prog.AddConstraint(
eq(
x[t + 1,
ind_x:ind_x + 2], x[t, ind_x:ind_x + 2] + h[t] *
x[t + 1, ind_x + 2:ind_x + 4])) # Backward euler
prog.AddConstraint(
eq(x[t + 1,
ind_x + 2:ind_x + 4], x[t, ind_x + 2:ind_x + 4] +
h[t] * np.array([0, -9.81])))
else:
AddDirectCollocationConstraint(
dir_coll_constr, np.array([[h[t]]]),
x[t, ind_x:ind_x + 4].reshape(-1, 1),
x[t + 1, ind_x:ind_x + 4].reshape(-1, 1),
u[t, 0:0].reshape(-1, 1), u[t + 1, 0:0].reshape(-1,
1), prog)
# Initial conditions
prog.AddLinearConstraint(eq(x[0, :], state_init))
# Final conditions
prog.AddLinearConstraint(eq(x[T, 0:14], state_final[0:14]))
# Quadrotor final conditions (full state)
for i in range(n_quadrotors):
ind = 6 * i
prog.AddLinearConstraint(
eq(x[T, ind:ind + 6], state_final[ind:ind + 6]))
# Ball final conditions (position only)
for i in range(n_balls):
ind = 6 * n_quadrotors + 4 * i
prog.AddLinearConstraint(
eq(x[T, ind:ind + 2], state_final[ind:ind + 2]))
# Input constraints
for i in range(n_quadrotors):
prog.AddLinearConstraint(ge(u[:, 2 * i], -20.0))
prog.AddLinearConstraint(le(u[:, 2 * i], 20.0))
prog.AddLinearConstraint(ge(u[:, 2 * i + 1], -20.0))
prog.AddLinearConstraint(le(u[:, 2 * i + 1], 20.0))
# Don't allow quadrotor to pitch more than 60 degrees
for i in range(n_quadrotors):
prog.AddLinearConstraint(ge(x[:, 6 * i + 2], -np.pi / 3))
prog.AddLinearConstraint(le(x[:, 6 * i + 2], np.pi / 3))
# Ball position constraints
# for i in range(n_balls):
# ind_i = 6*n_quadrotors + 4*i
# prog.AddLinearConstraint(ge(x[:,ind_i],-2.0))
# prog.AddLinearConstraint(le(x[:,ind_i], 2.0))
# prog.AddLinearConstraint(ge(x[:,ind_i+1],-3.0))
# prog.AddLinearConstraint(le(x[:,ind_i+1], 3.0))
# Impact constraint
quad_temp = Quadrotor2D()
for i in range(n_quadrotors):
for j in range(n_balls):
ind_q = 6 * i
ind_b = 6 * n_quadrotors + 4 * j
for t in range(T):
if np.any(
t == t_impact
): # If quad i and ball j impact at time t, add impact constraint
impact_indices = impact_combination[np.argmax(
t == t_impact)]
quad_ind, ball_ind = impact_indices[0], impact_indices[1]
if quad_ind == i and ball_ind == j:
# At impact, witness function == 0
prog.AddConstraint(lambda a: np.array([
CalcClosestDistanceQuadBall(a[0:3], a[3:5])
]).reshape(1, 1),
lb=np.zeros((1, 1)),
ub=np.zeros((1, 1)),
vars=np.concatenate(
(x[t, ind_q:ind_q + 3],
x[t,
ind_b:ind_b + 2])).reshape(
-1, 1))
# At impact, enforce discrete collision update for both ball and quadrotor
prog.AddConstraint(
CalcPostCollisionStateQuadBallResidual,
lb=np.zeros((6, 1)),
ub=np.zeros((6, 1)),
vars=np.concatenate(
(x[t, ind_q:ind_q + 6], x[t, ind_b:ind_b + 4],
x[t + 1, ind_q:ind_q + 6])).reshape(-1, 1))
prog.AddConstraint(
CalcPostCollisionStateBallQuadResidual,
lb=np.zeros((4, 1)),
ub=np.zeros((4, 1)),
vars=np.concatenate(
(x[t, ind_q:ind_q + 6], x[t, ind_b:ind_b + 4],
x[t + 1, ind_b:ind_b + 4])).reshape(-1, 1))
# rough constraints to enforce hitting center-ish of paddle
prog.AddLinearConstraint(
x[t, ind_q] - x[t, ind_b] >= -0.01)
prog.AddLinearConstraint(
x[t, ind_q] - x[t, ind_b] <= 0.01)
continue
# Everywhere else, witness function must be > 0
prog.AddConstraint(lambda a: np.array([
CalcClosestDistanceQuadBall(a[ind_q:ind_q + 3], a[
ind_b:ind_b + 2])
]).reshape(1, 1),
lb=np.zeros((1, 1)),
ub=np.inf * np.ones((1, 1)),
vars=x[t, :].reshape(-1, 1))
# Don't allow quadrotor collisions
# for t in range(T):
# for i in range(n_quadrotors):
# for j in range(i+1, n_quadrotors):
# prog.AddConstraint((x[t,6*i]-x[t,6*j])**2 + (x[t,6*i+1]-x[t,6*j+1])**2 >= 0.65**2)
# Quadrotors stay on their own side
# prog.AddLinearConstraint(ge(x[:, 0], 0.3))
# prog.AddLinearConstraint(le(x[:, 6], -0.3))
###############################################################################
# Set up initial guesses
initial_guess = np.empty(prog.num_vars())
# # initial guess for the time step
prog.SetDecisionVariableValueInVector(h, h_guess, initial_guess)
x_init[0, :] = state_init
prog.SetDecisionVariableValueInVector(x, x_guess, initial_guess)
prog.SetDecisionVariableValueInVector(u, u_guess, initial_guess)
solver = SnoptSolver()
print("Solving...")
result = solver.Solve(prog, initial_guess)
# print(GetInfeasibleConstraints(prog,result))
# be sure that the solution is optimal
assert result.is_success()
print(f'Solution found? {result.is_success()}.')
#################################################################################
# Extract results
# get optimal solution
h_opt = result.GetSolution(h)
x_opt = result.GetSolution(x)
u_opt = result.GetSolution(u)
time_breaks_opt = np.array([sum(h_opt[:t]) for t in range(T + 1)])
u_opt_poly = PiecewisePolynomial.ZeroOrderHold(time_breaks_opt, u_opt.T)
# x_opt_poly = PiecewisePolynomial.Cubic(time_breaks_opt, x_opt.T, False)
x_opt_poly = PiecewisePolynomial.FirstOrderHold(
time_breaks_opt, x_opt.T
) # Switch to first order hold instead of cubic because cubic was taking too long to create
#################################################################################
# Create list of K matrices for time varying LQR
context = quad_plant.CreateDefaultContext()
breaks = copy.copy(
time_breaks_opt) #np.linspace(0, x_opt_poly.end_time(), 100)
K_samples = np.zeros((breaks.size, 12 * n_quadrotors))
for i in range(n_quadrotors):
K = None
for j in range(breaks.size):
context.SetContinuousState(
x_opt_poly.value(breaks[j])[6 * i:6 * (i + 1)])
context.FixInputPort(
0,
u_opt_poly.value(breaks[j])[2 * i:2 * (i + 1)])
linear_system = FirstOrderTaylorApproximation(quad_plant, context)
A = linear_system.A()
B = linear_system.B()
try:
K, _, _ = control.lqr(A, B, Q, R)
except:
assert K is not None, "Failed to calculate initial K for quadrotor " + str(
i)
print("Warning: Failed to calculate K at timestep", j,
"for quadrotor", i, ". Using K from previous timestep")
K_samples[j, 12 * i:12 * (i + 1)] = K.reshape(-1)
K_poly = PiecewisePolynomial.ZeroOrderHold(breaks, K_samples.T)
return u_opt_poly, x_opt_poly, K_poly, h_opt
import numpy as np
from pydrake.all import Quaternion
from pydrake.all import MathematicalProgram, Solve, eq, le, ge, SolverOptions, SnoptSolver
import pdb
prog = MathematicalProgram()
x = prog.NewContinuousVariables(rows=1, name='x')
y = prog.NewContinuousVariables(rows=1, name='y')
slack = prog.NewContinuousVariables(rows=1, name="slack")
prog.AddConstraint(eq(x * y, slack))
prog.AddLinearConstraint(ge(x, 0))
prog.AddLinearConstraint(ge(y, 0))
prog.AddLinearConstraint(le(x, 1))
prog.AddLinearConstraint(le(y, 1))
prog.AddCost(1e5 * (slack**2)[0])
prog.AddCost(-(2 * x[0] + y[0]))
solver = SnoptSolver()
result = solver.Solve(prog)
print(
f"Success: {result.is_success()}, x = {result.GetSolution(x)}, y = {result.GetSolution(y)}, slack = {result.GetSolution(slack)}"
)
def optimize(self, vehs):
c = self.c
p = self.p
n_veh, L, N, dt, u_max = c.n_veh, c.circumference, c.n_opt, c.sim_step, c.u_max
L_veh = vehs[0].length
a, b, s0, T, v0, delta = p.a, p.b, p.s0, p.tau, p.v0, p.delta
vehs = vehs[::-1]
np.set_printoptions(linewidth=100)
# the controlled vehicle is now the first vehicle, positions are sorted from largest to smallest
print(f'Current positions {[veh.pos for veh in vehs]}')
v_init = [veh.speed for veh in vehs]
print(f'Current speeds {v_init}')
# spacing
s_init = [vehs[-1].pos - vehs[0].pos - L_veh] + [
veh_im1.pos - veh_i.pos - L_veh
for veh_im1, veh_i in zip(vehs[:-1], vehs[1:])
]
s_init = [s + L if s < 0 else s for s in s_init] # handle wrap
print(f'Current spacings {s_init}')
# can still follow current plan
if self.plan_index < len(self.plan):
accel = self.plan[self.plan_index]
self.plan_index += 1
return accel
print(f'Optimizing trajectory for {c.n_opt} steps')
## solve for equilibrium conditions (when all vehicles have same spacing and going at optimal speed)
import scipy.optimize
sf = L / n_veh - L_veh # equilibrium space
accel_fn = lambda v: a * (1 - (v / v0)**delta - ((s0 + v * T) / sf)**2)
sol = scipy.optimize.root(accel_fn, 0)
vf = sol.x.item() # equilibrium speed
sstarf = s0 + vf * T
print('Equilibrium speed', vf)
# nonconvex optimization
from pydrake.all import MathematicalProgram, SnoptSolver, eq, le, ge
# get guesses for solutions
v_guess = [np.mean(v_init)]
for _ in range(N):
v_guess.append(dt * accel_fn(v_guess[-1]))
s_guess = [sf] * (N + 1)
prog = MathematicalProgram()
v = prog.NewContinuousVariables(N + 1, n_veh, 'v') # speed
s = prog.NewContinuousVariables(N + 1, n_veh, 's') # spacing
flat = lambda x: x.reshape(-1)
# Guess
prog.SetInitialGuess(s, np.stack([s_guess] * n_veh, axis=1))
prog.SetInitialGuess(v, np.stack([v_guess] * n_veh, axis=1))
# initial conditions constraint
prog.AddLinearConstraint(eq(v[0], v_init))
prog.AddLinearConstraint(eq(s[0], s_init))
# velocity constraint
prog.AddLinearConstraint(ge(flat(v[1:]), 0))
prog.AddLinearConstraint(le(flat(v[1:]),
vf + 1)) # extra constraint to help solver
# spacing constraint
prog.AddLinearConstraint(ge(flat(s[1:]), s0))
prog.AddLinearConstraint(le(flat(s[1:]),
sf * 2)) # extra constraint to help solver
prog.AddLinearConstraint(eq(flat(s[1:].sum(axis=1)),
L - L_veh * n_veh))
# spacing update constraint
s_n = s[:-1, 1:] # s_i[n]
s_np1 = s[1:, 1:] # s_i[n + 1]
v_n = v[:-1, 1:] # v_i[n]
v_np1 = v[1:, 1:] # v_i[n + 1]
v_n_im1 = v[:-1, :-1] # v_{i - 1}[n]
v_np1_im1 = v[1:, :-1] # v_{i - 1}[n + 1]
prog.AddLinearConstraint(
eq(flat(s_np1 - s_n),
flat(0.5 * dt * (v_n_im1 + v_np1_im1 - v_n - v_np1))))
# handle position wrap for vehicle 1
prog.AddLinearConstraint(
eq(s[1:, 0] - s[:-1, 0],
0.5 * dt * (v[:-1, -1] + v[1:, -1] - v[:-1, 0] - v[1:, 0])))
# vehicle 0's action constraint
prog.AddLinearConstraint(ge(v[1:, 0], v[:-1, 0] - u_max * dt))
prog.AddLinearConstraint(le(v[1:, 0], v[:-1, 0] + u_max * dt))
# idm constraint
prog.AddConstraint(
eq((v_np1 - v_n - dt * a * (1 - (v_n / v0)**delta)) * s_n**2,
-dt * a * (s0 + v_n * T + v_n * (v_n - v_n_im1) /
(2 * np.sqrt(a * b)))**2))
if c.cost == 'mean':
prog.AddCost(-v.mean())
elif c.cost == 'target':
prog.AddCost(((v - vf)**2).mean() + ((s - sf)**2).mean())
solver = SnoptSolver()
result = solver.Solve(prog)
if not result.is_success():
accel = self.idm_backup.step(s_init[0], v_init[0], v_init[-1])
print(f'Optimization failed, using accel {accel} from IDM')
return accel
v_desired = result.GetSolution(v)
print('Planned speeds')
print(v_desired)
print('Planned spacings')
print(result.GetSolution(s))
a_desired = (v_desired[1:, 0] - v_desired[:-1, 0]
) / dt # we're optimizing the velocity of the 0th vehicle
self.plan = a_desired
self.plan_index = 1
return self.plan[0]
def compute_input(self, x, xd, initial_guess=None, tol=0.0):
prog = MathematicalProgram()
# Joint configuration states & Contact forces
q = prog.NewContinuousVariables(rows=self.T + 1,
cols=self.nq,
name='q')
v = prog.NewContinuousVariables(rows=self.T + 1,
cols=self.nq,
name='v')
u = prog.NewContinuousVariables(rows=self.T, cols=self.nu, name='u')
contact = prog.NewContinuousVariables(rows=self.T,
cols=self.nf,
name='lambda')
#
alpha = prog.NewContinuousVariables(rows=self.T, cols=2, name='alpha')
beta = prog.NewContinuousVariables(rows=self.T, cols=2, name='beta')
# Add Initial Condition Constraint
prog.AddConstraint(eq(q[0], np.array(x[0:3])))
prog.AddConstraint(eq(v[0], np.array(x[3:6])))
# Add Final Condition Constraint
prog.AddConstraint(eq(q[self.T], np.array(xd[0:3])))
prog.AddConstraint(eq(v[self.T], np.array(xd[3:6])))
# Add Dynamics Constraints
for t in range(self.T):
# Add Dynamics Constraints
prog.AddConstraint(
eq(q[t + 1], (q[t] + self.sim.params['h'] * v[t + 1])))
prog.AddConstraint(v[t + 1, 0] == (
v[t, 0] + self.sim.params['h'] *
(-self.sim.params['c'] * v[t, 0] - contact[t, 0] + u[t, 0])))
prog.AddConstraint(v[t + 1,
1] == (v[t, 1] + self.sim.params['h'] *
(-self.sim.params['c'] * v[t, 1] +
contact[t, 0] - contact[t, 1])))
prog.AddConstraint(v[t + 1, 2] == (
v[t, 2] + self.sim.params['h'] *
(-self.sim.params['c'] * v[t, 2] + contact[t, 1] + u[t, 1])))
# Add Contact Constraints
prog.AddConstraint(ge(alpha[t], 0))
prog.AddConstraint(ge(beta[t], 0))
prog.AddConstraint(alpha[t, 0] == contact[t, 0])
prog.AddConstraint(alpha[t, 1] == contact[t, 1])
prog.AddConstraint(
beta[t, 0] == (contact[t, 0] + self.sim.params['k'] *
(q[t, 1] - q[t, 0] - self.sim.params['d'])))
prog.AddConstraint(
beta[t, 1] == (contact[t, 1] + self.sim.params['k'] *
(q[t, 2] - q[t, 1] - self.sim.params['d'])))
# Complementarity constraints. Start with relaxed version and start constraining.
prog.AddConstraint(alpha[t, 0] * beta[t, 0] <= tol)
prog.AddConstraint(alpha[t, 1] * beta[t, 1] <= tol)
# Add Input Constraints and Contact Constraints
prog.AddConstraint(le(contact[t], self.contact_max))
prog.AddConstraint(ge(contact[t], -self.contact_max))
prog.AddConstraint(le(u[t], self.input_max))
prog.AddConstraint(ge(u[t], -self.input_max))
# Add Costs
prog.AddCost(u[t].dot(u[t]))
# Set Initial Guess as empty. Otherwise, start from last solver iteration.
if (type(initial_guess) == type(None)):
initial_guess = np.empty(prog.num_vars())
# Populate initial guess by linearly interpolating between initial
# and final states
#qinit = np.linspace(x[0:3], xd[0:3], self.T + 1)
qinit = np.tile(np.array(x[0:3]), (self.T + 1, 1))
vinit = np.tile(np.array(x[3:6]), (self.T + 1, 1))
uinit = np.tile(np.array([0, 0]), (self.T, 1))
finit = np.tile(np.array([0, 0]), (self.T, 1))
prog.SetDecisionVariableValueInVector(q, qinit, initial_guess)
prog.SetDecisionVariableValueInVector(v, vinit, initial_guess)
prog.SetDecisionVariableValueInVector(u, uinit, initial_guess)
prog.SetDecisionVariableValueInVector(contact, finit,
initial_guess)
# Solve the program
if (self.solver == "ipopt"):
solver = IpoptSolver()
elif (self.solver == "snopt"):
solver = SnoptSolver()
result = solver.Solve(prog, initial_guess)
if (self.solver == "ipopt"):
print("Ipopt Solver Status: ",
result.get_solver_details().status, ", meaning ",
result.get_solver_details().ConvertStatusToString())
elif (self.solver == "snopt"):
val = result.get_solver_details().info
status = self.snopt_status(val)
print("Snopt Solver Status: ",
result.get_solver_details().info, ", meaning ", status)
sol = result.GetSolution()
q_opt = result.GetSolution(q)
v_opt = result.GetSolution(v)
u_opt = result.GetSolution(u)
f_opt = result.GetSolution(contact)
return sol, q_opt, v_opt, u_opt, f_opt