import numpy as np
from pydrake.all import MathematicalProgram, IpoptSolver, ge, le, eq
from pydrake.symbolic import ToLatex
from IPython.display import display, Latex
prog = MathematicalProgram() # Instance of mathematical program
x = prog.NewContinuousVariables(1, 'x') # 1 element vector of decision variables named x
prog.AddQuadraticCost(x @ x) # Matrix product (since x is 1 element vector also known as 1 x 1 matrix)
prog.AddLinearConstraint(ge(x, 0.5 * np.ones_like(x))) # x >= 0.5 constraint
solver = IpoptSolver() # Internal-point method solver, default solver from Solve function is SNOPT which has more restricted licence
result = solver.Solve(prog) # Solve the mathematical program instance
assert result.is_success()
print(result.GetSolution(x))
prog = MathematicalProgram() # Instantiate new mathematical program
q_0 = np.float32([0, 0]) # Initial position q = (x, y)
q_0_dot = prog.NewContinuousVariables(2, '\dot{q_0}') # Decision variable for initial velocity q_dot = (x_dot, y_dot)
q_T = np.float32([3, 1]) # Final position q = (x, y)
t = 0.1 # Time of flight
g = 9.81 # Gravity
q_dot_dot = np.float32([0, -g]) # Gravity in full coordinates
# Physics constraint derived from equation of motion: m * q_dot_dot = m * g
# integrated two times with q_0 and q_dot_0 initial conditions
prog.AddConstraint(eq(q_0 + t * q_0_dot + t ** 2 * q_dot_dot / 2, q_T))
solver = IpoptSolver()
result = solver.Solve(prog)
assert result.is_success()
print(result.GetSolution(q_0_dot))
prog = MathematicalProgram()
q_0_dot = prog.NewContinuousVariables(2, '\dot{q_0}')
t = prog.NewContinuousVariables(1, 't') # Time of flight as decision variable on which we put the cost function J = t^2
# Initial guess for time of flight - in nonlinear case, the optimization is very sensitive to initial guess.
# In general, one can use motion capture of some performed trajectory, linear/polynomial interpolation between initial/goal state etc.
prog.SetInitialGuess(t, np.float32([1.0]))
q_0 = np.float32([0, 0])
q_T = np.float32([3, 1])
g = 9.81
q_dot_dot = np.float32([0, -g])
prog.AddConstraint(eq(q_0 + t * q_0_dot + t ** 2 * q_dot_dot / 2, q_T))
prog.AddQuadraticCost(t @ t) # Cost on time of flight. The shorter the better
# Note that optimal cost value (0) for t=0 violates equation of motiion constraint that (added just above).
# The optimizer tries it best to find the smallest t that still satisfies the constraints.
# Formulated like this, initial velocity can be arbitraly high and either local minima will be encountered
# or iterations limit for optimization. One can add more constraints (e.x. limit speed v = sqrt(x_dot^2 + y_dot^2)
# or add cost on speed.
solver = IpoptSolver()
result = solver.Solve(prog)
assert result.is_success()
print(f"q_dot_0 = {result.GetSolution(q_0_dot)}")
print(f"t_optimal = {result.GetSolution(t)}")
from functools import partial
from pydrake.all import sqrt, Solve
def canon_dynamics(state, c, g):
# State consist of position, q=(x, y), and velocity, q_dot=(x_dot, y_dot)
x, y, x_dot, y_dot = state
v_sqr = x_dot * x_dot + y_dot * y_dot
v = v_sqr ** 0.5
x_dot_dot = -c * (x_dot * v)
y_dot_dot = -c * (y_dot * v) - g
state_dot = np.array([x_dot, y_dot, x_dot_dot, y_dot_dot])
return state_dot
def initialize_state_trajectory(N):
# Initialize
traj_initial = np.zeros((N, 4), dtype=np.float32)
traj_initial[:, 0] = np.linspace(0, 3, N)
traj_initial[:, 1] = np.linspace(0, 1, N)
traj_initial[:-1, 2] = traj_initial[1:, 0] - traj_initial[:-1, 0]
traj_initial[:-1, 3] = traj_initial[1:, 1] - traj_initial[:-1, 1]
return traj_initial
c_no_drag = 0.0
c_drag = 0.4
g = 9.81
canon_dynamics_no_drag = partial(canon_dynamics, c=c_no_drag, g=g)
canon_dynamics_with_drag = partial(canon_dynamics, c=c_drag, g=g)
# Test to check the dynamics function
state_test = np.ones(4)
c_test = 0.7
state_test_dot = np.array([1., 1., -0.98994949, -10.79994949])
np.testing.assert_allclose(canon_dynamics(state_test, c_test, g), state_test_dot)
def solve_optimization(dynamics=canon_dynamics_no_drag, traj_initial=None, N=50):
prog = MathematicalProgram()
N = N
t = prog.NewContinuousVariables(1, 't')
prog.AddConstraint(ge(t, np.array([0.01])))
prog.SetInitialGuess(t, np.float32([1.0]))
h = t / N
traj = prog.NewContinuousVariables(N, 4, 'traj') # x, y, x_dot, y_dot trajectory
if traj_initial is None:
traj_initial = initialize_state_trajectory(N)
else:
traj_initial=traj_initial
prog.SetInitialGuess(traj, traj_initial)
# TODO: Add constraint on initial position (first two elements of traj[0])
prog.AddConstraint(eq(np.array([traj[0, 0],traj[0, 1]]), np.array([0, 0])))
# TODO: Add constraint on final position (first two elements of traj[0])
prog.AddConstraint(eq(np.array([traj[N-1, 0],traj[N-1, 1]]), np.array([3, 1])))
# TODO: Initialize derivative of trajectory
# Hint: dynamics takes state (x, y, x_dot, y_dot) and returns derivative
# initialize traj_dot so traj_dot[i] contains derivative at i'th step in trajectory
traj_dot = [dynamics(traj[i]) for i in range(N)]
# TODO: Add dynamics constraint
# Hint: Check 3.2 @ https://epubs.siam.org/doi/epdf/10.1137/16M1062569
# Hint: x_k is state at k=i, f_k is state_dot at k=i (dynamics at x_k)
for i in range(N - 1):
prog.AddConstraint(eq(traj[i+1], traj[i] + h * traj_dot[i]))
# TODO: Add cost on squared initial speed
# prog.AddCost(...)
prog.AddCost(traj[0, 2]**2 + traj[0, 3]**2)
solver = IpoptSolver()
result = solver.Solve(prog)
assert result.is_success()
return result.GetSolution(traj), result.GetSolution(t)
from matplotlib import pyplot as plt
traj_sol_no_drag, t_sol = solve_optimization(dynamics=canon_dynamics_no_drag)
xy = traj_sol_no_drag[:, :2]
plt.plot(xy[:, 0], xy[:, 1])
plt.plot(xy[:, 0], xy[:, 1], 'o');
plt.plot(traj_sol[:, 2])
plt.title('Horizontal velocity $\dot{x}$');
traj_sol, t_sol = solve_optimization(dynamics=canon_dynamics_with_drag, traj_initial=traj_sol)
xy = traj_sol[:, :2]
plt.plot(xy[:, 0], xy[:, 1])
plt.plot(xy[:, 0], xy[:, 1], 'o');
plt.plot(traj_sol[:, 2])
plt.title('Horizontal velocity $\dot{x}$');
# TODO: Solve the optimization, this time set initial velocity
# at some reasonable value and keep the rest of initial
# trajectory from initialize_state_trajectory
traj_initial = initialize_state_trajectory(50)
traj_initial[0, 2] = 10
traj_initial[0, 3] = 8
# Dla wyniku solvera dla równiania bez oporów
traj_sol, t_sol = solve_optimization(dynamics=canon_dynamics_with_drag, traj_initial=traj_sol_no_drag)
# Dla ustawionej prędkości
# traj_sol, t_sol = solve_optimization(dynamics=canon_dynamics_with_drag, traj_initial=traj_initial)
xy = traj_sol[:, :2]
plt.plot(xy[:, 0], xy[:, 1])
plt.plot(xy[:, 0], xy[:, 1], 'o');