SolutionResult)
plant = VanDerPolOscillator()
context = plant.CreateDefaultContext()
dircol = DirectCollocation(plant, context, num_time_samples=61,
minimum_timestep=0.01, maximum_timestep=0.5)
dircol.AddEqualTimeIntervalsConstraints()
# Initial state on the surface of section (and velocity > .1).
dircol.AddBoundingBoxConstraint([0., 0.1], [0., 10.], dircol.initial_state())
# Periodicity constraint.
# TODO(russt): Replace this with the vectorized version pending drake #8315.
dircol.AddLinearConstraint(
dircol.final_state()[0] == dircol.initial_state()[0])
dircol.AddLinearConstraint(
dircol.final_state()[1] == dircol.initial_state()[1])
samples = np.linspace(0, 2*math.pi, 10)
x_guess = np.vstack(([2*math.sin(t) for t in samples],
[2*math.cos(t) for t in samples]))
initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(samples, x_guess)
dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)
fig = plt.figure()
h, = plt.plot([], [], '.-')
plt.xlim((-2.5, 2.5))
plt.ylim((-3., 3.))
def direct_collocation_zhao_glider():
print("Running direct collocation")
plant = SlotineGlider()
context = plant.CreateDefaultContext()
N = 21
initial_guess = True
max_dt = 0.5
max_tf = N * max_dt
dircol = DirectCollocation(
plant,
context,
num_time_samples=N,
minimum_timestep=0.05,
maximum_timestep=max_dt,
)
# Constrain all timesteps, $h[k]$, to be equal, so the trajectory breaks are evenly distributed.
dircol.AddEqualTimeIntervalsConstraints()
# Add input constraints
u = dircol.input()
dircol.AddConstraintToAllKnotPoints(0 <= u[0])
dircol.AddConstraintToAllKnotPoints(u[0] <= 3)
dircol.AddConstraintToAllKnotPoints(-np.pi / 2 <= u[1])
dircol.AddConstraintToAllKnotPoints(u[1] <= np.pi / 2)
# Add state constraints
x = dircol.state()
min_speed = 5
dircol.AddConstraintToAllKnotPoints(x[0] >= min_speed)
min_height = 0.5
dircol.AddConstraintToAllKnotPoints(x[3] >= min_height)
# Add initial state
travel_angle = (3 / 2) * np.pi
h0 = 10
dir_vector = np.array([np.cos(travel_angle), np.sin(travel_angle)])
# Start at initial position
x0_pos = np.array([h0, 0, 0])
dircol.AddBoundingBoxConstraint(x0_pos, x0_pos,
dircol.initial_state()[3:6])
# Periodicity constraints
dircol.AddLinearConstraint(
dircol.final_state()[0] == dircol.initial_state()[0])
dircol.AddLinearConstraint(
dircol.final_state()[1] == dircol.initial_state()[1])
dircol.AddLinearConstraint(
dircol.final_state()[2] == dircol.initial_state()[2])
dircol.AddLinearConstraint(
dircol.final_state()[3] == dircol.initial_state()[3])
# Always end in right direction
# NOTE this assumes that we always are starting in origin
if travel_angle % np.pi == 0: # Travel along x-axis
dircol.AddConstraint(
dircol.final_state()[5] == dircol.initial_state()[5])
elif travel_angle % ((1 / 2) * np.pi) == 0: # Travel along y-axis
dircol.AddConstraint(
dircol.final_state()[4] == dircol.initial_state()[4])
else:
dircol.AddConstraint(
dircol.final_state()[5] == dircol.final_state()[4] *
np.tan(travel_angle))
# Maximize distance travelled in desired direction
p0 = dircol.initial_state()
p1 = dircol.final_state()
Q = 1
dist_travelled = np.array([p1[4], p1[5]]) # NOTE assume starting in origin
dircol.AddFinalCost(-(dir_vector.T.dot(dist_travelled)) * Q)
if True:
# Cost on input effort
R = 0.1
dircol.AddRunningCost(R * (u[0])**2 + R * u[1]**2)
# Initial guess is a straight line from x0 in direction
if initial_guess:
avg_vel_guess = 10 # Guess for initial velocity
x0_guess = np.array([avg_vel_guess, travel_angle, 0, h0, 0, 0])
guessed_total_dist_travelled = 200
xf_guess = np.array([
avg_vel_guess,
travel_angle,
0,
h0,
dir_vector[0] * guessed_total_dist_travelled,
dir_vector[1] * guessed_total_dist_travelled,
])
initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(
[0.0, 4.0], np.column_stack((x0_guess, xf_guess)))
dircol.SetInitialTrajectory(PiecewisePolynomial(),
initial_x_trajectory)
# Solve direct collocation
result = Solve(dircol)
assert result.is_success()
print("Found a solution!")
# PLOTTING
N_plot = 200
# Plot trajectory
x_trajectory = dircol.ReconstructStateTrajectory(result)
times = np.linspace(x_trajectory.start_time(), x_trajectory.end_time(),
N_plot)
x_knots = np.hstack([x_trajectory.value(t) for t in times])
z = x_knots[3, :]
x = x_knots[4, :]
y = x_knots[5, :]
plot_trj_3_wind(np.vstack((x, y, z)).T, dir_vector)
# Plot input
u_trajectory = dircol.ReconstructInputTrajectory(result)
u_knots = np.hstack([u_trajectory.value(t) for t in times])
plot_input_zhao_glider(times, u_knots.T)
plt.show()
return 0
