def min_time_traj_avoid_obs(self,
p0,
v0,
pf,
vf,
obstacles=None,
p_puck=None):
"""Minimum time trajectory while avoiding obstacles."""
x0 = np.array(np.concatenate((p0, v0), axis=0))
xf = np.concatenate((pf, vf), axis=0)
N = 20
prog = DirectCollocation(self.sys_c,
self.sys_c.CreateDefaultContext(),
N,
minimum_timestep=self.params.dt,
maximum_timestep=self.params.dt)
# Initial and final state
prog.AddBoundingBoxConstraint(x0, x0, prog.initial_state())
prog.AddQuadraticErrorCost(Q=np.eye(4),
x_desired=xf,
vars=prog.final_state())
u = prog.input()
prog.AddRunningCost(0.1 * u.dot(u))
prog.AddEqualTimeIntervalsConstraints()
## Input saturation
self.add_input_limits(prog)
# Arena constraints
self.add_arena_limits(prog)
prog.AddFinalCost(prog.time())
# Add non-linear constraints - will solve with SNOPT
# Avoid other players
if obstacles != None:
for p_obs in obstacles:
distance = prog.state()[0:2] - p_obs
prog.AddConstraintToAllKnotPoints(
distance.dot(distance) >= (2.0 *
self.params.player_radius)**2)
# avoid hitting the puck while generating a kicking trajectory
#if not p_puck.any(None):
# distance = prog.state()[0:2] - p_puck
# prog.AddConstraintToAllKnotPoints(distance.dot(distance) >= (self.params.player_radius + self.params.puck_radius)**2)
solver = SnoptSolver()
result = solver.Solve(prog)
solution_found = result.is_success()
if not solution_found:
print("Solution not found for intercepting_with_obs_avoidance")
u_traj = prog.ReconstructInputTrajectory(result)
u_values = u_traj.vector_values(u_traj.get_segment_times())
return solution_found, u_values
def compute_control(self, x0_p1, x0_p2, xf_p1, xf_p2, obstacles):
"""This is basically the single-agent MPC algorithm"""
prog = DirectCollocation(self.mpc_params.sys_two_players_c, self.mpc_params.sys_two_players_c.CreateDefaultContext(), self.mpc_params.N+1,
minimum_timestep=self.mpc_params.minT, maximum_timestep=self.mpc_params.maxT)
x0 = np.concatenate((x0_p1, x0_p2), axis=0)
prog.AddBoundingBoxConstraint(x0, x0, prog.initial_state())
x_des = np.concatenate((xf_p1, xf_p2))
Q = np.zeros((8,8))
Q[0:4,0:4] = self.mpc_params.Omega_N_max
Q[4:8,4:8] = self.mpc_params.Omega_N_max
prog.AddQuadraticErrorCost(Q, x_desired=x_des, vars=prog.final_state())
prog.AddEqualTimeIntervalsConstraints()
for obs_pos in obstacles: # both players should avoid the other players
for n in range(self.mpc_params.N):
x = prog.state()
prog.AddConstraintToAllKnotPoints((x[0:2]-obs_pos).dot(x[0:2]-obs_pos) >= (2.0*self.sim_params.player_radius)**2)
prog.AddConstraintToAllKnotPoints((x[4:6]-obs_pos).dot(x[4:6]-obs_pos) >= (2.0*self.sim_params.player_radius)**2)
# players should avoid each other
prog.AddConstraintToAllKnotPoints((x[0:2]-x[4:6]).dot(x[0:2]-x[4:6]) >= (2.0*self.sim_params.player_radius)**2)
# input constraints
for i in range(4):
prog.AddConstraintToAllKnotPoints(prog.input()[i] <= self.sim_params.input_limit)
prog.AddConstraintToAllKnotPoints(prog.input()[i] >= -self.sim_params.input_limit)
r = self.sim_params.player_radius
prog.AddConstraintToAllKnotPoints(prog.state()[0] + r <= self.sim_params.arena_limits_x / 2.0)
prog.AddConstraintToAllKnotPoints(prog.state()[0] - r >= -self.sim_params.arena_limits_x / 2.0)
prog.AddConstraintToAllKnotPoints(prog.state()[1] + r <= self.sim_params.arena_limits_y / 2.0)
prog.AddConstraintToAllKnotPoints(prog.state()[1] - r >= -self.sim_params.arena_limits_y / 2.0)
prog.AddConstraintToAllKnotPoints(prog.state()[4] + r <= self.sim_params.arena_limits_x / 2.0)
prog.AddConstraintToAllKnotPoints(prog.state()[4] - r >= -self.sim_params.arena_limits_x / 2.0)
prog.AddConstraintToAllKnotPoints(prog.state()[5] + r <= self.sim_params.arena_limits_y / 2.0)
prog.AddConstraintToAllKnotPoints(prog.state()[5] - r >= -self.sim_params.arena_limits_y / 2.0)
prog.AddFinalCost(prog.time())
if not self.prev_u is None and not self.prev_x is None:
prog.SetInitialTrajectory(traj_init_u=self.prev_u, traj_init_x=self.prev_x)
solver = SnoptSolver()
result = solver.Solve(prog)
u_traj = prog.ReconstructInputTrajectory(result)
x_traj = prog.ReconstructStateTrajectory(result)
self.prev_u = u_traj
self.prev_x = x_traj
u_vals = u_traj.vector_values(u_traj.get_segment_times())
x_vals = x_traj.vector_values(x_traj.get_segment_times())
return True, u_vals[0:2,0], u_vals[2:4,0]
def intercepting_with_obs_avoidance(self,
p0,
v0,
pf,
vf,
T,
obstacles=None,
p_puck=None):
"""Intercepting trajectory that avoids obs"""
x0 = np.array(np.concatenate((p0, v0), axis=0))
xf = np.concatenate((pf, vf), axis=0)
prog = MathematicalProgram()
# state and control inputs
N = int(T / self.params.dt) # number of time steps
state = prog.NewContinuousVariables(N + 1, 4, 'state')
cmd = prog.NewContinuousVariables(N, 2, 'input')
# Initial and final state
prog.AddLinearConstraint(eq(state[0], x0))
#prog.AddLinearConstraint(eq(state[-1], xf))
prog.AddQuadraticErrorCost(Q=10.0 * np.eye(4),
x_desired=xf,
vars=state[-1])
self.add_dynamics(prog, N, state, cmd)
## Input saturation
self.add_input_limits(prog, cmd, N)
# Arena constraints
self.add_arena_limits(prog, state, N)
for k in range(N):
prog.AddQuadraticCost(cmd[k].flatten().dot(cmd[k]))
# Add non-linear constraints - will solve with SNOPT
# Avoid other players
if obstacles != None:
for obs in obstacles:
self.avoid_other_player_nl(prog, state, obs, N)
# avoid hitting the puck while generating a kicking trajectory
if not p_puck.any(None):
self.avoid_puck_nl(prog, state, N, p_puck)
solver = SnoptSolver()
result = solver.Solve(prog)
solution_found = result.is_success()
if not solution_found:
print("Solution not found for intercepting_with_obs_avoidance")
u_values = result.GetSolution(cmd)
return solution_found, u_values.T
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 compute_control(self, x_des, sim_state, team_name, player_id):
prog = DirectCollocation(self.mpc_params.sys,
self.mpc_params.sys.CreateDefaultContext(),
self.mpc_params.N,
minimum_timestep=self.mpc_params.minT,
maximum_timestep=self.mpc_params.maxT)
pos0 = sim_state.get_player_pos(team_name, player_id)
vel0 = sim_state.get_player_vel(team_name, player_id)
x0 = np.concatenate((pos0, vel0), axis=0)
prog.AddBoundingBoxConstraint(x0, x0, prog.initial_state())
prog.AddQuadraticErrorCost(Q=self.mpc_params.Omega_N_max,
x_desired=x_des,
vars=prog.final_state())
obstacle_positions = self.get_obstacle_positions(
sim_state, team_name, player_id)
for obs_pos in obstacle_positions:
for n in range(self.mpc_params.N):
x = prog.state()
prog.AddConstraintToAllKnotPoints(
(x[0:2] - obs_pos).dot(x[0:2] - obs_pos) >= (
2.0 * self.sim_params.player_radius)**2)
prog.AddEqualTimeIntervalsConstraints()
self.add_input_limits(prog)
self.add_arena_limits(prog)
prog.AddFinalCost(prog.time())
if not self.prev_u is None and not self.prev_x is None:
prog.SetInitialTrajectory(traj_init_u=self.prev_u,
traj_init_x=self.prev_x)
solver = SnoptSolver()
result = solver.Solve(prog)
u_traj = prog.ReconstructInputTrajectory(result)
x_traj = prog.ReconstructStateTrajectory(result)
u_vals = u_traj.vector_values(u_traj.get_segment_times())
self.prev_u = u_traj
self.prev_x = x_traj
return u_vals[:, 0]
def nonlinear_optimization(env, x_ref, x_bar, u_bar, x0, N, dt):
''' Builds and solves single step linear optimization problem
with convex polygon obstacles and linearized dynamics
'''
prog, dvs = initialize_problem(n_x, n_u, N, x0)
prog = nonlinear_dynamics_constraints(prog, dvs, x_bar, u_bar, N, dt)
prog = add_input_state_constraints(prog, dvs, bounds)
prog, dvs = add_polyhedron_obstacle_constraints(
prog, dvs, env.obstacles, buffer, N)
prog = add_costs(prog, dvs, env, x_ref, Q, R, N)
# Use SNOPT solver
solver = SnoptSolver()
result = solver.Solve(prog)
assert(result.is_success), "Optimization Failed"
x, u, z, v = dvs
x_sol = result.GetSolution(x)
u_sol = result.GetSolution(u)
z_sol = [result.GetSolution(zj) for zj in z]
v_sol = [result.GetSolution(vj) for vj in v]
return result, (x_sol, u_sol, z_sol, v_sol), x_bar
def min_time_traj_dir_col(self, p0, v0, pf, vf):
"""generate minimum time trajectory while avoiding obs"""
N = 15
minT = self.params.dt / N
maxT = 5.0 / N
x0 = np.concatenate((p0, v0), axis=0)
xf = np.concatenate((pf, vf), axis=0)
prog = DirectCollocation(self.sys_c, self.sys_c.CreateDefaultContext(), num_time_samples=N,
minimum_timestep=minT,
maximum_timestep=maxT)
prog.AddBoundingBoxConstraint(x0, x0, prog.initial_state())
prog.AddEqualTimeIntervalsConstraints()
self.add_input_limits(prog)
self.add_arena_limits(prog)
prog.AddQuadraticErrorCost(Q=10.0*np.eye(4), x_desired=xf, vars=prog.final_state())
prog.AddFinalCost(prog.time())
solver = SnoptSolver()
result = solver.Solve(prog)
if not result.is_success():
print("Minimum time trajectory: optimization failed")
return False, np.zeros((2, 1))
# subsample trajectory accordingly
u_trajectory = prog.ReconstructInputTrajectory(result)
duration = u_trajectory.end_time() - u_trajectory.start_time()
if duration > self.params.dt:
times = np.linspace(u_trajectory.start_time(), u_trajectory.end_time(), (u_trajectory.end_time() - u_trajectory.start_time()) / self.params.dt )
else:
times = np.array([0])
u_values = np.empty((2, len(times)))
for i, t in enumerate(times):
u_values[:, i] = u_trajectory.value(t).flatten()
return result.is_success(), u_values
def min_time_bounce_kick_traj_dir_col(self, p0, v0, p0_puck, v0_puck, v_puck_desired):
"""DO NOT USE. NOT WORKING.
Minimum time trajectory + bounce kick off the wall."""
N = 15
minT = self.params.dt / N
maxT = 5.0 / N
x0 = np.concatenate((p0, v0), axis=0)
prog = DirectCollocation(self.sys_c, self.sys_c.CreateDefaultContext(), num_time_samples=N,
minimum_timestep=minT,
maximum_timestep=maxT)
prog.AddBoundingBoxConstraint(x0, x0, prog.initial_state())
prog.AddEqualTimeIntervalsConstraints()
self.add_final_state_constraint_elastic_collision(prog, p0_puck, v0_puck, v_puck_desired)
self.add_input_limits(prog)
self.add_arena_limits(prog)
# prog.AddQuadraticErrorCost(Q=10.0*np.eye(4), x_desired=xf, vars=prog.final_state())
pf = p0_puck - self.get_normalized_vector(v_puck_desired)*(self.params.puck_radius + self.params.player_radius)
prog.AddQuadraticErrorCost(Q=10.0*np.eye(2), x_desired=pf, vars=prog.final_state()[:2])
prog.AddFinalCost(prog.time())
solver = SnoptSolver()
result = solver.Solve(prog)
if not result.is_success():
print("Minimum time trajectory: optimization failed")
return False, np.zeros((2, 1))
u_trajectory = prog.ReconstructInputTrajectory(result)
times = np.linspace(u_trajectory.start_time(), u_trajectory.end_time(), (u_trajectory.end_time() - u_trajectory.start_time()) / self.params.dt )
u_values = np.empty((2, len(times)))
for i, t in enumerate(times):
u_values[:, i] = u_trajectory.value(t).flatten()
return result.is_success(), u_values
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 direct_transcription():
prog = MathematicalProgram()
N = 500
dt = 0.01
# Create decision variables
n_x = 6
n_u = 2
x = np.empty((n_x, N), dtype=Variable)
u = np.empty((n_u, N - 1), dtype=Variable)
for n in range(N - 1):
x[:, n] = prog.NewContinuousVariables(n_x, "x" + str(n))
u[:, n] = prog.NewContinuousVariables(n_u, "u" + str(n))
x[:, N - 1] = prog.NewContinuousVariables(n_x, "x" + str(N))
T = N - 1
# Add constraints
# x0 = np.array([10, -np.pi / 2, 0, 40, 0, 0])
# Slotine dynamics: x = [airspeed, heading, flight_path_angle, z, x, y]
for n in range(N - 1):
# Dynamics
prog.AddConstraint(
eq(x[:, n + 1],
x[:, n] + dt * continuous_dynamics(x[:, n], u[:, n])))
# Never below sea level
prog.AddConstraint(x[3, n + 1] >= 0 + 0.5)
# Always positive velocity
prog.AddConstraint(x[0, n + 1] >= 0)
# TODO use residuals
# Add periodic constraints
prog.AddConstraint(x[0, 0] - x[0, T] == 0)
prog.AddConstraint(x[1, 0] - x[1, T] == 0)
prog.AddConstraint(x[2, 0] - x[2, T] == 0)
prog.AddConstraint(x[3, 0] - x[3, T] == 0)
# Start at 20 meter height
prog.AddConstraint(x[4, 0] == 0)
prog.AddConstraint(x[5, 0] == 0)
h0 = 20
prog.AddConstraint(x[3, 0] == h0)
# Add cost function
p0 = x[4:6, 0]
p1 = x[4:6, T]
travel_angle = np.pi # TODO start along y direction
dir_vector = np.array([np.sin(travel_angle), np.cos(travel_angle)])
prog.AddCost(dir_vector.T.dot(p1 - p0)) # Maximize distance travelled
print("Initialized opt problem")
# Initial guess
V0 = 10
x_guess = np.zeros((n_x, N))
x_guess[:, 0] = np.array([V0, travel_angle, 0, h0, 0, 0])
for n in range(N - 1):
# Constant airspeed, heading and height
x_guess[0, n + 1] = V0
x_guess[1, n + 1] = travel_angle
x_guess[3, n + 1] = h0
# Interpolate position
avg_speed = 10 # conservative estimate
x_guess[4:6, n + 1] = dir_vector * avg_speed * n * dt
# Let the rest of the variables be initialized to zero
prog.SetInitialGuess(x, x_guess)
# solve mathematical program
solver = SnoptSolver()
result = solver.Solve(prog)
# be sure that the solution is optimal
# assert result.is_success()
# retrieve optimal solution
thrust_opt = result.GetSolution(x)
state_opt = result.GetSolution(u)
result = Solve(prog)
x_sol = result.GetSolution(x)
u_sol = result.GetSolution(u)
breakpoint()
# Slotine dynamics: x = [airspeed, heading, flight_path_angle, z, x, y]
z = x_sol[:, 3]
x = x_sol[:, 4]
y = x_sol[:, 5]
plot_trj_3_wind(np.vstack((x, y, z)).T, get_wind_field)
return
def optimize(self, vehs):
c = self.c
p = self.p
n_veh, L, N, dt, u_max = c.n_veh, c.circumference, c.n_opt, c.sim_step, c.u_max
L_veh = vehs[0].length
a, b, s0, T, v0, delta = p.a, p.b, p.s0, p.tau, p.v0, p.delta
vehs = vehs[::-1]
np.set_printoptions(linewidth=100)
# the controlled vehicle is now the first vehicle, positions are sorted from largest to smallest
print(f'Current positions {[veh.pos for veh in vehs]}')
v_init = [veh.speed for veh in vehs]
print(f'Current speeds {v_init}')
# spacing
s_init = [vehs[-1].pos - vehs[0].pos - L_veh] + [
veh_im1.pos - veh_i.pos - L_veh
for veh_im1, veh_i in zip(vehs[:-1], vehs[1:])
]
s_init = [s + L if s < 0 else s for s in s_init] # handle wrap
print(f'Current spacings {s_init}')
# can still follow current plan
if self.plan_index < len(self.plan):
accel = self.plan[self.plan_index]
self.plan_index += 1
return accel
print(f'Optimizing trajectory for {c.n_opt} steps')
## solve for equilibrium conditions (when all vehicles have same spacing and going at optimal speed)
import scipy.optimize
sf = L / n_veh - L_veh # equilibrium space
accel_fn = lambda v: a * (1 - (v / v0)**delta - ((s0 + v * T) / sf)**2)
sol = scipy.optimize.root(accel_fn, 0)
vf = sol.x.item() # equilibrium speed
sstarf = s0 + vf * T
print('Equilibrium speed', vf)
# nonconvex optimization
from pydrake.all import MathematicalProgram, SnoptSolver, eq, le, ge
# get guesses for solutions
v_guess = [np.mean(v_init)]
for _ in range(N):
v_guess.append(dt * accel_fn(v_guess[-1]))
s_guess = [sf] * (N + 1)
prog = MathematicalProgram()
v = prog.NewContinuousVariables(N + 1, n_veh, 'v') # speed
s = prog.NewContinuousVariables(N + 1, n_veh, 's') # spacing
flat = lambda x: x.reshape(-1)
# Guess
prog.SetInitialGuess(s, np.stack([s_guess] * n_veh, axis=1))
prog.SetInitialGuess(v, np.stack([v_guess] * n_veh, axis=1))
# initial conditions constraint
prog.AddLinearConstraint(eq(v[0], v_init))
prog.AddLinearConstraint(eq(s[0], s_init))
# velocity constraint
prog.AddLinearConstraint(ge(flat(v[1:]), 0))
prog.AddLinearConstraint(le(flat(v[1:]),
vf + 1)) # extra constraint to help solver
# spacing constraint
prog.AddLinearConstraint(ge(flat(s[1:]), s0))
prog.AddLinearConstraint(le(flat(s[1:]),
sf * 2)) # extra constraint to help solver
prog.AddLinearConstraint(eq(flat(s[1:].sum(axis=1)),
L - L_veh * n_veh))
# spacing update constraint
s_n = s[:-1, 1:] # s_i[n]
s_np1 = s[1:, 1:] # s_i[n + 1]
v_n = v[:-1, 1:] # v_i[n]
v_np1 = v[1:, 1:] # v_i[n + 1]
v_n_im1 = v[:-1, :-1] # v_{i - 1}[n]
v_np1_im1 = v[1:, :-1] # v_{i - 1}[n + 1]
prog.AddLinearConstraint(
eq(flat(s_np1 - s_n),
flat(0.5 * dt * (v_n_im1 + v_np1_im1 - v_n - v_np1))))
# handle position wrap for vehicle 1
prog.AddLinearConstraint(
eq(s[1:, 0] - s[:-1, 0],
0.5 * dt * (v[:-1, -1] + v[1:, -1] - v[:-1, 0] - v[1:, 0])))
# vehicle 0's action constraint
prog.AddLinearConstraint(ge(v[1:, 0], v[:-1, 0] - u_max * dt))
prog.AddLinearConstraint(le(v[1:, 0], v[:-1, 0] + u_max * dt))
# idm constraint
prog.AddConstraint(
eq((v_np1 - v_n - dt * a * (1 - (v_n / v0)**delta)) * s_n**2,
-dt * a * (s0 + v_n * T + v_n * (v_n - v_n_im1) /
(2 * np.sqrt(a * b)))**2))
if c.cost == 'mean':
prog.AddCost(-v.mean())
elif c.cost == 'target':
prog.AddCost(((v - vf)**2).mean() + ((s - sf)**2).mean())
solver = SnoptSolver()
result = solver.Solve(prog)
if not result.is_success():
accel = self.idm_backup.step(s_init[0], v_init[0], v_init[-1])
print(f'Optimization failed, using accel {accel} from IDM')
return accel
v_desired = result.GetSolution(v)
print('Planned speeds')
print(v_desired)
print('Planned spacings')
print(result.GetSolution(s))
a_desired = (v_desired[1:, 0] - v_desired[:-1, 0]
) / dt # we're optimizing the velocity of the 0th vehicle
self.plan = a_desired
self.plan_index = 1
return self.plan[0]
root_node_tf=torch.eye(4))
torch.random.manual_seed(42)
tree = grammar.sample_tree()
assert torch.isfinite(
tree.score(verbose=True)), "Sampled tree was infeasible."
draw_scene_tree_contents_meshcat(tree, zmq_url=vis.window.zmq_url)
draw_scene_tree_structure_meshcat(tree, zmq_url=vis.window.zmq_url)
@pytest.mark.skip(
reason="This test times out. Parsing this grammar is hard now. What changed?"
)
@pytest.mark.skipif(os.environ.get('GUROBI_PATH') is None
or not SnoptSolver().available(),
reason='This test relies on Gurobi and SNOPT.')
def test_parsing():
# Try to parse an example of this grammar.
grammar = SpatialSceneGrammar(root_node_type=Desk,
root_node_tf=torch.eye(4))
torch.random.manual_seed(42)
observed_tree = grammar.sample_tree(detach=True)
observed_nodes = observed_tree.get_observed_nodes()
inference_results = infer_mle_tree_with_mip(
grammar, observed_nodes, verbose=True, max_scene_extent_in_any_dir=10.)
mip_optimized_tree = get_optimized_tree_from_mip_results(inference_results)
refinement_results = optimize_scene_tree_with_nlp(
grammar,
mip_optimized_tree,
from spatial_scene_grammars.sampling import *
from .grammar import *
from torch.distributions import constraints
torch.set_default_tensor_type(torch.DoubleTensor)
@pytest.fixture(params=range(3))
def set_seed(request):
pyro.set_rng_seed(request.param)
# Proof-of-life of sampling routines
@pytest.mark.skipif(not SnoptSolver().available(),
reason='This test relies on Gurobi and SNOPT.')
@pytest.mark.parametrize("perturb_in_config_space", [True, False])
@pytest.mark.parametrize("do_hit_and_run_postprocess", [True, False])
def test_sampling(set_seed, perturb_in_config_space,
do_hit_and_run_postprocess):
grammar = SpatialSceneGrammar(root_node_type=NodeA,
root_node_tf=torch.eye(4))
tree = grammar.sample_tree(detach=True)
# Hack tree to be totally unobserved to get more code coverage
for node in tree:
node.observed = False
N_samples = 5
sampled_trees = do_fixed_structure_mcmc(
grammar,
0.7659655809383968
]))).evaluator().set_description("Final orientation constraint"))
(prog.AddLinearConstraint(
ge(dt,
0.0)).evaluator().set_description("Timestep greater than 0 constraint"))
(prog.AddConstraint(
np.sum(dt) == 1.0).evaluator().set_description("Total time constriant"))
for k in range(N):
prog.SetInitialGuess(w_axis[k], [0, 0, 1])
prog.SetInitialGuess(w_mag[k], [0])
prog.SetInitialGuess(q[k], [1., 0., 0., 0.])
prog.SetInitialGuess(dt[k], [1.0 / N])
# prog.AddCost((w_mag[k]*w_mag[k])[0])
# prog.AddQuadraticCost(Q=np.identity(1), b=np.zeros((1,)), c = 0.0, vars=np.reshape(w_mag[k], (1,1)))
solver = SnoptSolver()
result = solver.Solve(prog)
print(result.is_success())
if not result.is_success():
print("---------- INFEASIBLE ----------")
print(result.GetInfeasibleConstraintNames(prog))
print("--------------------")
print(f"Cost = {result.get_optimal_cost()}")
q_sol = result.GetSolution(q)
print(f"q_sol = {q_sol}")
w_axis_sol = result.GetSolution(w_axis)
print(f"w_axis_sol = {w_axis_sol}")
w_mag_sol = result.GetSolution(w_mag)
print(f"w_mag_sol = {w_mag_sol}")
dt_sol = result.GetSolution(dt)
print(f"dt_sol = {dt_sol}")
def test_sampling_baseline(set_seed):
vis = meshcat.Visualizer()
# Draw a random sample from the grammar and visualize it.
grammar = SpatialSceneGrammar(
root_node_type = DishBinBaseline,
root_node_tf = torch.eye(4)
)
torch.random.manual_seed(42)
tree = grammar.sample_tree()
assert torch.isfinite(tree.score(verbose=True)), "Sampled tree was infeasible."
@pytest.mark.skipif(os.environ.get('GUROBI_PATH') is None or not SnoptSolver().available(),
reason='This test relies on Gurobi and SNOPT.')
def test_parsing_mip(set_seed):
# Try to parse an example of this grammar.
grammar = SpatialSceneGrammar(
root_node_type = DishBin,
root_node_tf = torch.eye(4)
)
torch.random.manual_seed(42)
observed_tree = grammar.sample_tree(detach=True)
observed_nodes = observed_tree.get_observed_nodes()
inference_results = infer_mle_tree_with_mip(
grammar, observed_nodes, verbose=True,
max_scene_extent_in_any_dir=10.
)
def get_initial_guess(self, x_p1, p_goal, p_puck, obstacles):
"""This is basically the single-agent MPC algorithm"""
hit_dir = p_goal - p_puck
hit_dir = 6.0 * hit_dir / np.linalg.norm(hit_dir)
x_des = np.array([p_puck[0], p_puck[1], hit_dir[0], hit_dir[1]])
#x_des = np.array([1.0, 1.0, 0, 0])
print("x_des: {}, {}".format(x_des[0], x_des[1]))
print("x_des shape", x_des.shape)
print("zeros.shape", np.zeros(4).shape)
print("p_player", x_p1[0:2])
print("p_puck {}, {}".format(p_puck[0], p_puck[1]))
print("p_goal", p_goal)
prog = DirectCollocation(self.mpc_params.sys_c,
self.mpc_params.sys_c.CreateDefaultContext(),
self.mpc_params.N + 1,
minimum_timestep=self.mpc_params.minT,
maximum_timestep=self.mpc_params.maxT)
prog.AddBoundingBoxConstraint(x_p1, x_p1, prog.initial_state())
prog.AddQuadraticErrorCost(Q=self.mpc_params.Omega_N_max,
x_desired=x_des,
vars=prog.final_state())
prog.AddEqualTimeIntervalsConstraints()
# generate trajectory non in collision with puck
#for n in range(self.mpc_params.N):
# x = prog.state()
# eps = 0.1
# obs_pos = p_puck[0:2]
# prog.AddConstraintToAllKnotPoints((x[0:2]-obs_pos).dot(x[0:2]-obs_pos) >= (self.sim_params.player_radius + self.sim_params.puck_radius - eps)**2)
for obs_pos in obstacles:
for n in range(self.mpc_params.N):
x = prog.state()
prog.AddConstraintToAllKnotPoints(
(x[0:2] - obs_pos).dot(x[0:2] - obs_pos) >= (
2.0 * self.sim_params.player_radius)**2)
prog.AddConstraintToAllKnotPoints(
prog.input()[0] <= self.sim_params.input_limit)
prog.AddConstraintToAllKnotPoints(
prog.input()[0] >= -self.sim_params.input_limit)
prog.AddConstraintToAllKnotPoints(
prog.input()[1] <= self.sim_params.input_limit)
prog.AddConstraintToAllKnotPoints(
prog.input()[1] >= -self.sim_params.input_limit)
r = self.sim_params.player_radius
prog.AddConstraintToAllKnotPoints(
prog.state()[0] + r <= self.sim_params.arena_limits_x / 2.0)
prog.AddConstraintToAllKnotPoints(
prog.state()[0] - r >= -self.sim_params.arena_limits_x / 2.0)
prog.AddConstraintToAllKnotPoints(
prog.state()[1] + r <= self.sim_params.arena_limits_y / 2.0)
prog.AddConstraintToAllKnotPoints(
prog.state()[1] - r >= -self.sim_params.arena_limits_y / 2.0)
prog.AddFinalCost(prog.time())
if not self.prev_u is None and not self.prev_x is None:
prog.SetInitialTrajectory(traj_init_u=self.prev_u,
traj_init_x=self.prev_x)
solver = SnoptSolver()
result = solver.Solve(prog)
u_traj = prog.ReconstructInputTrajectory(result)
x_traj = prog.ReconstructStateTrajectory(result)
self.prev_u = u_traj
self.prev_x = x_traj
u_vals = u_traj.vector_values(u_traj.get_segment_times())
x_vals = x_traj.vector_values(x_traj.get_segment_times())
print(u_vals)
print(u_vals[:, 0])
return u_vals[:, 0]
def project_tree_to_feasibility(tree,
constraints=[],
jitter_q=None,
do_forward_sim=False,
zmq_url=None,
prefix="projection",
timestep=0.001,
T=1.):
# Mutates tree into tree with bodies in closest
# nonpenetrating configuration.
builder, mbp, sg, node_to_free_body_ids_map, body_id_to_node_map = \
compile_scene_tree_to_mbp_and_sg(tree, timestep=timestep)
mbp.Finalize()
# Connect visualizer if requested. Wrap carefully to keep it
# from spamming the console.
if zmq_url is not None:
with open(os.devnull, 'w') as devnull:
with contextlib.redirect_stdout(devnull):
visualizer = ConnectMeshcatVisualizer(builder,
sg,
zmq_url=zmq_url,
prefix=prefix)
diagram = builder.Build()
diagram_context = diagram.CreateDefaultContext()
mbp_context = diagram.GetMutableSubsystemContext(mbp, diagram_context)
q0 = mbp.GetPositions(mbp_context)
nq = len(q0)
if nq == 0:
logging.warn("Generated MBP had no positions.")
return tree
# Set up projection NLP.
ik = InverseKinematics(mbp, mbp_context)
q_dec = ik.q()
prog = ik.prog()
# It's always a projection, so we always have this
# Euclidean norm error between the optimized q and
# q0.
prog.AddQuadraticErrorCost(np.eye(nq), q0, q_dec)
# Nonpenetration constraint.
ik.AddMinimumDistanceConstraint(0.001)
# Other requested constraints.
for constraint in constraints:
constraint.add_to_ik_prog(tree, ik, mbp, mbp_context,
node_to_free_body_ids_map)
# Initial guess, which can be slightly randomized by request.
q_guess = q0
if jitter_q:
q_guess = q_guess + np.random.normal(0., jitter_q, size=q_guess.size)
prog.SetInitialGuess(q_dec, q_guess)
# Solve.
solver = SnoptSolver()
options = SolverOptions()
logfile = "/tmp/snopt.log"
os.system("rm %s" % logfile)
options.SetOption(solver.id(), "Print file", logfile)
options.SetOption(solver.id(), "Major feasibility tolerance", 1E-3)
options.SetOption(solver.id(), "Major optimality tolerance", 1E-3)
options.SetOption(solver.id(), "Major iterations limit", 300)
result = solver.Solve(prog, None, options)
if not result.is_success():
logging.warn("Projection failed.")
print("Logfile: ")
with open(logfile) as f:
print(f.read())
qf = result.GetSolution(q_dec)
mbp.SetPositions(mbp_context, qf)
# If forward sim is requested, do a quick forward sim to get to
# a statically stable config.
if do_forward_sim:
sim = Simulator(diagram, diagram_context)
sim.set_target_realtime_rate(1000.)
sim.AdvanceTo(T)
# Reload poses back into tree
free_bodies = mbp.GetFloatingBaseBodies()
for body_id, node in body_id_to_node_map.items():
if body_id not in free_bodies:
continue
node.tf = drake_tf_to_torch_tf(
mbp.GetFreeBodyPose(mbp_context, mbp.get_body(body_id)))
return tree
def trajopt_simulation(x0, xf, u_max=0.5): # numeric parameters time_interval = .1 time_steps = 400 # initialize optimization prog = MathematicalProgram() # optimization variables state = prog.NewContinuousVariables(time_steps + 1, 3, 'state') u = prog.NewContinuousVariables(time_steps, 1, 'u') # final position constraint # for residual in constraint_state_to_orbit(state[-2], state[-1], xf, 1/u_max, time_interval): # prog.AddConstraint(residual == 0) # initial position constraint prog.AddConstraint(eq(state[0],x0)) # discretized dynamics for t in range(time_steps): residuals = dubins_discrete_dynamics(state[t], state[t+1], u[t], time_interval) for residual in residuals: prog.AddConstraint(residual == 0) # control limits prog.AddConstraint(u[t][0] <= u_max) prog.AddConstraint(-u_max <= u[t][0]) # cost - increases cost if off the orbit p = state[t][:2] - xf[:2] v = (state[t+1][:2] - state[t][:2]) / time_interval # prog.AddCost(time_interval) # prog.AddCost(u[t][0]*u[t][0]*time_interval) # prog.AddCost(p.dot(v)*time_interval) for residual in constraint_state_to_orbit(state[t], state[t+1], xf, 1/u_max, time_interval): # prog.AddConstraint(residual >= 0) prog.AddQuadraticCost(residual) # prog.AddCost((p.dot(p))) # # initial guess state_guess = interpolate_dubins_state( np.array([0, 0, np.pi]), xf, time_steps, time_interval, 1/u_max ) prog.SetInitialGuess(state, state_guess) solver = SnoptSolver() result = solver.Solve(prog) assert result.is_success() # retrieve optimal solution u_opt = result.GetSolution(u).T state_opt = result.GetSolution(state).T return state_opt, u_opt
qdd_guess = np.hstack([qdd_guess_poly.value(t * h_guess) for t in range(T)]).T
prog.SetDecisionVariableValueInVector(q, q_guess, initial_guess)
prog.SetDecisionVariableValueInVector(qd, qd_guess, initial_guess)
prog.SetDecisionVariableValueInVector(qdd, qdd_guess, initial_guess)
# initial guess for the normal component of the stance-leg force
bodies = ['body', 'stance_leg', 'swing_leg']
mass = sum(compass_gait.GetBodyByName(body).default_mass() for body in bodies)
g = - compass_gait.gravity_field().gravity_vector()[-1]
weight = mass * g
prog.SetDecisionVariableValueInVector(f[:, 1], [weight] * T, initial_guess)
"""We can finally solve the problem! Be sure that the solver actually converged: you can check this by looking at the variable `result.is_success()` (printed below)."""
# solve mathematical program with initial guess
solver = SnoptSolver()
result = solver.Solve(prog, initial_guess)
# ensure solution is found
print(f'Solution found? {result.is_success()}.')
"""In the following cell we retrieve the optimal value of the decision variables."""
# get optimal solution
h_opt = result.GetSolution(h)
q_opt = result.GetSolution(q)
qd_opt = result.GetSolution(qd)
qdd_opt = result.GetSolution(qdd)
f_opt = result.GetSolution(f)
imp_opt = result.GetSolution(imp)
qd_post_opt = result.GetSolution(qd_post)
