Add tool to run a bunch of swerve simulations in parallel
This lets us grid spaces and search for problems visually and quickly.
Change-Id: Iebbc4150cba923406e2df6174f6c37ee22c70565
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index dcbf50a..c7355b5 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -279,3 +279,16 @@
"@pip//scipy",
],
)
+
+py_binary(
+ name = "multi_casadi_velocity_mpc",
+ srcs = ["multi_casadi_velocity_mpc.py"],
+ data = [":casadi_velocity_mpc"],
+ deps = [
+ ":dynamics",
+ "@pip//absl_py",
+ "@pip//matplotlib",
+ "@pip//numpy",
+ "@pip//pygobject",
+ ],
+)
diff --git a/frc971/control_loops/swerve/casadi_velocity_mpc.py b/frc971/control_loops/swerve/casadi_velocity_mpc.py
index d3bf370..fc4a4c6 100644
--- a/frc971/control_loops/swerve/casadi_velocity_mpc.py
+++ b/frc971/control_loops/swerve/casadi_velocity_mpc.py
@@ -1,6 +1,7 @@
#!/usr/bin/env python3
from frc971.control_loops.swerve import dynamics
+import pickle
import matplotlib.pyplot as pyplot
import matplotlib
from matplotlib import pylab
@@ -19,6 +20,8 @@
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.')
+flags.DEFINE_bool('pickle', False, 'Write optimization results.')
+flags.DEFINE_string('outputdir', None, 'Directory to write problem results to')
matplotlib.use("GTK3Agg")
@@ -370,6 +373,9 @@
def main(argv):
+ if FLAGS.outputdir:
+ os.chdir(FLAGS.outputdir)
+
module_velocity = 0.0
X_initial = numpy.zeros((25, 1))
@@ -403,10 +409,65 @@
results = solver.solve(mpc=mpc,
X_initial=X_initial,
R_goal=R_goal,
- debug=True)
+ debug=(FLAGS.pickle == False))
else:
return 0
+ if FLAGS.pickle:
+ with open('t.pkl', 'wb') as f:
+ pickle.dump(solver.t, f)
+ with open('X_plot.pkl', 'wb') as f:
+ pickle.dump(solver.X_plot, f)
+ with open('U_plot.pkl', 'wb') as f:
+ pickle.dump(solver.U_plot, f)
+
+ fig0, axs0 = pylab.subplots(2)
+ fig1, axs1 = pylab.subplots(2)
+
+ axs0[0].clear()
+ axs0[1].clear()
+
+ axs0[0].plot(solver.t, solver.X_plot[dynamics.STATE_VX, :], label="vx")
+ axs0[0].plot(solver.t, solver.X_plot[dynamics.STATE_VY, :], label="vy")
+ axs0[0].legend()
+
+ axs0[1].plot(solver.t, solver.U_plot[0, :], label="Is0")
+ axs0[1].plot(solver.t, solver.U_plot[1, :], label="Id0")
+ axs0[1].legend()
+
+ axs1[0].clear()
+ axs1[1].clear()
+
+ axs1[0].plot(solver.t,
+ solver.X_plot[dynamics.STATE_THETAS0, :],
+ label='steer0')
+ axs1[0].plot(solver.t,
+ solver.X_plot[dynamics.STATE_THETAS1, :],
+ label='steer1')
+ axs1[0].plot(solver.t,
+ solver.X_plot[dynamics.STATE_THETAS2, :],
+ label='steer2')
+ axs1[0].plot(solver.t,
+ solver.X_plot[dynamics.STATE_THETAS3, :],
+ label='steer3')
+ axs1[0].legend()
+ axs1[1].plot(solver.t,
+ solver.X_plot[dynamics.STATE_OMEGAS0, :],
+ label='steer_velocity0')
+ axs1[1].plot(solver.t,
+ solver.X_plot[dynamics.STATE_OMEGAS1, :],
+ label='steer_velocity1')
+ axs1[1].plot(solver.t,
+ solver.X_plot[dynamics.STATE_OMEGAS2, :],
+ label='steer_velocity2')
+ axs1[1].plot(solver.t,
+ solver.X_plot[dynamics.STATE_OMEGAS3, :],
+ label='steer_velocity3')
+ axs1[1].legend()
+
+ fig0.savefig('state.svg')
+ fig1.savefig('steer.svg')
+
if __name__ == '__main__':
app.run(main)
diff --git a/frc971/control_loops/swerve/multi_casadi_velocity_mpc.py b/frc971/control_loops/swerve/multi_casadi_velocity_mpc.py
new file mode 100644
index 0000000..8474108
--- /dev/null
+++ b/frc971/control_loops/swerve/multi_casadi_velocity_mpc.py
@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+from absl import app
+from frc971.control_loops.swerve import dynamics
+from absl import flags
+from matplotlib import pylab
+import matplotlib
+import sys, os, pickle
+from multiprocessing.pool import ThreadPool
+import numpy
+import pathlib, collections
+import subprocess
+import itertools
+import threading
+
+matplotlib.use("GTK3Agg")
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('outdir', '/tmp/swerve', "Directory to write results to.")
+
+workerid_lock = threading.Lock()
+workerid_global = 0
+workerid = threading.local()
+
+
+def set_workerid():
+ global workerid_global
+ with workerid_lock:
+ workerid.v = workerid_global
+ workerid_global += 1
+ print(f'Starting worker {workerid.v}')
+
+
+Object = lambda **kwargs: type("Object", (), kwargs)
+
+
+def solve_mpc(velocity):
+ filename = f'vx{velocity[0]}vy{velocity[1]}omega{velocity[2]}'
+ if FLAGS.outdir:
+ subdir = pathlib.Path(FLAGS.outdir) / filename
+ else:
+ subdir = pathlib.Path(filename)
+ subdir.mkdir(parents=True, exist_ok=True)
+
+ subprocess.check_call(args=[
+ sys.executable,
+ "frc971/control_loops/swerve/casadi_velocity_mpc",
+ f"--vx={velocity[0]}",
+ f"--vy={velocity[1]}",
+ f"--omega={velocity[2]}",
+ f"--outputdir={subdir.resolve()}",
+ "--pickle",
+ ])
+
+ with open(subdir / 't.pkl', 'rb') as f:
+ t = pickle.load(f)
+
+ with open(subdir / 'X_plot.pkl', 'rb') as f:
+ X_plot = pickle.load(f)
+
+ with open(subdir / 'U_plot.pkl', 'rb') as f:
+ U_plot = pickle.load(f)
+
+ return Object(t=t, X_plot=X_plot, U_plot=U_plot)
+
+
+def main(argv):
+ # Load a simple problem first so we compile with less system load. This
+ # makes it faster on a processor with frequency boosting.
+ subprocess.check_call(args=[
+ sys.executable,
+ "frc971/control_loops/swerve/casadi_velocity_mpc",
+ "--compileonly",
+ ])
+
+ # Try a bunch of goals now
+ vxy = numpy.array(
+ numpy.meshgrid(numpy.linspace(-1, 1, 9),
+ numpy.linspace(-1, 1, 9))).T.reshape(-1, 2)
+
+ velocity = numpy.hstack((vxy, numpy.zeros((vxy.shape[0], 1))))
+
+ with ThreadPool(initializer=set_workerid) as pool:
+ results = pool.starmap(solve_mpc, zip(velocity, ))
+
+ fig0, axs0 = pylab.subplots(2)
+ for r in results:
+ axs0[0].plot(r.X_plot[dynamics.STATE_VX, :],
+ r.X_plot[dynamics.STATE_VY, :],
+ label='trajectory')
+ axs0[1].plot(r.t, r.U_plot[0, :], label="Is0")
+ axs0[1].plot(r.t, r.U_plot[1, :], label="Id0")
+
+ axs0[0].legend()
+ axs0[1].legend()
+ pylab.show()
+
+
+if __name__ == '__main__':
+ app.run(main)