def add_final_state_constraint_elastic_collision(self, prog, p0_puck, v0_puck, v_puck_desired):
"""Utility function for bounce kick from the wall. Probably not a linear constraint."""
m1 = self.params.player_mass
m2 = self.params.puck_mass
p1 = prog.final_state()[:2] # var: player's final position
p2 = p0_puck
v1 = prog.final_state()[2:] # var: player's final velocity
v2 = v0_puck
# Final position constraint
pf = p0_puck - self.get_normalized_vector(v_puck_desired)*(self.params.puck_radius + self.params.player_radius)
prog.AddConstraint(eq(p1, pf))
# Final velocity constraint
v_puck_after_collision = v2 - 2*m1/(m1+m2)*(v2-v1).dot(p2-pf)/(p2-pf).dot(p2-pf)*(p2-pf) # perhaps not a nonlinear constraint, since dot product of v1 is required
prog.AddConstraint(eq(v_puck_after_collision, v_puck_desired))
def intercepting_with_obs_avoidance(self,
p0,
v0,
pf,
vf,
T,
obstacles=None,
p_puck=None):
"""Intercepting trajectory that avoids obs"""
x0 = np.array(np.concatenate((p0, v0), axis=0))
xf = np.concatenate((pf, vf), axis=0)
prog = MathematicalProgram()
# state and control inputs
N = int(T / self.params.dt) # number of time steps
state = prog.NewContinuousVariables(N + 1, 4, 'state')
cmd = prog.NewContinuousVariables(N, 2, 'input')
# Initial and final state
prog.AddLinearConstraint(eq(state[0], x0))
#prog.AddLinearConstraint(eq(state[-1], xf))
prog.AddQuadraticErrorCost(Q=10.0 * np.eye(4),
x_desired=xf,
vars=state[-1])
self.add_dynamics(prog, N, state, cmd)
## Input saturation
self.add_input_limits(prog, cmd, N)
# Arena constraints
self.add_arena_limits(prog, state, N)
for k in range(N):
prog.AddQuadraticCost(cmd[k].flatten().dot(cmd[k]))
# Add non-linear constraints - will solve with SNOPT
# Avoid other players
if obstacles != None:
for obs in obstacles:
self.avoid_other_player_nl(prog, state, obs, N)
# avoid hitting the puck while generating a kicking trajectory
if not p_puck.any(None):
self.avoid_puck_nl(prog, state, N, p_puck)
solver = SnoptSolver()
result = solver.Solve(prog)
solution_found = result.is_success()
if not solution_found:
print("Solution not found for intercepting_with_obs_avoidance")
u_values = result.GetSolution(cmd)
return solution_found, u_values.T
def set_initial_and_goal_position(prog, terrain, decision_variables):
# unpack only decision variables needed in this function
position_left, position_right = decision_variables[:2]
# initial position of the feet of the robot
foot_offset = np.array([0, .2])
center = terrain.get_stone_by_name('initial').center
initial_position_left = center
initial_position_right = center - foot_offset
# goal position of the feet of the robot
center = terrain.get_stone_by_name('goal').center
goal_position_left = center
goal_position_right = center - foot_offset
# enforce initial position of the feet
prog.AddLinearConstraint(eq(position_left[0], initial_position_left))
prog.AddLinearConstraint(eq(position_right[0], initial_position_right))
# enforce goal position of the feet
prog.AddLinearConstraint(eq(position_left[-1], goal_position_left))
prog.AddLinearConstraint(eq(position_right[-1], goal_position_right))
], q[k]).evaluator().set_description(f"q[{k}] unit quaternion constraint"))
# Impose unit length constraint on the rotation axis.
(prog.AddConstraint(lambda x: [x @ x], [1], [1],
w_axis[k]).evaluator().set_description(
f"w_axis[{k}] unit axis constraint"))
for k in range(1, N):
(prog.AddConstraint(
lambda q_qprev_v_dt: backward_euler(q_qprev_v_dt),
lb=[0.0] * 4,
ub=[0.0] * 4,
vars=np.concatenate([
q[k], q[k - 1], w_axis[k], w_mag[k], dt[k]
])).evaluator().set_description(f"q[{k}] backward euler constraint"))
(prog.AddLinearConstraint(eq(q[0], np.array(
[1.0, 0.0, 0.0,
0.0]))).evaluator().set_description("Initial orientation constraint"))
(prog.AddLinearConstraint(
eq(
q[-1],
np.array([
-0.2955511242573139, 0.25532186031279896, 0.5106437206255979,
0.7659655809383968
]))).evaluator().set_description("Final orientation constraint"))
(prog.AddLinearConstraint(
ge(dt,
0.0)).evaluator().set_description("Timestep greater than 0 constraint"))
(prog.AddConstraint(
np.sum(dt) == 1.0).evaluator().set_description("Total time constriant"))
for k in range(N):
def intercepting_with_obs_avoidance_bb(self,
p0,
v0,
pf,
vf,
T,
obstacles=None,
p_puck=None):
"""kick while avoiding obstacles using big M formulation and branch and bound"""
x0 = np.array(np.concatenate((p0, v0), axis=0))
xf = np.concatenate((pf, vf), axis=0)
prog = MathematicalProgram()
# state and control inputs
N = int(T / self.params.dt) # number of command steps
state = prog.NewContinuousVariables(N + 1, 4, 'state')
cmd = prog.NewContinuousVariables(N, 2, 'input')
# Initial and final state
prog.AddLinearConstraint(eq(state[0], x0))
prog.AddLinearConstraint(eq(state[-1], xf))
self.add_dynamics(prog, N, state, cmd)
## Input saturation
self.add_input_limits(prog, cmd, N)
# Arena constraints
self.add_arena_limits(prog, state, N)
# Add Mixed-integer constraints - will solve with BB
# avoid hitting the puck while generating a kicking trajectory
# MILP formulation with B&B solver
x_obs_puck = prog.NewBinaryVariables(rows=N + 1,
cols=2) # obs x_min, obs x_max
y_obs_puck = prog.NewBinaryVariables(rows=N + 1,
cols=2) # obs y_min, obs y_max
self.avoid_puck_bigm(prog, x_obs_puck, y_obs_puck, state, N, p_puck)
# Avoid other players
if obstacles != None:
x_obs_player = list()
y_obs_player = list()
for i, obs in enumerate(obstacles):
x_obs_player.append(prog.NewBinaryVariables(
rows=N + 1, cols=2)) # obs x_min, obs x_max
y_obs_player.append(prog.NewBinaryVariables(
rows=N + 1, cols=2)) # obs y_min, obs y_max
self.avoid_other_player_bigm(prog, x_obs_player[i],
y_obs_player[i], state, obs, N)
# Solve with simple b&b solver
for k in range(N):
prog.AddQuadraticCost(cmd[k].flatten().dot(cmd[k]))
bb = branch_and_bound.MixedIntegerBranchAndBound(
prog,
OsqpSolver().solver_id())
result = bb.Solve()
if result != result.kSolutionFound:
raise ValueError('Infeasible optimization problem.')
u_values = np.array([bb.GetSolution(u) for u in cmd]).T
solution_found = result.kSolutionFound
return solution_found, u_values
def add_dynamics(self, prog, N, state, acc):
"""Add model dynamics to the program as equality constraints for N steps"""
for k in range(N):
prog.AddLinearConstraint(
eq(state[k + 1],
np.matmul(self.A, state[k]) + np.matmul(self.B, acc[k])))
def CalcOutput(self, context, output):
# ============================ LOAD INPUTS ============================
# Torque inputs
tau_g = np.expand_dims(
np.array(self.GetInputPort("tau_g").Eval(context)), 1)
joint_centering_torque = np.expand_dims(
np.array(
self.GetInputPort("joint_centering_torque").Eval(context)), 1)
# Force inputs
# PROGRAMMING: Add zeros here?
F_GT = self.GetInputPort("F_GT").Eval(context)[0]
F_GN = self.GetInputPort("F_GN").Eval(context)[0]
# Positions
theta_L = self.GetInputPort("theta_L").Eval(context)[0]
theta_MX = self.GetInputPort("theta_MX").Eval(context)[0]
theta_MY = self.GetInputPort("theta_MY").Eval(context)[0]
theta_MZ = self.GetInputPort("theta_MZ").Eval(context)[0]
p_CT = self.GetInputPort("p_CT").Eval(context)[0]
p_CN = self.GetInputPort("p_CN").Eval(context)[0]
p_LT = self.GetInputPort("p_LT").Eval(context)[0]
p_LN = self.GetInputPort("p_LN").Eval(context)[0]
d_T = self.GetInputPort("d_T").Eval(context)[0]
d_N = self.GetInputPort("d_N").Eval(context)[0]
d_X = self.GetInputPort("d_X").Eval(context)[0]
p_MConM = np.array([self.GetInputPort("p_MConM").Eval(context)]).T
# Velocities
d_theta_L = self.GetInputPort("d_theta_L").Eval(context)[0]
d_theta_MX = self.GetInputPort("d_theta_MX").Eval(context)[0]
d_theta_MY = self.GetInputPort("d_theta_MY").Eval(context)[0]
d_theta_MZ = self.GetInputPort("d_theta_MZ").Eval(context)[0]
d_d_T = self.GetInputPort("d_d_T").Eval(context)[0]
d_d_N = self.GetInputPort("d_d_N").Eval(context)[0]
d_d_X = self.GetInputPort("d_d_X").Eval(context)[0]
# Manipulator inputs
J = np.array(self.GetInputPort("J").Eval(context)).reshape(
(6, self.nq))
J_translational = np.array(
self.GetInputPort("J_translational").Eval(context)).reshape(
(3, self.nq))
J_rotational = np.array(
self.GetInputPort("J_rotational").Eval(context)).reshape(
(3, self.nq))
Jdot_qdot = np.expand_dims(
np.array(self.GetInputPort("Jdot_qdot").Eval(context)), 1)
M = np.array(self.GetInputPort("M").Eval(context)).reshape(
(self.nq, self.nq))
Cv = np.expand_dims(np.array(self.GetInputPort("Cv").Eval(context)), 1)
# Other inputs
mu_S = self.GetInputPort("mu_S").Eval(context)[0]
hats_T = self.GetInputPort("hats_T").Eval(context)[0]
s_hat_X = self.GetInputPort("s_hat_X").Eval(context)[0]
# ============================= OTHER PREP ============================
if self.theta_MYd is None:
self.theta_MYd = theta_MY
if self.theta_MZd is None:
self.theta_MZd = theta_MZ
# Calculate desired values
dd_d_Td = 1000 * (self.d_Td - d_T) - 100 * d_d_T
dd_theta_Ld = 10 * (self.d_theta_Ld - d_theta_L)
a_MX_d = 100 * (self.d_Xd - d_X) - 10 * d_d_X
theta_MXd = theta_L
alpha_MXd = 100 * (theta_MXd - theta_MX) + 10 * (d_theta_L -
d_theta_MX)
alpha_MYd = 10 * (self.theta_MYd - theta_MY) - 5 * d_theta_MY
alpha_MZd = 10 * (self.theta_MZd - theta_MZ) - 5 * d_theta_MZ
dd_d_Nd = 0
# =========================== SOLVE PROGRAM ===========================
## 1. Define an instance of MathematicalProgram
prog = MathematicalProgram()
## 2. Add decision variables
# Contact values
F_NM = prog.NewContinuousVariables(1, 1, name="F_NM")
F_ContactMY = prog.NewContinuousVariables(1, 1, name="F_ContactMY")
F_ContactMZ = prog.NewContinuousVariables(1, 1, name="F_ContactMZ")
F_NL = prog.NewContinuousVariables(1, 1, name="F_NL")
if self.model_friction:
F_FMT = prog.NewContinuousVariables(1, 1, name="F_FMT")
F_FMX = prog.NewContinuousVariables(1, 1, name="F_FMX")
F_FLT = prog.NewContinuousVariables(1, 1, name="F_FLT")
F_FLX = prog.NewContinuousVariables(1, 1, name="F_FLX")
else:
F_FMT = np.array([[0]])
F_FMX = np.array([[0]])
F_FLT = np.array([[0]])
F_FLX = np.array([[0]])
F_ContactM_XYZ = np.array([F_FMX, F_ContactMY, F_ContactMZ])[:, :, 0]
# Object forces and torques
if not self.measure_joint_wrench:
F_OT = prog.NewContinuousVariables(1, 1, name="F_OT")
F_ON = prog.NewContinuousVariables(1, 1, name="F_ON")
tau_O = -self.sys_consts.k_J*theta_L \
- self.sys_consts.b_J*d_theta_L
# Control values
tau_ctrl = prog.NewContinuousVariables(self.nq, 1, name="tau_ctrl")
# Object accelerations
a_MX = prog.NewContinuousVariables(1, 1, name="a_MX")
a_MT = prog.NewContinuousVariables(1, 1, name="a_MT")
a_MY = prog.NewContinuousVariables(1, 1, name="a_MY")
a_MZ = prog.NewContinuousVariables(1, 1, name="a_MZ")
a_MN = prog.NewContinuousVariables(1, 1, name="a_MN")
a_LT = prog.NewContinuousVariables(1, 1, name="a_LT")
a_LN = prog.NewContinuousVariables(1, 1, name="a_LN")
alpha_MX = prog.NewContinuousVariables(1, 1, name="alpha_MX")
alpha_MY = prog.NewContinuousVariables(1, 1, name="alpha_MY")
alpha_MZ = prog.NewContinuousVariables(1, 1, name="alpha_MZ")
alpha_a_MXYZ = np.array(
[alpha_MX, alpha_MY, alpha_MZ, a_MX, a_MY, a_MZ])[:, :, 0]
# Derived accelerations
dd_theta_L = prog.NewContinuousVariables(1, 1, name="dd_theta_L")
dd_d_N = prog.NewContinuousVariables(1, 1, name="dd_d_N")
dd_d_T = prog.NewContinuousVariables(1, 1, name="dd_d_T")
ddq = prog.NewContinuousVariables(self.nq, 1, name="ddq")
## Constraints
# "set_description" calls gives us useful names for printing
prog.AddConstraint(eq(
self.sys_consts.m_L * a_LT, F_FLT + F_GT +
F_OT)).evaluator().set_description("Link tangential force balance")
prog.AddConstraint(
eq(self.sys_consts.m_L * a_LN, F_NL + F_GN +
F_ON)).evaluator().set_description("Link normal force balance")
prog.AddConstraint(eq(
self.sys_consts.I_L*dd_theta_L, \
(-self.sys_consts.w_L/2)*F_ON - (p_CN-p_LN) * F_FLT + \
(p_CT-p_LT)*F_NL + tau_O
)).evaluator().set_description("Link moment balance")
prog.AddConstraint(eq(
F_NL, -F_NM)).evaluator().set_description("3rd law normal forces")
if self.model_friction:
prog.AddConstraint(eq(F_FMT, -F_FLT)).evaluator().set_description(
"3rd law friction forces (T hat)")
prog.AddConstraint(eq(F_FMX, -F_FLX)).evaluator().set_description(
"3rd law friction forces (X hat)")
prog.AddConstraint(eq(
-dd_theta_L*(self.sys_consts.h_L/2+self.sys_consts.r) + \
d_theta_L**2*self.sys_consts.w_L/2 - a_LT + a_MT,
-dd_theta_L*d_N + dd_d_T - d_theta_L**2*d_T - 2*d_theta_L*d_d_N
)).evaluator().set_description("d_N derivative is derivative")
prog.AddConstraint(eq(
-dd_theta_L*self.sys_consts.w_L/2 - \
d_theta_L**2*self.sys_consts.h_L/2 - \
d_theta_L**2*self.sys_consts.r - a_LN + a_MN,
dd_theta_L*d_T + dd_d_N - d_theta_L**2*d_N + 2*d_theta_L*d_d_T
)).evaluator().set_description("d_N derivative is derivative")
prog.AddConstraint(eq(dd_d_N,
0)).evaluator().set_description("No penetration")
if self.model_friction:
prog.AddConstraint(
eq(F_FLT, mu_S * F_NL * self.sys_consts.mu *
hats_T)).evaluator().set_description(
"Friction relationship LT")
prog.AddConstraint(
eq(F_FLX, mu_S * F_NL * self.sys_consts.mu *
s_hat_X)).evaluator().set_description(
"Friction relationship LX")
if not self.measure_joint_wrench:
prog.AddConstraint(
eq(a_LT, -(self.sys_consts.w_L / 2) * d_theta_L**
2)).evaluator().set_description("Hinge constraint (T hat)")
prog.AddConstraint(eq(a_LN, (self.sys_consts.w_L / 2) *
dd_theta_L)).evaluator().set_description(
"Hinge constraint (N hat)")
for i in range(6):
lhs_i = alpha_a_MXYZ[i, 0]
assert not hasattr(lhs_i, "shape")
rhs_i = (Jdot_qdot + np.matmul(J, ddq))[i, 0]
assert not hasattr(rhs_i, "shape")
prog.AddConstraint(lhs_i == rhs_i).evaluator().set_description(
"Relate manipulator and end effector with joint " + \
"accelerations " + str(i))
tau_contact_trn = np.matmul(J_translational.T, F_ContactM_XYZ)
tau_contact_rot = np.matmul(J_rotational.T,
np.cross(p_MConM, F_ContactM_XYZ, axis=0))
tau_contact = tau_contact_trn + tau_contact_rot
tau_out = tau_ctrl - tau_g + joint_centering_torque
for i in range(self.nq):
M_ddq_row_i = (np.matmul(M, ddq) + Cv)[i, 0]
assert not hasattr(M_ddq_row_i, "shape")
tau_i = (tau_g + tau_contact + tau_out)[i, 0]
assert not hasattr(tau_i, "shape")
prog.AddConstraint(M_ddq_row_i == tau_i).evaluator(
).set_description("Manipulator equations " + str(i))
# Projection equations
prog.AddConstraint(
eq(a_MT,
np.cos(theta_L) * a_MY + np.sin(theta_L) * a_MZ))
prog.AddConstraint(
eq(a_MN, -np.sin(theta_L) * a_MY + np.cos(theta_L) * a_MZ))
prog.AddConstraint(
eq(F_FMT,
np.cos(theta_L) * F_ContactMY + np.sin(theta_L) * F_ContactMZ))
prog.AddConstraint(
eq(F_NM,
-np.sin(theta_L) * F_ContactMY + np.cos(theta_L) * F_ContactMZ))
prog.AddConstraint(dd_d_T[0,
0] == dd_d_Td).evaluator().set_description(
"Desired dd_d_Td constraint" + str(i))
prog.AddConstraint(dd_theta_L[0, 0] == dd_theta_Ld).evaluator(
).set_description("Desired a_LN constraint" + str(i))
prog.AddConstraint(a_MX[0, 0] == a_MX_d).evaluator().set_description(
"Desired a_MX constraint" + str(i))
prog.AddConstraint(alpha_MX[0, 0] == alpha_MXd).evaluator(
).set_description("Desired alpha_MX constraint" + str(i))
prog.AddConstraint(alpha_MY[0, 0] == alpha_MYd).evaluator(
).set_description("Desired alpha_MY constraint" + str(i))
prog.AddConstraint(alpha_MZ[0, 0] == alpha_MZd).evaluator(
).set_description("Desired alpha_MZ constraint" + str(i))
prog.AddConstraint(dd_d_N[0,
0] == dd_d_Nd).evaluator().set_description(
"Desired dd_d_N constraint" + str(i))
result = Solve(prog)
assert result.is_success()
tau_ctrl_result = []
for i in range(self.nq):
tau_ctrl_result.append(
result.GetSolution()[prog.FindDecisionVariableIndex(
tau_ctrl[i, 0])])
tau_ctrl_result = np.expand_dims(tau_ctrl_result, 1)
# ======================== UPDATE DEBUG VALUES ========================
self.debug["times"].append(context.get_time())
# control effort
self.debug["dd_d_Td"].append(dd_d_Td)
self.debug["dd_theta_Ld"].append(dd_theta_Ld)
self.debug["a_MX_d"].append(a_MX_d)
self.debug["alpha_MXd"].append(alpha_MXd)
self.debug["alpha_MYd"].append(alpha_MYd)
self.debug["alpha_MZd"].append(alpha_MZd)
self.debug["dd_d_Nd"].append(dd_d_Nd)
# decision variables
if self.model_friction:
self.debug["F_FMT"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(F_FMT[0,
0])])
self.debug["F_FMX"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(F_FMX[0,
0])])
self.debug["F_FLT"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(F_FLT[0,
0])])
self.debug["F_FLX"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(F_FLX[0,
0])])
else:
self.debug["F_FMT"].append(F_FMT)
self.debug["F_FMX"].append(F_FMX)
self.debug["F_FLT"].append(F_FLT)
self.debug["F_FLX"].append(F_FLX)
self.debug["F_NM"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(F_NM[0, 0])])
self.debug["F_ContactMY"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(
F_ContactMY[0, 0])])
self.debug["F_ContactMZ"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(
F_ContactMZ[0, 0])])
self.debug["F_NL"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(F_NL[0, 0])])
for i in range(self.nq):
self.debug["tau_ctrl_" + str(i)].append(
result.GetSolution()[prog.FindDecisionVariableIndex(
tau_ctrl[i, 0])])
self.debug["a_MX"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(a_MX[0, 0])])
self.debug["a_MT"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(a_MT[0, 0])])
self.debug["a_MY"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(a_MY[0, 0])])
self.debug["a_MZ"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(a_MZ[0, 0])])
self.debug["a_MN"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(a_MN[0, 0])])
self.debug["a_LT"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(a_LT[0, 0])])
self.debug["a_LN"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(a_LN[0, 0])])
self.debug["alpha_MX"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(alpha_MX[0,
0])])
self.debug["alpha_MY"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(alpha_MY[0,
0])])
self.debug["alpha_MZ"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(alpha_MZ[0,
0])])
self.debug["dd_theta_L"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(
dd_theta_L[0, 0])])
self.debug["dd_d_N"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(dd_d_N[0, 0])])
self.debug["dd_d_T"].append(
result.GetSolution()[prog.FindDecisionVariableIndex(dd_d_T[0, 0])])
for i in range(self.nq):
self.debug["ddq_" + str(i)].append(
result.GetSolution()[prog.FindDecisionVariableIndex(ddq[i,
0])])
self.debug["theta_MXd"].append(theta_MXd)
output.SetFromVector(tau_ctrl_result.flatten())
def direct_transcription():
prog = MathematicalProgram()
N = 500
dt = 0.01
# Create decision variables
n_x = 6
n_u = 2
x = np.empty((n_x, N), dtype=Variable)
u = np.empty((n_u, N - 1), dtype=Variable)
for n in range(N - 1):
x[:, n] = prog.NewContinuousVariables(n_x, "x" + str(n))
u[:, n] = prog.NewContinuousVariables(n_u, "u" + str(n))
x[:, N - 1] = prog.NewContinuousVariables(n_x, "x" + str(N))
T = N - 1
# Add constraints
# x0 = np.array([10, -np.pi / 2, 0, 40, 0, 0])
# Slotine dynamics: x = [airspeed, heading, flight_path_angle, z, x, y]
for n in range(N - 1):
# Dynamics
prog.AddConstraint(
eq(x[:, n + 1],
x[:, n] + dt * continuous_dynamics(x[:, n], u[:, n])))
# Never below sea level
prog.AddConstraint(x[3, n + 1] >= 0 + 0.5)
# Always positive velocity
prog.AddConstraint(x[0, n + 1] >= 0)
# TODO use residuals
# Add periodic constraints
prog.AddConstraint(x[0, 0] - x[0, T] == 0)
prog.AddConstraint(x[1, 0] - x[1, T] == 0)
prog.AddConstraint(x[2, 0] - x[2, T] == 0)
prog.AddConstraint(x[3, 0] - x[3, T] == 0)
# Start at 20 meter height
prog.AddConstraint(x[4, 0] == 0)
prog.AddConstraint(x[5, 0] == 0)
h0 = 20
prog.AddConstraint(x[3, 0] == h0)
# Add cost function
p0 = x[4:6, 0]
p1 = x[4:6, T]
travel_angle = np.pi # TODO start along y direction
dir_vector = np.array([np.sin(travel_angle), np.cos(travel_angle)])
prog.AddCost(dir_vector.T.dot(p1 - p0)) # Maximize distance travelled
print("Initialized opt problem")
# Initial guess
V0 = 10
x_guess = np.zeros((n_x, N))
x_guess[:, 0] = np.array([V0, travel_angle, 0, h0, 0, 0])
for n in range(N - 1):
# Constant airspeed, heading and height
x_guess[0, n + 1] = V0
x_guess[1, n + 1] = travel_angle
x_guess[3, n + 1] = h0
# Interpolate position
avg_speed = 10 # conservative estimate
x_guess[4:6, n + 1] = dir_vector * avg_speed * n * dt
# Let the rest of the variables be initialized to zero
prog.SetInitialGuess(x, x_guess)
# solve mathematical program
solver = SnoptSolver()
result = solver.Solve(prog)
# be sure that the solution is optimal
# assert result.is_success()
# retrieve optimal solution
thrust_opt = result.GetSolution(x)
state_opt = result.GetSolution(u)
result = Solve(prog)
x_sol = result.GetSolution(x)
u_sol = result.GetSolution(u)
breakpoint()
# Slotine dynamics: x = [airspeed, heading, flight_path_angle, z, x, y]
z = x_sol[:, 3]
x = x_sol[:, 4]
y = x_sol[:, 5]
plot_trj_3_wind(np.vstack((x, y, z)).T, get_wind_field)
return
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
def compute_control(self, x_p1, x_p2, x_puck, p_goal, obstacles):
"""Solve for initial velocity for the puck to bounce the wall once and hit the goal."""
"""Currently does not work"""
# initialize program
N = self.mpc_params.N # length of receding horizon
prog = MathematicalProgram()
# State and input variables
x1 = prog.NewContinuousVariables(N + 1, 4,
name='p1_state') # state of player 1
u1 = prog.NewContinuousVariables(
N, 2, name='p1_input') # inputs for player 1
xp = prog.NewContinuousVariables(
N + 1, 4, name='puck_state') # state of the puck
# Slack variables
t1_kick = prog.NewContinuousVariables(
N + 1, name='kick_time'
) # slack variables that captures if player 1 is kicking or not at the given time step
# Defined as continous, but treated as mixed integer. 1 when kicking
v1_kick = prog.NewContinuousVariables(
N + 1, 2, name='v1_kick') # velocity of player after kick to puck
vp_kick = prog.NewContinuousVariables(
N + 1, 2, name='vp_kick'
) # velocity of puck after being kicked by the player
dist = prog.NewContinuousVariables(
N + 1, name='dist_puck_player') # distance between puck and player
cost = prog.NewContinuousVariables(
N + 1, name='cost') # slack variable to monitor cost
# Compute distance between puck and player as slack variable.
for k in range(N + 1):
r1 = self.sim_params.player_radius
rp = self.sim_params.puck_radius
prog.AddConstraint(
dist[k] == (x1[k, 0:2] - xp[k, 0:2]).dot(x1[k, 0:2] -
xp[k, 0:2]) -
(r1 + rp)**2)
#### COST and final states
# TODO: find adequate final velocity
x_puck_des = np.concatenate(
(p_goal, np.zeros(2)), axis=0) # desired position and vel for puck
for k in range(N + 1):
Q_dist_puck_goal = 10 * np.eye(2)
q_dist_puck_player = 0.1
e1 = x_puck_des[0:2] - xp[k, 0:2]
e2 = xp[k, 0:2] - x1[k, 0:2]
prog.AddConstraint(
cost[k] == e1.dot(np.matmul(Q_dist_puck_goal, e1)) +
q_dist_puck_player * dist[k])
prog.AddCost(cost[k])
#prog.AddQuadraticErrorCost(Q=self.mpc_params.Q_puck, x_desired=x_puck_des, vars=xp[k]) # puck in the goal
#prog.AddQuadraticErrorCost(Q=np.eye(2), x_desired=x_puck_des[0:2], vars=x1[k, 0:2])
#prog.AddQuadraticErrorCost(Q=10*np.eye(2), x_desired=np.array([2, 2]), vars=x1[k, 0:2]) # TEST: control position of the player instead of puck
#for k in range(N):
# prog.AddQuadraticCost(1e-2*u1[k].dot(u1[k])) # be wise on control effort
# Initial states for puck and player
prog.AddBoundingBoxConstraint(x_p1, x_p1,
x1[0]) # Intial state for player 1
prog.AddBoundingBoxConstraint(x_puck, x_puck,
xp[0]) # Initial state for puck
# Compute elastic collision for every possible timestep.
for k in range(N + 1):
v_puck_bfr = xp[k, 2:4]
v_player_bfr = x1[k, 2:4]
v_puck_aft, v_player_aft = self.get_reset_velocities(
v_puck_bfr, v_player_bfr)
prog.AddConstraint(eq(vp_kick[k], v_puck_aft))
prog.AddConstraint(eq(v1_kick[k], v_player_aft))
# Use slack variable to activate guard condition based on distance.
M = 15**2
for k in range(N + 1):
prog.AddLinearConstraint(dist[k] <= M * (1 - t1_kick[k]))
prog.AddLinearConstraint(dist[k] >= -t1_kick[k] * M)
# Hybrid dynamics for player
#for k in range(N):
# A = self.mpc_params.A_player
# B = self.mpc_params.B_player
# x1_kick = np.concatenate((x1[k][0:2], v1_kick[k]), axis=0) # The state of the player after it kicks the puck
# x1_next = np.matmul(A, (1 - t1_kick[k])*x1[k] + t1_kick[k]*x1_kick) + np.matmul(B, u1[k])
# prog.AddConstraint(eq(x1[k+1], x1_next))
# Assuming player dynamics are not affected by collision
for k in range(N):
A = self.mpc_params.A_player
B = self.mpc_params.B_player
x1_next = np.matmul(A, x1[k]) + np.matmul(B, u1[k])
prog.AddConstraint(eq(x1[k + 1], x1_next))
# Hybrid dynamics for puck_mass
for k in range(N):
A = self.mpc_params.A_puck
xp_kick = np.concatenate(
(xp[k][0:2], vp_kick[k]),
axis=0) # State of the puck after it gets kicked
xp_next = np.matmul(A, (1 - t1_kick[k]) * xp[k] +
t1_kick[k] * xp_kick)
prog.AddConstraint(eq(xp[k + 1], xp_next))
# Generate trajectory that is not in direct collision with the puck
for k in range(N + 1):
eps = 0.1
prog.AddConstraint(dist[k] >= -eps)
# Input and arena constraint
self.add_input_limits(prog, u1, N)
self.add_arena_limits(prog, x1, N)
self.add_arena_limits(prog, xp, N)
# fake mixed-integer constraint
#for k in range(N+1):
# prog.AddConstraint(t1_kick[k]*(1-t1_kick[k])==0)
# Hot-start
guess_u1, guess_x1 = self.get_initial_guess(x_p1, p_goal, x_puck[0:2])
prog.SetInitialGuess(x1, guess_x1)
prog.SetInitialGuess(u1, guess_u1)
if not self.prev_xp is None:
prog.SetInitialGuess(xp, self.prev_xp)
#prog.SetInitialGuess(t1_kick, np.ones_like(t1_kick))
# solve for the periods
# solver = SnoptSolver()
#result = solver.Solve(prog)
#if not result.is_success():
# print("Unable to find solution.")
# save for hot-start
#self.prev_xp = result.GetSolution(xp)
#if True:
# print("x1:{}".format(result.GetSolution(x1)))
# print("u1: {}".format( result.GetSolution(u1)))
# print("xp: {}".format( result.GetSolution(xp)))
# print('dist{}'.format(result.GetSolution(dist)))
# print("t1_kick:{}".format(result.GetSolution(t1_kick)))
# print("v1_kick:{}".format(result.GetSolution(v1_kick)))
# print("vp_kick:{}".format(result.GetSolution(vp_kick)))
# print("cost:{}".format(result.GetSolution(cost)))
# return whether successful, and the initial player velocity
#u1_opt = result.GetSolution(u1)
return True, guess_u1[0, :], np.zeros((2))
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):
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='lambda1')
z = prog.NewBinaryVariables(rows=self.T, cols=self.nf, name='z')
# 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])))
prog.AddConstraint(z[0, 0] == 0)
prog.AddConstraint(z[0, 1] == 0)
# 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 with big M = self.contact
prog.AddConstraint(ge(contact[t], 0))
prog.AddConstraint(contact[t, 0] + self.sim.params['k'] *
(q[t, 1] - q[t, 0] - self.sim.params['d']) >= 0)
prog.AddConstraint(contact[t, 1] + self.sim.params['k'] *
(q[t, 2] - q[t, 1] - self.sim.params['d']) >= 0)
# Mixed Integer Constraints
M = self.contact_max
prog.AddConstraint(contact[t, 0] <= M)
prog.AddConstraint(contact[t, 1] <= M)
prog.AddConstraint(contact[t, 0] <= M * z[t, 0])
prog.AddConstraint(contact[t, 1] <= M * z[t, 1])
prog.AddConstraint(
contact[t, 0] + self.sim.params['k'] *
(q[t, 1] - q[t, 0] - self.sim.params['d']) <= M *
(1 - z[t, 0]))
prog.AddConstraint(
contact[t, 1] + self.sim.params['k'] *
(q[t, 2] - q[t, 1] - self.sim.params['d']) <= M *
(1 - z[t, 1]))
prog.AddConstraint(z[t, 0] + z[t, 1] == 1)
# Add Input Constraints. Contact Constraints already enforced in big-M
# 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_id = IpoptSolver().solver_id()
elif (self.solver == "snopt"):
solver_id = SnoptSolver().solver_id()
elif (self.solver == "osqp"):
solver_id = OsqpSolver().solver_id()
elif (self.solver == "mosek"):
solver_id = MosekSolver().solver_id()
elif (self.solver == "gurobi"):
solver_id = GurobiSolver().solver_id()
solver = MixedIntegerBranchAndBound(prog, solver_id)
#result = solver.Solve(prog, initial_guess)
result = solver.Solve()
if result != result.kSolutionFound:
raise ValueError('Infeasible optimization problem')
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
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
def trajopt_simulation(x0, xf, u_max=0.5): # numeric parameters time_interval = .1 time_steps = 400 # initialize optimization prog = MathematicalProgram() # optimization variables state = prog.NewContinuousVariables(time_steps + 1, 3, 'state') u = prog.NewContinuousVariables(time_steps, 1, 'u') # final position constraint # for residual in constraint_state_to_orbit(state[-2], state[-1], xf, 1/u_max, time_interval): # prog.AddConstraint(residual == 0) # initial position constraint prog.AddConstraint(eq(state[0],x0)) # discretized dynamics for t in range(time_steps): residuals = dubins_discrete_dynamics(state[t], state[t+1], u[t], time_interval) for residual in residuals: prog.AddConstraint(residual == 0) # control limits prog.AddConstraint(u[t][0] <= u_max) prog.AddConstraint(-u_max <= u[t][0]) # cost - increases cost if off the orbit p = state[t][:2] - xf[:2] v = (state[t+1][:2] - state[t][:2]) / time_interval # prog.AddCost(time_interval) # prog.AddCost(u[t][0]*u[t][0]*time_interval) # prog.AddCost(p.dot(v)*time_interval) for residual in constraint_state_to_orbit(state[t], state[t+1], xf, 1/u_max, time_interval): # prog.AddConstraint(residual >= 0) prog.AddQuadraticCost(residual) # prog.AddCost((p.dot(p))) # # initial guess state_guess = interpolate_dubins_state( np.array([0, 0, np.pi]), xf, time_steps, time_interval, 1/u_max ) prog.SetInitialGuess(state, state_guess) solver = SnoptSolver() result = solver.Solve(prog) assert result.is_success() # retrieve optimal solution u_opt = result.GetSolution(u).T state_opt = result.GetSolution(state).T return state_opt, u_opt
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
imp = prog.NewContinuousVariables(nf, name='imp')
# generalized velocity after the heel strike
# (if "mirrored", this velocity must coincide with the
# initial velocity qd[0] to ensure periodicity)
qd_post = prog.NewContinuousVariables(nq, name='qd_post')
"""Here are part of the constraints of the optimization problem."""
# lower and upper bound on the time steps for all t
prog.AddBoundingBoxConstraint([h_min] * T, [h_max] * T, h)
# link the configurations, velocities, and accelerations
# uses implicit Euler method, https://en.wikipedia.org/wiki/Backward_Euler_method
for t in range(T):
prog.AddConstraint(eq(q[t+1], q[t] + h[t] * qd[t+1]))
prog.AddConstraint(eq(qd[t+1], qd[t] + h[t] * qdd[t]))
# manipulator equations for all t (implicit Euler)
for t in range(T):
vars = np.concatenate((q[t+1], qd[t+1], qdd[t], f[t]))
prog.AddConstraint(manipulator_equations, lb=[0]*nq, ub=[0]*nq, vars=vars)
# velocity reset across heel strike
# see http://underactuated.mit.edu/multibody.html#impulsive_collision
vars = np.concatenate((q[-1], qd[-1], qd_post, imp))
prog.AddConstraint(reset_velocity_heelstrike, lb=[0]*(nq+nf), ub=[0]*(nq+nf), vars=vars)
# mirror initial and final configuration
# see "The Walking Cycle" section of this notebook
prog.AddLinearConstraint(eq(q[0], - q[-1]))
