def do_single_basic_pend_traj_opt(should_init):
plant = PendulumPlant()
context = plant.CreateDefaultContext()
kNumTimeSamples = 15
kMinimumTimeStep = 0.01
kMaximumTimeStep = 0.2
dircol = DirectCollocation(plant, context, kNumTimeSamples, kMinimumTimeStep, kMaximumTimeStep)
dircol.AddEqualTimeIntervalsConstraints()
kTorqueLimit = 3.0 # N*m.
u = dircol.input()
dircol.AddConstraintToAllKnotPoints(-kTorqueLimit <= u[0])
dircol.AddConstraintToAllKnotPoints(u[0] <= kTorqueLimit)
initial_state = PendulumState()
initial_state.set_theta(0.0)
initial_state.set_thetadot(0.0)
dircol.AddBoundingBoxConstraint(initial_state.get_value(), initial_state.get_value(), dircol.initial_state())
# dircol.AddLinearConstraint(dircol.initial_state() == initial_state.get_value())
final_state = PendulumState()
final_state.set_theta(math.pi)
final_state.set_thetadot(0.0)
dircol.AddBoundingBoxConstraint(final_state.get_value(), final_state.get_value(), dircol.final_state())
# dircol.AddLinearConstraint(dircol.final_state() == final_state.get_value())
R = 10 # Cost on input "effort".
dircol.AddRunningCost(R*u[0]**2)
if should_init:
initial_x_trajectory = PiecewisePolynomial.FirstOrderHold([0., 4.], [initial_state.get_value(), final_state.get_value()])
dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)
def cb(decision_vars):
global cb_counter
cb_counter += 1
if cb_counter % 10 != 1:
return
# Get the total cost
all_costs = dircol.EvalBindings(dircol.GetAllCosts(), decision_vars)
# :all_constraints = dircol.EvalBindings(dircol.GetAllConstraints(), decision_vars)
# Get the total cost of the constraints.
# Additionally, the number and extent of any constraint violations.
violated_constraint_count = 0
violated_constraint_cost = 0
constraint_cost = 0
for constraint in dircol.GetAllConstraints():
val = dircol.EvalBinding(constraint, decision_vars)
# TODO: switch to DoCheckSatisfied...
nudge = 1e-1 # This much constraint violation is not considered bad...
lb = constraint.evaluator().lower_bound()
ub = constraint.evaluator().upper_bound()
good_lb = np.all( np.less_equal(lb, val+nudge) )
good_ub = np.all( np.greater_equal(ub, val-nudge) )
if not good_lb or not good_ub:
# print("val,lb,ub: ", val, lb, ub)
violated_constraint_count += 1
violated_constraint_cost += np.sum(val)
constraint_cost += np.sum(val)
print("total cost: {:.2f} | constraint {:.2f}, bad {}, {:.2f}".format(
sum(all_costs), constraint_cost, violated_constraint_count, violated_constraint_cost))
# dircol.AddVisualizationCallback(cb, np.array(dircol.decision_variables()))
result = dircol.Solve()
assert(result == SolutionResult.kSolutionFound)
sol_costs = np.hstack([dircol.EvalBindingAtSolution(cost) for cost in dircol.GetAllCosts()])
sol_constraints = np.hstack([dircol.EvalBindingAtSolution(constraint) for constraint in dircol.GetAllConstraints()])
dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)
fig = plt.figure()
h, = plt.plot([], [], '.-')
plt.xlim((-2.5, 2.5))
plt.ylim((-3., 3.))
def draw_trajectory(t, x):
h.set_xdata(x[0, :])
h.set_ydata(x[1, :])
fig.canvas.draw()
if plt.get_backend() == u'MacOSX':
plt.pause(1e-10)
dircol.AddStateTrajectoryCallback(draw_trajectory)
result = dircol.Solve()
assert(result == SolutionResult.kSolutionFound)
x_trajectory = dircol.ReconstructStateTrajectory()
x_knots = np.hstack([x_trajectory.value(t) for t in
np.linspace(x_trajectory.start_time(),
x_trajectory.end_time(), 100)])
plt.plot(x_knots[0, :], x_knots[1, :])
plt.show()
def _solve_traj_opt(self, initial_state, constrain_final_state=True, duration_bounds=None, d=0.0, verbose=False):
'''Finds a trajectory from an initial state, optionally to a final state.
Args:
initial_state (tuple): the initial state
final_state (tuple): the final state (default to None, final state unconstrained)
duration (tuple): the min and max duration of the trajectory (default to None,
no duration constraints)
d (float): constant disturbance force
verbose (bool): enables/disables verbose output
Returns:
pydrake.trajectories.PiecewisePolynomial: the planned trajectory
float: the cost of the planned trajectory
Raises:
RuntimeError: raised if the optimization fails
'''
print("Initial state: {}\nFinal state: {}\nMin duration: {} s\nMax duration: {} s".format(
initial_state, constrain_final_state, duration_bounds[0], duration_bounds[1]))
traj_opt = DirectCollocation(self.plant, self.context,
self.opt_params['num_time_samples'],
self.opt_params['minimum_timestep'],
self.opt_params['maximum_timestep'])
traj_opt.AddEqualTimeIntervalsConstraints()
# Add bounds on the total duration of the trajectory
if duration_bounds:
traj_opt.AddDurationBounds(duration_bounds[0], duration_bounds[1])
# TODO make input limits a paramter
limits_low = [-15., -15.]
limits_upp = [15., 15.]
x = traj_opt.state()
u = traj_opt.input()
t = traj_opt.time()
# TODO assuming disturbance is at the last index
for i in range(len(u) - 1):
traj_opt.AddConstraintToAllKnotPoints(limits_low[i] <= u[i])
traj_opt.AddConstraintToAllKnotPoints(u[i] <= limits_upp[i])
traj_opt.AddConstraintToAllKnotPoints(u[len(u) - 1] == d)
for signed_dist_func in self.signed_dist_funcs:
traj_opt.AddConstraintToAllKnotPoints(signed_dist_func(x) >= 0)
traj_opt.AddRunningCost(traj_opt.timestep(0) * self.running_cost(x, u, t))
traj_opt.AddFinalCost(self.final_cost(x, u, t))
# Add initial and final state constraints
traj_opt.AddBoundingBoxConstraint(initial_state, initial_state,
traj_opt.initial_state())
if self.final_state_constraint and constrain_final_state:
traj_opt.AddConstraint(self.final_state_constraint(traj_opt.final_state()) == 0)
# # TODO this is redundant with the final state equality constraint above
# if final_state:
# traj_opt.AddBoundingBoxConstraint(final_state, final_state,
# traj_opt.final_state())
# initial_x_trajectory = PiecewisePolynomial.FirstOrderHold([0., 0.4 * 21],
# np.column_stack((initial_state,
# final_state)))
# traj_opt.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)
# else:
# initial_x_trajectory = PiecewisePolynomial.FirstOrderHold([0., 0.4 * 21],
# np.column_stack((initial_state,
# initial_state)))
# traj_opt.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)
initial_x_trajectory = PiecewisePolynomial.FirstOrderHold([0., 0.4 * 21],
np.column_stack((initial_state,
initial_state)))
traj_opt.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)
result = traj_opt.Solve()
if result != SolutionResult.kSolutionFound:
raise RuntimeError('Direct collocation failed from initial state {}!'.format(initial_state))
state_samples = traj_opt.GetStateSamples()
input_samples = traj_opt.GetInputSamples()
time_samples = traj_opt.GetSampleTimes()
# for debugging
hs = [time_samples[i+1] - time_samples[i] for i in range(len(time_samples)) if i < len(time_samples) - 1]
#print(hs)
total_cost = 0.
for k in range(state_samples.shape[1]):
total_cost += (hs[0] *
self.running_cost(state_samples[:, k],
input_samples[:, k],
time_samples[k]))
if verbose:
for i, phi in enumerate(self.signed_dist_funcs):
print("\tsigned dist {}: {}".format(i, signed_dist_func(state_samples[:, k])))
if verbose:
print("Total cost is {}".format(total_cost))
u_traj = traj_opt.ReconstructInputTrajectory()
times = np.linspace(u_traj.start_time(), u_traj.end_time(), 100)
u_lookup = np.vectorize(lambda t: u_traj.value(t)[0])
u_values = u_lookup(times)
plt.figure()
plt.plot(times, u_values)
plt.xlabel('time (seconds)')
plt.ylabel('force (Newtons)')
plt.show()
return traj_opt.ReconstructStateTrajectory(), total_cost
