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 q_guess = np.hstack([q_guess_poly.value(t * h_guess) for t in range(T + 1)]).T qd_guess = np.hstack([qd_guess_poly.value(t * h_guess) for t in range(T + 1)]).T qdd_guess = np.hstack([qdd_guess_poly.value(t * h_guess) for t in range(T)]).T prog.SetDecisionVariableValueInVector(q, q_guess, initial_guess)
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 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