Use flags in casadi_velocity_mpc
This sets us up to do batch runs easily.
Change-Id: I0677b2f36ca091b930aeb56e6e48f862901f9c7a
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index cbc4e66..dcbf50a 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -219,6 +219,7 @@
],
deps = [
":dynamics",
+ "@pip//absl_py",
"@pip//casadi",
"@pip//matplotlib",
"@pip//numpy",
diff --git a/frc971/control_loops/swerve/casadi_velocity_mpc.py b/frc971/control_loops/swerve/casadi_velocity_mpc.py
index f5e746b..d3bf370 100644
--- a/frc971/control_loops/swerve/casadi_velocity_mpc.py
+++ b/frc971/control_loops/swerve/casadi_velocity_mpc.py
@@ -1,4 +1,4 @@
-#!/usr/bin/python3
+#!/usr/bin/env python3
from frc971.control_loops.swerve import dynamics
import matplotlib.pyplot as pyplot
@@ -9,6 +9,16 @@
import scipy
import casadi
import os, sys
+from absl import flags
+from absl import app
+
+FLAGS = flags.FLAGS
+flags.DEFINE_bool('compileonly', False,
+ 'If true, load casadi, don\'t compile it')
+flags.DEFINE_float('vx', 1.0, 'Goal velocity in m/s in x')
+flags.DEFINE_float('vy', 0.0, 'Goal velocity in m/s in y')
+flags.DEFINE_float('omega', 0.0, 'Goal velocity in m/s in omega')
+flags.DEFINE_float('duration', 0.5, 'Time to simulate in seconds.')
matplotlib.use("GTK3Agg")
@@ -78,6 +88,7 @@
Xn = casadi.horzcat(*Xn_variables)
U = casadi.horzcat(*U_variables)
+ # printme(number) is the debug.
Xk_begin = casadi.horzcat(self.X0, Xn)
Xk_end = self.next_X.map(self.N, "thread")(Xk_begin, U)
J = casadi.sum2(self.cost.map(self.N, "thread")(Xk_end, U, self.R))
@@ -140,7 +151,6 @@
"debug": True,
"equality": equality,
}
- self.solver = casadi.nlpsol('solver', 'fatrop', prob, options)
else:
options = {
"jit": True,
@@ -154,8 +164,10 @@
options["ipopt"] = {
"print_level": 12,
}
- self.solver = casadi.nlpsol('solver', 'ipopt', prob, options)
+ self.solver = casadi.nlpsol('solver', solver, prob, options)
+
+ # TODO(austin): Vary the number of sub steps to be more short term and fewer long term?
def make_physics(self):
X0 = casadi.MX.sym('X0', dynamics.NUM_VELOCITY_STATES)
U = casadi.MX.sym('U', 8)
@@ -249,112 +261,152 @@
(i - 1):(8 + dynamics.NUM_VELOCITY_STATES) * i]
-mpc = MPC(solver='fatrop') if not full_debug else MPC(solver='ipopt')
+class Solver(object):
-R_goal = numpy.zeros((3, 1))
-R_goal[0, 0] = 0.0
-R_goal[1, 0] = -1.0
-R_goal[2, 0] = 0.0
+ def __init__(self):
+ self.iterations = int(round(FLAGS.duration / 0.005))
-module_velocity = 0.0
+ self.X_plot = numpy.zeros((25, self.iterations))
+ self.U_plot = numpy.zeros((8, self.iterations))
+ self.t = []
-X_initial = numpy.zeros((25, 1))
-# All the wheels are spinning at the speed needed to hit 1 m/s
-X_initial[dynamics.STATE_THETAS0, 0] = 0.0
-X_initial[dynamics.STATE_OMEGAS0, 0] = module_velocity
+ def solve(self, mpc, X_initial, R_goal, debug=False):
+ X = X_initial.copy()
-X_initial[dynamics.STATE_THETAS1, 0] = 0.0
-X_initial[dynamics.STATE_OMEGAS1, 0] = module_velocity
+ if debug:
+ pyplot.ion()
+ fig0, axs0 = pylab.subplots(2)
+ fig1, axs1 = pylab.subplots(2)
-X_initial[dynamics.STATE_THETAS2, 0] = 0.0
-X_initial[dynamics.STATE_OMEGAS2, 0] = module_velocity
-
-X_initial[dynamics.STATE_THETAS3, 0] = 0.0
-X_initial[dynamics.STATE_OMEGAS3, 0] = module_velocity
-
-# Robot is moving at 0 m/s
-X_initial[dynamics.STATE_VX, 0] = 0.0
-X_initial[dynamics.STATE_VY, 0] = 0.0
-# No angular velocity
-X_initial[dynamics.STATE_OMEGA, 0] = 0.0
-
-iterations = 100
-
-X_plot = numpy.zeros((25, iterations))
-U_plot = numpy.zeros((8, iterations))
-t = []
-
-X = X_initial.copy()
-
-pyplot.ion()
-
-fig0, axs0 = pylab.subplots(2)
-fig1, axs1 = pylab.subplots(2)
-last_time = time.time()
-
-seed = ([0, 0] * 4 + list(dynamics.to_velocity_state(X))) * (mpc.N -
- 1) + [0, 0] * 4
-
-overall_time = 0
-for i in range(iterations):
- t.append(i * mpc.dt)
- print("Current X at", i * mpc.dt, X.transpose())
- print("Goal R at", i * mpc.dt, R_goal)
- start_time = time.time()
- sol = mpc.solve(
- # TODO(austin): Is this better or worse than constraints on the initial state for convergence?
- p=numpy.vstack((dynamics.to_velocity_state(X), R_goal)),
- seed=seed)
- end_time = time.time()
- print(f"Took {end_time - start_time} seconds to solve.")
- overall_time += end_time - start_time
-
- X_plot[:, i] = X[:, 0]
-
- U = mpc.unpack_u(sol, 0)
- seed = (
- list(sol['x'].full().flatten()[8 + dynamics.NUM_VELOCITY_STATES:]) +
- list(sol['x'].full().flatten()[-(8 + dynamics.NUM_VELOCITY_STATES):]))
- U_plot[:, i] = U
-
- print('x(0):', X.transpose())
- for j in range(mpc.N):
- print(f'u({j}): ', mpc.unpack_u(sol, j))
- print(f'x({j+1}): ', mpc.unpack_x(sol, j + 1))
-
- result = scipy.integrate.solve_ivp(
- lambda t, x: mpc.wrapped_swerve_physics(x, U).flatten(), [0, mpc.dt],
- X.flatten())
- X[:, 0] = result.y[:, -1]
-
- if time.time() > last_time + 2 or i == iterations - 1:
- axs0[0].clear()
- axs0[1].clear()
-
- axs0[0].plot(t, X_plot[dynamics.STATE_VX, 0:i + 1], label="vx")
- axs0[0].plot(t, X_plot[dynamics.STATE_VY, 0:i + 1], label="vy")
- axs0[0].legend()
-
- axs0[1].plot(t, U_plot[0, 0:i + 1], label="Is0")
- axs0[1].plot(t, U_plot[1, 0:i + 1], label="Id0")
- axs0[1].legend()
-
- axs1[0].clear()
- axs1[1].clear()
-
- axs1[0].plot(t, X_plot[0, 0:i + 1], label='steer0')
- axs1[0].plot(t, X_plot[4, 0:i + 1], label='steer1')
- axs1[0].plot(t, X_plot[8, 0:i + 1], label='steer2')
- axs1[0].plot(t, X_plot[12, 0:i + 1], label='steer3')
- axs1[0].legend()
- axs1[1].plot(t, X_plot[2, 0:i + 1], label='steer_velocity0')
- axs1[1].plot(t, X_plot[6, 0:i + 1], label='steer_velocity1')
- axs1[1].plot(t, X_plot[10, 0:i + 1], label='steer_velocity2')
- axs1[1].plot(t, X_plot[14, 0:i + 1], label='steer_velocity3')
- axs1[1].legend()
- pyplot.pause(0.0001)
last_time = time.time()
-print(f"Tool {overall_time} seconds overall to solve.")
+ seed = ([0, 0] * 4 +
+ list(dynamics.to_velocity_state(X))) * (mpc.N - 1) + [0, 0] * 4
-pyplot.pause(-1)
+ overall_time = 0
+ for i in range(self.iterations):
+ self.t.append(i * mpc.dt)
+ print("Current X at", i * mpc.dt, X.transpose())
+ print("Goal R at", i * mpc.dt, R_goal)
+ start_time = time.time()
+ sol = mpc.solve(
+ # TODO(austin): Is this better or worse than constraints on the initial state for convergence?
+ p=numpy.vstack((dynamics.to_velocity_state(X), R_goal)),
+ seed=seed)
+ end_time = time.time()
+ print(f"Took {end_time - start_time} seconds to solve.")
+ overall_time += end_time - start_time
+
+ self.X_plot[:, i] = X[:, 0]
+
+ U = mpc.unpack_u(sol, 0)
+ seed = (list(
+ sol['x'].full().flatten()[8 + dynamics.NUM_VELOCITY_STATES:]) +
+ list(sol['x'].full().flatten()
+ [-(8 + dynamics.NUM_VELOCITY_STATES):]))
+ self.U_plot[:, i] = U
+
+ print('x(0):', X.transpose())
+ for j in range(mpc.N):
+ print(f'u({j}): ', mpc.unpack_u(sol, j))
+ print(f'x({j+1}): ', mpc.unpack_x(sol, j + 1))
+
+ result = scipy.integrate.solve_ivp(
+ lambda t, x: mpc.wrapped_swerve_physics(x, U).flatten(),
+ [0, mpc.dt], X.flatten())
+ X[:, 0] = result.y[:, -1]
+
+ if time.time() > last_time + 2 or i == self.iterations - 1:
+ if debug:
+ axs0[0].clear()
+ axs0[1].clear()
+
+ axs0[0].plot(self.t,
+ self.X_plot[dynamics.STATE_VX, 0:i + 1],
+ label="vx")
+ axs0[0].plot(self.t,
+ self.X_plot[dynamics.STATE_VY, 0:i + 1],
+ label="vy")
+ axs0[0].legend()
+
+ axs0[1].plot(self.t, self.U_plot[0, 0:i + 1], label="Is0")
+ axs0[1].plot(self.t, self.U_plot[1, 0:i + 1], label="Id0")
+ axs0[1].legend()
+
+ axs1[0].clear()
+ axs1[1].clear()
+
+ axs1[0].plot(self.t,
+ self.X_plot[dynamics.STATE_THETAS0, 0:i + 1],
+ label='steer0')
+ axs1[0].plot(self.t,
+ self.X_plot[dynamics.STATE_THETAS1, 0:i + 1],
+ label='steer1')
+ axs1[0].plot(self.t,
+ self.X_plot[dynamics.STATE_THETAS2, 0:i + 1],
+ label='steer2')
+ axs1[0].plot(self.t,
+ self.X_plot[dynamics.STATE_THETAS3, 0:i + 1],
+ label='steer3')
+ axs1[0].legend()
+ axs1[1].plot(self.t,
+ self.X_plot[dynamics.STATE_OMEGAS0, 0:i + 1],
+ label='steer_velocity0')
+ axs1[1].plot(self.t,
+ self.X_plot[dynamics.STATE_OMEGAS1, 0:i + 1],
+ label='steer_velocity1')
+ axs1[1].plot(self.t,
+ self.X_plot[dynamics.STATE_OMEGAS2, 0:i + 1],
+ label='steer_velocity2')
+ axs1[1].plot(self.t,
+ self.X_plot[dynamics.STATE_OMEGAS3, 0:i + 1],
+ label='steer_velocity3')
+ axs1[1].legend()
+ pyplot.pause(0.0001)
+
+ last_time = time.time()
+
+ print(f"Tool {overall_time} seconds overall to solve.")
+
+
+def main(argv):
+ module_velocity = 0.0
+
+ X_initial = numpy.zeros((25, 1))
+ # All the wheels are spinning at the speed needed to hit 1 m/s
+ X_initial[dynamics.STATE_THETAS0, 0] = 0.0
+ X_initial[dynamics.STATE_OMEGAS0, 0] = module_velocity
+
+ X_initial[dynamics.STATE_THETAS1, 0] = 0.0
+ X_initial[dynamics.STATE_OMEGAS1, 0] = module_velocity
+
+ X_initial[dynamics.STATE_THETAS2, 0] = 0.0
+ X_initial[dynamics.STATE_OMEGAS2, 0] = module_velocity
+
+ X_initial[dynamics.STATE_THETAS3, 0] = 0.0
+ X_initial[dynamics.STATE_OMEGAS3, 0] = module_velocity
+
+ # Robot is moving at 0 m/s
+ X_initial[dynamics.STATE_VX, 0] = 0.0
+ X_initial[dynamics.STATE_VY, 0] = 0.0
+ # No angular velocity
+ X_initial[dynamics.STATE_OMEGA, 0] = 0.0
+
+ R_goal = numpy.zeros((3, 1))
+ R_goal[0, 0] = FLAGS.vx
+ R_goal[1, 0] = FLAGS.vy
+ R_goal[2, 0] = FLAGS.omega
+
+ mpc = MPC(solver='fatrop') if not full_debug else MPC(solver='ipopt')
+ solver = Solver()
+ if not FLAGS.compileonly:
+ results = solver.solve(mpc=mpc,
+ X_initial=X_initial,
+ R_goal=R_goal,
+ debug=True)
+ else:
+ return 0
+
+
+if __name__ == '__main__':
+ app.run(main)