diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 513bf1b..2c3ad50 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -227,11 +227,15 @@
     main = "physics_test.py",
     target_compatible_with = ["@platforms//cpu:x86_64"],
     deps = [
+        ":casadi_velocity_mpc_lib",
         ":dynamics",
         ":jax_dynamics",
         ":physics_test_utils",
+        "@pip//absl_py",
         "@pip//casadi",
+        "@pip//matplotlib",
         "@pip//numpy",
+        "@pip//pygobject",
         "@pip//scipy",
     ],
 )
@@ -248,12 +252,28 @@
     main = "physics_test.py",
     target_compatible_with = ["@platforms//cpu:x86_64"],
     deps = [
+        ":casadi_velocity_mpc_lib",
         ":dynamics",
         ":jax_dynamics",
         ":physics_test_utils",
+        "@pip//absl_py",
+        "@pip//casadi",
+        "@pip//matplotlib",
+        "@pip//numpy",
+        "@pip//pygobject",
+        "@pip//scipy",
+    ],
+)
+
+py_binary(
+    name = "casadi_velocity_mpc_lib",
+    srcs = [
+        "casadi_velocity_mpc_lib.py",
+    ],
+    deps = [
+        ":dynamics",
         "@pip//casadi",
         "@pip//numpy",
-        "@pip//scipy",
     ],
 )
 
@@ -263,9 +283,8 @@
         "casadi_velocity_mpc.py",
     ],
     deps = [
-        ":dynamics",
+        ":casadi_velocity_mpc_lib",
         "@pip//absl_py",
-        "@pip//casadi",
         "@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 517de5a..d62f7ed 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
+from frc971.control_loops.swerve.casadi_velocity_mpc_lib import MPC
 import pickle
 import matplotlib.pyplot as pyplot
 import matplotlib
@@ -23,296 +24,10 @@
 flags.DEFINE_bool('pickle', False, 'Write optimization results.')
 flags.DEFINE_string('outputdir', None, 'Directory to write problem results to')
 
-matplotlib.use("GTK3Agg")
-
 # Full print level on ipopt. Piping to a file and using a filter or search method is suggested
 # grad_x prints out the gradient at each iteration in the following sequence: U0, X1, U1, etc.
-full_debug = False
-
-
-class MPC(object):
-
-    def __init__(self, solver='fatrop'):
-        self.fdot = dynamics.swerve_full_dynamics(
-            casadi.SX.sym("X", dynamics.NUM_STATES, 1),
-            casadi.SX.sym("U", 8, 1))
-        self.velocity_fdot = dynamics.velocity_swerve_physics(
-            casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
-            casadi.SX.sym("U", 8, 1))
-
-        self.wrapped_swerve_physics = lambda X, U: numpy.array(self.fdot(X, U))
-
-        self.dt = 0.005
-
-        # TODO(austin): Do we need a disturbance torque per module to account for friction?
-        # Do it only in the observer/post?
-
-        self.force = [
-            dynamics.F(i, casadi.SX.sym("X", 25, 1), casadi.SX.sym("U", 8, 1))
-            for i in range(4)
-        ]
-        self.force_vel = [
-            dynamics.F_vel(i,
-                           casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
-                           casadi.SX.sym("U", 8, 1)) for i in range(4)
-        ]
-        self.slip_angle = [
-            dynamics.slip_angle(i, casadi.SX.sym("X", 25, 1),
-                                casadi.SX.sym("U", 8, 1)) for i in range(4)
-        ]
-        self.mounting_location = [
-            dynamics.mounting_location(
-                i, casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
-                casadi.SX.sym("U", 8, 1)) for i in range(4)
-        ]
-        self.torque_cross = self.torque_cross_func(casadi.SX.sym("r", 2, 1),
-                                                   casadi.SX.sym("F", 2, 1))
-        self.force_cross = self.force_cross_func(casadi.SX.sym("Tau", 1, 1),
-                                                 casadi.SX.sym("r", 2, 1))
-        self.next_X = self.make_physics()
-        self.cost = self.make_cost()
-
-        self.N = 200
-
-        # Start with an empty nonlinear program.
-        self.w = []
-        self.lbw = []
-        self.ubw = []
-        J = 0
-        self.g = []
-        self.lbg = []
-        self.ubg = []
-
-        self.X0 = casadi.MX.sym('X0', dynamics.NUM_VELOCITY_STATES)
-
-        # We care about the linear and angular velocities only.
-        self.R = casadi.MX.sym('R', 3)
-
-        # Make Xn and U for each step.  fatrop wants us to interleave the control variables and
-        # states so that it can produce a banded/structured problem which it can solve a lot
-        # faster.
-        Xn_variables = []
-        U_variables = []
-        for i in range(self.N):
-            U_variables.append(casadi.MX.sym(f'U{i}', 8))
-
-            if i == 0:
-                continue
-
-            Xn_variables.append(
-                casadi.MX.sym(f'X{i}', dynamics.NUM_VELOCITY_STATES))
-
-        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))
-
-        # Put U and Xn interleaved into w to go fast.
-        for i in range(self.N):
-            self.w += [U_variables[i]]
-            self.ubw += [100] * 8
-            self.lbw += [-100] * 8
-
-            if i == self.N - 1:
-                continue
-
-            self.w += [Xn_variables[i]]
-            self.ubw += [casadi.inf] * dynamics.NUM_VELOCITY_STATES
-            self.lbw += [-casadi.inf] * dynamics.NUM_VELOCITY_STATES
-
-        self.g += [
-            casadi.reshape(Xn - Xk_end[:, 0:(self.N - 1)],
-                           dynamics.NUM_VELOCITY_STATES * (self.N - 1), 1)
-        ]
-
-        self.lbg += [0] * dynamics.NUM_VELOCITY_STATES * (self.N - 1)
-        self.ubg += [0] * dynamics.NUM_VELOCITY_STATES * (self.N - 1)
-
-        prob = {
-            'f': J,
-            # lbx <= x <= ubx
-            'x': casadi.vertcat(*self.w),
-            # lbg <= g(x, p) <= ubg
-            'g': casadi.vertcat(*self.g),
-            # Input parameters (initial position + goal)
-            'p': casadi.vertcat(self.X0, self.R),
-        }
-
-        compiler = "ccache clang"
-        flags = ["-O3"]
-        jit_options = {
-            "flags": flags,
-            "verbose": False,
-            "compiler": compiler,
-            "temp_suffix": False,
-        }
-
-        if solver == 'fatrop':
-            equality = [
-                True
-                for _ in range(dynamics.NUM_VELOCITY_STATES * (self.N - 1))
-            ]
-            options = {
-                "jit": True,
-                "jit_cleanup": False,
-                "jit_temp_suffix": False,
-                "compiler": "shell",
-                "jit_options": jit_options,
-                "structure_detection": "auto",
-                "fatrop": {
-                    "tol": 1e-7
-                },
-                "debug": True,
-                "equality": equality,
-            }
-        else:
-            options = {
-                "jit": True,
-                "jit_cleanup": False,
-                "jit_temp_suffix": False,
-                "compiler": "shell",
-                "jit_options": jit_options,
-            }
-            if full_debug:
-                options["jit"] = False
-                options["ipopt"] = {
-                    "print_level": 12,
-                }
-
-        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)
-
-        X = X0
-        M = 2  # RK4 steps per interval
-        DT = self.dt / M
-
-        for j in range(M):
-            k1 = self.velocity_fdot(X, U)
-            k2 = self.velocity_fdot(X + DT / 2 * k1, U)
-            k3 = self.velocity_fdot(X + DT / 2 * k2, U)
-            k4 = self.velocity_fdot(X + DT * k3, U)
-            X = X + DT / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
-
-        return casadi.Function("next_X", [X0, U], [X])
-
-    def make_cost(self):
-        # TODO(austin): tune cost fn?
-        # Do we want to penalize slipping tires?
-        # Do we want to penalize powers unevenly applied across the motors?
-        # Need to do some simulations to see what works well.
-
-        X = casadi.MX.sym('X', dynamics.NUM_VELOCITY_STATES)
-        U = casadi.MX.sym('U', 8)
-        R = casadi.MX.sym('R', 3)
-
-        J = 0
-        vnorm = casadi.sqrt(R[0]**2.0 + R[1]**2.0)
-
-        vnormx = casadi.if_else(vnorm > 0.0001, R[0] / vnorm, 1.0)
-        vnormy = casadi.if_else(vnorm > 0.0001, R[1] / vnorm, 0.0)
-
-        vperpx = casadi.if_else(vnorm > 0.0001, -vnormy, 0.0)
-        vperpy = casadi.if_else(vnorm > 0.0001, vnormx, 1.0)
-
-        # TODO(austin): Do we want to do something more special for 0?
-
-        # cost velocity a lot more in the perpendicular direction to allow tire to spin up
-        # we also only want to get moving in the correct direction as fast as possible
-        J += 75 * ((R[0] - X[dynamics.VELOCITY_STATE_VX]) * vnormx +
-                   (R[1] - X[dynamics.VELOCITY_STATE_VY]) * vnormy)**2.0
-
-        J += 1500 * ((R[0] - X[dynamics.VELOCITY_STATE_VX]) * vperpx +
-                     (R[1] - X[dynamics.VELOCITY_STATE_VY]) * vperpy)**2.0
-
-        J += 1000 * (R[2] - X[dynamics.VELOCITY_STATE_OMEGA])**2.0
-
-        kSteerPositionGain = 0
-        kSteerVelocityGain = 0.10
-        J += kSteerPositionGain * (X[dynamics.VELOCITY_STATE_THETAS0])**2.0
-        J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS0])**2.0
-
-        J += kSteerPositionGain * (X[dynamics.VELOCITY_STATE_THETAS1])**2.0
-        J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS1])**2.0
-
-        J += kSteerPositionGain * (X[dynamics.VELOCITY_STATE_THETAS2])**2.0
-        J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS2])**2.0
-
-        J += kSteerPositionGain * (X[dynamics.VELOCITY_STATE_THETAS3])**2.0
-        J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS3])**2.0
-
-        # cost variance of the force by a tire and the expected average force and torque on it
-        total_force = self.force_vel[0](X, U)
-        total_torque = self.torque_cross(self.mounting_location[0](X, U),
-                                         self.force_vel[0](X, U))
-        for i in range(3):
-            total_force += self.force_vel[i + 1](X, U)
-            total_torque += self.torque_cross(
-                self.mounting_location[i + 1](X, U), self.force_vel[i + 1](X,
-                                                                           U))
-
-        total_force /= 4
-        total_torque /= 4
-        for i in range(4):
-            f_diff = (total_force +
-                      self.force_cross(total_torque, self.mounting_location[i]
-                                       (X, U))) - self.force_vel[i](X, U)
-            J += 0.01 * (f_diff[0, 0]**2.0 + f_diff[1, 0]**2.0)
-
-        # TODO(austin): Don't penalize torque steering current.
-        for i in range(4):
-            Is = U[2 * i + 0]
-            Id = U[2 * i + 1]
-            # Steer, cost it a lot less than drive to be more agressive in steering
-            J += ((Is + dynamics.STEER_CURRENT_COUPLING_FACTOR * Id)**
-                  2.0) / 100000.0
-            # Drive
-            J += Id * Id / 1000.0
-
-        return casadi.Function("Jn", [X, U, R], [J])
-
-    def torque_cross_func(self, r, F):
-        result = casadi.SX.sym('Tau', 1, 1)
-        result[0, 0] = r[0, 0] * F[1, 0] - r[1, 0] * F[0, 0]
-        return casadi.Function('Tau', [r, F], [result])
-
-    def force_cross_func(self, Tau, r):
-        result = casadi.SX.sym('F', 2, 1)
-        result[0, 0] = -r[1, 0] * Tau[0, 0] / casadi.norm_2(r)**2.0
-        result[1, 0] = r[0, 0] * Tau[0, 0] / casadi.norm_2(r)**2.0
-        return casadi.Function('F', [Tau, r], [result])
-
-    def solve(self, p, seed=None):
-        if seed is None:
-            seed = []
-
-            for i in range(self.N):
-                seed += [0, 0] * 4
-                if i < self.N - 1:
-                    seed += list(p[:dynamics.NUM_VELOCITY_STATES, 0])
-
-        return self.solver(x0=seed,
-                           lbx=self.lbw,
-                           ubx=self.ubw,
-                           lbg=self.lbg,
-                           ubg=self.ubg,
-                           p=casadi.DM(p))
-
-    def unpack_u(self, sol, i):
-        return sol['x'].full().flatten()[
-            (8 + dynamics.NUM_VELOCITY_STATES) *
-            i:((8 + dynamics.NUM_VELOCITY_STATES) * i + 8)]
-
-    def unpack_x(self, sol, i):
-        return sol['x'].full().flatten(
-        )[8 + (8 + dynamics.NUM_VELOCITY_STATES) *
-          (i - 1):(8 + dynamics.NUM_VELOCITY_STATES) * i]
+flags.DEFINE_bool('full_debug', False,
+                  'If true, turn on all the debugging in the solver.')
 
 
 class Solver(object):
@@ -423,10 +138,12 @@
 
                 last_time = time.time()
 
-        print(f"Tool {overall_time} seconds overall to solve.")
+        print(f"Took {overall_time} seconds overall to solve.")
 
 
 def main(argv):
+    matplotlib.use("GTK3Agg")
+
     if FLAGS.outputdir:
         os.chdir(FLAGS.outputdir)
 
@@ -457,7 +174,7 @@
     R_goal[1, 0] = FLAGS.vy
     R_goal[2, 0] = FLAGS.omega
 
-    mpc = MPC(solver='fatrop') if not full_debug else MPC(solver='ipopt')
+    mpc = MPC(solver='fatrop') if not FLAGS.full_debug else MPC(solver='ipopt')
     solver = Solver()
     if not FLAGS.compileonly:
         results = solver.solve(mpc=mpc,
diff --git a/frc971/control_loops/swerve/casadi_velocity_mpc_lib.py b/frc971/control_loops/swerve/casadi_velocity_mpc_lib.py
new file mode 100644
index 0000000..f464268
--- /dev/null
+++ b/frc971/control_loops/swerve/casadi_velocity_mpc_lib.py
@@ -0,0 +1,280 @@
+#!/usr/bin/env python3
+
+from frc971.control_loops.swerve import dynamics
+import casadi
+import numpy
+
+
+class MPC(object):
+
+    def __init__(self, solver='fatrop', jit=True):
+        self.fdot = dynamics.swerve_full_dynamics(
+            casadi.SX.sym("X", dynamics.NUM_STATES, 1),
+            casadi.SX.sym("U", 8, 1))
+        self.velocity_fdot = dynamics.velocity_swerve_physics(
+            casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
+            casadi.SX.sym("U", 8, 1))
+
+        self.wrapped_swerve_physics = lambda X, U: numpy.array(self.fdot(X, U))
+
+        self.dt = 0.005
+
+        # TODO(austin): Do we need a disturbance torque per module to account for friction?
+        # Do it only in the observer/post?
+
+        self.force = [
+            dynamics.F(i, casadi.SX.sym("X", 25, 1), casadi.SX.sym("U", 8, 1))
+            for i in range(4)
+        ]
+        self.force_vel = [
+            dynamics.F_vel(i,
+                           casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
+                           casadi.SX.sym("U", 8, 1)) for i in range(4)
+        ]
+        self.slip_angle = [
+            dynamics.slip_angle(i, casadi.SX.sym("X", 25, 1),
+                                casadi.SX.sym("U", 8, 1)) for i in range(4)
+        ]
+        self.rotated_mounting_location = [
+            dynamics.rotated_mounting_location(
+                i, casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
+                casadi.SX.sym("U", 8, 1)) for i in range(4)
+        ]
+        self.torque_cross = self.torque_cross_func(casadi.SX.sym("r", 2, 1),
+                                                   casadi.SX.sym("F", 2, 1))
+        self.force_cross = self.force_cross_func(casadi.SX.sym("Tau", 1, 1),
+                                                 casadi.SX.sym("r", 2, 1))
+        self.next_X = self.make_physics()
+        self.cost = self.make_cost()
+
+        self.N = 200
+
+        # Start with an empty nonlinear program.
+        self.w = []
+        self.lbw = []
+        self.ubw = []
+        J = 0
+        self.g = []
+        self.lbg = []
+        self.ubg = []
+
+        self.X0 = casadi.MX.sym('X0', dynamics.NUM_VELOCITY_STATES)
+
+        # We care about the linear and angular velocities only.
+        self.R = casadi.MX.sym('R', 3)
+
+        # Make Xn and U for each step.  fatrop wants us to interleave the control variables and
+        # states so that it can produce a banded/structured problem which it can solve a lot
+        # faster.
+        Xn_variables = []
+        U_variables = []
+        for i in range(self.N):
+            U_variables.append(casadi.MX.sym(f'U{i}', 8))
+
+            if i == 0:
+                continue
+
+            Xn_variables.append(
+                casadi.MX.sym(f'X{i}', dynamics.NUM_VELOCITY_STATES))
+
+        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))
+
+        # Put U and Xn interleaved into w to go fast.
+        for i in range(self.N):
+            self.w += [U_variables[i]]
+            self.ubw += [100] * 8
+            self.lbw += [-100] * 8
+
+            if i == self.N - 1:
+                continue
+
+            self.w += [Xn_variables[i]]
+            self.ubw += [casadi.inf] * dynamics.NUM_VELOCITY_STATES
+            self.lbw += [-casadi.inf] * dynamics.NUM_VELOCITY_STATES
+
+        self.g += [
+            casadi.reshape(Xn - Xk_end[:, 0:(self.N - 1)],
+                           dynamics.NUM_VELOCITY_STATES * (self.N - 1), 1)
+        ]
+
+        self.lbg += [0] * dynamics.NUM_VELOCITY_STATES * (self.N - 1)
+        self.ubg += [0] * dynamics.NUM_VELOCITY_STATES * (self.N - 1)
+
+        prob = {
+            'f': J,
+            # lbx <= x <= ubx
+            'x': casadi.vertcat(*self.w),
+            # lbg <= g(x, p) <= ubg
+            'g': casadi.vertcat(*self.g),
+            # Input parameters (initial position + goal)
+            'p': casadi.vertcat(self.X0, self.R),
+        }
+
+        compiler = "ccache clang"
+        flags = ["-O3"]
+        jit_options = {
+            "flags": flags,
+            "verbose": True,
+            "compiler": compiler,
+            "temp_suffix": False,
+        }
+
+        if solver == 'fatrop':
+            equality = [
+                True
+                for _ in range(dynamics.NUM_VELOCITY_STATES * (self.N - 1))
+            ]
+            options = {
+                "jit": jit,
+                "jit_cleanup": False,
+                "jit_temp_suffix": False,
+                "compiler": "shell",
+                "jit_options": jit_options,
+                "structure_detection": "auto",
+                "fatrop": {
+                    "tol": 1e-7
+                },
+                "debug": True,
+                "equality": equality,
+            }
+        else:
+            options = {
+                "jit": jit,
+                "jit_cleanup": False,
+                "jit_temp_suffix": False,
+                "compiler": "shell",
+                "jit_options": jit_options,
+            }
+            if FLAGS.full_debug:
+                options["jit"] = False
+                options["ipopt"] = {
+                    "print_level": 12,
+                }
+
+        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)
+
+        X = X0
+        M = 2  # RK4 steps per interval
+        DT = self.dt / M
+
+        for j in range(M):
+            k1 = self.velocity_fdot(X, U)
+            k2 = self.velocity_fdot(X + DT / 2 * k1, U)
+            k3 = self.velocity_fdot(X + DT / 2 * k2, U)
+            k4 = self.velocity_fdot(X + DT * k3, U)
+            X = X + DT / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
+
+        return casadi.Function("next_X", [X0, U], [X])
+
+    def make_cost(self):
+        # TODO(austin): tune cost fn?
+        # Need to do some simulations to see what works well.
+
+        X = casadi.MX.sym('X', dynamics.NUM_VELOCITY_STATES)
+        U = casadi.MX.sym('U', 8)
+        R = casadi.MX.sym('R', 3)
+
+        J = 0
+        vnorm = casadi.sqrt(R[0]**2.0 + R[1]**2.0)
+
+        vnormx = casadi.if_else(vnorm > 0.0001, R[0] / vnorm, 1.0)
+        vnormy = casadi.if_else(vnorm > 0.0001, R[1] / vnorm, 0.0)
+
+        vperpx = casadi.if_else(vnorm > 0.0001, -vnormy, 0.0)
+        vperpy = casadi.if_else(vnorm > 0.0001, vnormx, 1.0)
+
+        # TODO(austin): Do we want to do something more special for 0?
+
+        # cost velocity a lot more in the perpendicular direction to allow tire to spin up
+        # we also only want to get moving in the correct direction as fast as possible
+        J += 75 * ((R[0] - X[dynamics.VELOCITY_STATE_VX]) * vnormx +
+                   (R[1] - X[dynamics.VELOCITY_STATE_VY]) * vnormy)**2.0
+
+        J += 1500 * ((R[0] - X[dynamics.VELOCITY_STATE_VX]) * vperpx +
+                     (R[1] - X[dynamics.VELOCITY_STATE_VY]) * vperpy)**2.0
+
+        J += 1000 * (R[2] - X[dynamics.VELOCITY_STATE_OMEGA])**2.0
+
+        kSteerVelocityGain = 0.10
+        J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS0])**2.0
+        J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS1])**2.0
+        J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS2])**2.0
+        J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS3])**2.0
+
+        # cost variance of the force by a tire and the expected average force and torque on it
+        total_force = self.force_vel[0](X, U)
+        total_torque = self.torque_cross(
+            self.rotated_mounting_location[0](X, U), self.force_vel[0](X, U))
+        for i in range(3):
+            total_force += self.force_vel[i + 1](X, U)
+            total_torque += self.torque_cross(
+                self.rotated_mounting_location[i + 1](X, U),
+                self.force_vel[i + 1](X, U))
+
+        total_force /= 4
+        total_torque /= 4
+        for i in range(4):
+            f_diff = (total_force + self.force_cross(
+                total_torque, self.rotated_mounting_location[i]
+                (X, U))) - self.force_vel[i](X, U)
+            J += 0.01 * (f_diff[0, 0]**2.0 + f_diff[1, 0]**2.0)
+
+        # TODO(austin): Don't penalize torque steering current.
+        for i in range(4):
+            Is = U[2 * i + 0]
+            Id = U[2 * i + 1]
+            # Steer, cost it a lot less than drive to be more agressive in steering
+            J += ((Is + dynamics.STEER_CURRENT_COUPLING_FACTOR * Id)**
+                  2.0) / 100000.0
+            # Drive
+            J += Id * Id / 1000.0
+
+        return casadi.Function("Jn", [X, U, R], [J])
+
+    def torque_cross_func(self, r, F):
+        result = casadi.SX.sym('Tau', 1, 1)
+        result[0, 0] = r[0, 0] * F[1, 0] - r[1, 0] * F[0, 0]
+        return casadi.Function('Tau', [r, F], [result])
+
+    def force_cross_func(self, Tau, r):
+        result = casadi.SX.sym('F', 2, 1)
+        result[0, 0] = -r[1, 0] * Tau[0, 0] / casadi.norm_2(r)**2.0
+        result[1, 0] = r[0, 0] * Tau[0, 0] / casadi.norm_2(r)**2.0
+        return casadi.Function('F', [Tau, r], [result])
+
+    def solve(self, p, seed=None):
+        if seed is None:
+            seed = []
+
+            for i in range(self.N):
+                seed += [0, 0] * 4
+                if i < self.N - 1:
+                    seed += list(p[:dynamics.NUM_VELOCITY_STATES, 0])
+
+        return self.solver(x0=seed,
+                           lbx=self.lbw,
+                           ubx=self.ubw,
+                           lbg=self.lbg,
+                           ubg=self.ubg,
+                           p=casadi.DM(p))
+
+    def unpack_u(self, sol, i):
+        return sol['x'].full().flatten()[
+            (8 + dynamics.NUM_VELOCITY_STATES) *
+            i:((8 + dynamics.NUM_VELOCITY_STATES) * i + 8)]
+
+    def unpack_x(self, sol, i):
+        return sol['x'].full().flatten(
+        )[8 + (8 + dynamics.NUM_VELOCITY_STATES) *
+          (i - 1):(8 + dynamics.NUM_VELOCITY_STATES) * i]
diff --git a/frc971/control_loops/swerve/generate_physics.cc b/frc971/control_loops/swerve/generate_physics.cc
index 7ea3db5..de52d79 100644
--- a/frc971/control_loops/swerve/generate_physics.cc
+++ b/frc971/control_loops/swerve/generate_physics.cc
@@ -70,6 +70,7 @@
 // State per module.
 struct Module {
   DenseMatrix mounting_location;
+  DenseMatrix rotated_mounting_location;
 
   RCP<const Symbol> Is;
 
@@ -938,6 +939,14 @@
         },
         &result_py);
 
+    DefineVector2dVelocityFunction(
+        "rotated_mounting_location",
+        "Returns the mounting location of wheel in field aligned coordinates",
+        [](const Module &m, int dimension) {
+          return ccode(*m.rotated_mounting_location.get(dimension, 0));
+        },
+        &result_py);
+
     DefineScalarFunction(
         "Ms", "Returns the self aligning moment of the ith wheel",
         [this](const Module &m) {
@@ -1189,17 +1198,18 @@
                     DenseMatrix(2, 1, {result.full.Fwx, result.Fwy}),
                     result.full.F);
 
-    DenseMatrix rotated_mounting_location = DenseMatrix(2, 1);
+    result.rotated_mounting_location = DenseMatrix(2, 1);
     mul_dense_dense(R(theta_), result.mounting_location,
-                    rotated_mounting_location);
-    result.full.torque = force_cross(rotated_mounting_location, result.full.F);
+                    result.rotated_mounting_location);
+    result.full.torque =
+        force_cross(result.rotated_mounting_location, result.full.F);
 
     result.direct.F = DenseMatrix(2, 1);
     mul_dense_dense(R(add(theta_, result.thetas)),
                     DenseMatrix(2, 1, {result.direct.Fwx, result.Fwy}),
                     result.direct.F);
     result.direct.torque =
-        force_cross(rotated_mounting_location, result.direct.F);
+        force_cross(result.rotated_mounting_location, result.direct.F);
 
     VLOG(1);
     VLOG(1) << "full torque = " << result.full.torque->__str__();
diff --git a/frc971/control_loops/swerve/jax_dynamics.py b/frc971/control_loops/swerve/jax_dynamics.py
index 6e1305d..19a48de 100644
--- a/frc971/control_loops/swerve/jax_dynamics.py
+++ b/frc971/control_loops/swerve/jax_dynamics.py
@@ -352,3 +352,73 @@
         X[STATE_VY],
         X[STATE_OMEGA],
     ])
+
+
+def mpc_cost(coefficients: CoefficientsType, X, U, goal):
+    J = 0
+
+    rnorm = jax.numpy.linalg.norm(goal[0:2])
+
+    vnorm = jax.lax.select(rnorm > 0.0001, goal[0:2] / rnorm,
+                           jax.numpy.array([1.0, 0.0]))
+    vperp = jax.lax.select(rnorm > 0.0001,
+                           jax.numpy.array([-vnorm[1], vnorm[0]]),
+                           jax.numpy.array([0.0, 1.0]))
+
+    velocity_error = goal[0:2] - X[VELOCITY_STATE_VX:VELOCITY_STATE_VY + 1]
+
+    # TODO(austin): Do we want to do something more special for 0?
+
+    J += 75 * (jax.numpy.dot(velocity_error, vnorm)**2.0)
+    J += 1500 * (jax.numpy.dot(velocity_error, vperp)**2.0)
+    J += 1000 * (goal[2] - X[VELOCITY_STATE_OMEGA])**2.0
+
+    kSteerVelocityGain = 0.10
+    J += kSteerVelocityGain * (X[VELOCITY_STATE_OMEGAS0])**2.0
+    J += kSteerVelocityGain * (X[VELOCITY_STATE_OMEGAS1])**2.0
+    J += kSteerVelocityGain * (X[VELOCITY_STATE_OMEGAS2])**2.0
+    J += kSteerVelocityGain * (X[VELOCITY_STATE_OMEGAS3])**2.0
+
+    mounting_locations = jax.numpy.array(
+        [[coefficients.robot_width / 2.0, coefficients.robot_width / 2.0],
+         [-coefficients.robot_width / 2.0, coefficients.robot_width / 2.0],
+         [-coefficients.robot_width / 2.0, -coefficients.robot_width / 2.0],
+         [coefficients.robot_width / 2.0, -coefficients.robot_width / 2.0]])
+
+    Rtheta = R(X[VELOCITY_STATE_THETA])
+    _, F0, torque0 = velocity_module_physics(coefficients, Rtheta, 0,
+                                             mounting_locations[0], X, U)
+    _, F1, torque1 = velocity_module_physics(coefficients, Rtheta, 1,
+                                             mounting_locations[1], X, U)
+    _, F2, torque2 = velocity_module_physics(coefficients, Rtheta, 2,
+                                             mounting_locations[2], X, U)
+    _, F3, torque3 = velocity_module_physics(coefficients, Rtheta, 3,
+                                             mounting_locations[3], X, U)
+
+    forces = [F0, F1, F2, F3]
+
+    F = (F0 + F1 + F2 + F3)
+    torque = (torque0 + torque1 + torque2 + torque3)
+
+    def force_cross(torque, r):
+        r_squared_norm = jax.numpy.inner(r, r)
+
+        return jax.numpy.array(
+            [-r[1] * torque / r_squared_norm, r[0] * torque / r_squared_norm])
+
+    # TODO(austin): Are these penalties reasonable?  Do they give us a decent time constant?
+    for i in range(4):
+        desired_force = F / 4.0 + force_cross(
+            torque / 4.0, Rtheta @ mounting_locations[i, :])
+        force_error = desired_force - forces[i]
+        J += 0.01 * jax.numpy.inner(force_error, force_error)
+
+    for i in range(4):
+        Is = U[2 * i + 0]
+        Id = U[2 * i + 1]
+        # Steer
+        J += ((Is + STEER_CURRENT_COUPLING_FACTOR * Id)**2.0) / 100000.0
+        # Drive
+        J += (Id**2.0) / 1000.0
+
+    return J
diff --git a/frc971/control_loops/swerve/physics_test.py b/frc971/control_loops/swerve/physics_test.py
index 1da2a31..bc366fa 100644
--- a/frc971/control_loops/swerve/physics_test.py
+++ b/frc971/control_loops/swerve/physics_test.py
@@ -18,6 +18,7 @@
 from frc971.control_loops.swerve import dynamics
 from frc971.control_loops.swerve import nocaster_dynamics
 from frc971.control_loops.swerve import physics_test_utils as utils
+from frc971.control_loops.swerve import casadi_velocity_mpc_lib
 from frc971.control_loops.swerve import jax_dynamics
 from frc971.control_loops.swerve.cpp_dynamics import swerve_dynamics as cpp_dynamics
 
@@ -707,8 +708,56 @@
         Xdot_rot = self.swerve_full_dynamics(X_rot, steer_I, skip_compare=True)
 
         self.assertGreater(Xdot[dynamics.STATE_OMEGA, 0], 0.0)
-        self.assertAlmostEquals(Xdot[dynamics.STATE_OMEGA, 0],
-                                Xdot_rot[dynamics.STATE_OMEGA, 0])
+        self.assertAlmostEqual(Xdot[dynamics.STATE_OMEGA, 0],
+                               Xdot_rot[dynamics.STATE_OMEGA, 0])
+
+    def test_cost_equality(self):
+        """Tests that the casadi and jax cost functions match."""
+        mpc = casadi_velocity_mpc_lib.MPC(jit=False)
+        cost = mpc.make_cost()
+
+        for i in range(10):
+            X = numpy.random.uniform(size=(dynamics.NUM_VELOCITY_STATES, ))
+            U = numpy.random.uniform(low=-10, high=10, size=(8, ))
+            R = numpy.random.uniform(low=-1, high=1, size=(3, ))
+
+            J = numpy.array(cost(X, U, R))[0, 0]
+            jax_J = jax_dynamics.mpc_cost(self.coefficients, X, U, R)
+
+            self.assertAlmostEqual(J, jax_J)
+
+        R = jax.numpy.array([0.0, 0.0, 1.0])
+
+        # Now try spinning in place and make sure the cost doesn't change.
+        # This tells us if we got our rotations right.
+        steer_I = numpy.array([(i % 2) * 20 for i in range(8)])
+
+        X = utils.state_vector(velocity=numpy.array([[0.0], [0.0]]),
+                               omega=0.0,
+                               module_angles=[
+                                   3 * numpy.pi / 4.0, -3 * numpy.pi / 4.0,
+                                   -numpy.pi / 4.0, numpy.pi / 4.0
+                               ],
+                               drive_wheel_velocity=1.0)
+
+        jax_J_orig = jax_dynamics.mpc_cost(self.coefficients,
+                                           self.to_velocity_state(X)[:, 0],
+                                           steer_I, R)
+
+        X_rotated = utils.state_vector(velocity=numpy.array([[0.0], [0.0]]),
+                                       omega=0.0,
+                                       theta=numpy.pi,
+                                       module_angles=[
+                                           3 * numpy.pi / 4.0,
+                                           -3 * numpy.pi / 4.0,
+                                           -numpy.pi / 4.0, numpy.pi / 4.0
+                                       ],
+                                       drive_wheel_velocity=1.0)
+        jax_J_rotated = jax_dynamics.mpc_cost(
+            self.coefficients,
+            self.to_velocity_state(X_rotated)[:, 0], steer_I, R)
+
+        self.assertAlmostEqual(jax_J_orig, jax_J_rotated)
 
     def test_cpp_consistency(self):
         """Tests that the C++ physics are consistent with the Python physics."""
diff --git a/frc971/control_loops/swerve/velocity_controller/physics.py b/frc971/control_loops/swerve/velocity_controller/physics.py
index d6515fe..b0da0b4 100644
--- a/frc971/control_loops/swerve/velocity_controller/physics.py
+++ b/frc971/control_loops/swerve/velocity_controller/physics.py
@@ -1,7 +1,7 @@
 import jax, numpy
 from functools import partial
 from absl import logging
-from frc971.control_loops.swerve import dynamics, jax_dynamics
+from frc971.control_loops.swerve import jax_dynamics
 from frc971.control_loops.python import controls
 from flax.typing import PRNGKey
 
diff --git a/frc971/control_loops/swerve/velocity_controller/train.py b/frc971/control_loops/swerve/velocity_controller/train.py
index f5ba092..b1ad9a8 100644
--- a/frc971/control_loops/swerve/velocity_controller/train.py
+++ b/frc971/control_loops/swerve/velocity_controller/train.py
@@ -20,7 +20,6 @@
 from functools import partial
 import flashbax
 from jax.experimental.ode import odeint
-from frc971.control_loops.swerve import dynamics
 import orbax.checkpoint
 from frc971.control_loops.swerve.velocity_controller.model import *
 from frc971.control_loops.swerve.velocity_controller.physics import *