def __init__(self, rtree, q_nom, control_period=0.001):
self.rtree = rtree
self.nq = rtree.get_num_positions()
self.nv = rtree.get_num_velocities()
self.nu = rtree.get_num_actuators()
dim = 3 # 3D
nd = 4 # for friction cone approx, hard coded for now
self.nc = 4 # number of contacts; TODO don't hard code
self.nf = self.nc * nd # number of contact force variables
self.ncf = 2 # number of loop constraint forces; TODO don't hard code
self.neps = self.nc * dim # number of slack variables for contact
self.q_nom = q_nom
self.com_des = rtree.centerOfMass(self.rtree.doKinematics(q_nom))
self.u_des = CassieFixedPointTorque()
self.lfoot = rtree.FindBody("toe_left")
self.rfoot = rtree.FindBody("toe_right")
self.springs = 2 + np.array([
rtree.FindIndexOfChildBodyOfJoint("knee_spring_left_fixed"),
rtree.FindIndexOfChildBodyOfJoint("knee_spring_right_fixed"),
rtree.FindIndexOfChildBodyOfJoint("ankle_spring_joint_left"),
rtree.FindIndexOfChildBodyOfJoint("ankle_spring_joint_right")
])
umin = np.zeros(self.nu)
umax = np.zeros(self.nu)
ii = 0
for motor in rtree.actuators:
umin[ii] = motor.effort_limit_min
umax[ii] = motor.effort_limit_max
ii += 1
slack_limit = 10.0
self.initialized = False
#------------------------------------------------------------
# Add Decision Variables ------------------------------------
prog = MathematicalProgram()
qddot = prog.NewContinuousVariables(self.nq, "joint acceleration")
u = prog.NewContinuousVariables(self.nu, "input")
bar = prog.NewContinuousVariables(self.ncf, "four bar forces")
beta = prog.NewContinuousVariables(self.nf, "friction forces")
eps = prog.NewContinuousVariables(self.neps, "slack variable")
nparams = prog.num_vars()
#------------------------------------------------------------
# Problem Constraints ---------------------------------------
self.con_u_lim = prog.AddBoundingBoxConstraint(umin, umax,
u).evaluator()
self.con_fric_lim = prog.AddBoundingBoxConstraint(0, 100000,
beta).evaluator()
self.con_slack_lim = prog.AddBoundingBoxConstraint(
-slack_limit, slack_limit, eps).evaluator()
bar_con = np.zeros((self.ncf, self.nq))
self.con_4bar = prog.AddLinearEqualityConstraint(
bar_con, np.zeros(self.ncf), qddot).evaluator()
if self.nc > 0:
dyn_con = np.zeros(
(self.nq, self.nq + self.nu + self.ncf + self.nf))
dyn_vars = np.concatenate((qddot, u, bar, beta))
self.con_dyn = prog.AddLinearEqualityConstraint(
dyn_con, np.zeros(self.nq), dyn_vars).evaluator()
foot_con = np.zeros((self.neps, self.nq + self.neps))
foot_vars = np.concatenate((qddot, eps))
self.con_foot = prog.AddLinearEqualityConstraint(
foot_con, np.zeros(self.neps), foot_vars).evaluator()
else:
dyn_con = np.zeros((self.nq, self.nq + self.nu + self.ncf))
dyn_vars = np.concatenate((qddot, u, bar))
self.con_dyn = prog.AddLinearEqualityConstraint(
dyn_con, np.zeros(self.nq), dyn_vars).evaluator()
#------------------------------------------------------------
# Problem Costs ---------------------------------------------
self.Kp_com = 50
self.Kd_com = 2.0 * sqrt(self.Kp_com)
# self.Kp_qdd = 10*np.eye(self.nq)#np.diag(np.diag(H)/np.max(np.diag(H)))
self.Kp_qdd = np.diag(
np.concatenate(([10, 10, 10, 5, 5, 5], np.zeros(self.nq - 6))))
self.Kd_qdd = 1.0 * np.sqrt(self.Kp_qdd)
self.Kp_qdd[self.springs, self.springs] = 0
self.Kd_qdd[self.springs, self.springs] = 0
com_ddot_des = np.zeros(dim)
qddot_des = np.zeros(self.nq)
self.w_com = 0
self.w_qdd = np.zeros(self.nq)
self.w_qdd[:6] = 10
self.w_qdd[8:10] = 1
self.w_qdd[3:6] = 1000
# self.w_qdd[self.springs] = 0
self.w_u = 0.001 * np.ones(self.nu)
self.w_u[2:4] = 0.01
self.w_slack = 0.1
self.cost_qdd = prog.AddQuadraticErrorCost(np.diag(self.w_qdd),
qddot_des,
qddot).evaluator()
self.cost_u = prog.AddQuadraticErrorCost(np.diag(self.w_u), self.u_des,
u).evaluator()
self.cost_slack = prog.AddQuadraticErrorCost(
self.w_slack * np.eye(self.neps), np.zeros(self.neps),
eps).evaluator()
# self.cost_com = prog.AddQuadraticCost().evaluator()
# self.cost_qdd = prog.AddQuadraticCost(
# 2*np.diag(self.w_qdd), -2*np.matmul(qddot_des, np.diag(self.w_qdd)), qddot).evaluator()
# self.cost_u = prog.AddQuadraticCost(
# 2*np.diag(self.w_u), -2*np.matmul(self.u_des, np.diag(self.w_u)) , u).evaluator()
# self.cost_slack = prog.AddQuadraticCost(
# 2*self.w_slack*np.eye(self.neps), np.zeros(self.neps), eps).evaluator()
REG = 1e-8
self.cost_reg = prog.AddQuadraticErrorCost(
REG * np.eye(nparams), np.zeros(nparams),
prog.decision_variables()).evaluator()
# self.cost_reg = prog.AddQuadraticCost(2*REG*np.eye(nparams),
# np.zeros(nparams), prog.decision_variables()).evaluator()
#------------------------------------------------------------
# Solver settings -------------------------------------------
self.solver = GurobiSolver()
# self.solver = OsqpSolver()
prog.SetSolverOption(SolverType.kGurobi, "Method", 2)
#------------------------------------------------------------
# Save MathProg Variables -----------------------------------
self.qddot = qddot
self.u = u
self.bar = bar
self.beta = beta
self.eps = eps
self.prog = prog
def compute_trajectory_to_other_world(self, state_initial, minimum_time,
maximum_time):
'''
Your mission is to implement this function.
A successful implementation of this function will compute a dynamically feasible trajectory
which satisfies these criteria:
- Efficiently conserve fuel
- Reach "orbit" of the far right world
- Approximately obey the dynamic constraints
- Begin at the state_initial provided
- Take no more than maximum_time, no less than minimum_time
The above are defined more precisely in the provided notebook.
Please note there are two return args.
:param: state_initial: :param state_initial: numpy array of length 4, see rocket_dynamics for documentation
:param: minimum_time: float, minimum time allowed for trajectory
:param: maximum_time: float, maximum time allowed for trajectory
:return: three return args separated by commas:
trajectory, input_trajectory, time_array
trajectory: a 2d array with N rows, and 4 columns. See simulate_states_over_time for more documentation.
input_trajectory: a 2d array with N-1 row, and 2 columns. See simulate_states_over_time for more documentation.
time_array: an array with N rows.
'''
# length of horizon (excluding init state)
N = 50
trajectory = np.zeros((N + 1, 4))
input_trajectory = np.ones((N, 2)) * 10.0
### My implementation of Direct Transcription
# (states and control are all decision vars using Euler integration)
mp = MathematicalProgram()
# let trajectory duration be a decision var
total_time = mp.NewContinuousVariables(1, "total_time")
dt = total_time[0] / N
# create the control decision var (m*N) and state decision var (n*[N+1])
idx = 0
u_list = mp.NewContinuousVariables(2, "u_{}".format(idx))
state_list = mp.NewContinuousVariables(4, "state_{}".format(idx))
state_list = np.vstack(
(state_list,
mp.NewContinuousVariables(4, "state_{}".format(idx + 1))))
for idx in range(1, N):
u_list = np.vstack(
(u_list, mp.NewContinuousVariables(2, "u_{}".format(idx))))
state_list = np.vstack(
(state_list,
mp.NewContinuousVariables(4, "state_{}".format(idx + 1))))
### Constraints
# set init state as constraint on stage 0 decision vars
for state_idx in range(4):
mp.AddLinearConstraint(
state_list[0, state_idx] == state_initial[state_idx])
# interstage equality constraint on state vars via Euler integration
# note: Direct Collocation could improve accuracy for same computation
for idx in range(1, N + 1):
state_new = state_list[idx - 1, :] + dt * self.rocket_dynamics(
state_list[idx - 1, :], u_list[idx - 1, :])
for state_idx in range(4):
mp.AddConstraint(state_list[idx,
state_idx] == state_new[state_idx])
# constraint on time
mp.AddLinearConstraint(total_time[0] <= maximum_time)
mp.AddLinearConstraint(total_time[0] >= minimum_time)
# constraint on final state distance (squared)to second planet
final_dist_to_world_2_sq = (
self.world_2_position[0] - state_list[-1, 0])**2 + (
self.world_2_position[1] - state_list[-1, 1])**2
mp.AddConstraint(final_dist_to_world_2_sq <= 0.25)
# constraint on final state speed (squared
final_speed_sq = state_list[-1, 2]**2 + state_list[-1, 3]**2
mp.AddConstraint(final_speed_sq <= 1)
### Cost
# equal cost on vertical/horizontal accels, weight shouldn't matter since this is the only cost
mp.AddQuadraticCost(1 * u_list[:, 0].dot(u_list[:, 0]))
mp.AddQuadraticCost(1 * u_list[:, 1].dot(u_list[:, 1]))
### Solve and parse
result = Solve(mp)
trajectory = result.GetSolution(state_list)
input_trajectory = result.GetSolution(u_list)
tf = result.GetSolution(total_time)
time_array = np.linspace(0, tf[0], N + 1)
print "optimization successful: ", result.is_success()
print "total num decision vars (x: (N+1)*4, u: 2N, total_time: 1): {}".format(
mp.num_vars())
print "solver used: ", result.get_solver_id().name()
print "optimal trajectory time: {:.2f} sec".format(tf[0])
return trajectory, input_trajectory, time_array
class Optimization:
def __init__(self, config):
self.physical_params = config["physical_params"]
self.T = config["T"]
self.dt = config["dt"]
self.initial_state = config["xinit"]
self.goal_state = config["xgoal"]
self.ellipse_arc = config["ellipse_arc"]
self.deviation_cost = config["deviation_cost"]
self.Qf = config["Qf"]
self.limits = config["limits"]
self.n_state = 6
self.n_nominal_forces = 4
self.tire_model = config["tire_model"]
self.initial_guess_config = config["initial_guess_config"]
self.puddle_model = config["puddle_model"]
self.force_constraint = config["force_constraint"]
self.visualize_initial_guess = config["visualize_initial_guess"]
self.dynamics = Dynamics(self.physical_params.lf,
self.physical_params.lr,
self.physical_params.m,
self.physical_params.Iz, self.dt)
def build_program(self):
self.prog = MathematicalProgram()
# Declare variables.
state = self.prog.NewContinuousVariables(rows=self.T + 1,
cols=self.n_state,
name='state')
nominal_forces = self.prog.NewContinuousVariables(
rows=self.T, cols=self.n_nominal_forces, name='nominal_forces')
steers = self.prog.NewContinuousVariables(rows=self.T, name="steers")
slip_ratios = self.prog.NewContinuousVariables(rows=self.T,
cols=2,
name="slip_ratios")
self.state = state
self.nominal_forces = nominal_forces
self.steers = steers
self.slip_ratios = slip_ratios
# Set the initial state.
xinit = pack_state_vector(self.initial_state)
for i, s in enumerate(xinit):
self.prog.AddConstraint(state[0, i] == s)
# Constrain xdot to always be at least some value to prevent numerical issues with optimizer.
for i in range(self.T + 1):
s = unpack_state_vector(state[i])
self.prog.AddConstraint(s["xdot"] >= self.limits.min_xdot)
# Constrain slip ratio to be at least a certain magnitude.
if i != self.T:
self.prog.AddConstraint(
slip_ratios[i,
0]**2.0 >= self.limits.min_slip_ratio_mag**2.0)
self.prog.AddConstraint(
slip_ratios[i,
1]**2.0 >= self.limits.min_slip_ratio_mag**2.0)
# Control rate limits.
for i in range(self.T - 1):
ddelta = self.dt * (steers[i + 1] - steers[i])
f_dkappa = self.dt * (slip_ratios[i + 1, 0] - slip_ratios[i, 0])
r_dkappa = self.dt * (slip_ratios[i + 1, 1] - slip_ratios[i, 1])
self.prog.AddConstraint(ddelta <= self.limits.max_ddelta)
self.prog.AddConstraint(ddelta >= -self.limits.max_ddelta)
self.prog.AddConstraint(f_dkappa <= self.limits.max_dkappa)
self.prog.AddConstraint(f_dkappa >= -self.limits.max_dkappa)
self.prog.AddConstraint(r_dkappa <= self.limits.max_dkappa)
self.prog.AddConstraint(r_dkappa >= -self.limits.max_dkappa)
# Control value limits.
for i in range(self.T):
self.prog.AddConstraint(steers[i] <= self.limits.max_delta)
self.prog.AddConstraint(steers[i] >= -self.limits.max_delta)
# Add dynamics constraints by constraining residuals to be zero.
for i in range(self.T):
residuals = self.dynamics.nominal_dynamics(state[i], state[i + 1],
nominal_forces[i],
steers[i])
for r in residuals:
self.prog.AddConstraint(r == 0.0)
# Add the cost function.
self.add_cost(state)
# Add a different force constraint depending on the configuration.
if self.force_constraint == ForceConstraint.NO_PUDDLE:
self.add_no_puddle_force_constraints(
state, nominal_forces, steers,
self.tire_model.get_deterministic_model(), slip_ratios)
elif self.force_constraint == ForceConstraint.MEAN_CONSTRAINED:
self.add_mean_constrained()
elif self.force_constraint == ForceConstraint.CHANCE_CONSTRAINED:
self.add_chance_constraints()
else:
raise NotImplementedError("ForceConstraint type invalid.")
return
def constant_input_initial_guess(self, state, steers, slip_ratios,
nominal_forces):
# Guess the slip ratio as the minimum allowed value.
gslip_ratios = np.tile(
np.array([
self.limits.min_slip_ratio_mag, self.limits.min_slip_ratio_mag
]), (self.T, 1))
# Guess the slip angle as a linearly ramping steer that then becomes constant.
# This is done by creating an array of values corresponding to end_delta then
# filling in the initial ramp up phase.
gsteers = self.initial_guess_config["end_delta"] * np.ones(self.T)
igc = self.initial_guess_config
for i in range(igc["ramp_steps"]):
gsteers[i] = (i / igc["ramp_steps"]) * (igc["end_delta"] -
igc["start_delta"])
# Create empty arrays for state and forces.
gstate = np.empty((self.T + 1, self.n_state))
gstate[0] = pack_state_vector(self.initial_state)
gforces = np.empty((self.T, 4))
all_slips = self.T * [None]
# Simulate the dynamics for the initial guess differently depending on the force
# constraint being used.
if self.force_constraint == ForceConstraint.NO_PUDDLE:
tire_model = self.tire_model.get_deterministic_model()
for i in range(self.T):
s = unpack_state_vector(gstate[i])
# Simulate the dynamics for one time step.
gstate[i + 1], forces, slips = self.dynamics.simulate(
gstate[i], gsteers[i], gslip_ratios[i, 0],
gslip_ratios[i, 1], tire_model, self.dt)
# Store the results.
gforces[i] = pack_force_vector(forces)
all_slips[i] = slips
elif self.force_constraint == ForceConstraint.MEAN_CONSTRAINED or self.force_constraint == ForceConstraint.CHANCE_CONSTRAINED:
# mean model is a function that maps (slip_ratio, slip_angle, x, y) -> (E[Fx], E[Fy])
mean_model = self.tire_model.get_mean_model(
self.puddle_model.get_mean_fun())
for i in range(self.T):
# Update the tire model based off the conditions of the world
# at (x, y)
s = unpack_state_vector(gstate[i])
tire_model = lambda slip_ratio, slip_angle: mean_model(
slip_ratio, slip_angle, s["X"], s["Y"])
# Simulate the dynamics for one time step.
gstate[i + 1], forces, slips = self.dynamics.simulate(
gstate[i], gsteers[i], gslip_ratios[i, 0],
gslip_ratios[i, 1], tire_model, self.dt)
# Store the results.
gforces[i] = pack_force_vector(forces)
all_slips[i] = slips
else:
raise NotImplementedError(
"Invalid value for self.force_constraint")
# Declare array for the initial guess and set the values.
initial_guess = np.empty(self.prog.num_vars())
self.prog.SetDecisionVariableValueInVector(state, gstate,
initial_guess)
self.prog.SetDecisionVariableValueInVector(steers, gsteers,
initial_guess)
self.prog.SetDecisionVariableValueInVector(slip_ratios, gslip_ratios,
initial_guess)
self.prog.SetDecisionVariableValueInVector(nominal_forces, gforces,
initial_guess)
if self.visualize_initial_guess:
# TODO: reorganize visualizations
psis = gstate[:, 2]
xs = gstate[:, 4]
ys = gstate[:, 5]
fig, ax = plt.subplots()
if self.force_constraint != ForceConstraint.NO_PUDDLE:
plot_puddles(ax, self.puddle_model)
plot_ellipse_arc(ax, self.ellipse_arc)
plot_planned_trajectory(ax, xs, ys, psis, gsteers,
self.physical_params)
# Plot the slip ratios/angles
plot_slips(all_slips)
plot_forces(gforces)
if self.force_constraint == ForceConstraint.NO_PUDDLE:
generate_animation(xs,
ys,
psis,
gsteers,
self.physical_params,
self.dt,
puddle_model=None)
else:
generate_animation(xs,
ys,
psis,
gsteers,
self.physical_params,
self.dt,
puddle_model=self.puddle_model)
return initial_guess
def solve_program(self):
# Generate initial guess
initial_guess = self.constant_input_initial_guess(
self.state, self.steers, self.slip_ratios, self.nominal_forces)
# Solve the problem.
solver = SnoptSolver()
result = solver.Solve(self.prog, initial_guess)
solver_details = result.get_solver_details()
print("Exit flag: " + str(solver_details.info))
self.visualize_results(result)
def visualize_results(self, result):
state_res = result.GetSolution(self.state)
psis = state_res[:, 2]
xs = state_res[:, 4]
ys = state_res[:, 5]
steers = result.GetSolution(self.steers)
fig, ax = plt.subplots()
plot_ellipse_arc(ax, self.ellipse_arc)
plot_puddles(ax, self.puddle_model)
plot_planned_trajectory(ax, xs, ys, psis, steers, self.physical_params)
generate_animation(xs,
ys,
psis,
steers,
self.physical_params,
self.dt,
puddle_model=self.puddle_model)
plt.show()
def add_cost(self, state):
# Add the final state cost function.
diff_state = pack_state_vector(self.goal_state) - state[-1]
self.prog.AddQuadraticCost(diff_state.T @ self.Qf @ diff_state)
# Get the approx distance function for the ellipse.
fun = self.ellipse_arc.approx_dist_fun()
for i in range(self.T):
s = unpack_state_vector(state[i])
self.prog.AddCost(self.deviation_cost * fun(s["X"], s["Y"]))
def add_mean_constrained(self):
# mean model is a function that maps (slip_ratio, slip_angle, x, y) -> (E[Fx], E[Fy])
mean_model = self.tire_model.get_mean_model(
self.puddle_model.get_mean_fun())
for i in range(self.T):
# Get slip angles and slip ratios.
s = unpack_state_vector(self.state[i])
F = unpack_force_vector(self.nominal_forces[i])
# get the tire model at this position in space.
tire_model = lambda slip_ratio, slip_angle: mean_model(
slip_ratio, slip_angle, s["X"], s["Y"])
# Unpack values
delta = self.steers[i]
alpha_f, alpha_r = self.dynamics.slip_angles(
s["xdot"], s["ydot"], s["psidot"], delta)
kappa_f = self.slip_ratios[i, 0]
kappa_r = self.slip_ratios[i, 1]
# Compute expected forces.
E_Ffx, E_Ffy = tire_model(kappa_f, alpha_f)
E_Frx, E_Fry = tire_model(kappa_r, alpha_r)
# Constrain these expected force values to be equal to the nominal
# forces in the optimization.
self.prog.AddConstraint(E_Ffx - F["f_long"] == 0.0)
self.prog.AddConstraint(E_Ffy - F["f_lat"] == 0.0)
self.prog.AddConstraint(E_Frx - F["r_long"] == 0.0)
self.prog.AddConstraint(E_Fry - F["r_lat"] == 0.0)
def add_chance_constrained(self):
pass
def add_no_puddle_force_constraints(self, state, forces, steers,
pacejka_model, slip_ratios):
"""
Args:
prog:
state:
forces:
steers:
pacejka_model: function with signature (slip_ratio, slip_angle) using pydrake.math
"""
for i in range(self.T):
# Get slip angles and slip ratios.
s = unpack_state_vector(state[i])
F = unpack_force_vector(forces[i])
delta = steers[i]
alpha_f, alpha_r = self.dynamics.slip_angles(
s["xdot"], s["ydot"], s["psidot"], delta)
kappa_f = slip_ratios[i, 0]
kappa_r = slip_ratios[i, 1]
Ffx, Ffy = pacejka_model(kappa_f, alpha_f)
Frx, Fry = pacejka_model(kappa_r, alpha_r)
# Constrain the values from the pacejka model to be equal
# to the desired values in the optimization.
self.prog.AddConstraint(Ffx - F["f_long"] == 0.0)
self.prog.AddConstraint(Ffy - F["f_lat"] == 0.0)
self.prog.AddConstraint(Frx - F["r_long"] == 0.0)
self.prog.AddConstraint(Fry - F["r_lat"] == 0.0)
def __init__(self,
rtree,
q_nom,
control_period=0.001,
print_period=1.0,
sim=True,
cost_pub=False):
LeafSystem.__init__(self)
self.rtree = rtree
self.nq = rtree.get_num_positions()
self.nv = rtree.get_num_velocities()
self.nu = rtree.get_num_actuators()
self.cost_pub = cost_pub
dim = 3 # 3D
nd = 4 # for friction cone approx, hard coded for now
self.nc = 4 # number of contacts; TODO don't hard code
self.nf = self.nc * nd # number of contact force variables
self.ncf = 2 # number of loop constraint forces; TODO don't hard code
self.neps = self.nc * dim # number of slack variables for contact
if sim:
self.sim = True
self._DeclareInputPort(PortDataType.kVectorValued,
self.nq + self.nv)
self._DeclareDiscreteState(self.nu)
self._DeclareVectorOutputPort(BasicVector(self.nu),
self.CopyStateOutSim)
self.print_period = print_period
self.last_print_time = -print_period
self.last_v = np.zeros(3)
self.lcmgl = lcmgl("Balancing-plan", true_lcm.LCM())
# self.lcmgl = None
else:
self.sim = False
self._DeclareInputPort(PortDataType.kVectorValued, 1)
self._DeclareInputPort(PortDataType.kVectorValued,
kCassiePositions)
self._DeclareInputPort(PortDataType.kVectorValued,
kCassieVelocities)
self._DeclareInputPort(PortDataType.kVectorValued, 2)
self._DeclareDiscreteState(kCassieActuators * 8 + 1)
self._DeclareVectorOutputPort(
BasicVector(1),
lambda c, o: self.CopyStateOut(c, o, 8 * kCassieActuators, 1))
self._DeclareVectorOutputPort(
BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
c, o, 0, kCassieActuators)) #torque
self._DeclareVectorOutputPort(
BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
c, o, kCassieActuators, kCassieActuators)) #motor_pos
self._DeclareVectorOutputPort(
BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
c, o, 2 * kCassieActuators, kCassieActuators)) #motor_vel
self._DeclareVectorOutputPort(
BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
c, o, 3 * kCassieActuators, kCassieActuators)) #kp
self._DeclareVectorOutputPort(
BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
c, o, 4 * kCassieActuators, kCassieActuators)) #kd
self._DeclareVectorOutputPort(
BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
c, o, 5 * kCassieActuators, kCassieActuators)) #ki
self._DeclareVectorOutputPort(
BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
c, o, 6 * kCassieActuators, kCassieActuators)) #leak
self._DeclareVectorOutputPort(
BasicVector(kCassieActuators), lambda c, o: self.CopyStateOut(
c, o, 7 * kCassieActuators, kCassieActuators)) #clamp
self._DeclarePeriodicDiscreteUpdate(0.0005)
self.q_nom = q_nom
self.com_des = rtree.centerOfMass(self.rtree.doKinematics(q_nom))
self.u_des = CassieFixedPointTorque()
self.lfoot = rtree.FindBody("toe_left")
self.rfoot = rtree.FindBody("toe_right")
self.springs = 2 + np.array([
rtree.FindIndexOfChildBodyOfJoint("knee_spring_left_fixed"),
rtree.FindIndexOfChildBodyOfJoint("knee_spring_right_fixed"),
rtree.FindIndexOfChildBodyOfJoint("ankle_spring_joint_left"),
rtree.FindIndexOfChildBodyOfJoint("ankle_spring_joint_right")
])
umin = np.zeros(self.nu)
umax = np.zeros(self.nu)
ii = 0
for motor in rtree.actuators:
umin[ii] = motor.effort_limit_min
umax[ii] = motor.effort_limit_max
ii += 1
slack_limit = 10.0
self.initialized = False
#------------------------------------------------------------
# Add Decision Variables ------------------------------------
prog = MathematicalProgram()
qddot = prog.NewContinuousVariables(self.nq, "joint acceleration")
u = prog.NewContinuousVariables(self.nu, "input")
bar = prog.NewContinuousVariables(self.ncf, "four bar forces")
beta = prog.NewContinuousVariables(self.nf, "friction forces")
eps = prog.NewContinuousVariables(self.neps, "slack variable")
nparams = prog.num_vars()
#------------------------------------------------------------
# Problem Constraints ---------------------------------------
self.con_u_lim = prog.AddBoundingBoxConstraint(umin, umax,
u).evaluator()
self.con_fric_lim = prog.AddBoundingBoxConstraint(0, 100000,
beta).evaluator()
self.con_slack_lim = prog.AddBoundingBoxConstraint(
-slack_limit, slack_limit, eps).evaluator()
bar_con = np.zeros((self.ncf, self.nq))
self.con_4bar = prog.AddLinearEqualityConstraint(
bar_con, np.zeros(self.ncf), qddot).evaluator()
if self.nc > 0:
dyn_con = np.zeros(
(self.nq, self.nq + self.nu + self.ncf + self.nf))
dyn_vars = np.concatenate((qddot, u, bar, beta))
self.con_dyn = prog.AddLinearEqualityConstraint(
dyn_con, np.zeros(self.nq), dyn_vars).evaluator()
foot_con = np.zeros((self.neps, self.nq + self.neps))
foot_vars = np.concatenate((qddot, eps))
self.con_foot = prog.AddLinearEqualityConstraint(
foot_con, np.zeros(self.neps), foot_vars).evaluator()
else:
dyn_con = np.zeros((self.nq, self.nq + self.nu + self.ncf))
dyn_vars = np.concatenate((qddot, u, bar))
self.con_dyn = prog.AddLinearEqualityConstraint(
dyn_con, np.zeros(self.nq), dyn_vars).evaluator()
#------------------------------------------------------------
# Problem Costs ---------------------------------------------
self.Kp_com = 50
self.Kd_com = 2.0 * sqrt(self.Kp_com)
# self.Kp_qdd = 10*np.eye(self.nq)#np.diag(np.diag(H)/np.max(np.diag(H)))
self.Kp_qdd = np.diag(
np.concatenate(([10, 10, 10, 5, 5, 5], np.zeros(self.nq - 6))))
self.Kd_qdd = 1.0 * np.sqrt(self.Kp_qdd)
self.Kp_qdd[self.springs, self.springs] = 0
self.Kd_qdd[self.springs, self.springs] = 0
com_ddot_des = np.zeros(dim)
qddot_des = np.zeros(self.nq)
self.w_com = 0
self.w_qdd = np.zeros(self.nq)
self.w_qdd[:6] = 10
self.w_qdd[8:10] = 1
self.w_qdd[3:6] = 1000
# self.w_qdd[self.springs] = 0
self.w_u = 0.001 * np.ones(self.nu)
self.w_u[2:4] = 0.01
self.w_slack = 0.1
self.cost_qdd = prog.AddQuadraticErrorCost(np.diag(self.w_qdd),
qddot_des,
qddot).evaluator()
self.cost_u = prog.AddQuadraticErrorCost(np.diag(self.w_u), self.u_des,
u).evaluator()
self.cost_slack = prog.AddQuadraticErrorCost(
self.w_slack * np.eye(self.neps), np.zeros(self.neps),
eps).evaluator()
# self.cost_com = prog.AddQuadraticCost().evaluator()
# self.cost_qdd = prog.AddQuadraticCost(
# 2*np.diag(self.w_qdd), -2*np.matmul(qddot_des, np.diag(self.w_qdd)), qddot).evaluator()
# self.cost_u = prog.AddQuadraticCost(
# 2*np.diag(self.w_u), -2*np.matmul(self.u_des, np.diag(self.w_u)) , u).evaluator()
# self.cost_slack = prog.AddQuadraticCost(
# 2*self.w_slack*np.eye(self.neps), np.zeros(self.neps), eps).evaluator()
REG = 1e-8
self.cost_reg = prog.AddQuadraticErrorCost(
REG * np.eye(nparams), np.zeros(nparams),
prog.decision_variables()).evaluator()
# self.cost_reg = prog.AddQuadraticCost(2*REG*np.eye(nparams),
# np.zeros(nparams), prog.decision_variables()).evaluator()
#------------------------------------------------------------
# Solver settings -------------------------------------------
self.solver = GurobiSolver()
# self.solver = OsqpSolver()
prog.SetSolverOption(SolverType.kGurobi, "Method", 2)
#------------------------------------------------------------
# Save MathProg Variables -----------------------------------
self.qddot = qddot
self.u = u
self.bar = bar
self.beta = beta
self.eps = eps
self.prog = prog
For the time steps `h` we just initialize them to their maximal value `h_max` (somewhat an arbitrary decision, but it works).
For the robot configuration `q`, we interpolate between the initial value `q0_guess` and the final value `- q0_guess`.
In our implementation, the value given below for `q0_guess` made the optimization converge.
But, if you find the need, feel free to tweak this parameter.
The initial guess for the velocity and the acceleration is obtained by differentiating the one for the position.
The normal force `f` at the stance foot is equal to the total `weight` of the robot.
All the other optimization variables are initialized at zero.
(Note that, if the initial guess for a variable is not specified, the default value is zero.)
"""
# vector of the initial guess
initial_guess = np.empty(prog.num_vars())
# initial guess for the time step
h_guess = h_max
prog.SetDecisionVariableValueInVector(h, [h_guess] * T, initial_guess)
# linear interpolation of the configuration
q0_guess = np.array([0, 0, .15, -.3])
q_guess_poly = PiecewisePolynomial.FirstOrderHold(
[0, T * h_guess],
np.column_stack((q0_guess, - q0_guess))
)
qd_guess_poly = q_guess_poly.derivative()
qdd_guess_poly = q_guess_poly.derivative()
# set initial guess for configuration, velocity, and acceleration
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 solve(self, quad_start_q, quad_final_q, time_used):
mp = MathematicalProgram()
# We want to solve this for a certain number of knot points
N = 40 # num knot points
time_increment = time_used / (N + 1)
dt = time_increment
time_array = np.arange(0.0, time_used, time_increment)
quad_u = mp.NewContinuousVariables(2, "u_0")
quad_q = mp.NewContinuousVariables(6, "quad_q_0")
for i in range(1, N):
u = mp.NewContinuousVariables(2, "u_%d" % i)
quad = mp.NewContinuousVariables(6, "quad_q_%d" % i)
quad_u = np.vstack((quad_u, u))
quad_q = np.vstack((quad_q, quad))
for i in range(N):
mp.AddLinearConstraint(quad_u[i][0] <= 3.0) # force
mp.AddLinearConstraint(quad_u[i][0] >= 0.0) # force
mp.AddLinearConstraint(quad_u[i][1] <= 10.0) # torque
mp.AddLinearConstraint(quad_u[i][1] >= -10.0) # torque
mp.AddLinearConstraint(quad_q[i][0] <= 1000.0) # pos x
mp.AddLinearConstraint(quad_q[i][0] >= -1000.0)
mp.AddLinearConstraint(quad_q[i][1] <= 1000.0) # pos y
mp.AddLinearConstraint(quad_q[i][1] >= -1000.0)
mp.AddLinearConstraint(
quad_q[i][2] <= 60.0 * np.pi / 180.0) # pos theta
mp.AddLinearConstraint(quad_q[i][2] >= -60.0 * np.pi / 180.0)
mp.AddLinearConstraint(quad_q[i][3] <= 1000.0) # vel x
mp.AddLinearConstraint(quad_q[i][3] >= -1000.0)
mp.AddLinearConstraint(quad_q[i][4] <= 1000.0) # vel y
mp.AddLinearConstraint(quad_q[i][4] >= -1000.0)
mp.AddLinearConstraint(quad_q[i][5] <= 10.0) # vel theta
mp.AddLinearConstraint(quad_q[i][5] >= -10.0)
for i in range(1, N):
quad_q_dyn_feasible = self.dynamics(quad_q[i - 1, :],
quad_u[i - 1, :], dt)
# Direct transcription constraints on states to dynamics
for j in range(6):
quad_state_err = (quad_q[i][j] - quad_q_dyn_feasible[j])
eps = 0.001
mp.AddConstraint(quad_state_err <= eps)
mp.AddConstraint(quad_state_err >= -eps)
# Initial and final quad and ball states
for j in range(6):
mp.AddLinearConstraint(quad_q[0][j] == quad_start_q[j])
mp.AddLinearConstraint(quad_q[-1][j] == quad_final_q[j])
# Quadratic cost on the control input
R_force = 10.0
R_torque = 100.0
Q_quad_x = 100.0
Q_quad_y = 100.0
mp.AddQuadraticCost(R_force * quad_u[:, 0].dot(quad_u[:, 0]))
mp.AddQuadraticCost(R_torque * quad_u[:, 1].dot(quad_u[:, 1]))
mp.AddQuadraticCost(Q_quad_x * quad_q[:, 0].dot(quad_q[:, 1]))
mp.AddQuadraticCost(Q_quad_y * quad_q[:, 1].dot(quad_q[:, 1]))
# Solve the optimization
print "Number of decision vars: ", mp.num_vars()
# mp.SetSolverOption(SolverType.kSnopt, "Major iterations limit", 100000)
print "Solve: ", mp.Solve()
quad_traj = mp.GetSolution(quad_q)
input_traj = mp.GetSolution(quad_u)
return (quad_traj, input_traj, time_array)
