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
import matplotlib.pyplot as plt from pydrake.all import (DirectCollocation, FloatingBaseType, PiecewisePolynomial, RigidBodyTree, RigidBodyPlant, SolutionResult) from pydrake.examples.acrobot import AcrobotPlant from underactuated import (FindResource, PlanarRigidBodyVisualizer) plant = AcrobotPlant() context = plant.CreateDefaultContext() dircol = DirectCollocation(plant, context, num_time_samples=21, minimum_timestep=0.2, maximum_timestep=0.5) dircol.AddEqualTimeIntervalsConstraints() dircol.AddDurationBounds(0.5, 4) # Add input limits. torque_limit = 8.0 # N*m. u = dircol.input() dircol.AddConstraintToAllKnotPoints(-torque_limit <= u[0]) dircol.AddConstraintToAllKnotPoints(u[0] <= torque_limit) initial_state = (0., 0., 0., 0.) dircol.AddBoundingBoxConstraint(initial_state, initial_state, dircol.initial_state()) # More elegant version is blocked on drake #8315: # dircol.AddLinearConstraint(dircol.initial_state() == initial_state) final_state = (math.pi, 0., 0., 0.) dircol.AddBoundingBoxConstraint(final_state, final_state,