diff --git a/WORKSPACE b/WORKSPACE
index ad59314..6fd635e 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -702,14 +702,16 @@
 cc_library(
     name = 'api-cpp',
     visibility = ['//visibility:public'],
-    hdrs = glob(['ctre/phoenix6/**/*.hpp']),
+    hdrs = glob(['ctre/phoenix6/**/*.hpp', 'ctre/unit/**/*.h']),
     includes = ["."],
-    deps = ["@//third_party/allwpilib/wpimath"],
+    deps = ["@//third_party/allwpilib/wpimath",
+            "@ctre_phoenix6_tools_headers//:tools",
+            ],
 )
 """,
-    sha256 = "1b9295ac868d619d0263f285b372a22297798e45c8fdeb83ca99154a097c5c98",
+    sha256 = "67fdd9a4bf275c69666bbc8bf38312eea8a85bf9fda3901d02af3a18136ffb3e",
     urls = [
-        "https://maven.ctr-electronics.com/release/com/ctre/phoenix6/api-cpp/24.1.0/api-cpp-24.1.0-headers.zip",
+        "https://maven.ctr-electronics.com/release/com/ctre/phoenix6/api-cpp/24.50.0-alpha-2/api-cpp-24.50.0-alpha-2-headers.zip",
     ],
 )
 
@@ -731,9 +733,9 @@
     target_compatible_with = ['@//tools/platforms/hardware:roborio'],
 )
 """,
-    sha256 = "51f4921f8995e3e80ba347a90f6fa5c0cb7d0e438923a1a736554f263e2d5b8e",
+    sha256 = "00da14f437cfeb2c344674dee59fe70477a3a0d612eb93a1441188dc94a5136f",
     urls = [
-        "https://maven.ctr-electronics.com/release/com/ctre/phoenix6/api-cpp/24.1.0/api-cpp-24.1.0-linuxathena.zip",
+        "https://maven.ctr-electronics.com/release/com/ctre/phoenix6/api-cpp/24.50.0-alpha-2/api-cpp-24.50.0-alpha-2-linuxathena.zip",
     ],
 )
 
@@ -743,12 +745,12 @@
 cc_library(
     name = 'tools',
     visibility = ['//visibility:public'],
-    hdrs = glob(['ctre/**/*.h', 'ctre/**/*.hpp']),
+    hdrs = glob(['ctre/**/*.h', 'ctre/phoenix/**/*.hpp', 'ctre/phoenix6/**/*.hpp']),
 )
 """,
-    sha256 = "72be3e50205e2546736361e3aba9c0de5d5318b263c195600f9b558c3d92cb9e",
+    sha256 = "77624291ec19a03c9e068347ef800e435782444722793d56c9e39f6108da33a8",
     urls = [
-        "https://maven.ctr-electronics.com/release/com/ctre/phoenix6/tools/24.1.0/tools-24.1.0-headers.zip",
+        "https://maven.ctr-electronics.com/release/com/ctre/phoenix6/tools/24.50.0-alpha-2/tools-24.50.0-alpha-2-headers.zip",
     ],
 )
 
@@ -770,9 +772,9 @@
     target_compatible_with = ['@//tools/platforms/hardware:roborio'],
 )
 """,
-    sha256 = "76d30cb8c0988eb6accab10e77ceae492782edb64c06f7df628b8c161cf0220a",
+    sha256 = "217bcd6aecb71224fd32725f857da58e61a6ea451ca07f85fb55e49f83dbf113",
     urls = [
-        "https://maven.ctr-electronics.com/release/com/ctre/phoenix6/tools/24.1.0/tools-24.1.0-linuxathena.zip",
+        "https://maven.ctr-electronics.com/release/com/ctre/phoenix6/tools/24.50.0-alpha-2/tools-24.50.0-alpha-2-linuxathena.zip",
     ],
 )
 
@@ -782,12 +784,13 @@
 cc_library(
     name = 'api-cpp',
     visibility = ['//visibility:public'],
-    hdrs = glob(['ctre/phoenix/**/*.h']),
+    hdrs = glob(['ctre/phoenix/**/*.h', 'ctre/unit/**/*.h']),
+    includes = ["."]
 )
 """,
-    sha256 = "d2790e2495a59a3b14ab4b0fd03c2e25292d6874168f5642165acb0d1bab2f99",
+    sha256 = "4fba20441ea1d61f8487897682c723a48ee2c2741946eab4062b4248434c0afc",
     urls = [
-        "https://maven.ctr-electronics.com/release/com/ctre/phoenix/api-cpp/5.33.0/api-cpp-5.33.0-headers.zip",
+        "https://maven.ctr-electronics.com/release/com/ctre/phoenix/api-cpp/5.34.0-alpha-1/api-cpp-5.34.0-alpha-1-headers.zip",
     ],
 )
 
@@ -824,9 +827,9 @@
     hdrs = glob(['ctre/phoenix/**/*.h']),
 )
 """,
-    sha256 = "56cb8272d623721560eea360730d2508f4f365e2ac1060aec5fbd546cffc0771",
+    sha256 = "cd39f32341037a7ec1074710044b332db6ca607dfd548b5f951c75f7a8506ab5",
     urls = [
-        "https://maven.ctr-electronics.com/release/com/ctre/phoenix/cci/5.33.0/cci-5.33.0-headers.zip",
+        "https://maven.ctr-electronics.com/release/com/ctre/phoenix/cci/5.34.0-alpha-1/cci-5.34.0-alpha-1-headers.zip",
     ],
 )
 
@@ -867,6 +870,7 @@
     target_compatible_with = ['@platforms//cpu:arm64'],
 )
 
+
 # TODO(max): Use cc_import once they add a defines property.
 # See: https://github.com/bazelbuild/bazel/issues/19753
 cc_library(
@@ -891,9 +895,9 @@
     ],
 )
 """,
-    sha256 = "635b19f019d1283749fb23a95ad9e13ddd9d5013cff4c838303d898e7d9556bd",
+    sha256 = "0f1312f39eacc490fb253198c2d0e61e48ae00eff6a87cfd362358b1ad36a930",
     urls = [
-        "https://software.frc971.org/Build-Dependencies/phoenix6_24.2.0_5.4.2024.tar.gz",
+        "https://software.frc971.org/Build-Dependencies/phoenix6_24.50.0-alpha-2_arm64-2024.10.26.tar.gz",
     ],
 )
 
diff --git a/aos/containers/error_list.h b/aos/containers/error_list.h
index 7ceb1fe..762d2ad 100644
--- a/aos/containers/error_list.h
+++ b/aos/containers/error_list.h
@@ -5,6 +5,7 @@
 
 #include <algorithm>
 
+#include "absl/log/check.h"
 #include "flatbuffers/buffer.h"
 #include "flatbuffers/flatbuffer_builder.h"
 #include "flatbuffers/vector.h"
@@ -128,7 +129,11 @@
       flatbuffers::FlatBufferBuilder *fbb) const {
     return fbb->CreateVector(array_.data(), array_.size());
   }
-};  // namespace aos
+  template <typename StaticBuilder>
+  void ToStaticFlatbuffer(StaticBuilder *vector_builder) const {
+    CHECK(vector_builder->FromData(array_.data(), array_.size()));
+  }
+};
 
 }  // namespace aos
 
diff --git a/aos/events/logging/config_remapper.cc b/aos/events/logging/config_remapper.cc
index 763cbc3..edc2326 100644
--- a/aos/events/logging/config_remapper.cc
+++ b/aos/events/logging/config_remapper.cc
@@ -566,7 +566,9 @@
 
     // Add the schema if it doesn't exist.
     if (schema_map.find(c->type()->string_view()) == schema_map.end()) {
-      CHECK(c->has_schema());
+      if (!c->has_schema()) {
+        LOG(FATAL) << "Could not find schema for " << c->type()->string_view();
+      }
       schema_map.insert(std::make_pair(c->type()->string_view(),
                                        RecursiveCopyFlatBuffer(c->schema())));
     }
diff --git a/frc971/control_loops/dlqr.h b/frc971/control_loops/dlqr.h
index 57aa88b..d31f492 100644
--- a/frc971/control_loops/dlqr.h
+++ b/frc971/control_loops/dlqr.h
@@ -41,7 +41,9 @@
          ::Eigen::Matrix<double, kM, kN> *K,
          ::Eigen::Matrix<double, kN, kN> *S) {
   *K = ::Eigen::Matrix<double, kM, kN>::Zero();
-  *S = ::Eigen::Matrix<double, kN, kN>::Zero();
+  if (S != nullptr) {
+    *S = ::Eigen::Matrix<double, kN, kN>::Zero();
+  }
   // Discrete (not continuous)
   char DICO = 'D';
   // B and R are provided instead of G.
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 2c3ad50..6be2604 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -14,8 +14,17 @@
 )
 
 static_flatbuffer(
+    name = "swerve_drivetrain_joystick_goal_fbs",
+    srcs = ["swerve_drivetrain_joystick_goal.fbs"],
+)
+
+static_flatbuffer(
     name = "swerve_drivetrain_goal_fbs",
     srcs = ["swerve_drivetrain_goal.fbs"],
+    deps = [
+        ":swerve_drivetrain_joystick_goal_fbs",
+        "//frc971/math:matrix_fbs",
+    ],
 )
 
 static_flatbuffer(
@@ -292,6 +301,33 @@
     ],
 )
 
+py_binary(
+    name = "experience_collector",
+    srcs = [
+        "experience_collector.py",
+    ],
+    deps = [
+        ":casadi_velocity_mpc_lib",
+        ":jax_dynamics",
+        "//frc971/control_loops/swerve/velocity_controller:physics",
+        "@pip//absl_py",
+        "@pip//matplotlib",
+        "@pip//numpy",
+        "@pip//pygobject",
+        "@pip//scipy",
+        "@pip//tensorflow",
+    ],
+)
+
+py_binary(
+    name = "multi_experience_collector",
+    srcs = ["multi_experience_collector.py"],
+    data = [":experience_collector"],
+    deps = [
+        "@pip//absl_py",
+    ],
+)
+
 py_library(
     name = "physics_test_utils",
     srcs = [
@@ -371,3 +407,92 @@
         "@com_google_absl//absl/log:check",
     ],
 )
+
+cc_library(
+    name = "linearization_utils",
+    hdrs = ["linearization_utils.h"],
+)
+
+cc_library(
+    name = "linearized_controller",
+    hdrs = ["linearized_controller.h"],
+    deps = [
+        ":eigen_dynamics",
+        ":linearization_utils",
+        "//frc971/control_loops:c2d",
+        "//frc971/control_loops:dlqr",
+        "//frc971/control_loops:jacobian",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+    ],
+)
+
+cc_test(
+    name = "linearized_controller_test",
+    srcs = ["linearized_controller_test.cc"],
+    deps = [
+        ":linearized_controller",
+        "//aos/testing:googletest",
+    ],
+)
+
+cc_library(
+    name = "auto_diff_jacobian",
+    hdrs = ["auto_diff_jacobian.h"],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        "@com_google_ceres_solver//:ceres",
+    ],
+)
+
+cc_test(
+    name = "auto_diff_jacobian_test",
+    srcs = ["auto_diff_jacobian_test.cc"],
+    deps = [
+        ":auto_diff_jacobian",
+        "//aos/testing:googletest",
+    ],
+)
+
+cc_library(
+    name = "simplified_dynamics",
+    hdrs = ["simplified_dynamics.h"],
+    deps = [
+        ":auto_diff_jacobian",
+        ":eigen_dynamics",
+        ":motors",
+        "//aos/util:math",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+    ],
+)
+
+cc_test(
+    name = "simplified_dynamics_test",
+    srcs = ["simplified_dynamics_test.cc"],
+    deps = [
+        ":simplified_dynamics",
+        "//aos/testing:googletest",
+        "//aos/time",
+        "//frc971/control_loops:jacobian",
+        "@com_google_absl//absl/log",
+    ],
+)
+
+cc_library(
+    name = "inverse_kinematics",
+    hdrs = ["inverse_kinematics.h"],
+    deps = [
+        ":simplified_dynamics",
+        "//aos/util:math",
+    ],
+)
+
+cc_test(
+    name = "inverse_kinematics_test",
+    srcs = ["inverse_kinematics_test.cc"],
+    deps = [
+        ":inverse_kinematics",
+        "//aos/testing:googletest",
+    ],
+)
diff --git a/frc971/control_loops/swerve/auto_diff_jacobian.h b/frc971/control_loops/swerve/auto_diff_jacobian.h
new file mode 100644
index 0000000..2e7947a
--- /dev/null
+++ b/frc971/control_loops/swerve/auto_diff_jacobian.h
@@ -0,0 +1,62 @@
+#ifndef FRC971_CONTROL_LOOPS_SWERVE_AUTO_DIFF_JACOBIAN_H_
+#define FRC971_CONTROL_LOOPS_SWERVE_AUTO_DIFF_JACOBIAN_H_
+#include "include/ceres/tiny_solver.h"
+#include "include/ceres/tiny_solver_autodiff_function.h"
+
+namespace frc971::control_loops::swerve {
+// Class to conveniently scope a function that makes use of Ceres'
+// autodifferentiation methods to calculate the jacobian of the provided method.
+// Template parameters:
+// Scalar: scalar type to use (typically double or float; this is used for
+//    allowing you to control what precision you use; you cannot use this
+//    with ceres Jets because this has to use Jets internally).
+// Function: The type of the function itself. A Function f must be callable as
+//    Eigen::Matrix<Scalar, kNumOutputs, 1> = f(Eigen::Matrix<Scalar,
+//    kNumInputs, 1>{});
+template <typename Scalar, typename Function, size_t kNumInputs,
+          size_t kNumOutputs>
+class AutoDiffJacobian {
+ public:
+  // Calculates the jacobian of the provided method, function, at the provided
+  // input X.
+  static Eigen::Matrix<Scalar, kNumOutputs, kNumInputs> Jacobian(
+      const Function &function, const Eigen::Matrix<Scalar, kNumInputs, 1> &X) {
+    AutoDiffCeresFunctor ceres_functor(function);
+    TinySolverFunctor tiny_solver(ceres_functor);
+    // residual is unused, it's just a place to store the evaluated function at
+    // the current state/input.
+    Eigen::Matrix<Scalar, kNumOutputs, 1> residual;
+    Eigen::Matrix<Scalar, kNumOutputs, kNumInputs> jacobian;
+    tiny_solver(X.data(), residual.data(), jacobian.data());
+    return jacobian;
+  }
+
+ private:
+  // Borrow the TinySolver's auto-differentiation execution for use here to
+  // calculate the linearized dynamics. We aren't actually doing any solving,
+  // just letting it do the jacobian calculation for us.
+  // As such, construct a "residual" function whose residuals are just the
+  // derivative of the state and whose parameters are the stacked state + input.
+  class AutoDiffCeresFunctor {
+   public:
+    AutoDiffCeresFunctor(const Function &function) : function_(function) {}
+    template <typename ScalarT>
+    bool operator()(const ScalarT *const parameters,
+                    ScalarT *const residuals) const {
+      const Eigen::Map<const Eigen::Matrix<ScalarT, kNumInputs, 1>>
+          eigen_parameters(parameters);
+      Eigen::Map<Eigen::Matrix<ScalarT, kNumOutputs, 1>> eigen_residuals(
+          residuals);
+      eigen_residuals = function_(eigen_parameters);
+      return true;
+    }
+
+   private:
+    const Function &function_;
+  };
+  typedef ceres::TinySolverAutoDiffFunction<AutoDiffCeresFunctor, kNumOutputs,
+                                            kNumInputs, Scalar>
+      TinySolverFunctor;
+};
+}  // namespace frc971::control_loops::swerve
+#endif  // FRC971_CONTROL_LOOPS_SWERVE_AUTO_DIFF_JACOBIAN_H_
diff --git a/frc971/control_loops/swerve/auto_diff_jacobian_test.cc b/frc971/control_loops/swerve/auto_diff_jacobian_test.cc
new file mode 100644
index 0000000..d04726c
--- /dev/null
+++ b/frc971/control_loops/swerve/auto_diff_jacobian_test.cc
@@ -0,0 +1,23 @@
+#include "frc971/control_loops/swerve/auto_diff_jacobian.h"
+
+#include <functional>
+
+#include "absl/log/log.h"
+#include "gtest/gtest.h"
+
+namespace frc971::control_loops::swerve::testing {
+struct TestFunction {
+  template <typename Scalar>
+  Eigen::Matrix<Scalar, 3, 1> operator()(
+      const Eigen::Map<const Eigen::Matrix<Scalar, 2, 1>> X) const {
+    return Eigen::Matrix<double, 3, 2>{{1, 2}, {3, 4}, {5, 6}} * X;
+  }
+};
+
+TEST(AutoDiffJacobianTest, EvaluatesJacobian) {
+  EXPECT_EQ((AutoDiffJacobian<double, TestFunction, 2, 3>::Jacobian(
+                TestFunction{}, Eigen::Vector2d::Zero())),
+            (Eigen::Matrix<double, 3, 2>{{1, 2}, {3, 4}, {5, 6}}));
+}
+
+}  // namespace frc971::control_loops::swerve::testing
diff --git a/frc971/control_loops/swerve/casadi_velocity_mpc.py b/frc971/control_loops/swerve/casadi_velocity_mpc.py
index d62f7ed..54e33f1 100644
--- a/frc971/control_loops/swerve/casadi_velocity_mpc.py
+++ b/frc971/control_loops/swerve/casadi_velocity_mpc.py
@@ -24,11 +24,6 @@
 flags.DEFINE_bool('pickle', False, 'Write optimization results.')
 flags.DEFINE_string('outputdir', None, 'Directory to write problem results to')
 
-# 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.
-flags.DEFINE_bool('full_debug', False,
-                  'If true, turn on all the debugging in the solver.')
-
 
 class Solver(object):
 
diff --git a/frc971/control_loops/swerve/casadi_velocity_mpc_lib.py b/frc971/control_loops/swerve/casadi_velocity_mpc_lib.py
index f464268..e358422 100644
--- a/frc971/control_loops/swerve/casadi_velocity_mpc_lib.py
+++ b/frc971/control_loops/swerve/casadi_velocity_mpc_lib.py
@@ -3,11 +3,19 @@
 from frc971.control_loops.swerve import dynamics
 import casadi
 import numpy
+from absl import flags
+
+FLAGS = flags.FLAGS
+
+# 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.
+flags.DEFINE_bool('full_debug', False,
+                  'If true, turn on all the debugging in the solver.')
 
 
 class MPC(object):
 
-    def __init__(self, solver='fatrop', jit=True):
+    def __init__(self, solver='fatrop', jit=True, N=200):
         self.fdot = dynamics.swerve_full_dynamics(
             casadi.SX.sym("X", dynamics.NUM_STATES, 1),
             casadi.SX.sym("U", 8, 1))
@@ -47,7 +55,7 @@
         self.next_X = self.make_physics()
         self.cost = self.make_cost()
 
-        self.N = 200
+        self.N = N
 
         # Start with an empty nonlinear program.
         self.w = []
diff --git a/frc971/control_loops/swerve/experience_collector.py b/frc971/control_loops/swerve/experience_collector.py
new file mode 100644
index 0000000..b4c17fb
--- /dev/null
+++ b/frc971/control_loops/swerve/experience_collector.py
@@ -0,0 +1,201 @@
+#!/usr/bin/env python3
+import os, sys
+
+# Setup XLA first.
+os.environ['XLA_FLAGS'] = ' '.join([
+    # Teach it where to find CUDA
+    '--xla_gpu_cuda_data_dir=/usr/lib/cuda',
+    # Use up to 20 cores
+    #'--xla_force_host_platform_device_count=6',
+    # Dump XLA to /tmp/foo to aid debugging
+    #'--xla_dump_to=/tmp/foo',
+    #'--xla_gpu_enable_command_buffer='
+    # Dump sharding
+    #"--xla_dump_to=/tmp/foo",
+    #"--xla_dump_hlo_pass_re=spmd|propagation"
+])
+os.environ['JAX_PLATFORMS'] = 'cpu'
+os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
+os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '.50'
+
+from absl import flags
+from absl import app
+from absl import logging
+import pickle
+import numpy
+from frc971.control_loops.swerve import dynamics
+from frc971.control_loops.swerve.casadi_velocity_mpc_lib import MPC
+import jax
+import tensorflow as tf
+from frc971.control_loops.swerve.velocity_controller.physics import SwerveProblem
+from frc971.control_loops.swerve import jax_dynamics
+
+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_integer('seed', 0, 'Seed for random initial state.')
+
+flags.DEFINE_bool('save_plots', True,
+                  'If true, save plots for each run as well.')
+flags.DEFINE_string('outputdir', None, 'Directory to write problem results to')
+flags.DEFINE_bool('quiet', False, 'If true, print a lot less')
+
+flags.DEFINE_integer('num_solutions', 100,
+                     'Number of random problems to solve.')
+flags.DEFINE_integer('horizon', 200, 'Horizon to solve for')
+
+try:
+    from matplotlib import pylab
+except ModuleNotFoundError:
+    pass
+
+
+def collect_experience(problem, mpc, rng):
+    X_initial = numpy.array(problem.random_states(rng,
+                                                  dimensions=1)).transpose()
+
+    R_goal = numpy.zeros((3, 1))
+    R_goal[0, 0] = FLAGS.vx
+    R_goal[1, 0] = FLAGS.vy
+    R_goal[2, 0] = FLAGS.omega
+
+    solution = mpc.solve(p=numpy.vstack((X_initial, R_goal)))
+    sys.stderr.flush()
+    sys.stdout.flush()
+
+    # Solver doesn't solve for the last state.  So we get N-1 states back.
+    experience = {
+        'observations1': numpy.zeros((mpc.N - 1, problem.num_states)),
+        'observations2': numpy.zeros((mpc.N - 1, problem.num_states)),
+        'actions': numpy.zeros((mpc.N - 1, problem.num_outputs)),
+        'rewards': numpy.zeros((mpc.N - 1, 1)),
+        'goals': numpy.zeros((mpc.N - 1, problem.num_goals)),
+    }
+
+    if not FLAGS.quiet:
+        print('x(0):', X_initial.transpose())
+
+    logging.info('Finished solving')
+    X_prior = X_initial.squeeze()
+    for j in range(mpc.N - 1):
+        if not FLAGS.quiet:
+            print(f'u({j}): ', mpc.unpack_u(solution, j))
+            print(f'x({j+1}): ', mpc.unpack_x(solution, j + 1))
+        experience['observations1'][j, :] = X_prior
+        X_prior = mpc.unpack_x(solution, j + 1)
+        experience['observations2'][j, :] = X_prior
+        experience['actions'][j, :] = mpc.unpack_u(solution, j)
+        experience['goals'][j, :] = R_goal[:, 0]
+
+    logging.info('Finished all but reward')
+    for j in range(mpc.N - 1):
+        experience['rewards'][j, :] = problem.reward(
+            X=X_prior,
+            U=mpc.unpack_u(solution, j),
+            goal=R_goal[:, 0],
+        )
+    sys.stderr.flush()
+    sys.stdout.flush()
+
+    return experience
+
+
+def save_experience(problem, mpc, experience, experience_number):
+    with open(f'experience_{experience_number}.pkl', 'wb') as f:
+        pickle.dump(experience, f)
+
+    if not FLAGS.save_plots:
+        return
+
+    fig0, axs0 = pylab.subplots(3)
+    fig1, axs1 = pylab.subplots(2)
+
+    axs0[0].clear()
+    axs0[1].clear()
+    axs0[2].clear()
+
+    t = problem.dt * numpy.array(list(range(mpc.N - 1)))
+
+    X_plot = experience['observations1']
+    U_plot = experience['actions']
+
+    axs0[0].plot(t, X_plot[:, dynamics.VELOCITY_STATE_VX], label="vx")
+    axs0[0].plot(t, X_plot[:, dynamics.VELOCITY_STATE_VY], label="vy")
+    axs0[0].plot(t, X_plot[:, dynamics.VELOCITY_STATE_OMEGA], label="omega")
+    axs0[0].legend()
+
+    axs0[1].plot(t, U_plot[:, 0], label="Is0")
+    axs0[1].plot(t, U_plot[:, 2], label="Is1")
+    axs0[1].plot(t, U_plot[:, 4], label="Is2")
+    axs0[1].plot(t, U_plot[:, 6], label="Is3")
+    axs0[1].legend()
+
+    axs0[2].plot(t, U_plot[:, 1], label="Id0")
+    axs0[2].plot(t, U_plot[:, 3], label="Id1")
+    axs0[2].plot(t, U_plot[:, 5], label="Id2")
+    axs0[2].plot(t, U_plot[:, 7], label="Id3")
+    axs0[2].legend()
+
+    axs1[0].clear()
+    axs1[1].clear()
+
+    axs1[0].plot(t, X_plot[:, dynamics.VELOCITY_STATE_THETAS0], label='steer0')
+    axs1[0].plot(t, X_plot[:, dynamics.VELOCITY_STATE_THETAS1], label='steer1')
+    axs1[0].plot(t, X_plot[:, dynamics.VELOCITY_STATE_THETAS2], label='steer2')
+    axs1[0].plot(t, X_plot[:, dynamics.VELOCITY_STATE_THETAS3], label='steer3')
+    axs1[0].legend()
+    axs1[1].plot(t,
+                 X_plot[:, dynamics.VELOCITY_STATE_OMEGAS0],
+                 label='steer_velocity0')
+    axs1[1].plot(t,
+                 X_plot[:, dynamics.VELOCITY_STATE_OMEGAS1],
+                 label='steer_velocity1')
+    axs1[1].plot(t,
+                 X_plot[:, dynamics.VELOCITY_STATE_OMEGAS2],
+                 label='steer_velocity2')
+    axs1[1].plot(t,
+                 X_plot[:, dynamics.VELOCITY_STATE_OMEGAS3],
+                 label='steer_velocity3')
+    axs1[1].legend()
+
+    fig0.savefig(f'state_{experience_number}.svg')
+    fig1.savefig(f'steer_{experience_number}.svg')
+
+    # Free the memory associated with the figures.
+    fig0.clf()
+    fig1.clf()
+    pylab.close(fig0)
+    pylab.close(fig1)
+
+
+def main(argv):
+    if FLAGS.outputdir:
+        os.chdir(FLAGS.outputdir)
+
+    # Hide any GPUs from TensorFlow. Otherwise it might reserve memory.
+    tf.config.experimental.set_visible_devices([], 'GPU')
+    rng = jax.random.key(FLAGS.seed)
+
+    physics_constants = jax_dynamics.Coefficients()
+    problem = SwerveProblem(physics_constants)
+    mpc = MPC(solver='ipopt', N=(FLAGS.horizon + 1))
+
+    if FLAGS.compileonly:
+        return
+
+    for i in range(FLAGS.num_solutions):
+        rng, rng_init = jax.random.split(rng)
+        experience = collect_experience(problem, mpc, rng_init)
+        logging.info('Solved problem %d', i)
+
+        save_experience(problem, mpc, experience, i)
+        logging.info('Wrote problem %d', i)
+
+
+if __name__ == '__main__':
+    app.run(main)
diff --git a/frc971/control_loops/swerve/generate_physics.cc b/frc971/control_loops/swerve/generate_physics.cc
index de52d79..8a7fe16 100644
--- a/frc971/control_loops/swerve/generate_physics.cc
+++ b/frc971/control_loops/swerve/generate_physics.cc
@@ -339,7 +339,7 @@
     result_h.emplace_back(
         "inline constexpr size_t kNumVelocityStates = "
         "static_cast<size_t>(VelocityStates::kNumStates);");
-    result_h.emplace_back("struct Inputs {");
+    result_h.emplace_back("struct InputStates {");
     result_h.emplace_back("enum States {");
     result_h.emplace_back("  kIs0 = 0,");
     result_h.emplace_back("  kId0 = 1,");
@@ -354,7 +354,7 @@
     result_h.emplace_back("};");
     result_h.emplace_back(
         "inline constexpr size_t kNumInputs = "
-        "static_cast<size_t>(Inputs::kNumInputs);");
+        "static_cast<size_t>(InputStates::kNumInputs);");
     result_h.emplace_back("");
     result_h.emplace_back("// Returns the derivative of our state vector");
     result_h.emplace_back(
@@ -365,29 +365,6 @@
     result_h.emplace_back(
         "    Eigen::Ref<const Eigen::Matrix<double, kNumInputs, 1>> U);");
     result_h.emplace_back("");
-    result_h.emplace_back(
-        "Eigen::Matrix<double, kNumVelocityStates, 1> ToVelocityState(");
-    result_h.emplace_back(
-        "    Eigen::Ref<const Eigen::Matrix<double, kNumFullDynamicsStates, "
-        "1>> X);");
-    result_h.emplace_back("");
-    result_h.emplace_back(
-        "Eigen::Matrix<double, kNumFullDynamicsStates, 1> FromVelocityState(");
-    result_h.emplace_back(
-        "    Eigen::Ref<const Eigen::Matrix<double, kNumVelocityStates, 1>> "
-        "X);");
-    result_h.emplace_back("");
-    result_h.emplace_back(
-        "inline Eigen::Matrix<double, kNumVelocityStates, 1> VelocityPhysics(");
-    result_h.emplace_back(
-        "    Eigen::Ref<const Eigen::Matrix<double, kNumVelocityStates, 1>> "
-        "X,");
-    result_h.emplace_back(
-        "    Eigen::Ref<const Eigen::Matrix<double, kNumInputs, 1>> U) {");
-    result_h.emplace_back(
-        "  return ToVelocityState(SwervePhysics(FromVelocityState(X), U));");
-    result_h.emplace_back("}");
-    result_h.emplace_back("");
     result_h.emplace_back("}  // namespace frc971::control_loops::swerve");
     result_h.emplace_back("");
     result_h.emplace_back(absl::Substitute("#endif  // $0_", include_guard));
@@ -401,40 +378,6 @@
     result_cc.emplace_back("namespace frc971::control_loops::swerve {");
     result_cc.emplace_back("");
     result_cc.emplace_back(
-        "Eigen::Matrix<double, kNumVelocityStates, 1> ToVelocityState(");
-    result_cc.emplace_back(
-        "    Eigen::Ref<const Eigen::Matrix<double, kNumFullDynamicsStates, "
-        "1>> X) {");
-    result_cc.emplace_back(
-        "    Eigen::Matrix<double, kNumVelocityStates, 1> velocity;");
-    const std::vector<std::string_view> velocity_states = {
-        "kThetas0", "kOmegas0", "kThetas1", "kOmegas1", "kThetas2", "kOmegas2",
-        "kThetas3", "kOmegas3", "kTheta",   "kVx",      "kVy",      "kOmega"};
-    for (const std::string_view velocity_state : velocity_states) {
-      result_cc.emplace_back(absl::StrFormat(
-          "  velocity(VelocityStates::%s) = X(FullDynamicsStates::%s);",
-          velocity_state, velocity_state));
-    }
-    result_cc.emplace_back("  return velocity;");
-    result_cc.emplace_back("}");
-    result_cc.emplace_back("");
-    result_cc.emplace_back(
-        "Eigen::Matrix<double, kNumFullDynamicsStates, 1> FromVelocityState(");
-    result_cc.emplace_back(
-        "    Eigen::Ref<const Eigen::Matrix<double, kNumVelocityStates, 1>> X) "
-        "{");
-    result_cc.emplace_back(
-        "    Eigen::Matrix<double, kNumFullDynamicsStates, 1> full;");
-    result_cc.emplace_back("    full.setZero();");
-    for (const std::string_view velocity_state : velocity_states) {
-      result_cc.emplace_back(absl::StrFormat(
-          "  full(FullDynamicsStates::%s) = X(VelocityStates::%s);",
-          velocity_state, velocity_state));
-    }
-    result_cc.emplace_back("  return full;");
-    result_cc.emplace_back("}");
-    result_cc.emplace_back("");
-    result_cc.emplace_back(
         "Eigen::Matrix<double, kNumFullDynamicsStates, 1> SwervePhysics(");
     result_cc.emplace_back(
         "    Eigen::Ref<const Eigen::Matrix<double, kNumFullDynamicsStates, "
diff --git a/frc971/control_loops/swerve/inverse_kinematics.h b/frc971/control_loops/swerve/inverse_kinematics.h
new file mode 100644
index 0000000..a4e43e1
--- /dev/null
+++ b/frc971/control_loops/swerve/inverse_kinematics.h
@@ -0,0 +1,86 @@
+#ifndef FRC971_CONTROL_LOOPS_SWERVE_INVERSE_KINEMATICS_H_
+#define FRC971_CONTROL_LOOPS_SWERVE_INVERSE_KINEMATICS_H_
+#include "aos/util/math.h"
+#include "frc971/control_loops/swerve/simplified_dynamics.h"
+namespace frc971::control_loops::swerve {
+// Class to do straightforwards inverse kinematics of a swerve drivebase. This
+// is meant largely as a sanity-check/initializer for more sophisticated
+// methods. This calculates which directions the modules must be pointed to
+// cause them to be pointed directly along the direction of motion of the
+// drivebase. Accounting for slip angles and the such must be done as part of
+// more sophisticated inverse dynamics.
+template <typename Scalar>
+class InverseKinematics {
+ public:
+  using ModuleParams = SimplifiedDynamics<Scalar>::ModuleParams;
+  using Parameters = SimplifiedDynamics<Scalar>::Parameters;
+  using States = SimplifiedDynamics<Scalar>::States;
+  using State = SimplifiedDynamics<Scalar>::template VelocityState<Scalar>;
+  InverseKinematics(const Parameters &params) : params_(params) {}
+
+  // Uses kVx, kVy, kTheta, and kOmega from the input goal state for the
+  // absolute kinematics. Also uses the specified theta values to bias theta
+  // output values towards the current state (i.e., if the module 0 theta is
+  // currently 0 and we are asked to drive straight backwards, this will prefer
+  // a theta of zero rather than a theta of pi).
+  State Solve(const State &goal) {
+    State result = goal;
+    for (size_t module_index = 0; module_index < params_.modules.size();
+         ++module_index) {
+      SolveModule(goal, params_.modules[module_index].position,
+                  &result(States::kThetas0 + 2 * module_index),
+                  &result(States::kOmegas0 + 2 * module_index));
+    }
+    return result;
+  }
+
+  void SolveModule(const State &goal,
+                   const Eigen::Matrix<Scalar, 2, 1> &module_position,
+                   Scalar *module_theta, Scalar *module_omega) {
+    const Scalar vx = goal(States::kVx);
+    const Scalar vy = goal(States::kVy);
+    const Scalar omega = goal(States::kOmega);
+    // module_velocity_in_robot_frame = R(-theta) * robot_vel +
+    //     omega.cross(module_position);
+    // module_vel_x = (cos(-theta) * vx - sin(-theta) * vy) - omega * module_y
+    // module_vel_y = (sin(-theta) * vx + cos(-theta) * vy) + omega * module_x
+    // module_theta = atan2(module_vel_y, module_vel_x)
+    // module_omega = datan2(module_vel_y, module_vel_x) / dt
+    // datan2(y, x) / dt = (x * dy/dt - y * dx / dt) / (x^2 + y^2)
+    // robot accelerations are assumed to be zero.
+    // dmodule_vel_x / dt = (sin(-theta) * vx + cos(-theta) * vy) * omega
+    // dmodule_vel_y / dt = (-cos(-theta) * vx + sin(-theta) * vy) * omega
+    const Scalar ctheta = std::cos(-goal(States::kTheta));
+    const Scalar stheta = std::sin(-goal(States::kTheta));
+    const Scalar module_vel_x =
+        (ctheta * vx - stheta * vy) - omega * module_position.y();
+    const Scalar module_vel_y =
+        (stheta * vx + ctheta * vy) + omega * module_position.x();
+    const Scalar nominal_module_theta = atan2(module_vel_y, module_vel_x);
+    // If the current module theta is more than 90 deg from the desired theta,
+    // flip the desired theta by 180 deg.
+    if (std::abs(aos::math::DiffAngle(nominal_module_theta, *module_theta)) >
+        M_PI_2) {
+      *module_theta = aos::math::NormalizeAngle(nominal_module_theta + M_PI);
+    } else {
+      *module_theta = nominal_module_theta;
+    }
+    const Scalar module_accel_x = (stheta * vx + ctheta * vy) * omega;
+    const Scalar module_accel_y = (-ctheta * vx + stheta * vy) * omega;
+    const Scalar module_vel_norm_squared =
+        (module_vel_x * module_vel_x + module_vel_y * module_vel_y);
+    if (module_vel_norm_squared < 1e-5) {
+      // Prevent poor conditioning of module velocities at near-zero speeds.
+      *module_omega = 0.0;
+    } else {
+      *module_omega =
+          (module_vel_x * module_accel_y - module_vel_y * module_accel_x) /
+          module_vel_norm_squared;
+    }
+  }
+
+ private:
+  Parameters params_;
+};
+}  // namespace frc971::control_loops::swerve
+#endif  // FRC971_CONTROL_LOOPS_SWERVE_INVERSE_KINEMATICS_H_
diff --git a/frc971/control_loops/swerve/inverse_kinematics_test.cc b/frc971/control_loops/swerve/inverse_kinematics_test.cc
new file mode 100644
index 0000000..407d540
--- /dev/null
+++ b/frc971/control_loops/swerve/inverse_kinematics_test.cc
@@ -0,0 +1,150 @@
+#include "frc971/control_loops/swerve/inverse_kinematics.h"
+
+#include "gtest/gtest.h"
+
+namespace frc971::control_loops::swerve::testing {
+class InverseKinematicsTest : public ::testing::Test {
+ protected:
+  typedef double Scalar;
+  using State = InverseKinematics<Scalar>::State;
+  using States = InverseKinematics<Scalar>::States;
+  using ModuleParams = InverseKinematics<Scalar>::ModuleParams;
+  using Parameters = InverseKinematics<Scalar>::Parameters;
+  static ModuleParams MakeModule(const Eigen::Matrix<Scalar, 2, 1> &position) {
+    return ModuleParams{.position = position,
+                        .slip_angle_coefficient = 0.0,
+                        .slip_angle_alignment_coefficient = 0.0,
+                        .steer_motor = KrakenFOC(),
+                        .drive_motor = KrakenFOC(),
+                        .steer_ratio = 1.0,
+                        .drive_ratio = 1.0};
+  }
+  static Parameters MakeParams() {
+    return {.mass = 1.0,
+            .moment_of_inertia = 1.0,
+            .modules = {
+                MakeModule({1.0, 1.0}),
+                MakeModule({-1.0, 1.0}),
+                MakeModule({-1.0, -1.0}),
+                MakeModule({1.0, -1.0}),
+            }};
+  }
+
+  InverseKinematicsTest() : inverse_kinematics_(MakeParams()) {}
+
+  struct Goal {
+    Scalar vx;
+    Scalar vy;
+    Scalar omega;
+    Scalar theta;
+  };
+
+  void CheckState(
+      const Goal &goal, const std::array<Scalar, 4> &expected_thetas,
+      const std::optional<Eigen::Vector4d> &expected_omegas = std::nullopt) {
+    State goal_state = State::Zero();
+    goal_state(States::kVx) = goal.vx;
+    goal_state(States::kVy) = goal.vy;
+    goal_state(States::kOmega) = goal.omega;
+    goal_state(States::kTheta) = goal.theta;
+    SCOPED_TRACE(goal_state.bottomRows<4>().transpose());
+    const State nominal_state = inverse_kinematics_.Solve(goal_state);
+    // Now, calculate the numerical derivative of the state and validate that it
+    // matches expectations.
+    const Scalar kDt = 1e-5;
+    const Scalar dtheta = kDt * goal.omega;
+    goal_state(States::kTheta) += dtheta / 2.0;
+    const State state_eps_pos = inverse_kinematics_.Solve(goal_state);
+    goal_state(States::kTheta) -= dtheta;
+    const State state_eps_neg = inverse_kinematics_.Solve(goal_state);
+    const State state_derivative = (state_eps_pos - state_eps_neg) / kDt;
+    for (size_t module_index = 0; module_index < 4; ++module_index) {
+      SCOPED_TRACE(module_index);
+      const int omega_idx = States::kOmegas0 + 2 * module_index;
+      const int theta_idx = States::kThetas0 + 2 * module_index;
+      EXPECT_NEAR(nominal_state(omega_idx), state_derivative(theta_idx), 1e-10);
+      EXPECT_NEAR(nominal_state(theta_idx), expected_thetas[module_index],
+                  1e-10);
+      if (expected_omegas.has_value()) {
+        EXPECT_NEAR(nominal_state(omega_idx),
+                    expected_omegas.value()(module_index), 1e-10);
+      }
+    }
+  }
+
+  InverseKinematics<Scalar> inverse_kinematics_;
+};
+
+// Tests that if we are driving straight with no yaw that we get sane
+// kinematics.
+TEST_F(InverseKinematicsTest, StraightDrivingNoYaw) {
+  // Sanity-check zero-speed operation.
+  CheckState({.vx = 0.0, .vy = 0.0, .omega = 0.0, .theta = 0.0},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 1.0, .vy = 0.0, .omega = 0.0, .theta = 0.0},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+  // Reverse should prefer to bias the modules towards [-pi/2, pi/2] due to
+  // hysteresis from starting the modules at thetas of 0.
+  CheckState({.vx = -1.0, .vy = 0.0, .omega = 0.0, .theta = 0.0},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 0.0, .vy = 1.0, .omega = 0.0, .theta = 0.0},
+             {M_PI_2, M_PI_2, M_PI_2, M_PI_2}, Eigen::Vector4d::Zero());
+  // For module hysteresis, this is a corner case where we are exactly 90 deg
+  // from the current value; the exact result is unimportant.
+  CheckState({.vx = 0.0, .vy = -1.0, .omega = 0.0, .theta = 0.0},
+             {-M_PI_2, -M_PI_2, -M_PI_2, -M_PI_2}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 1.0, .vy = 1.0, .omega = 0.0, .theta = 0.0},
+             {M_PI_4, M_PI_4, M_PI_4, M_PI_4}, Eigen::Vector4d::Zero());
+  // Reverse should prefer to bias the modules towards [-pi/2, pi/2] due to
+  // hysteresis from starting the modules at thetas of 0.
+  CheckState({.vx = -1.0, .vy = -1.0, .omega = 0.0, .theta = 0.0},
+             {M_PI_4, M_PI_4, M_PI_4, M_PI_4}, Eigen::Vector4d::Zero());
+}
+
+// Tests that if we are driving straight with non-zero yaw that we get sane
+// kinematics.
+TEST_F(InverseKinematicsTest, StraightDrivingYawed) {
+  CheckState({.vx = 1.0, .vy = 0.0, .omega = 0.0, .theta = 0.1},
+             {-0.1, -0.1, -0.1, -0.1}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 1.0, .vy = 0.0, .omega = 0.0, .theta = M_PI_2},
+             {-M_PI_2, -M_PI_2, -M_PI_2, -M_PI_2}, Eigen::Vector4d::Zero());
+  CheckState({.vx = 1.0, .vy = 0.0, .omega = 0.0, .theta = M_PI},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 0.0, .vy = 1.0, .omega = 0.0, .theta = M_PI_2},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+  // Reverse should prefer to bias the modules towards [-pi/2, pi/2] due to
+  // hysteresis from starting the modules at thetas of 0.
+  CheckState({.vx = 0.0, .vy = -1.0, .omega = 0.0, .theta = M_PI_2},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+}
+
+// Tests that we can spin in place.
+TEST_F(InverseKinematicsTest, SpinInPlace) {
+  CheckState({.vx = 0.0, .vy = 0.0, .omega = 1.0, .theta = 0.0},
+             {-M_PI_4, M_PI_4, -M_PI_4, M_PI_4}, Eigen::Vector4d::Zero());
+  // And changing the current theta should not matter.
+  CheckState({.vx = 0.0, .vy = 0.0, .omega = 1.0, .theta = 1.0},
+             {-M_PI_4, M_PI_4, -M_PI_4, M_PI_4}, Eigen::Vector4d::Zero());
+}
+
+// Tests that if we are spinning while moving that we correctly calculate module
+// yaw rates.
+TEST_F(InverseKinematicsTest, SpinWhileMoving) {
+  // Set up a situation where we are driving straight forwards, with a
+  // yaw rate of 1 rad / sec.
+  // The modules are all at radii of sqrt(2), so the contribution from both
+  // the translational and rotational velocities should be equal; in this case
+  // each module will have an angle that is an equal combination of straight
+  // forwards (0 deg) and some 45 deg offset.
+  CheckState({.vx = std::sqrt(static_cast<Scalar>(2)),
+              .vy = 0.0,
+              .omega = 1.0,
+              .theta = 0.0},
+             {3 * M_PI_4 / 2, -3 * M_PI_4 / 2, -M_PI_4 / 2, M_PI_4 / 2});
+}
+}  // namespace frc971::control_loops::swerve::testing
diff --git a/frc971/control_loops/swerve/jax_dynamics.py b/frc971/control_loops/swerve/jax_dynamics.py
index 19a48de..58d5fcf 100644
--- a/frc971/control_loops/swerve/jax_dynamics.py
+++ b/frc971/control_loops/swerve/jax_dynamics.py
@@ -354,6 +354,7 @@
     ])
 
 
+@jax.jit
 def mpc_cost(coefficients: CoefficientsType, X, U, goal):
     J = 0
 
diff --git a/frc971/control_loops/swerve/linearization_utils.h b/frc971/control_loops/swerve/linearization_utils.h
new file mode 100644
index 0000000..ee1275e
--- /dev/null
+++ b/frc971/control_loops/swerve/linearization_utils.h
@@ -0,0 +1,13 @@
+#ifndef FRC971_CONTROL_LOOPS_SWERVE_LINEARIZATION_UTILS_H_
+#define FRC971_CONTROL_LOOPS_SWERVE_LINEARIZATION_UTILS_H_
+
+namespace frc971::control_loops::swerve {
+template <typename State, typename Input>
+struct DynamicsInterface {
+  virtual ~DynamicsInterface() {}
+  // To be overridden by the implementation; returns the derivative of the state
+  // at the given state with the provided control input.
+  virtual State operator()(const State &X, const Input &U) const = 0;
+};
+}  // namespace frc971::control_loops::swerve
+#endif  // FRC971_CONTROL_LOOPS_SWERVE_LINEARIZATION_UTILS_H_
diff --git a/frc971/control_loops/swerve/linearized_controller.h b/frc971/control_loops/swerve/linearized_controller.h
new file mode 100644
index 0000000..1df9640
--- /dev/null
+++ b/frc971/control_loops/swerve/linearized_controller.h
@@ -0,0 +1,126 @@
+#include <memory>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include <Eigen/Dense>
+
+#include "frc971/control_loops/c2d.h"
+#include "frc971/control_loops/dlqr.h"
+#include "frc971/control_loops/jacobian.h"
+#include "frc971/control_loops/swerve/dynamics.h"
+#include "frc971/control_loops/swerve/linearization_utils.h"
+
+namespace frc971::control_loops::swerve {
+
+// Provides a simple LQR controller that takes a non-linear system, linearizes
+// the dynamics at each timepoint, recalculates the LQR gains for those
+// dynamics, and calculates the relevant feedback inputs to provide.
+template <int NStates, typename Scalar = double>
+class LinearizedController {
+ public:
+  typedef Eigen::Matrix<Scalar, NStates, 1> State;
+  typedef Eigen::Matrix<Scalar, NStates, NStates> StateSquare;
+  typedef Eigen::Matrix<Scalar, kNumInputs, 1> Input;
+  typedef Eigen::Matrix<Scalar, kNumInputs, kNumInputs> InputSquare;
+  typedef Eigen::Matrix<Scalar, NStates, kNumInputs> BMatrix;
+  typedef DynamicsInterface<State, Input> Dynamics;
+
+  struct Parameters {
+    // State cost matrix.
+    StateSquare Q;
+    // Input cost matrix.
+    InputSquare R;
+    // period at which the controller is called.
+    std::chrono::nanoseconds dt;
+    // The dynamics to use.
+    // TODO(james): I wrote this before creating the auto-differentiation
+    // functions; we should swap to the auto-differentiation, since the
+    // numerical linearization is one of the bigger timesinks in this controller
+    // right now.
+    std::unique_ptr<Dynamics> dynamics;
+  };
+
+  // Represents the linearized dynamics of the system.
+  struct LinearDynamics {
+    StateSquare A;
+    BMatrix B;
+  };
+
+  // Debug information for a given cycle of the controller.
+  struct ControllerDebug {
+    // Feedforward input which we provided.
+    Input U_ff;
+    // Calculated feedback input to provide.
+    Input U_feedback;
+    Eigen::Matrix<Scalar, kNumInputs, NStates> feedback_contributions;
+  };
+
+  struct ControllerResult {
+    // Control input to provide to the robot.
+    Input U;
+    ControllerDebug debug;
+  };
+
+  LinearizedController(Parameters params) : params_(std::move(params)) {}
+
+  // Runs the controller for a given iteration, relinearizing the dynamics about
+  // the provided current state X, attempting to control the robot to the
+  // desired goal state.
+  // The U_ff input will be added into the returned control input.
+  ControllerResult RunController(const State &X, const State &goal,
+                                 Input U_ff) {
+    auto start_time = aos::monotonic_clock::now();
+    // TODO(james): Swap this to the auto-diff methods; this is currently about
+    // a third of the total time spent in this method when run on the roborio.
+    const struct LinearDynamics continuous_dynamics =
+        LinearizeDynamics(X, U_ff);
+    auto linearization_time = aos::monotonic_clock::now();
+    struct LinearDynamics discrete_dynamics;
+    frc971::controls::C2D(continuous_dynamics.A, continuous_dynamics.B,
+                          params_.dt, &discrete_dynamics.A,
+                          &discrete_dynamics.B);
+    auto c2d_time = aos::monotonic_clock::now();
+    VLOG(2) << "Controllability of dynamics (ideally should be " << NStates
+            << "): "
+            << frc971::controls::Controllability(discrete_dynamics.A,
+                                                 discrete_dynamics.B);
+    Eigen::Matrix<Scalar, kNumInputs, NStates> K;
+    Eigen::Matrix<Scalar, NStates, NStates> S;
+    // TODO(james): Swap this to a cheaper DARE solver; we should probably just
+    // do something like we do in Trajectory::CalculatePathGains for the tank
+    // spline controller where we approximate the infinite-horizon DARE solution
+    // by doing a finite-horizon LQR.
+    // Currently the dlqr call represents ~60% of the time spent in the
+    // RunController() method.
+    frc971::controls::dlqr(discrete_dynamics.A, discrete_dynamics.B, params_.Q,
+                           params_.R, &K, &S);
+    auto dlqr_time = aos::monotonic_clock::now();
+    const Input U_feedback = K * (goal - X);
+    const Input U = U_ff + U_feedback;
+    Eigen::Matrix<Scalar, kNumInputs, NStates> feedback_contributions;
+    for (int state_idx = 0; state_idx < NStates; ++state_idx) {
+      feedback_contributions.col(state_idx) =
+          K.col(state_idx) * (goal - X)(state_idx);
+    }
+    VLOG(2) << "linearization time "
+            << aos::time::DurationInSeconds(linearization_time - start_time)
+            << " c2d time "
+            << aos::time::DurationInSeconds(c2d_time - linearization_time)
+            << " dlqr time "
+            << aos::time::DurationInSeconds(dlqr_time - c2d_time);
+    return {.U = U,
+            .debug = {.U_ff = U_ff,
+                      .U_feedback = U_feedback,
+                      .feedback_contributions = feedback_contributions}};
+  }
+
+  LinearDynamics LinearizeDynamics(const State &X, const Input &U) {
+    return {.A = NumericalJacobianX(*params_.dynamics, X, U),
+            .B = NumericalJacobianU(*params_.dynamics, X, U)};
+  }
+
+ private:
+  const Parameters params_;
+};
+
+}  // namespace frc971::control_loops::swerve
diff --git a/frc971/control_loops/swerve/linearized_controller_test.cc b/frc971/control_loops/swerve/linearized_controller_test.cc
new file mode 100644
index 0000000..d5e7aee
--- /dev/null
+++ b/frc971/control_loops/swerve/linearized_controller_test.cc
@@ -0,0 +1,87 @@
+#include "frc971/control_loops/swerve/linearized_controller.h"
+
+#include "gtest/gtest.h"
+
+namespace frc971::control_loops::swerve::test {
+
+class LinearizedControllerTest : public ::testing::Test {
+ protected:
+  typedef LinearizedController<2> Controller;
+  typedef Controller::State State;
+  typedef Controller::Input Input;
+  struct LinearDynamics : public Controller::Dynamics {
+    Eigen::Vector2d operator()(
+        const Eigen::Vector2d &X,
+        const Eigen::Matrix<double, 8, 1> &U) const override {
+      return Eigen::Matrix2d{{0.0, 1.0}, {0.0, -0.01}} * X +
+             Eigen::Matrix<double, 2, 8>{
+                 {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0},
+                 {1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}} *
+                 U;
+    }
+  };
+  static Eigen::Matrix<double, kNumInputs, kNumInputs> MakeR() {
+    Eigen::Matrix<double, kNumInputs, kNumInputs> R;
+    R.setIdentity();
+    return R;
+  }
+  LinearizedControllerTest()
+      : controller_({.Q = Eigen::Matrix<double, 2, 2>{{1.0, 0.0}, {0.0, 1.0}},
+                     .R = MakeR(),
+                     .dt = std::chrono::milliseconds(10),
+                     .dynamics = std::make_unique<LinearDynamics>()}) {}
+
+  Controller controller_;
+};
+
+// Sanity check that the dynamics linearization is working correctly.
+TEST_F(LinearizedControllerTest, LinearizedDynamics) {
+  auto dynamics = controller_.LinearizeDynamics(State::Zero(), Input::Zero());
+  EXPECT_EQ(0.0, dynamics.A(0, 0));
+  EXPECT_EQ(0.0, dynamics.A(1, 0));
+  EXPECT_EQ(1.0, dynamics.A(0, 1));
+  EXPECT_EQ(-0.01, dynamics.A(1, 1));
+  // All elements of B except for (1, 0) should be exactly 0.
+  EXPECT_EQ(1.0, dynamics.B(1, 0));
+  EXPECT_EQ(1.0, dynamics.B.norm());
+}
+
+// Confirm that the generated LQR controller is able to generate correct
+// inputs when state and goal are at zero.
+TEST_F(LinearizedControllerTest, ControllerResultAtZero) {
+  auto result =
+      controller_.RunController(State::Zero(), State::Zero(), Input::Zero());
+  EXPECT_EQ(0.0, result.U.norm());
+  EXPECT_EQ(0.0, result.debug.U_ff.norm());
+  EXPECT_EQ(0.0, result.debug.U_feedback.norm());
+}
+
+// Confirm that the generated LQR controller is able to generate correct
+// inputs when state is zero and the goal is non-zero.
+TEST_F(LinearizedControllerTest, ControlToZero) {
+  auto result = controller_.RunController(State::Zero(), State{{1.0}, {0.0}},
+                                          Input::Zero());
+  EXPECT_LT(0.0, result.U(0, 0));
+  // All other U inputs should be zero.
+  EXPECT_EQ(0.0, result.U.bottomRows<7>().norm());
+  EXPECT_EQ(0.0, result.debug.U_ff.norm());
+  EXPECT_EQ(0.0,
+            (result.U - (result.debug.U_ff + result.debug.U_feedback)).norm());
+}
+
+// Confirm that the generated LQR controller is able to pass through the
+// feedforwards when we have no difference between the goal and the current
+// state.
+TEST_F(LinearizedControllerTest, ControlToNonzeroState) {
+  const State state{{1.0}, {1.0}};
+  auto result = controller_.RunController(
+      state, state,
+      Input{{1.0}, {0.0}, {0.0}, {0.0}, {0.0}, {0.0}, {0.0}, {0.0}});
+  EXPECT_EQ(1.0, result.U(0, 0));
+  // All other U inputs should be zero.
+  EXPECT_EQ(0.0, result.U.bottomRows<7>().norm());
+  EXPECT_EQ(0.0, result.debug.U_feedback.norm());
+  EXPECT_EQ(0.0,
+            (result.U - (result.debug.U_ff + result.debug.U_feedback)).norm());
+}
+}  // namespace frc971::control_loops::swerve::test
diff --git a/frc971/control_loops/swerve/multi_experience_collector.py b/frc971/control_loops/swerve/multi_experience_collector.py
new file mode 100644
index 0000000..d0551e1
--- /dev/null
+++ b/frc971/control_loops/swerve/multi_experience_collector.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+from absl import app
+from absl import flags
+import sys
+from multiprocessing.pool import ThreadPool
+import pathlib
+import subprocess
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('outdir', '/tmp/swerve', "Directory to write results to.")
+flags.DEFINE_integer('num_actors', 20, 'Number of actors to run in parallel.')
+flags.DEFINE_integer('num_solutions', 100,
+                     'Number of random problems to solve.')
+
+
+def collect_experience(agent_number):
+    filename = f'{agent_number}'
+    if FLAGS.outdir:
+        subdir = pathlib.Path(FLAGS.outdir) / filename
+    else:
+        subdir = pathlib.Path(filename)
+    subdir.mkdir(parents=True, exist_ok=True)
+
+    with open(f'{subdir.resolve()}/log', 'w') as output:
+        subprocess.check_call(
+            args=[
+                sys.executable,
+                "frc971/control_loops/swerve/experience_collector",
+                f"--seed={agent_number}",
+                f"--outputdir={subdir.resolve()}",
+                "--quiet",
+                f"--num_solutions={FLAGS.num_solutions}",
+            ],
+            stdout=output,
+            stderr=output,
+        )
+
+
+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/experience_collector",
+        "--compileonly",
+    ])
+
+    # Try a bunch of goals now
+    with ThreadPool(FLAGS.num_actors) as pool:
+        pool.starmap(collect_experience, zip(range(FLAGS.num_actors), ))
+
+
+if __name__ == '__main__':
+    app.run(main)
diff --git a/frc971/control_loops/swerve/simplified_dynamics.h b/frc971/control_loops/swerve/simplified_dynamics.h
new file mode 100644
index 0000000..9afe636
--- /dev/null
+++ b/frc971/control_loops/swerve/simplified_dynamics.h
@@ -0,0 +1,372 @@
+#ifndef FRC971_CONTROL_LOOPS_SWERVE_SIMPLIFIED_DYNAMICS_H_
+#define FRC971_CONTROL_LOOPS_SWERVE_SIMPLIFIED_DYNAMICS_H_
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+
+#include "aos/util/math.h"
+#include "frc971/control_loops/swerve/auto_diff_jacobian.h"
+#include "frc971/control_loops/swerve/dynamics.h"
+#include "frc971/control_loops/swerve/motors.h"
+
+namespace frc971::control_loops::swerve {
+
+// Provides a simplified set of physics representing a swerve drivetrain.
+// Broadly speaking, these dynamics model:
+// * Standard motors on the drive and steer axes, with no coupling between
+//   the motors.
+// * Assume that the steer direction of each module is only influenced by
+//   the inertia of the motor rotor plus some small aligning force.
+// * Assume that torque from the drive motor is transferred directly to the
+//   carpet.
+// * Assume that a lateral force on the wheel is generated proportional to the
+//   slip angle of the wheel.
+//
+// This class is templated on a Scalar that is used to determine whether you
+// want to use double or single-precision floats for the calculations here.
+//
+// Several individual methods on this class are also templated on a LocalScalar
+// type. This is provided to allow those methods to be called with ceres Jets to
+// do autodifferentiation within various solvers/jacobian calculators.
+template <typename Scalar = double>
+class SimplifiedDynamics {
+ public:
+  struct ModuleParams {
+    // Module position relative to the center of mass of the robot.
+    Eigen::Matrix<Scalar, 2, 1> position;
+    // Coefficient dictating how much sideways force is generated at a given
+    // slip angle. Units are effectively Newtons / radian of slip.
+    Scalar slip_angle_coefficient;
+    // Coefficient dicating how much the steer wheel is forced into alignment
+    // by the motion of the wheel over the ground (i.e., we are assuming that
+    // if you push the robot along it will cause the wheels to eventually
+    // align with the direction of motion).
+    // In radians / sec^2 / radians of slip angle.
+    Scalar slip_angle_alignment_coefficient;
+    // Parameters for the steer and drive motors.
+    Motor steer_motor;
+    Motor drive_motor;
+    // radians of module steering = steer_ratio * radians of motor shaft
+    Scalar steer_ratio;
+    // meters of driving = drive_ratio * radians of motor shaft
+    Scalar drive_ratio;
+  };
+  struct Parameters {
+    // Mass of the robot, in kg.
+    Scalar mass;
+    // Moment of inertia of the robot about the yaw axis, in kg * m^2.
+    Scalar moment_of_inertia;
+    // Note: While this technically would support an arbitrary number of
+    // modules, the statically-sized state vectors do limit us to 4 modules
+    // currently, and it should not be counted on that other pieces of code will
+    // be able to support non-4-module swerves.
+    std::vector<ModuleParams> modules;
+  };
+  enum States {
+    // Thetas* and Omegas* are the yaw and yaw rate of the indicated modules.
+    // (note that we do not actually need to track drive speed per module,
+    // as with current control we have the ability to directly command torque
+    // to those motors; however, if we wished to account for the saturation
+    // limits on the motor, then we would need to have access to those states,
+    // although they can be fully derived from the robot vx, vy, theta, and
+    // omega).
+    kThetas0 = 0,
+    kOmegas0 = 1,
+    kThetas1 = 2,
+    kOmegas1 = 3,
+    kThetas2 = 4,
+    kOmegas2 = 5,
+    kThetas3 = 6,
+    kOmegas3 = 7,
+    // Robot yaw, in radians.
+    kTheta = 8,
+    // Robot speed in the global frame, in meters / sec.
+    kVx = 9,
+    kVy = 10,
+    // Robot yaw rate, in radians / sec.
+    kOmega = 11,
+    kNumVelocityStates = 12,
+    // Augmented states for doing position control.
+    // Robot X position in the global frame.
+    kX = 12,
+    // Robot Y position in the global frame.
+    kY = 13,
+    kNumPositionStates = 14,
+  };
+  using Inputs = InputStates::States;
+
+  template <typename ScalarT = Scalar>
+  using VelocityState = Eigen::Matrix<ScalarT, kNumVelocityStates, 1>;
+  template <typename ScalarT = Scalar>
+  using PositionState = Eigen::Matrix<ScalarT, kNumPositionStates, 1>;
+  template <typename ScalarT = Scalar>
+  using VelocityStateSquare =
+      Eigen::Matrix<ScalarT, kNumVelocityStates, kNumVelocityStates>;
+  template <typename ScalarT = Scalar>
+  using PositionStateSquare =
+      Eigen::Matrix<ScalarT, kNumPositionStates, kNumPositionStates>;
+  template <typename ScalarT = Scalar>
+  using PositionBMatrix =
+      Eigen::Matrix<ScalarT, kNumPositionStates, kNumInputs>;
+  template <typename ScalarT = Scalar>
+  using Input = Eigen::Matrix<ScalarT, kNumInputs, 1>;
+
+  SimplifiedDynamics(const Parameters &params) : params_(params) {
+    for (size_t module_index = 0; module_index < params_.modules.size();
+         ++module_index) {
+      module_dynamics_.emplace_back(params_, module_index);
+    }
+  }
+
+  // Returns the derivative of state for the given state and input.
+  template <typename LocalScalar>
+  PositionState<LocalScalar> Dynamics(const PositionState<LocalScalar> &state,
+                                      const Input<LocalScalar> &input) const {
+    PositionState<LocalScalar> Xdot = PositionState<LocalScalar>::Zero();
+
+    for (const ModuleDynamics &module : module_dynamics_) {
+      Xdot += module.PartialDynamics(state, input);
+    }
+
+    // And finally catch the global states:
+    Xdot(kX) = state(kVx);
+    Xdot(kY) = state(kVy);
+    Xdot(kTheta) = state(kOmega);
+
+    return Xdot;
+  }
+
+  template <typename LocalScalar>
+  VelocityState<LocalScalar> VelocityDynamics(
+      const VelocityState<LocalScalar> &state,
+      const Input<LocalScalar> &input) const {
+    PositionState<LocalScalar> input_state = PositionState<LocalScalar>::Zero();
+    input_state.template topRows<kNumVelocityStates>() = state;
+    return Dynamics(input_state, input).template topRows<kNumVelocityStates>();
+  }
+
+  std::pair<PositionStateSquare<>, PositionBMatrix<>> LinearizedDynamics(
+      const PositionState<> &state, const Input<> &input) {
+    DynamicsFunctor functor(*this);
+    Eigen::Matrix<Scalar, kNumPositionStates + kNumInputs, 1> parameters;
+    parameters.template topRows<kNumPositionStates>() = state;
+    parameters.template bottomRows<kNumInputs>() = input;
+    const Eigen::Matrix<Scalar, kNumPositionStates,
+                        kNumPositionStates + kNumInputs>
+        jacobian =
+            AutoDiffJacobian<Scalar, DynamicsFunctor,
+                             kNumPositionStates + kNumInputs,
+                             kNumPositionStates>::Jacobian(functor, parameters);
+    return {
+        jacobian.template block<kNumPositionStates, kNumPositionStates>(0, 0),
+        jacobian.template block<kNumPositionStates, kNumInputs>(
+            0, kNumPositionStates)};
+  }
+
+ private:
+  // Wrapper to provide an operator() for the dynamisc class that allows it to
+  // be used by the auto-differentiation code.
+  class DynamicsFunctor {
+   public:
+    DynamicsFunctor(const SimplifiedDynamics &dynamics) : dynamics_(dynamics) {}
+
+    template <typename LocalScalar>
+    Eigen::Matrix<LocalScalar, kNumPositionStates, 1> operator()(
+        const Eigen::Map<const Eigen::Matrix<
+            LocalScalar, kNumPositionStates + kNumInputs, 1>>
+            input) const {
+      return dynamics_.Dynamics(
+          PositionState<LocalScalar>(
+              input.template topRows<kNumPositionStates>()),
+          Input<LocalScalar>(input.template bottomRows<kNumInputs>()));
+    }
+
+   private:
+    const SimplifiedDynamics &dynamics_;
+  };
+
+  // Represents the dynamics of an individual module.
+  class ModuleDynamics {
+   public:
+    ModuleDynamics(const Parameters &robot_params, const size_t module_index)
+        : robot_params_(robot_params), module_index_(module_index) {
+      CHECK_LT(module_index_, robot_params_.modules.size());
+    }
+
+    // This returns the portions of the derivative of state that are due to the
+    // individual module. The result from this function should be able to be
+    // naively summed with the dynamics for each other module plus some global
+    // dynamics (which take care of that e.g. xdot = vx) and give you the
+    // overall dynamics of the system.
+    template <typename LocalScalar>
+    PositionState<LocalScalar> PartialDynamics(
+        const PositionState<LocalScalar> &state,
+        const Input<LocalScalar> &input) const {
+      PositionState<LocalScalar> Xdot = PositionState<LocalScalar>::Zero();
+
+      Xdot(ThetasIdx()) = state(OmegasIdx());
+
+      // Steering dynamics for an individual module assume ~zero friction,
+      // and thus ~the only inertia is from the motor rotor itself.
+      // torque_motor = stall_torque / stall_current * current
+      // accel_motor = torque_motor / motor_inertia
+      // accel_steer = accel_motor * steer_ratio
+      const Motor &steer_motor = module_params().steer_motor;
+      const LocalScalar steer_motor_accel =
+          input(IsIdx()) *
+          static_cast<Scalar>(
+              module_params().steer_ratio * steer_motor.stall_torque /
+              (steer_motor.stall_current * steer_motor.motor_inertia));
+
+      // For the impacts of the modules on the overall robot
+      // dynamics (X, Y, and theta acceleration), we calculate the forces
+      // generated by the module and then apply them. These forces come from
+      // two effects in this model:
+      // 1. Slip angle of the module (dependent on the current robot velocity &
+      //    module steer angle).
+      // 2. Drive torque from the module (dependent on the current drive
+      //    current and module steer angle).
+      // We assume no torque is generated from e.g. the wheel resisting the
+      // steering motion.
+      //
+      // clang-format off
+       //
+       // For slip angle we have:
+       // wheel_velocity = R(-theta - theta_steer) * (
+       //    robot_vel + omega.cross(R(theta) * module_position))
+       // slip_angle = -atan2(wheel_velocity)
+       // slip_force = slip_angle_coefficient * slip_angle
+       // slip_force_direction = theta + theta_steer + pi / 2
+       // force_x = slip_force * cos(slip_force_direction)
+       // force_y = slip_force * sin(slip_force_direction)
+       // accel_* = force_* / mass
+       // # And now calculate torque from slip angle.
+       // torque_vec = module_position.cross([slip_force * cos(theta_steer + pi / 2),
+       //                                     slip_force * sin(theta_steer + pi / 2),
+       //                                     0.0])
+       // torque_vec = module_position.cross([slip_force * -sin(theta_steer),
+       //                                     slip_force * cos(theta_steer),
+       //                                     0.0])
+       // robot_torque = torque_vec.z()
+       //
+       // For drive torque we have:
+       // drive_force = (drive_current * stall_torque / stall_current) / drive_ratio
+       // drive_force_direction = theta + theta_steer
+       // force_x = drive_force * cos(drive_force_direction)
+       // force_y = drive_force * sin(drive_force_direction)
+       // torque_vec = drive_force * module_position.cross([cos(theta_steer),
+       //                                                   sin(theta_steer),
+       //                                                   0.0])
+       // torque = torque_vec.z()
+       //
+      // clang-format on
+
+      const Eigen::Matrix<Scalar, 3, 1> module_position{
+          {module_params().position.x()},
+          {module_params().position.y()},
+          {0.0}};
+      const LocalScalar theta = state(kTheta);
+      const LocalScalar theta_steer = state(ThetasIdx());
+      const Eigen::Matrix<LocalScalar, 3, 1> wheel_velocity_in_global_frame =
+          Eigen::Matrix<LocalScalar, 3, 1>(state(kVx), state(kVy),
+                                           static_cast<LocalScalar>(0.0)) +
+          (Eigen::Matrix<LocalScalar, 3, 1>(static_cast<LocalScalar>(0.0),
+                                            static_cast<LocalScalar>(0.0),
+                                            state(kOmega))
+               .cross(Eigen::AngleAxis<LocalScalar>(
+                          theta, Eigen::Matrix<LocalScalar, 3, 1>::UnitZ()) *
+                      module_position));
+      const Eigen::Matrix<LocalScalar, 3, 1> wheel_velocity_in_wheel_frame =
+          Eigen::AngleAxis<LocalScalar>(
+              -theta - theta_steer, Eigen::Matrix<LocalScalar, 3, 1>::UnitZ()) *
+          wheel_velocity_in_global_frame;
+      // The complicated dynamics use some obnoxious-looking functions to
+      // try to approximate how the slip angle behaves a low speeds to better
+      // condition the dynamics. Because I couldn't be bothered to copy those
+      // dynamics, instead just bias the slip angle to zero at low speeds.
+      const LocalScalar wheel_speed = wheel_velocity_in_wheel_frame.norm();
+      const Scalar start_speed = 0.1;
+      const LocalScalar heading_truth_proportion = -expm1(
+          /*arbitrary large number=*/static_cast<Scalar>(-100.0) *
+          (wheel_speed - start_speed));
+      const LocalScalar wheel_heading =
+          (wheel_speed < start_speed)
+              ? static_cast<LocalScalar>(0.0)
+              : heading_truth_proportion *
+                    atan2(wheel_velocity_in_wheel_frame.y(),
+                          wheel_velocity_in_wheel_frame.x());
+
+      // We wrap slip_angle with a sin() not because there is actually a sin()
+      // in the real math but rather because we need to smoothly and correctly
+      // handle slip angles between pi / 2 and 3 * pi / 2.
+      const LocalScalar slip_angle = sin(-wheel_heading);
+      const LocalScalar slip_force =
+          module_params().slip_angle_coefficient * slip_angle;
+      const LocalScalar slip_force_direction =
+          theta + theta_steer + static_cast<Scalar>(M_PI_2);
+      const Eigen::Matrix<LocalScalar, 3, 1> slip_force_vec =
+          slip_force * UnitYawVector<LocalScalar>(slip_force_direction);
+      const LocalScalar slip_torque =
+          module_position
+              .cross(slip_force *
+                     UnitYawVector<LocalScalar>(theta_steer +
+                                                static_cast<Scalar>(M_PI_2)))
+              .z();
+
+      // drive torque calculations
+      const Motor &drive_motor = module_params().drive_motor;
+      const LocalScalar drive_force =
+          input(IdIdx()) * static_cast<Scalar>(drive_motor.stall_torque /
+                                               drive_motor.stall_current /
+                                               module_params().drive_ratio);
+      const Eigen::Matrix<LocalScalar, 3, 1> drive_force_vec =
+          drive_force * UnitYawVector<LocalScalar>(theta + theta_steer);
+      const LocalScalar drive_torque =
+          drive_force *
+          module_position.cross(UnitYawVector<LocalScalar>(theta_steer)).z();
+      // We add in an aligning force on the wheels primarily to help provide a
+      // bit of impetus to the controllers/solvers to discourage aggressive
+      // slip angles. If we do not include this, then the dynamics make it look
+      // like there are no losses to using extremely aggressive slip angles.
+      const LocalScalar wheel_alignment_accel =
+          -module_params().slip_angle_alignment_coefficient * slip_angle;
+
+      Xdot(OmegasIdx()) = steer_motor_accel + wheel_alignment_accel;
+      // Sum up all the forces.
+      Xdot(kVx) =
+          (slip_force_vec.x() + drive_force_vec.x()) / robot_params_.mass;
+      Xdot(kVy) =
+          (slip_force_vec.y() + drive_force_vec.y()) / robot_params_.mass;
+      Xdot(kOmega) =
+          (slip_torque + drive_torque) / robot_params_.moment_of_inertia;
+
+      return Xdot;
+    }
+
+   private:
+    template <typename LocalScalar>
+    Eigen::Matrix<LocalScalar, 3, 1> UnitYawVector(LocalScalar yaw) const {
+      return Eigen::Matrix<LocalScalar, 3, 1>{
+          {static_cast<LocalScalar>(cos(yaw))},
+          {static_cast<LocalScalar>(sin(yaw))},
+          {static_cast<LocalScalar>(0.0)}};
+    }
+    size_t ThetasIdx() const { return kThetas0 + 2 * module_index_; }
+    size_t OmegasIdx() const { return kOmegas0 + 2 * module_index_; }
+    size_t IsIdx() const { return Inputs::kIs0 + 2 * module_index_; }
+    size_t IdIdx() const { return Inputs::kId0 + 2 * module_index_; }
+
+    const ModuleParams &module_params() const {
+      return robot_params_.modules[module_index_];
+    }
+
+    const Parameters robot_params_;
+
+    const size_t module_index_;
+  };
+
+  Parameters params_;
+  std::vector<ModuleDynamics> module_dynamics_;
+};
+
+}  // namespace frc971::control_loops::swerve
+#endif  // FRC971_CONTROL_LOOPS_SWERVE_SIMPLIFIED_DYNAMICS_H_
diff --git a/frc971/control_loops/swerve/simplified_dynamics_test.cc b/frc971/control_loops/swerve/simplified_dynamics_test.cc
new file mode 100644
index 0000000..894a85b
--- /dev/null
+++ b/frc971/control_loops/swerve/simplified_dynamics_test.cc
@@ -0,0 +1,244 @@
+#include "frc971/control_loops/swerve/simplified_dynamics.h"
+
+#include <functional>
+
+#include "absl/log/log.h"
+#include "gtest/gtest.h"
+
+#include "aos/time/time.h"
+#include "frc971/control_loops/jacobian.h"
+
+namespace frc971::control_loops::swerve::testing {
+class SimplifiedDynamicsTest : public ::testing::Test {
+ protected:
+  using Dynamics = SimplifiedDynamics<double>;
+  using PositionState = Dynamics::PositionState<double>;
+  using States = Dynamics::States;
+  using Inputs = Dynamics::Inputs;
+  using Input = Dynamics::Input<double>;
+  using ModuleParams = Dynamics::ModuleParams;
+  using Parameters = Dynamics::Parameters;
+  static ModuleParams MakeModule(const Eigen::Vector2d &position,
+                                 bool wheel_alignment) {
+    return ModuleParams{
+        .position = position,
+        .slip_angle_coefficient = 200.0,
+        .slip_angle_alignment_coefficient = wheel_alignment ? 1.0 : 0.0,
+        .steer_motor = KrakenFOC(),
+        .drive_motor = KrakenFOC(),
+        .steer_ratio = 0.1,
+        .drive_ratio = 0.01};
+  }
+  static Parameters MakeParams(bool wheel_alignment) {
+    return {.mass = 60,
+            .moment_of_inertia = 2,
+            .modules = {
+                MakeModule({1.0, 1.0}, wheel_alignment),
+                MakeModule({-1.0, 1.0}, wheel_alignment),
+                MakeModule({-1.0, -1.0}, wheel_alignment),
+                MakeModule({1.0, -1.0}, wheel_alignment),
+            }};
+  }
+  SimplifiedDynamicsTest() : dynamics_(MakeParams(false)) {}
+
+  PositionState ValidateDynamics(const PositionState &state,
+                                 const Input &input) {
+    const PositionState Xdot = dynamics_.Dynamics(state, input);
+    // Sanity check simple invariants:
+    EXPECT_EQ(Xdot(Dynamics::kX), state(Dynamics::kVx));
+    EXPECT_EQ(Xdot(Dynamics::kY), state(Dynamics::kVy));
+    EXPECT_EQ(Xdot(Dynamics::kTheta), state(Dynamics::kOmega));
+
+    // Check that the dynamics linearization produces numbers that match numeric
+    // differentiation of the dynamics.
+    aos::monotonic_clock::time_point start_time = aos::monotonic_clock::now();
+    const auto linearized_dynamics = dynamics_.LinearizedDynamics(state, input);
+    const aos::monotonic_clock::duration auto_diff_time =
+        aos::monotonic_clock::now() - start_time;
+
+    start_time = aos::monotonic_clock::now();
+    auto numerical_A = NumericalJacobianX(
+        std::bind(&Dynamics::Dynamics<double>, &dynamics_,
+                  std::placeholders::_1, std::placeholders::_2),
+        state, input);
+    auto numerical_B = NumericalJacobianU(
+        std::bind(&Dynamics::Dynamics<double>, &dynamics_,
+                  std::placeholders::_1, std::placeholders::_2),
+        state, input);
+    const aos::monotonic_clock::duration numerical_time =
+        aos::monotonic_clock::now() - start_time;
+    VLOG(1) << "Autodifferentiation took " << auto_diff_time
+            << " while numerical differentiation took " << numerical_time;
+    EXPECT_LT((numerical_A - linearized_dynamics.first).norm(), 1e-6)
+        << "Numerical result:\n"
+        << numerical_A << "\nAuto-diff result:\n"
+        << linearized_dynamics.first;
+    EXPECT_LT((numerical_B - linearized_dynamics.second).norm(), 1e-6)
+        << "Numerical result:\n"
+        << numerical_B << "\nAuto-diff result:\n"
+        << linearized_dynamics.second;
+    return Xdot;
+  }
+
+  Dynamics dynamics_;
+};
+
+// Test that if all states and inputs are at zero that the robot won't move.
+TEST_F(SimplifiedDynamicsTest, ZeroIsZero) {
+  EXPECT_EQ(PositionState::Zero(),
+            ValidateDynamics(PositionState::Zero(), Input::Zero()));
+}
+
+// Test that if we are travelling straight forwards with zero inputs that we
+// just coast.
+TEST_F(SimplifiedDynamicsTest, CoastForwards) {
+  PositionState state = PositionState::Zero();
+  state(States::kVx) = 1.0;
+  PositionState expected = PositionState::Zero();
+  expected(States::kX) = 1.0;
+  EXPECT_EQ(expected, ValidateDynamics(state, Input::Zero()));
+}
+
+// Tests that we can accelerate the robot.
+TEST_F(SimplifiedDynamicsTest, AccelerateStraight) {
+  // Check that the drive currents behave as anticipated and accelerate the
+  // robot.
+  Input input{{0.0}, {1.0}, {0.0}, {1.0}, {0.0}, {1.0}, {0.0}, {1.0}};
+  PositionState state = PositionState::Zero();
+  state(States::kVx) = 0.0;
+  PositionState result = ValidateDynamics(state, input);
+  EXPECT_EQ(result.norm(), result(States::kVx));
+  EXPECT_LT(0.1, result(States::kVx));
+}
+
+// Test that if we are driving straight sideways (so our wheel are at 90
+// degrees) that we experience a force slowing us down.
+TEST_F(SimplifiedDynamicsTest, ForceWheelsSideways) {
+  PositionState state = PositionState::Zero();
+  state(States::kVy) = 1.0;
+  PositionState result = ValidateDynamics(state, Input::Zero());
+  EXPECT_LT(result.topRows<States::kVy>().norm(), 1e-10)
+      << ": All derivatives prior to the vy state should be ~zero.\n"
+      << result;
+  EXPECT_LT(result(States::kVy), -1.0)
+      << ": expected non-trivial deceleration.";
+  EXPECT_EQ(result(States::kOmega), 0.0);
+}
+
+// Tests that we can make the robot spin in place by orienting all the wheels
+TEST_F(SimplifiedDynamicsTest, SpinInPlaceNoSlip) {
+  PositionState state = PositionState::Zero();
+  state(States::kThetas0) = 3.0 * M_PI / 4.0;
+  state(States::kThetas1) = 5.0 * M_PI / 4.0;
+  state(States::kThetas2) = 7.0 * M_PI / 4.0;
+  state(States::kThetas3) = 1.0 * M_PI / 4.0;
+  state(States::kOmega) = 1.0;
+  PositionState result = ValidateDynamics(state, Input::Zero());
+  EXPECT_NEAR(result.norm(), 1.0, 1e-10)
+      << ": Only non-zero state should be kTheta, which should be exactly 1.0.";
+  EXPECT_EQ(result(States::kTheta), 1.0);
+
+  // Sanity check that when we then apply drive torque to the wheels that that
+  // apins the robot more.
+  Input input{{0.0}, {1.0}, {0.0}, {1.0}, {0.0}, {1.0}, {0.0}, {1.0}};
+  result = ValidateDynamics(state, input);
+  EXPECT_EQ(result(States::kTheta), 1.0);
+  EXPECT_LT(1.0, result(States::kOmega));
+  // Everything else should be ~zero.
+  result(States::kTheta) = 0.0;
+  result(States::kOmega) = 0.0;
+  EXPECT_LT(result.norm(), 1e-10);
+}
+
+// Tests that we can spin in place when skid-steering (i.e., all wheels stay
+// pointed straight, but we still attempt to spint he robot).
+TEST_F(SimplifiedDynamicsTest, SpinInPlaceSkidSteer) {
+  PositionState state = PositionState::Zero();
+  state(States::kThetas0) = -M_PI;
+  state(States::kThetas1) = -M_PI;
+  state(States::kThetas2) = 0.0;
+  state(States::kThetas3) = 0.0;
+  state(States::kOmega) = 1.0;
+  PositionState result = ValidateDynamics(state, Input::Zero());
+  EXPECT_EQ(result(States::kTheta), 1.0);
+  EXPECT_LT(result(States::kOmega), -1.0)
+      << "We should be aggressively decelerrating when slipping wheels.";
+  // Everything else should be ~zero.
+  result(States::kTheta) = 0.0;
+  result(States::kOmega) = 0.0;
+  EXPECT_LT(result.norm(), 1e-10);
+
+  // Sanity check that when we then apply drive torque to the wheels that that
+  // we can counteract the spin.
+  Input input{{0.0}, {100.0}, {0.0}, {100.0}, {0.0}, {100.0}, {0.0}, {100.0}};
+  result = ValidateDynamics(state, input);
+  EXPECT_EQ(result(States::kTheta), 1.0);
+  EXPECT_LT(1.0, result(States::kOmega));
+  // Everything else should be ~zero.
+  result(States::kTheta) = 0.0;
+  result(States::kOmega) = 0.0;
+  EXPECT_LT(result.norm(), 1e-10);
+}
+
+// Tests that we can spin in place when skid-steering backwards (ensures that
+// slip angle calculations and the such handle the sign changes correctly).
+TEST_F(SimplifiedDynamicsTest, SpinInPlaceSkidSteerBackwards) {
+  PositionState state = PositionState::Zero();
+  state(States::kThetas0) = 0.0;
+  state(States::kThetas1) = 0.0;
+  state(States::kThetas2) = M_PI;
+  state(States::kThetas3) = M_PI;
+  state(States::kOmega) = 1.0;
+  PositionState result = ValidateDynamics(state, Input::Zero());
+  EXPECT_EQ(result(States::kTheta), 1.0);
+  EXPECT_LT(result(States::kOmega), -1.0)
+      << "We should be aggressively decelerrating when slipping wheels.";
+  // Everything else should be ~zero.
+  result(States::kTheta) = 0.0;
+  result(States::kOmega) = 0.0;
+  EXPECT_LT(result.norm(), 1e-10);
+
+  // Sanity check that when we then apply drive torque to the wheels that that
+  // we can counteract the spin.
+  Input input{{0.0}, {-100.0}, {0.0}, {-100.0},
+              {0.0}, {-100.0}, {0.0}, {-100.0}};
+  result = ValidateDynamics(state, input);
+  EXPECT_EQ(result(States::kTheta), 1.0);
+  EXPECT_LT(1.0, result(States::kOmega));
+  // Everything else should be ~zero.
+  result(States::kTheta) = 0.0;
+  result(States::kOmega) = 0.0;
+  EXPECT_LT(result.norm(), 1e-10);
+}
+
+// Test that if we turn on the wheel alignment forces that it results in forces
+// that cause the wheels to align to straight over time.
+TEST_F(SimplifiedDynamicsTest, WheelAlignmentForces) {
+  dynamics_ = Dynamics(MakeParams(true));
+  PositionState state = PositionState::Zero();
+  state(States::kThetas0) = -0.1;
+  state(States::kThetas1) = -0.1;
+  state(States::kThetas2) = -0.1;
+  state(States::kThetas3) = -0.1;
+  state(States::kVx) = 1.0;
+  PositionState result = ValidateDynamics(state, Input::Zero());
+  EXPECT_LT(1e-2, result(States::kOmegas0));
+  EXPECT_LT(1e-2, result(States::kOmegas1));
+  EXPECT_LT(1e-2, result(States::kOmegas2));
+  EXPECT_LT(1e-2, result(States::kOmegas3));
+}
+
+// Do some fuzz testing of the jacobian calculations.
+TEST_F(SimplifiedDynamicsTest, Fuzz) {
+  for (size_t state_index = 0; state_index < States::kNumPositionStates;
+       ++state_index) {
+    SCOPED_TRACE(state_index);
+    PositionState state = PositionState::Zero();
+    for (const double value : {-1.0, 0.0, 1.0}) {
+      SCOPED_TRACE(value);
+      state(state_index) = value;
+      ValidateDynamics(state, Input::Zero());
+    }
+  }
+}
+}  // namespace frc971::control_loops::swerve::testing
diff --git a/frc971/control_loops/swerve/swerve_drivetrain_goal.fbs b/frc971/control_loops/swerve/swerve_drivetrain_goal.fbs
index eab0d0d..7cbb11c 100644
--- a/frc971/control_loops/swerve/swerve_drivetrain_goal.fbs
+++ b/frc971/control_loops/swerve/swerve_drivetrain_goal.fbs
@@ -1,3 +1,6 @@
+include "frc971/math/matrix.fbs";
+include "frc971/control_loops/swerve/swerve_drivetrain_joystick_goal.fbs";
+
 namespace frc971.control_loops.swerve;
 
 // States what translation control type goal we will care about
@@ -22,11 +25,20 @@
   translation_speed:double (id: 3);
 }
 
+attribute "static_length";
+
+table LinearVelocityGoal {
+  state:frc971.fbs.Matrix (id: 0);
+  input:frc971.fbs.Matrix (id: 1);
+}
+
 table Goal {
     front_left_goal:SwerveModuleGoal (id: 0);
     front_right_goal:SwerveModuleGoal (id: 1);
     back_left_goal:SwerveModuleGoal (id: 2);
     back_right_goal:SwerveModuleGoal (id: 3);
+    linear_velocity_goal:LinearVelocityGoal (id: 4);
+    joystick_goal:JoystickGoal (id: 5);
 }
 
 root_type Goal;
diff --git a/frc971/control_loops/swerve/swerve_drivetrain_joystick_goal.fbs b/frc971/control_loops/swerve/swerve_drivetrain_joystick_goal.fbs
new file mode 100644
index 0000000..bf44cdd
--- /dev/null
+++ b/frc971/control_loops/swerve/swerve_drivetrain_joystick_goal.fbs
@@ -0,0 +1,9 @@
+namespace frc971.control_loops.swerve;
+
+table JoystickGoal {
+  vx:float (id: 0);
+  vy:float (id: 1);
+  omega:float (id: 2);
+}
+
+root_type JoystickGoal;
diff --git a/frc971/control_loops/swerve/velocity_controller/BUILD b/frc971/control_loops/swerve/velocity_controller/BUILD
index 7d376bd..a0078b6 100644
--- a/frc971/control_loops/swerve/velocity_controller/BUILD
+++ b/frc971/control_loops/swerve/velocity_controller/BUILD
@@ -44,6 +44,8 @@
         "@pip//matplotlib",
         "@pip//numpy",
         "@pip//tensorflow",
+        "@pip//tensorflow_probability",
+        "@pip//tf_keras",
     ],
 )
 
@@ -59,6 +61,30 @@
 )
 
 py_binary(
+    name = "plot",
+    srcs = [
+        "model.py",
+        "plot.py",
+    ],
+    deps = [
+        ":experience_buffer",
+        ":physics",
+        "//frc971/control_loops/swerve:jax_dynamics",
+        "@pip//absl_py",
+        "@pip//flashbax",
+        "@pip//flax",
+        "@pip//jax",
+        "@pip//jaxtyping",
+        "@pip//matplotlib",
+        "@pip//numpy",
+        "@pip//pygobject",
+        "@pip//tensorflow",
+        "@pip//tensorflow_probability",
+        "@pip//tf_keras",
+    ],
+)
+
+py_binary(
     name = "lqr_plot",
     srcs = [
         "lqr_plot.py",
diff --git a/frc971/control_loops/swerve/velocity_controller/lqr_plot.py b/frc971/control_loops/swerve/velocity_controller/lqr_plot.py
index bd6674c..a63e100 100644
--- a/frc971/control_loops/swerve/velocity_controller/lqr_plot.py
+++ b/frc971/control_loops/swerve/velocity_controller/lqr_plot.py
@@ -47,9 +47,26 @@
 
 # Container for the data.
 Data = collections.namedtuple('Data', [
-    't', 'X', 'X_lqr', 'U', 'U_lqr', 'cost', 'cost_lqr', 'q1_grid', 'q2_grid',
-    'q_grid', 'target_q_grid', 'lqr_grid', 'pi_grid_U', 'lqr_grid_U', 'grid_X',
-    'grid_Y', 'reward', 'reward_lqr', 'step'
+    't',
+    'X',
+    'X_lqr',
+    'U',
+    'U_lqr',
+    'cost',
+    'cost_lqr',
+    'q1_grid',
+    'q2_grid',
+    'q_grid',
+    'target_q_grid',
+    'lqr_grid',
+    'pi_grid_U',
+    'lqr_grid_U',
+    'grid_X',
+    'grid_Y',
+    'reward',
+    'reward_lqr',
+    'step',
+    'stdev',
 ])
 
 FLAGS = absl.flags.FLAGS
@@ -87,8 +104,13 @@
     rng = jax.random.key(0)
     rng, init_rng = jax.random.split(rng)
 
-    state = create_train_state(init_rng, problem, FLAGS.q_learning_rate,
-                               FLAGS.pi_learning_rate)
+    state = create_train_state(
+        init_rng,
+        problem,
+        q_learning_rate=FLAGS.q_learning_rate,
+        pi_learning_rate=FLAGS.pi_learning_rate,
+        alpha_learning_rate=FLAGS.alpha_learning_rate,
+    )
 
     state = restore_checkpoint(state, FLAGS.workdir)
     if step is not None and state.step == step:
@@ -156,12 +178,23 @@
 
     def compute_pi_U(X, Y):
         x = jax.numpy.array([X, Y])
-        U, _, _, _ = state.pi_apply(rng,
-                                    state.params,
-                                    observation=state.problem.unwrap_angles(x),
-                                    R=goal,
-                                    deterministic=True)
-        return U[0]
+        U, _, _ = state.pi_apply(rng,
+                                 state.params,
+                                 observation=state.problem.unwrap_angles(x),
+                                 R=goal,
+                                 deterministic=True)
+        return U[0] * problem.action_limit
+
+    def compute_pi_stdev(X, Y):
+        x = jax.numpy.array([X, Y])
+        _, _, std = state.pi_apply(rng,
+                                   state.params,
+                                   observation=state.problem.unwrap_angles(x),
+                                   R=goal,
+                                   deterministic=True)
+        return std[0]
+
+    std_grid = jax.vmap(jax.vmap(compute_pi_stdev))(grid_X, grid_Y)
 
     lqr_cost_U = jax.vmap(jax.vmap(compute_lqr_U))(grid_X, grid_Y)
     pi_cost_U = jax.vmap(jax.vmap(compute_pi_U))(grid_X, grid_Y)
@@ -173,22 +206,25 @@
         X, X_lqr, data, params = val
         t = data.t.at[i].set(i * problem.dt)
 
-        U, _, _, _ = state.pi_apply(rng,
-                                    params,
-                                    observation=state.problem.unwrap_angles(X),
-                                    R=goal,
-                                    deterministic=True)
+        normalized_U, _, _ = state.pi_apply(
+            rng,
+            params,
+            observation=state.problem.unwrap_angles(X),
+            R=goal,
+            deterministic=True)
         U_lqr = problem.F @ (goal - X_lqr)
 
         cost = jax.numpy.minimum(
             state.q1_apply(params,
                            observation=state.problem.unwrap_angles(X),
                            R=goal,
-                           action=U),
+                           action=normalized_U),
             state.q2_apply(params,
                            observation=state.problem.unwrap_angles(X),
                            R=goal,
-                           action=U))
+                           action=normalized_U))
+
+        U = normalized_U * problem.action_limit
 
         U_plot = data.U.at[i, :].set(U)
         U_lqr_plot = data.U_lqr.at[i, :].set(U_lqr)
@@ -200,8 +236,9 @@
         X = problem.A @ X + problem.B @ U
         X_lqr = problem.A @ X_lqr + problem.B @ U_lqr
 
-        reward = data.reward - state.problem.cost(X, U, goal)
-        reward_lqr = data.reward_lqr - state.problem.cost(X_lqr, U_lqr, goal)
+        reward = data.reward + state.problem.reward(X, normalized_U, goal)
+        reward_lqr = data.reward_lqr + state.problem.reward(
+            X_lqr, U_lqr / problem.action_limit, goal)
 
         return X, X_lqr, data._replace(
             t=t,
@@ -242,6 +279,7 @@
             reward=0.0,
             reward_lqr=0.0,
             step=state.step,
+            stdev=std_grid,
         ), X, X_lqr, state.params)
 
     logging.info('Finished integrating, reward of %f, lqr reward of %f',
@@ -265,6 +303,7 @@
         lqr_grid_U=numpy.array(data.lqr_grid_U),
         grid_X=numpy.array(data.grid_X),
         grid_Y=numpy.array(data.grid_Y),
+        stdev=numpy.array(data.stdev),
         reward=float(data.reward),
         reward_lqr=float(data.reward_lqr),
         step=data.step,
@@ -317,9 +356,10 @@
 
         self.Ufig = pyplot.figure(figsize=pyplot.figaspect(0.5))
         self.Uax = [
-            self.Ufig.add_subplot(1, 3, 1, projection='3d'),
-            self.Ufig.add_subplot(1, 3, 2, projection='3d'),
-            self.Ufig.add_subplot(1, 3, 3, projection='3d'),
+            self.Ufig.add_subplot(2, 2, 1, projection='3d'),
+            self.Ufig.add_subplot(2, 2, 2, projection='3d'),
+            self.Ufig.add_subplot(2, 2, 3, projection='3d'),
+            self.Ufig.add_subplot(2, 2, 4, projection='3d'),
         ]
 
         self.last_trajectory_step = 0
@@ -331,6 +371,7 @@
             return
         self.last_trajectory_step = data.step
         logging.info('Updating trajectory plots')
+        self.fig0.suptitle(f'Step {data.step}')
 
         # Put data in the trajectory plots.
         self.x.set_data(data.t, data.X[:, 0])
@@ -362,6 +403,7 @@
             return
         logging.info('Updating cost plots')
         self.last_cost_step = data.step
+        self.costfig.suptitle(f'Step {data.step}')
         # Put data in the cost plots.
         if hasattr(self, 'costsurf'):
             for surf in self.costsurf:
@@ -398,6 +440,7 @@
             return
         self.last_U_step = data.step
         logging.info('Updating U plots')
+        self.Ufig.suptitle(f'Step {data.step}')
         # Put data in the controller plots.
         if hasattr(self, 'Usurf'):
             for surf in self.Usurf:
@@ -407,6 +450,7 @@
             (data.lqr_grid_U, 'lqr'),
             (data.pi_grid_U, 'pi'),
             ((data.lqr_grid_U - data.pi_grid_U), 'error'),
+            (data.stdev, 'stdev'),
         ]
 
         self.Usurf = [
diff --git a/frc971/control_loops/swerve/velocity_controller/main.py b/frc971/control_loops/swerve/velocity_controller/main.py
index a0d248d..391b0fe 100644
--- a/frc971/control_loops/swerve/velocity_controller/main.py
+++ b/frc971/control_loops/swerve/velocity_controller/main.py
@@ -39,6 +39,9 @@
 
 flags.DEFINE_string('workdir', None, 'Directory to store model data.')
 
+flags.DEFINE_bool('swerve', True,
+                  'If true, train the swerve model, otherwise do the turret.')
+
 
 def main(argv):
     if len(argv) > 1:
@@ -64,7 +67,12 @@
         FLAGS.workdir,
     )
 
-    problem = physics.TurretProblem()
+    if FLAGS.swerve:
+        physics_constants = jax_dynamics.Coefficients()
+        problem = physics.SwerveProblem(physics_constants)
+    else:
+        problem = physics.TurretProblem()
+
     state = train.train(FLAGS.workdir, problem)
 
 
diff --git a/frc971/control_loops/swerve/velocity_controller/model.py b/frc971/control_loops/swerve/velocity_controller/model.py
index 1394463..7f01eef 100644
--- a/frc971/control_loops/swerve/velocity_controller/model.py
+++ b/frc971/control_loops/swerve/velocity_controller/model.py
@@ -14,9 +14,12 @@
 from jax.experimental import mesh_utils
 from jax.sharding import Mesh, PartitionSpec, NamedSharding
 from frc971.control_loops.swerve import jax_dynamics
-from frc971.control_loops.swerve import dynamics
 from frc971.control_loops.swerve.velocity_controller import physics
 from frc971.control_loops.swerve.velocity_controller import experience_buffer
+from tensorflow_probability.substrates import jax as tfp
+
+tfd = tfp.distributions
+tfb = tfp.bijectors
 
 from flax.typing import PRNGKey
 
@@ -29,6 +32,12 @@
 )
 
 absl.flags.DEFINE_float(
+    'alpha_learning_rate',
+    default=0.004,
+    help='Training learning rate for entropy.',
+)
+
+absl.flags.DEFINE_float(
     'q_learning_rate',
     default=0.002,
     help='Training learning rate.',
@@ -52,6 +61,13 @@
     help='Fraction of --pi_learning_rate to reduce by by the end.',
 )
 
+absl.flags.DEFINE_float(
+    'target_entropy_scalar',
+    default=1.0,
+    help=
+    'Target entropy scalar for use when using automatic temperature adjustment.',
+)
+
 absl.flags.DEFINE_integer(
     'replay_size',
     default=2000000,
@@ -64,6 +80,30 @@
     help='Batch size for learning Q and pi',
 )
 
+absl.flags.DEFINE_boolean(
+    'skip_layer',
+    default=False,
+    help='If true, add skip layer connections to the Q network.',
+)
+
+absl.flags.DEFINE_boolean(
+    'rmsnorm',
+    default=False,
+    help='If true, use rmsnorm instead of layer norm.',
+)
+
+absl.flags.DEFINE_boolean(
+    'dreamer_solver',
+    default=False,
+    help='If true, use the solver from dreamer v3 instead of adam.',
+)
+
+absl.flags.DEFINE_float(
+    'initial_logalpha',
+    default=0.0,
+    help='The initial value to set logalpha to.',
+)
+
 HIDDEN_WEIGHTS = 256
 
 LOG_STD_MIN = -20
@@ -127,36 +167,25 @@
         if rng is None:
             rng = self.make_rng('pi')
 
-        # Grab a random sample
-        random_sample = jax.random.normal(rng, shape=std.shape)
+        pi_distribution = tfd.TransformedDistribution(
+            distribution=tfd.Normal(loc=mu, scale=std),
+            bijector=tfb.Tanh(),
+        )
 
         if deterministic:
             # We are testing the optimal policy, just use the mean.
-            pi_action = mu
+            pi_action = flax.linen.activation.tanh(mu)
         else:
-            # Use the reparameterization trick.  Adjust the unit gausian with
-            # something we can solve for to get the desired noise.
-            pi_action = random_sample * std + mu
+            pi_action = pi_distribution.sample(shape=(1, ), seed=rng)
 
-        logp_pi = gaussian_likelihood(random_sample, log_std)
-        # Adjustment to log prob
-        # NOTE: This formula is a little bit magic. To get an understanding of where it
-        # comes from, check out the original SAC paper (arXiv 1801.01290) and look in
-        # appendix C. This is a more numerically-stable equivalent to Eq 21.
-        delta = (2.0 * (jax.numpy.log(2.0) - pi_action -
-                        flax.linen.softplus(-2.0 * pi_action)))
+        logp_pi = pi_distribution.log_prob(pi_action)
 
-        if len(delta.shape) > 1:
-            delta = jax.numpy.sum(delta, keepdims=True, axis=1)
+        if len(logp_pi.shape) > 1:
+            logp_pi = jax.numpy.sum(logp_pi, keepdims=True, axis=1)
         else:
-            delta = jax.numpy.sum(delta, keepdims=True)
+            logp_pi = jax.numpy.sum(logp_pi, keepdims=True)
 
-        logp_pi = logp_pi - delta
-
-        # Now, saturate the action to the limit using tanh
-        pi_action = self.action_limit * flax.linen.activation.tanh(pi_action)
-
-        return pi_action, logp_pi, self.action_limit * std, random_sample
+        return pi_action, logp_pi, self.action_limit * std
 
 
 class MLPQFunction(nn.Module):
@@ -173,10 +202,23 @@
         # Estimate Q with a simple multi layer dense network.
         x = jax.numpy.hstack((observation, R, action))
         for i, hidden_size in enumerate(self.hidden_sizes):
+            # Add d2rl skip layer connections if requested.
+            # Idea from D2RL: https://arxiv.org/pdf/2010.09163.
+            if FLAGS.skip_layer and i != 0:
+                x = jax.numpy.hstack((x, observation, R, action))
+
             x = nn.Dense(
                 name=f'denselayer{i}',
                 features=hidden_size,
             )(x)
+
+            if FLAGS.rmsnorm:
+                # Idea from Dreamerv3: https://arxiv.org/pdf/2301.04104v2.
+                x = nn.RMSNorm(name=f'rmsnorm{i}')(x)
+            else:
+                # Layernorm also improves stability.
+                # Idea from RLPD: https://arxiv.org/pdf/2302.02948.
+                x = nn.LayerNorm(name=f'layernorm{i}')(x)
             x = self.activation(x)
 
         x = nn.Dense(
@@ -369,7 +411,7 @@
             q_opt_state=q_opt_state,
             alpha_tx=alpha_tx,
             alpha_opt_state=alpha_opt_state,
-            target_entropy=-problem.num_states,
+            target_entropy=-problem.num_outputs * FLAGS.target_entropy_scalar,
             mesh=mesh,
             sharding=sharding,
             replicated_sharding=replicated_sharding,
@@ -380,14 +422,15 @@
 
 
 def create_train_state(rng: PRNGKey, problem: Problem, q_learning_rate,
-                       pi_learning_rate):
+                       pi_learning_rate, alpha_learning_rate):
     """Creates initial `TrainState`."""
-    pi = SquashedGaussianMLPActor(activation=nn.activation.gelu,
+    pi = SquashedGaussianMLPActor(activation=nn.activation.silu,
                                   action_space=problem.num_outputs,
                                   action_limit=problem.action_limit)
     # We want q1 and q2 to have different network architectures so they pick up differnet things.
-    q1 = MLPQFunction(activation=nn.activation.gelu, hidden_sizes=[128, 256])
-    q2 = MLPQFunction(activation=nn.activation.gelu, hidden_sizes=[256, 128])
+    # SiLu is used in DreamerV3 so we use it: https://arxiv.org/pdf/2301.04104v2.
+    q1 = MLPQFunction(activation=nn.activation.silu, hidden_sizes=[128, 256])
+    q2 = MLPQFunction(activation=nn.activation.silu, hidden_sizes=[256, 128])
 
     @jax.jit
     def init_params(rng):
@@ -412,7 +455,7 @@
         )['params']
 
         if FLAGS.alpha < 0.0:
-            logalpha = 0.0
+            logalpha = FLAGS.initial_logalpha
         else:
             logalpha = jax.numpy.log(FLAGS.alpha)
 
@@ -423,9 +466,14 @@
             'logalpha': logalpha,
         }
 
-    pi_tx = optax.sgd(learning_rate=pi_learning_rate)
-    q_tx = optax.sgd(learning_rate=q_learning_rate)
-    alpha_tx = optax.sgd(learning_rate=q_learning_rate)
+    if FLAGS.dreamer_solver:
+        pi_tx = create_dreamer_solver(learning_rate=pi_learning_rate)
+        q_tx = create_dreamer_solver(learning_rate=q_learning_rate)
+        alpha_tx = create_dreamer_solver(learning_rate=alpha_learning_rate)
+    else:
+        pi_tx = optax.adam(learning_rate=pi_learning_rate)
+        q_tx = optax.adam(learning_rate=q_learning_rate)
+        alpha_tx = optax.adam(learning_rate=alpha_learning_rate)
 
     result = TrainState.create(
         problem=problem,
@@ -441,6 +489,89 @@
     return result
 
 
+# Solver from dreamer v3: https://arxiv.org/pdf/2301.04104v2.
+# TODO(austin): How many of these pieces are actually in optax already?
+def scale_by_rms(beta=0.999, eps=1e-8):
+
+    def init_fn(params):
+        nu = jax.tree_util.tree_map(
+            lambda t: jax.numpy.zeros_like(t, jax.numpy.float32), params)
+        step = jax.numpy.zeros((), jax.numpy.int32)
+        return (step, nu)
+
+    def update_fn(updates, state, params=None):
+        step, nu = state
+        step = optax.safe_int32_increment(step)
+        nu = jax.tree_util.tree_map(
+            lambda v, u: beta * v + (1 - beta) * (u * u), nu, updates)
+        nu_hat = optax.bias_correction(nu, beta, step)
+        updates = jax.tree_util.tree_map(
+            lambda u, v: u / (jax.numpy.sqrt(v) + eps), updates, nu_hat)
+        return updates, (step, nu)
+
+    return optax.GradientTransformation(init_fn, update_fn)
+
+
+def scale_by_agc(clip=0.03, pmin=1e-3):
+
+    def init_fn(params):
+        return ()
+
+    def update_fn(updates, state, params=None):
+
+        def fn(param, update):
+            unorm = jax.numpy.linalg.norm(update.flatten(), 2)
+            pnorm = jax.numpy.linalg.norm(param.flatten(), 2)
+            upper = clip * jax.numpy.maximum(pmin, pnorm)
+            return update * (1 / jax.numpy.maximum(1.0, unorm / upper))
+
+        updates = jax.tree_util.tree_map(fn, params, updates)
+        return updates, ()
+
+    return optax.GradientTransformation(init_fn, update_fn)
+
+
+def scale_by_momentum(beta=0.9, nesterov=False):
+
+    def init_fn(params):
+        mu = jax.tree_util.tree_map(
+            lambda t: jax.numpy.zeros_like(t, jax.numpy.float32), params)
+        step = jax.numpy.zeros((), jax.numpy.int32)
+        return (step, mu)
+
+    def update_fn(updates, state, params=None):
+        step, mu = state
+        step = optax.safe_int32_increment(step)
+        mu = optax.update_moment(updates, mu, beta, 1)
+        if nesterov:
+            mu_nesterov = optax.update_moment(updates, mu, beta, 1)
+            mu_hat = optax.bias_correction(mu_nesterov, beta, step)
+        else:
+            mu_hat = optax.bias_correction(mu, beta, step)
+        return mu_hat, (step, mu)
+
+    return optax.GradientTransformation(init_fn, update_fn)
+
+
+def create_dreamer_solver(
+    learning_rate,
+    agc: float = 0.3,
+    pmin: float = 1e-3,
+    beta1: float = 0.9,
+    beta2: float = 0.999,
+    eps: float = 1e-20,
+    nesterov: bool = False,
+) -> optax.base.GradientTransformation:
+    # From dreamer v3.
+    return optax.chain(
+        # Adaptive gradient clipping.
+        scale_by_agc(agc, pmin),
+        scale_by_rms(beta2, eps),
+        scale_by_momentum(beta1, nesterov),
+        optax.scale_by_learning_rate(learning_rate),
+    )
+
+
 def create_learning_rate_fn(
     base_learning_rate: float,
     final_learning_rate: float,
diff --git a/frc971/control_loops/swerve/velocity_controller/physics.py b/frc971/control_loops/swerve/velocity_controller/physics.py
index b0da0b4..accd0d6 100644
--- a/frc971/control_loops/swerve/velocity_controller/physics.py
+++ b/frc971/control_loops/swerve/velocity_controller/physics.py
@@ -55,13 +55,17 @@
     def random_states(self, rng: PRNGKey, dimensions=None):
         raise NotImplemented("random_states not implemented")
 
-    def random_actions(self, rng: PRNGKey, dimensions=None):
+    def random_actions(self,
+                       rng: PRNGKey,
+                       X: jax.typing.ArrayLike,
+                       goal: jax.typing.ArrayLike,
+                       dimensions=None):
         """Produces a uniformly random action in the action space."""
         return jax.random.uniform(
             rng,
             (dimensions or FLAGS.num_agents, self.num_outputs),
-            minval=-self.action_limit,
-            maxval=self.action_limit,
+            minval=-1.0,
+            maxval=1.0,
         )
 
     def random_goals(self, rng: PRNGKey, dimensions=None):
@@ -94,12 +98,13 @@
         A_continuous = jax.numpy.array([[0., 1.], [0., -36.85154548]])
         B_continuous = jax.numpy.array([[0.], [56.08534375]])
 
+        U = U * self.action_limit
         return A_continuous @ X + B_continuous @ U
 
-    def cost(self, X: jax.typing.ArrayLike, U: jax.typing.ArrayLike,
-             goal: jax.typing.ArrayLike):
-        return (X - goal).T @ jax.numpy.array(
-            self.Q) @ (X - goal) + U.T @ jax.numpy.array(self.R) @ U
+    def reward(self, X: jax.typing.ArrayLike, U: jax.typing.ArrayLike,
+               goal: jax.typing.ArrayLike):
+        return -(X - goal).T @ jax.numpy.array(
+            self.Q) @ (X - goal) - U.T @ jax.numpy.array(self.R) @ U
 
     def random_states(self, rng: PRNGKey, dimensions=None):
         rng1, rng2 = jax.random.split(rng)
@@ -120,3 +125,123 @@
                                maxval=0.1),
             jax.numpy.zeros((dimensions or FLAGS.num_agents, 1)),
         ))
+
+
+class SwerveProblem(Problem):
+
+    def __init__(self, coefficients: jax_dynamics.CoefficientsType):
+        super().__init__(num_states=jax_dynamics.NUM_VELOCITY_STATES,
+                         num_unwrapped_states=17,
+                         num_outputs=8,
+                         num_goals=3,
+                         action_limit=40.0)
+
+        self.coefficients = coefficients
+
+    def random_actions(self,
+                       rng: PRNGKey,
+                       X: jax.typing.ArrayLike,
+                       goal: jax.typing.ArrayLike,
+                       dimensions=None):
+        """Produces a uniformly random action in the action space."""
+        return jax.random.uniform(
+            rng,
+            (dimensions or FLAGS.num_agents, self.num_outputs),
+            minval=-1.0,
+            maxval=1.0,
+        )
+
+    def unwrap_angles(self, X: jax.typing.ArrayLike):
+        return jax.numpy.stack([
+            jax.numpy.cos(X[..., jax_dynamics.VELOCITY_STATE_THETAS0]),
+            jax.numpy.sin(X[..., jax_dynamics.VELOCITY_STATE_THETAS0]),
+            X[..., jax_dynamics.VELOCITY_STATE_OMEGAS0],
+            jax.numpy.cos(X[..., jax_dynamics.VELOCITY_STATE_THETAS1]),
+            jax.numpy.sin(X[..., jax_dynamics.VELOCITY_STATE_THETAS1]),
+            X[..., jax_dynamics.VELOCITY_STATE_OMEGAS1],
+            jax.numpy.cos(X[..., jax_dynamics.VELOCITY_STATE_THETAS2]),
+            jax.numpy.sin(X[..., jax_dynamics.VELOCITY_STATE_THETAS2]),
+            X[..., jax_dynamics.VELOCITY_STATE_OMEGAS2],
+            jax.numpy.cos(X[..., jax_dynamics.VELOCITY_STATE_THETAS3]),
+            jax.numpy.sin(X[..., jax_dynamics.VELOCITY_STATE_THETAS3]),
+            X[..., jax_dynamics.VELOCITY_STATE_OMEGAS3],
+            jax.numpy.cos(X[..., jax_dynamics.VELOCITY_STATE_THETA]),
+            jax.numpy.sin(X[..., jax_dynamics.VELOCITY_STATE_THETA]),
+            X[..., jax_dynamics.VELOCITY_STATE_VX],
+            X[..., jax_dynamics.VELOCITY_STATE_VY],
+            X[..., jax_dynamics.VELOCITY_STATE_OMEGA],
+        ],
+                               axis=-1)
+
+    def xdot(self, X: jax.typing.ArrayLike, U: jax.typing.ArrayLike):
+        return jax_dynamics.velocity_dynamics(self.coefficients, X,
+                                              self.action_limit * U)
+
+    def reward(self, X: jax.typing.ArrayLike, U: jax.typing.ArrayLike,
+               goal: jax.typing.ArrayLike):
+        return -jax_dynamics.mpc_cost(coefficients=self.coefficients,
+                                      X=X,
+                                      U=self.action_limit * U,
+                                      goal=goal)
+
+    def random_states(self, rng: PRNGKey, dimensions=None):
+        rng, rng1, rng2, rng3, rng4, rng5, rng6, rng7, rng8, rng9, rng10, rng11 = jax.random.split(
+            rng, num=12)
+
+        return jax.numpy.hstack((
+            # VELOCITY_STATE_THETAS0 = 0
+            self._random_angle(rng1, dimensions),
+            # VELOCITY_STATE_OMEGAS0 = 1
+            self._random_module_velocity(rng2, dimensions),
+            # VELOCITY_STATE_THETAS1 = 2
+            self._random_angle(rng3, dimensions),
+            # VELOCITY_STATE_OMEGAS1 = 3
+            self._random_module_velocity(rng4, dimensions),
+            # VELOCITY_STATE_THETAS2 = 4
+            self._random_angle(rng5, dimensions),
+            # VELOCITY_STATE_OMEGAS2 = 5
+            self._random_module_velocity(rng6, dimensions),
+            # VELOCITY_STATE_THETAS3 = 6
+            self._random_angle(rng7, dimensions),
+            # VELOCITY_STATE_OMEGAS3 = 7
+            self._random_module_velocity(rng8, dimensions),
+            # VELOCITY_STATE_THETA = 8
+            self._random_angle(rng9, dimensions),
+            # VELOCITY_STATE_VX = 9
+            # VELOCITY_STATE_VY = 10
+            self._random_robot_velocity(rng10, dimensions),
+            # VELOCITY_STATE_OMEGA = 11
+            self._random_robot_angular_velocity(rng11, dimensions),
+        ))
+
+    def random_goals(self, rng: PRNGKey, dimensions=None):
+        """Produces a random goal in the goal space."""
+        return jax.numpy.hstack((
+            jax.random.uniform(rng, (dimensions or FLAGS.num_agents, 1),
+                               minval=1.0,
+                               maxval=1.0),
+            jax.numpy.zeros((dimensions or FLAGS.num_agents, 2)),
+        ))
+
+    MODULE_VELOCITY = 1.0
+    ROBOT_ANGULAR_VELOCITY = 0.5
+
+    def _random_angle(self, rng: PRNGKey, dimensions=None):
+        return jax.random.uniform(rng, (dimensions or FLAGS.num_agents, 1),
+                                  minval=-0.1 * jax.numpy.pi,
+                                  maxval=0.1 * jax.numpy.pi)
+
+    def _random_module_velocity(self, rng: PRNGKey, dimensions=None):
+        return jax.random.uniform(rng, (dimensions or FLAGS.num_agents, 1),
+                                  minval=-self.MODULE_VELOCITY,
+                                  maxval=self.MODULE_VELOCITY)
+
+    def _random_robot_velocity(self, rng: PRNGKey, dimensions=None):
+        return jax.random.uniform(rng, (dimensions or FLAGS.num_agents, 2),
+                                  minval=0.9,
+                                  maxval=1.1)
+
+    def _random_robot_angular_velocity(self, rng: PRNGKey, dimensions=None):
+        return jax.random.uniform(rng, (dimensions or FLAGS.num_agents, 1),
+                                  minval=-self.ROBOT_ANGULAR_VELOCITY,
+                                  maxval=self.ROBOT_ANGULAR_VELOCITY)
diff --git a/frc971/control_loops/swerve/velocity_controller/plot.py b/frc971/control_loops/swerve/velocity_controller/plot.py
new file mode 100644
index 0000000..502a3e2
--- /dev/null
+++ b/frc971/control_loops/swerve/velocity_controller/plot.py
@@ -0,0 +1,430 @@
+#!/usr/bin/env python3
+
+import os
+
+os.environ['XLA_FLAGS'] = ' '.join([
+    # Teach it where to find CUDA
+    '--xla_gpu_cuda_data_dir=/usr/lib/cuda',
+    # Use up to 20 cores
+    '--xla_force_host_platform_device_count=20',
+    # Dump XLA to /tmp/foo to aid debugging
+    #'--xla_dump_to=/tmp/foo',
+    #'--xla_gpu_enable_command_buffer='
+])
+
+os.environ['JAX_PLATFORMS'] = 'cpu'
+# Don't pre-allocate memory
+os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
+
+import absl
+from absl import logging
+from matplotlib.animation import FuncAnimation
+import matplotlib
+import numpy
+import scipy
+import time
+
+matplotlib.use("gtk3agg")
+
+from matplotlib import pylab
+from matplotlib import pyplot
+from flax.training import checkpoints
+import tensorflow as tf
+from jax.experimental.ode import odeint
+import threading
+import collections
+
+from frc971.control_loops.swerve.velocity_controller.model import *
+from frc971.control_loops.swerve.velocity_controller.physics import *
+
+# Container for the data.
+Data = collections.namedtuple('Data', [
+    't', 'X', 'U', 'logp_pi', 'cost', 'q1_grid', 'q2_grid', 'q_grid',
+    'target_q_grid', 'pi_grid_U', 'grid_X', 'grid_Y', 'reward', 'step',
+    'rewards'
+])
+
+FLAGS = absl.flags.FLAGS
+
+absl.flags.DEFINE_string('workdir', None, 'Directory to store model data.')
+
+absl.flags.DEFINE_integer(
+    'horizon',
+    default=100,
+    help='MPC horizon',
+)
+
+numpy.set_printoptions(linewidth=200, )
+
+absl.flags.DEFINE_float(
+    'alpha',
+    default=0.2,
+    help='Entropy.  If negative, automatically solve for it.',
+)
+
+
+def restore_checkpoint(state: TrainState, workdir: str):
+    return checkpoints.restore_checkpoint(workdir, state)
+
+
+dt = 0.005
+
+
+def X0Full():
+    module_theta = 0.0
+    module_omega = 0.0
+    theta = 0.0
+    vx = 1.0
+    dtheta = 0.05
+    vy = 0.0
+    drive_omega = jax.numpy.hypot(vx, vy) / jax_dynamics.WHEEL_RADIUS
+    omega = 0.0
+    return jax.numpy.array([
+        module_theta + dtheta,
+        0.0,
+        module_omega,
+        drive_omega,
+        module_theta + dtheta,
+        0.0,
+        module_omega,
+        drive_omega,
+        module_theta - dtheta,
+        0.0,
+        module_omega,
+        drive_omega,
+        module_theta - dtheta,
+        0.0,
+        module_omega,
+        drive_omega,
+        0.0,
+        0.0,
+        theta,
+        vx,
+        vy,
+        omega,
+        0.0,
+        0.0,
+        0.0,
+    ])
+
+
+def generate_data(step=None):
+    grid_X = numpy.arange(-1, 1, 0.1)
+    grid_Y = numpy.arange(-10, 10, 0.1)
+    grid_X, grid_Y = numpy.meshgrid(grid_X, grid_Y)
+    grid_X = jax.numpy.array(grid_X)
+    grid_Y = jax.numpy.array(grid_Y)
+    # Load the training state.
+    physics_constants = jax_dynamics.Coefficients()
+    problem = physics.SwerveProblem(physics_constants)
+
+    rng = jax.random.key(0)
+    rng, init_rng = jax.random.split(rng)
+
+    state = create_train_state(init_rng,
+                               problem,
+                               FLAGS.q_learning_rate,
+                               FLAGS.pi_learning_rate,
+                               alpha_learning_rate=0.001)
+
+    state = restore_checkpoint(state, FLAGS.workdir)
+
+    X = X0Full()
+    X_lqr = X.copy()
+    goal = jax.numpy.array([1.0, 0.0, 0.0])
+
+    logging.info('X: %s', X)
+    logging.info('goal: %s', goal)
+    logging.debug('params: %s', state.params)
+
+    # Now simulate the robot, accumulating up things to plot.
+    def loop(i, val):
+        X, data, params = val
+        t = data.t.at[i].set(i * problem.dt)
+
+        U, logp_pi, std = state.pi_apply(
+            rng,
+            params,
+            observation=state.problem.unwrap_angles(
+                jax_dynamics.to_velocity_state(X)),
+            R=goal,
+            deterministic=True)
+
+        logp_pi = logp_pi * jax.numpy.exp(params['logalpha'])
+
+        jax.debug.print('mu: {mu} std: {std}', mu=U, std=std)
+
+        step_reward = state.problem.q_reward(jax_dynamics.to_velocity_state(X),
+                                             U, goal)
+        reward = data.reward + step_reward
+
+        cost = jax.numpy.minimum(
+            state.q1_apply(params,
+                           observation=state.problem.unwrap_angles(
+                               jax_dynamics.to_velocity_state(X)),
+                           R=goal,
+                           action=U),
+            state.q2_apply(params,
+                           observation=state.problem.unwrap_angles(
+                               jax_dynamics.to_velocity_state(X)),
+                           R=goal,
+                           action=U))
+
+        U = U * problem.action_limit
+        U_plot = data.U.at[i, :].set(U)
+        rewards = data.rewards.at[i, :].set(step_reward)
+        X_plot = data.X.at[i, :].set(X)
+        cost_plot = data.cost.at[i, :].set(cost)
+        logp_pi_plot = data.logp_pi.at[i, :].set(logp_pi)
+
+        # TODO(austin): I'd really like to visualize the slip angle per wheel.
+        # Maybe also the force deviation, etc.
+        # I think that would help enormously in figuring out how good a specific solution is.
+
+        def fn(X, t):
+            return jax_dynamics.full_dynamics(problem.coefficients, X,
+                                              U).flatten()
+
+        X = odeint(fn, X, jax.numpy.array([0, problem.dt]))
+
+        return X[1, :], data._replace(
+            t=t,
+            U=U_plot,
+            X=X_plot,
+            logp_pi=logp_pi_plot,
+            cost=cost_plot,
+            reward=reward,
+            rewards=rewards,
+        ), params
+
+    # Do it.
+    @jax.jit
+    def integrate(data, X, params):
+        return jax.lax.fori_loop(0, FLAGS.horizon, loop, (X, data, params))
+
+    X, data, params = integrate(
+        Data(
+            t=jax.numpy.zeros((FLAGS.horizon, )),
+            X=jax.numpy.zeros((FLAGS.horizon, jax_dynamics.NUM_STATES)),
+            U=jax.numpy.zeros((FLAGS.horizon, state.problem.num_outputs)),
+            logp_pi=jax.numpy.zeros((FLAGS.horizon, 1)),
+            rewards=jax.numpy.zeros((FLAGS.horizon, 1)),
+            cost=jax.numpy.zeros((FLAGS.horizon, 1)),
+            q1_grid=jax.numpy.zeros(grid_X.shape),
+            q2_grid=jax.numpy.zeros(grid_X.shape),
+            q_grid=jax.numpy.zeros(grid_X.shape),
+            target_q_grid=jax.numpy.zeros(grid_X.shape),
+            pi_grid_U=jax.numpy.zeros(grid_X.shape),
+            grid_X=grid_X,
+            grid_Y=grid_Y,
+            reward=0.0,
+            step=state.step,
+        ), X, state.params)
+
+    logging.info('Reward: %s', float(data.reward))
+
+    # Convert back to numpy for plotting.
+    return Data(
+        t=numpy.array(data.t),
+        X=numpy.array(data.X),
+        U=numpy.array(data.U),
+        logp_pi=numpy.array(data.logp_pi),
+        cost=numpy.array(data.cost),
+        q1_grid=numpy.array(data.q1_grid),
+        q2_grid=numpy.array(data.q2_grid),
+        q_grid=numpy.array(data.q_grid),
+        target_q_grid=numpy.array(data.target_q_grid),
+        pi_grid_U=numpy.array(data.pi_grid_U),
+        grid_X=numpy.array(data.grid_X),
+        grid_Y=numpy.array(data.grid_Y),
+        rewards=numpy.array(data.rewards),
+        reward=float(data.reward),
+        step=data.step,
+    )
+
+
+class Plotter(object):
+
+    def __init__(self, data):
+        # Make all the plots and axis.
+        self.fig0, self.axs0 = pylab.subplots(3)
+        self.fig0.supxlabel('Seconds')
+
+        self.vx, = self.axs0[0].plot([], [], label="vx")
+        self.vy, = self.axs0[0].plot([], [], label="vy")
+        self.omega, = self.axs0[0].plot([], [], label="omega")
+        self.axs0[0].set_ylabel('Velocity')
+        self.axs0[0].legend()
+        self.axs0[0].grid()
+
+        self.steer0, = self.axs0[1].plot([], [], label="Steer0")
+        self.steer1, = self.axs0[1].plot([], [], label="Steer1")
+        self.steer2, = self.axs0[1].plot([], [], label="Steer2")
+        self.steer3, = self.axs0[1].plot([], [], label="Steer3")
+        self.axs0[1].set_ylabel('Amps')
+        self.axs0[1].legend()
+        self.axs0[1].grid()
+
+        self.drive0, = self.axs0[2].plot([], [], label="Drive0")
+        self.drive1, = self.axs0[2].plot([], [], label="Drive1")
+        self.drive2, = self.axs0[2].plot([], [], label="Drive2")
+        self.drive3, = self.axs0[2].plot([], [], label="Drive3")
+        self.axs0[2].set_ylabel('Amps')
+        self.axs0[2].legend()
+        self.axs0[2].grid()
+
+        self.fig1, self.axs1 = pylab.subplots(3)
+        self.fig1.supxlabel('Seconds')
+
+        self.theta0, = self.axs1[0].plot([], [], label='steer position0')
+        self.theta1, = self.axs1[0].plot([], [], label='steer position1')
+        self.theta2, = self.axs1[0].plot([], [], label='steer position2')
+        self.theta3, = self.axs1[0].plot([], [], label='steer position3')
+        self.axs1[0].set_ylabel('Radians')
+        self.axs1[0].legend()
+        self.omega0, = self.axs1[1].plot([], [], label='steer velocity0')
+        self.omega1, = self.axs1[1].plot([], [], label='steer velocity1')
+        self.omega2, = self.axs1[1].plot([], [], label='steer velocity2')
+        self.omega3, = self.axs1[1].plot([], [], label='steer velocity3')
+        self.axs1[1].set_ylabel('Radians/second')
+        self.axs1[1].legend()
+
+        self.logp_axis = self.axs1[2].twinx()
+        self.cost, = self.axs1[2].plot([], [], label='cost')
+        self.reward, = self.axs1[2].plot([], [], label='reward')
+        self.axs1[2].set_ylabel('Radians/second')
+        self.axs1[2].legend()
+
+        self.logp_pi, = self.logp_axis.plot([], [],
+                                            label='logp_pi',
+                                            color='C2')
+        self.logp_axis.set_ylabel('log(liklihood)*alpha')
+        self.logp_axis.legend()
+
+        self.last_robot_step = 0
+        self.last_steer_step = 0
+
+    def update_robot_plot(self, data):
+        if data.step is not None and data.step == self.last_robot_step:
+            return
+        self.last_robot_step = data.step
+        logging.info('Updating robot plots')
+        self.fig0.suptitle(f'Step {data.step}')
+
+        self.vx.set_data(data.t, data.X[:, jax_dynamics.STATE_VX])
+        self.vy.set_data(data.t, data.X[:, jax_dynamics.STATE_VY])
+        self.omega.set_data(data.t, data.X[:, jax_dynamics.STATE_OMEGA])
+
+        self.axs0[0].relim()
+        self.axs0[0].autoscale_view()
+
+        self.steer0.set_data(data.t, data.U[:, 0])
+        self.steer1.set_data(data.t, data.U[:, 2])
+        self.steer2.set_data(data.t, data.U[:, 4])
+        self.steer3.set_data(data.t, data.U[:, 6])
+        self.axs0[1].relim()
+        self.axs0[1].autoscale_view()
+
+        self.drive0.set_data(data.t, data.U[:, 1])
+        self.drive1.set_data(data.t, data.U[:, 3])
+        self.drive2.set_data(data.t, data.U[:, 5])
+        self.drive3.set_data(data.t, data.U[:, 7])
+        self.axs0[2].relim()
+        self.axs0[2].autoscale_view()
+
+        return (self.vx, self.vy, self.omega, self.steer0, self.steer1,
+                self.steer2, self.steer3, self.drive0, self.drive1,
+                self.drive2, self.drive3)
+
+    def update_steer_plot(self, data):
+        if data.step == self.last_steer_step:
+            return
+        self.last_steer_step = data.step
+        logging.info('Updating steer plots')
+        self.fig1.suptitle(f'Step {data.step}')
+
+        self.theta0.set_data(data.t, data.X[:, jax_dynamics.STATE_THETAS0])
+        self.theta1.set_data(data.t, data.X[:, jax_dynamics.STATE_THETAS1])
+        self.theta2.set_data(data.t, data.X[:, jax_dynamics.STATE_THETAS2])
+        self.theta3.set_data(data.t, data.X[:, jax_dynamics.STATE_THETAS3])
+        self.axs1[0].relim()
+        self.axs1[0].autoscale_view()
+
+        self.omega0.set_data(data.t, data.X[:, jax_dynamics.STATE_OMEGAS0])
+        self.omega1.set_data(data.t, data.X[:, jax_dynamics.STATE_OMEGAS1])
+        self.omega2.set_data(data.t, data.X[:, jax_dynamics.STATE_OMEGAS2])
+        self.omega3.set_data(data.t, data.X[:, jax_dynamics.STATE_OMEGAS3])
+        self.axs1[1].relim()
+        self.axs1[1].autoscale_view()
+
+        self.cost.set_data(data.t, data.cost)
+        self.reward.set_data(data.t, data.rewards)
+        self.logp_pi.set_data(data.t, data.logp_pi)
+        self.axs1[2].relim()
+        self.axs1[2].autoscale_view()
+        self.logp_axis.relim()
+        self.logp_axis.autoscale_view()
+
+        return (self.theta0, self.theta1, self.theta2, self.theta3,
+                self.omega0, self.omega1, self.omega2, self.omega3, self.cost,
+                self.logp_pi, self.reward)
+
+
+def main(argv):
+    if len(argv) > 1:
+        raise absl.app.UsageError('Too many command-line arguments.')
+
+    tf.config.experimental.set_visible_devices([], 'GPU')
+
+    lock = threading.Lock()
+
+    # Load data.
+    data = generate_data()
+
+    plotter = Plotter(data)
+
+    # Event for shutting down the thread.
+    shutdown = threading.Event()
+
+    # Thread to grab new data periodically.
+    def do_update():
+        while True:
+            nonlocal data
+
+            my_data = generate_data(data.step)
+
+            if my_data is not None:
+                with lock:
+                    data = my_data
+
+            if shutdown.wait(timeout=3):
+                return
+
+    update_thread = threading.Thread(target=do_update)
+    update_thread.start()
+
+    # Now, update each of the plots every second with the new data.
+    def update0(frame):
+        with lock:
+            my_data = data
+
+        return plotter.update_robot_plot(my_data)
+
+    def update1(frame):
+        with lock:
+            my_data = data
+
+        return plotter.update_steer_plot(my_data)
+
+    animation0 = FuncAnimation(plotter.fig0, update0, interval=1000)
+    animation1 = FuncAnimation(plotter.fig1, update1, interval=1000)
+
+    pyplot.show()
+
+    shutdown.set()
+    update_thread.join()
+
+
+if __name__ == '__main__':
+    absl.flags.mark_flags_as_required(['workdir'])
+    absl.app.run(main)
diff --git a/frc971/control_loops/swerve/velocity_controller/train.py b/frc971/control_loops/swerve/velocity_controller/train.py
index cf54bc4..34d904a 100644
--- a/frc971/control_loops/swerve/velocity_controller/train.py
+++ b/frc971/control_loops/swerve/velocity_controller/train.py
@@ -1,3 +1,6 @@
+# Machine learning based on Soft Actor Critic(SAC) which was initially proposed in https://arxiv.org/pdf/1801.01290.
+# Our implementation was heavily based on OpenAI's spinning up reference implementation https://spinningup.openai.com/en/latest/algorithms/sac.html.
+
 import absl
 import time
 import collections
@@ -35,7 +38,7 @@
 )
 
 absl.flags.DEFINE_integer(
-    'start_steps',
+    'random_sample_steps',
     default=10000,
     help='Number of steps to randomly sample before using the policy',
 )
@@ -76,6 +79,13 @@
     help='If true, explode on any NaNs found, and print them.',
 )
 
+absl.flags.DEFINE_bool(
+    'maximum_entropy_q',
+    default=True,
+    help=
+    'If false, do not add the maximum entropy term to the bellman backup for Q.',
+)
+
 
 def save_checkpoint(state: TrainState, workdir: str):
     """Saves a checkpoint in the workdir."""
@@ -123,10 +133,10 @@
     R = data['goals']
 
     # Compute the ending actions from the current network.
-    action2, logp_pi2, _, _ = state.pi_apply(rng=rng,
-                                             params=params,
-                                             observation=observations2,
-                                             R=R)
+    action2, logp_pi2, _ = state.pi_apply(rng=rng,
+                                          params=params,
+                                          observation=observations2,
+                                          R=R)
 
     # Compute target network Q values
     q1_pi_target = state.q1_apply(state.target_params,
@@ -142,8 +152,13 @@
     alpha = jax.numpy.exp(params['logalpha'])
 
     # Now we can compute the Bellman backup
-    bellman_backup = jax.lax.stop_gradient(rewards + FLAGS.gamma *
-                                           (q_pi_target - alpha * logp_pi2))
+    # Max entropy SAC is based on https://arxiv.org/pdf/1812.05905.
+    if FLAGS.maximum_entropy_q:
+        bellman_backup = jax.lax.stop_gradient(
+            rewards + FLAGS.gamma * (q_pi_target - alpha * logp_pi2))
+    else:
+        bellman_backup = jax.lax.stop_gradient(rewards +
+                                               FLAGS.gamma * q_pi_target)
 
     # Compute the starting Q values from the Q network being optimized.
     q1 = state.q1_apply(params, observation=observations1, R=R, action=actions)
@@ -156,19 +171,6 @@
 
 
 @jax.jit
-def compute_batched_loss_q(state: TrainState, rng: PRNGKey, params,
-                           data: ArrayLike):
-
-    def bound_compute_loss_q(rng, data):
-        return compute_loss_q(state, rng, params, data)
-
-    return jax.vmap(bound_compute_loss_q)(
-        jax.random.split(rng, FLAGS.num_agents),
-        data,
-    ).mean()
-
-
-@jax.jit
 def compute_loss_pi(state: TrainState, rng: PRNGKey, params, data: ArrayLike):
     """Computes the Soft Actor-Critic loss for pi."""
     observations1 = data['observations1']
@@ -176,10 +178,10 @@
     # TODO(austin): We've got differentiable policy and differentiable physics.  Can we use those here?  Have Q learn the future, not the current step?
 
     # Compute the action
-    pi, logp_pi, _, _ = state.pi_apply(rng=rng,
-                                       params=params,
-                                       observation=observations1,
-                                       R=R)
+    pi, logp_pi, _ = state.pi_apply(rng=rng,
+                                    params=params,
+                                    observation=observations1,
+                                    R=R)
     q1_pi = state.q1_apply(jax.lax.stop_gradient(params),
                            observation=observations1,
                            R=R,
@@ -199,6 +201,32 @@
 
 
 @jax.jit
+def compute_loss_alpha(state: TrainState, rng: PRNGKey, params,
+                       data: ArrayLike):
+    """Computes the Soft Actor-Critic loss for alpha."""
+    observations1 = data['observations1']
+    R = data['goals']
+    pi, logp_pi, _ = jax.lax.stop_gradient(
+        state.pi_apply(rng=rng, params=params, observation=observations1, R=R))
+
+    return (-jax.numpy.exp(params['logalpha']) *
+            (logp_pi + state.target_entropy)).mean(), logp_pi.mean()
+
+
+@jax.jit
+def compute_batched_loss_q(state: TrainState, rng: PRNGKey, params,
+                           data: ArrayLike):
+
+    def bound_compute_loss_q(rng, data):
+        return compute_loss_q(state, rng, params, data)
+
+    return jax.vmap(bound_compute_loss_q)(
+        jax.random.split(rng, FLAGS.num_agents),
+        data,
+    ).mean()
+
+
+@jax.jit
 def compute_batched_loss_pi(state: TrainState, rng: PRNGKey, params,
                             data: ArrayLike):
 
@@ -212,33 +240,21 @@
 
 
 @jax.jit
-def compute_loss_alpha(state: TrainState, rng: PRNGKey, params,
-                       data: ArrayLike):
-    """Computes the Soft Actor-Critic loss for alpha."""
-    observations1 = data['observations1']
-    R = data['goals']
-    pi, logp_pi, _, _ = jax.lax.stop_gradient(
-        state.pi_apply(rng=rng, params=params, R=R, observation=observations1))
-
-    return (-jax.numpy.exp(params['logalpha']) *
-            (logp_pi + state.target_entropy)).mean()
-
-
-@jax.jit
 def compute_batched_loss_alpha(state: TrainState, rng: PRNGKey, params,
                                data: ArrayLike):
 
     def bound_compute_loss_alpha(rng, data):
         return compute_loss_alpha(state, rng, params, data)
 
-    return jax.vmap(bound_compute_loss_alpha)(
+    loss, entropy = jax.vmap(bound_compute_loss_alpha)(
         jax.random.split(rng, FLAGS.num_agents),
         data,
-    ).mean()
+    )
+    return (loss.mean(), entropy.mean())
 
 
 @jax.jit
-def train_step(state: TrainState, data, action_data, update_rng: PRNGKey,
+def train_step(state: TrainState, data, update_rng: PRNGKey,
                step: int) -> TrainState:
     """Updates the parameters for Q, Pi, target Q, and alpha."""
     update_rng, q_grad_rng = jax.random.split(update_rng)
@@ -253,16 +269,11 @@
 
     state = state.q_apply_gradients(step=step, grads=q_grads)
 
-    update_rng, pi_grad_rng = jax.random.split(update_rng)
-
     # Update pi
+    update_rng, pi_grad_rng = jax.random.split(update_rng)
     pi_grad_fn = jax.value_and_grad(lambda params: compute_batched_loss_pi(
-        state, pi_grad_rng, params, action_data))
+        state, pi_grad_rng, params, data))
     pi_loss, pi_grads = pi_grad_fn(state.params)
-
-    print_nan(step, pi_loss)
-    print_nan(step, pi_grads)
-
     state = state.pi_apply_gradients(step=step, grads=pi_grads)
 
     update_rng, alpha_grad_rng = jax.random.split(update_rng)
@@ -271,15 +282,18 @@
         # Update alpha
         alpha_grad_fn = jax.value_and_grad(
             lambda params: compute_batched_loss_alpha(state, alpha_grad_rng,
-                                                      params, data))
-        alpha_loss, alpha_grads = alpha_grad_fn(state.params)
+                                                      params, data),
+            has_aux=True,
+        )
+        (alpha_loss, entropy), alpha_grads = alpha_grad_fn(state.params)
         print_nan(step, alpha_loss)
         print_nan(step, alpha_grads)
         state = state.alpha_apply_gradients(step=step, grads=alpha_grads)
     else:
+        entropy = 0.0
         alpha_loss = 0.0
 
-    return state, q_loss, pi_loss, alpha_loss
+    return state, q_loss, pi_loss, alpha_loss, entropy
 
 
 @jax.jit
@@ -304,20 +318,22 @@
 
         def true_fn(i):
             # We are at the beginning of the process, pick a random action.
-            return state.problem.random_actions(action_rng, FLAGS.num_agents)
+            return state.problem.random_actions(action_rng,
+                                                X=observation,
+                                                goal=R,
+                                                dimensions=FLAGS.num_agents)
 
         def false_fn(i):
             # We are past the beginning of the process, use the trained network.
-            pi_action, logp_pi, std, random_sample = state.pi_apply(
-                rng=action_rng,
-                params=state.params,
-                observation=observation,
-                R=R,
-                deterministic=False)
+            pi_action, _, _ = state.pi_apply(rng=action_rng,
+                                             params=state.params,
+                                             observation=observation,
+                                             R=R,
+                                             deterministic=False)
             return pi_action
 
         pi_action = jax.lax.cond(
-            step <= FLAGS.start_steps,
+            step <= FLAGS.random_sample_steps,
             true_fn,
             false_fn,
             i,
@@ -328,11 +344,9 @@
             lambda o, pi: state.problem.integrate_dynamics(o, pi),
             in_axes=(0, 0))(observation, pi_action)
 
-        # Soft Actor-Critic is designed to maximize reward.  LQR minimizes
-        # cost.  There is nothing which assumes anything about the sign of
-        # the reward, so use the negative of the cost.
-        reward = -jax.vmap(state.problem.cost)(
-            X=observation2, U=pi_action, goal=R)
+        reward = jax.vmap(state.problem.reward)(X=observation2,
+                                                U=pi_action,
+                                                goal=R)
 
         replay_buffer_state = state.replay_buffer.add(
             replay_buffer_state, {
@@ -357,10 +371,8 @@
                      step: int):
     rng, sample_rng = jax.random.split(rng)
 
-    action_data = state.replay_buffer.sample(replay_buffer_state, sample_rng)
-
     def update_iteration(i, val):
-        rng, state, q_loss, pi_loss, alpha_loss, replay_buffer_state, action_data = val
+        rng, state, q_loss, pi_loss, alpha_loss, replay_buffer_state, entropy = val
         rng, sample_rng, update_rng = jax.random.split(rng, 3)
 
         batch = state.replay_buffer.sample(replay_buffer_state, sample_rng)
@@ -368,22 +380,18 @@
         print_nan(i, replay_buffer_state)
         print_nan(i, batch)
 
-        state, q_loss, pi_loss, alpha_loss = train_step(
-            state,
-            data=batch.experience,
-            action_data=batch.experience,
-            update_rng=update_rng,
-            step=i)
+        state, q_loss, pi_loss, alpha_loss, entropy = train_step(
+            state, data=batch.experience, update_rng=update_rng, step=i)
 
-        return rng, state, q_loss, pi_loss, alpha_loss, replay_buffer_state, action_data
+        return rng, state, q_loss, pi_loss, alpha_loss, replay_buffer_state, entropy
 
-    rng, state, q_loss, pi_loss, alpha_loss, replay_buffer_state, action_data = jax.lax.fori_loop(
+    rng, state, q_loss, pi_loss, alpha_loss, replay_buffer_state, entropy = jax.lax.fori_loop(
         step, step + FLAGS.horizon + 1, update_iteration,
-        (rng, state, 0.0, 0.0, 0.0, replay_buffer_state, action_data))
+        (rng, state, 0.0, 0.0, 0.0, replay_buffer_state, 0))
 
     state = state.target_apply_gradients(step=state.step)
 
-    return rng, state, q_loss, pi_loss, alpha_loss
+    return rng, state, q_loss, pi_loss, alpha_loss, entropy
 
 
 def train(workdir: str, problem: Problem) -> train_state.TrainState:
@@ -396,6 +404,8 @@
     run['hparams'] = {
         'q_learning_rate': FLAGS.q_learning_rate,
         'pi_learning_rate': FLAGS.pi_learning_rate,
+        'alpha_learning_rate': FLAGS.alpha_learning_rate,
+        'random_sample_steps': FLAGS.random_sample_steps,
         'batch_size': FLAGS.batch_size,
         'horizon': FLAGS.horizon,
         'warmup_steps': FLAGS.warmup_steps,
@@ -422,11 +432,11 @@
         problem,
         q_learning_rate=q_learning_rate,
         pi_learning_rate=pi_learning_rate,
+        alpha_learning_rate=FLAGS.alpha_learning_rate,
     )
     state = restore_checkpoint(state, workdir)
 
-    state_sharding = nn.get_sharding(state, state.mesh)
-    logging.info(state_sharding)
+    logging.debug(nn.get_sharding(state, state.mesh))
 
     replay_buffer_state = state.replay_buffer.init({
         'observations1':
@@ -438,12 +448,10 @@
         'rewards':
         jax.numpy.zeros((1, )),
         'goals':
-        jax.numpy.zeros((problem.num_states, )),
+        jax.numpy.zeros((problem.num_goals, )),
     })
 
-    replay_buffer_state_sharding = nn.get_sharding(replay_buffer_state,
-                                                   state.mesh)
-    logging.info(replay_buffer_state_sharding)
+    logging.debug(nn.get_sharding(replay_buffer_state, state.mesh))
 
     # Number of gradients to accumulate before doing decent.
     update_after = FLAGS.batch_size // FLAGS.num_agents
@@ -461,24 +469,25 @@
         )
 
         def nop(rng, state, replay_buffer_state, step):
-            return rng, state.update_step(step=step), 0.0, 0.0, 0.0
+            return rng, state.update_step(step=step), 0.0, 0.0, 0.0, 0.0
 
         # Train
-        rng, state, q_loss, pi_loss, alpha_loss = jax.lax.cond(
+        rng, state, q_loss, pi_loss, alpha_loss, entropy = jax.lax.cond(
             step >= update_after, update_gradients, nop, rng, state,
             replay_buffer_state, step)
 
-        return state, replay_buffer_state, rng, q_loss, pi_loss, alpha_loss
+        return state, replay_buffer_state, rng, q_loss, pi_loss, alpha_loss, entropy
 
+    last_time = time.time()
     for step in range(0, FLAGS.steps, FLAGS.horizon):
-        state, replay_buffer_state, rng, q_loss, pi_loss, alpha_loss = train_loop(
+        state, replay_buffer_state, rng, q_loss, pi_loss, alpha_loss, entropy = train_loop(
             state, replay_buffer_state, rng, step)
 
         if FLAGS.debug_nan and has_nan(state.params):
             logging.fatal('Nan params, aborting')
 
         logging.info(
-            'Step %s: q_loss=%s, pi_loss=%s, alpha_loss=%s, q_learning_rate=%s, pi_learning_rate=%s, alpha=%s',
+            'Step %s: q_loss=%s, pi_loss=%s, alpha_loss=%s, q_learning_rate=%s, pi_learning_rate=%s, alpha=%s, entropy=%s, random=%s',
             step,
             q_loss,
             pi_loss,
@@ -486,6 +495,8 @@
             q_learning_rate(step),
             pi_learning_rate(step),
             jax.numpy.exp(state.params['logalpha']),
+            entropy,
+            step <= FLAGS.random_sample_steps,
         )
 
         run.track(
@@ -493,12 +504,14 @@
                 'q_loss': float(q_loss),
                 'pi_loss': float(pi_loss),
                 'alpha_loss': float(alpha_loss),
-                'alpha': float(jax.numpy.exp(state.params['logalpha']))
+                'alpha': float(jax.numpy.exp(state.params['logalpha'])),
+                'entropy': entropy,
             },
             step=step)
 
-        if step % 1000 == 0 and step > update_after:
+        if time.time() > last_time + 3.0 and step > update_after:
             # TODO(austin): Simulate a rollout and accumulate the reward.  How good are we doing?
             save_checkpoint(state, workdir)
+            last_time = time.time()
 
     return state
diff --git a/frc971/imu_fdcan/README.md b/frc971/imu_fdcan/README.md
index 8e82f73..7bf7328 100644
--- a/frc971/imu_fdcan/README.md
+++ b/frc971/imu_fdcan/README.md
@@ -23,7 +23,8 @@
     * The main code lives in [`Dual_IMU/Core/Src`](/Dual_IMU/Core/Src/). Make sure your changes happen inside sections marked `/* USER CODE BEGIN ... */` `/* USER CODE END ... */`. Code outside these markers will be overwritten by CubeIDE when generating code after changes to the `.ioc` file.
 3) Build + Run:
     * Option 1: Open CubeIDE GUI to build, debug, or run.
-    * Option 2: 
+    <!-- TODO(sindy): fix this build script -->
+    * Option 2 (DO NOT USE. NOT SAFE) 
         1) SSH onto the build server. 
         2) Run `bazel build -c opt --config=cortex-m4f-imu //frc971/imu_fdcan/Dual_IMU/Core:main.elf`. The output .elf file should be in bazel-bin/frc971/imu_fdcan/Dual_IMU/Core.
         3) (If deploying code locally) Move file to local directory. For example: `scp <username>@build.frc971.org:<path/to/main.elf> <local/path/to/save/file/`. A good spot to put this locally is ./Dual_IMU/Debug/.
diff --git a/frc971/wpilib/talonfx.cc b/frc971/wpilib/talonfx.cc
index 4338cfd..0da7603 100644
--- a/frc971/wpilib/talonfx.cc
+++ b/frc971/wpilib/talonfx.cc
@@ -58,11 +58,12 @@
 
 void TalonFX::WriteConfigs() {
   ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-  current_limits.StatorCurrentLimit = stator_current_limit_;
+  current_limits.StatorCurrentLimit =
+      units::current::ampere_t{stator_current_limit_};
   current_limits.StatorCurrentLimitEnable = true;
-  current_limits.SupplyCurrentLimit = supply_current_limit_;
+  current_limits.SupplyCurrentLimit =
+      units::current::ampere_t{supply_current_limit_};
   current_limits.SupplyCurrentLimitEnable = true;
-  current_limits.SupplyTimeThreshold = 0.0;
 
   ctre::phoenix6::configs::MotorOutputConfigs output_configs;
   output_configs.NeutralMode = neutral_mode_;
diff --git a/frc971/zeroing/continuous_absolute_encoder.cc b/frc971/zeroing/continuous_absolute_encoder.cc
index 3381833..733d46d 100644
--- a/frc971/zeroing/continuous_absolute_encoder.cc
+++ b/frc971/zeroing/continuous_absolute_encoder.cc
@@ -165,4 +165,14 @@
   return builder.Finish();
 }
 
+void ContinuousAbsoluteEncoderZeroingEstimator::GetEstimatorState(
+    AbsoluteEncoderEstimatorStateStatic *fbs) const {
+  errors_.ToStaticFlatbuffer(fbs->add_errors());
+
+  fbs->set_error(error_);
+  fbs->set_zeroed(zeroed_);
+  fbs->set_position(position_);
+  fbs->set_absolute_position(filtered_absolute_encoder_);
+}
+
 }  // namespace frc971::zeroing
diff --git a/frc971/zeroing/continuous_absolute_encoder.h b/frc971/zeroing/continuous_absolute_encoder.h
index 4994280..5e700ee 100644
--- a/frc971/zeroing/continuous_absolute_encoder.h
+++ b/frc971/zeroing/continuous_absolute_encoder.h
@@ -47,6 +47,8 @@
   virtual flatbuffers::Offset<State> GetEstimatorState(
       flatbuffers::FlatBufferBuilder *fbb) const override;
 
+  void GetEstimatorState(AbsoluteEncoderEstimatorStateStatic *fbs) const;
+
  private:
   struct PositionStruct {
     PositionStruct(const AbsolutePosition &position_buffer)
diff --git a/frc971/zeroing/zeroing.h b/frc971/zeroing/zeroing.h
index 68ef38f..01c2c61 100644
--- a/frc971/zeroing/zeroing.h
+++ b/frc971/zeroing/zeroing.h
@@ -10,7 +10,7 @@
 #include "flatbuffers/flatbuffers.h"
 
 #include "frc971/constants.h"
-#include "frc971/control_loops/control_loops_generated.h"
+#include "frc971/control_loops/control_loops_static.h"
 
 // TODO(pschrader): Flag an error if encoder index pulse is not n revolutions
 // away from the last one (i.e. got extra counts from noise, etc..)
diff --git a/scouting/scouting_test.cy.js b/scouting/scouting_test.cy.js
index 990e69f..7157880 100644
--- a/scouting/scouting_test.cy.js
+++ b/scouting/scouting_test.cy.js
@@ -119,7 +119,7 @@
       ' Ended Match; stageType: kHARMONY, trapNote: false, spotlight: false '
     );
   // Ensure that the penalties action is only submitted once.
-  cy.get('#review_data li').contains('Penalties').should('have.length', 1);
+  cy.get('#review_data li:contains("Penalties")').its('length').should('eq', 1);
 
   clickButton('Submit');
   headerShouldBe(teamNumber + ' Success ');
diff --git a/y2020/wpilib_interface.cc b/y2020/wpilib_interface.cc
index 2213c16..8c8b480 100644
--- a/y2020/wpilib_interface.cc
+++ b/y2020/wpilib_interface.cc
@@ -123,9 +123,11 @@
 void WriteConfigs(ctre::phoenix6::hardware::TalonFX *talon,
                   double stator_current_limit, double supply_current_limit) {
   ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-  current_limits.StatorCurrentLimit = stator_current_limit;
+  current_limits.StatorCurrentLimit =
+      units::current::ampere_t{stator_current_limit};
   current_limits.StatorCurrentLimitEnable = true;
-  current_limits.SupplyCurrentLimit = supply_current_limit;
+  current_limits.SupplyCurrentLimit =
+      units::current::ampere_t{supply_current_limit};
   current_limits.SupplyCurrentLimitEnable = true;
 
   ctre::phoenix6::configs::TalonFXConfiguration configuration;
@@ -454,7 +456,8 @@
       ::std::unique_ptr<::ctre::phoenix6::hardware::TalonFX> t) {
     climber_falcon_ = ::std::move(t);
     ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-    current_limits.SupplyCurrentLimit = Values::kClimberSupplyCurrentLimit();
+    current_limits.SupplyCurrentLimit =
+        units::current::ampere_t{Values::kClimberSupplyCurrentLimit()};
     current_limits.SupplyCurrentLimitEnable = true;
 
     ctre::phoenix6::configs::TalonFXConfiguration configuration;
diff --git a/y2021_bot3/wpilib_interface.cc b/y2021_bot3/wpilib_interface.cc
index 4006736..f3c956e 100644
--- a/y2021_bot3/wpilib_interface.cc
+++ b/y2021_bot3/wpilib_interface.cc
@@ -119,9 +119,11 @@
 void WriteConfigs(ctre::phoenix6::hardware::TalonFX *talon,
                   double stator_current_limit, double supply_current_limit) {
   ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-  current_limits.StatorCurrentLimit = stator_current_limit;
+  current_limits.StatorCurrentLimit =
+      units::current::ampere_t{stator_current_limit};
   current_limits.StatorCurrentLimitEnable = true;
-  current_limits.SupplyCurrentLimit = supply_current_limit;
+  current_limits.SupplyCurrentLimit =
+      units::current::ampere_t{supply_current_limit};
   current_limits.SupplyCurrentLimitEnable = true;
 
   ctre::phoenix6::configs::TalonFXConfiguration configuration;
@@ -249,7 +251,8 @@
       ::std::unique_ptr<::ctre::phoenix6::hardware::TalonFX> t) {
     climber_falcon_ = ::std::move(t);
     ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-    current_limits.SupplyCurrentLimit = Values::kClimberSupplyCurrentLimit();
+    current_limits.SupplyCurrentLimit =
+        units::current::ampere_t{Values::kClimberSupplyCurrentLimit()};
     current_limits.SupplyCurrentLimitEnable = true;
 
     ctre::phoenix6::configs::TalonFXConfiguration configuration;
diff --git a/y2022/wpilib_interface.cc b/y2022/wpilib_interface.cc
index 1f82e91..469c93b 100644
--- a/y2022/wpilib_interface.cc
+++ b/y2022/wpilib_interface.cc
@@ -129,9 +129,11 @@
 void WriteConfigs(ctre::phoenix6::hardware::TalonFX *talon,
                   double stator_current_limit, double supply_current_limit) {
   ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-  current_limits.StatorCurrentLimit = stator_current_limit;
+  current_limits.StatorCurrentLimit =
+      units::current::ampere_t{stator_current_limit};
   current_limits.StatorCurrentLimitEnable = true;
-  current_limits.SupplyCurrentLimit = supply_current_limit;
+  current_limits.SupplyCurrentLimit =
+      units::current::ampere_t{supply_current_limit};
   current_limits.SupplyCurrentLimitEnable = true;
 
   ctre::phoenix6::configs::TalonFXConfiguration configuration;
@@ -483,10 +485,10 @@
     for (auto &falcon : {catapult_falcon_1_can_, catapult_falcon_2_can_}) {
       ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
       current_limits.StatorCurrentLimit =
-          Values::kIntakeRollerStatorCurrentLimit();
+          units::current::ampere_t{Values::kIntakeRollerStatorCurrentLimit()};
       current_limits.StatorCurrentLimitEnable = true;
       current_limits.SupplyCurrentLimit =
-          Values::kIntakeRollerSupplyCurrentLimit();
+          units::current::ampere_t{Values::kIntakeRollerSupplyCurrentLimit()};
       current_limits.SupplyCurrentLimitEnable = true;
 
       ctre::phoenix6::configs::TalonFXConfiguration configuration;
diff --git a/y2023/wpilib_interface.cc b/y2023/wpilib_interface.cc
index bc14311..536f7bb 100644
--- a/y2023/wpilib_interface.cc
+++ b/y2023/wpilib_interface.cc
@@ -166,11 +166,11 @@
     inverted_ = invert;
 
     ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-    current_limits.StatorCurrentLimit =
-        constants::Values::kDrivetrainStatorCurrentLimit();
+    current_limits.StatorCurrentLimit = units::current::ampere_t{
+        constants::Values::kDrivetrainStatorCurrentLimit()};
     current_limits.StatorCurrentLimitEnable = true;
-    current_limits.SupplyCurrentLimit =
-        constants::Values::kDrivetrainSupplyCurrentLimit();
+    current_limits.SupplyCurrentLimit = units::current::ampere_t{
+        constants::Values::kDrivetrainSupplyCurrentLimit()};
     current_limits.SupplyCurrentLimitEnable = true;
 
     ctre::phoenix6::configs::MotorOutputConfigs output_configs;
@@ -196,11 +196,11 @@
 
   void WriteRollerConfigs() {
     ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-    current_limits.StatorCurrentLimit =
-        constants::Values::kRollerStatorCurrentLimit();
+    current_limits.StatorCurrentLimit = units::current::ampere_t{
+        constants::Values::kRollerStatorCurrentLimit()};
     current_limits.StatorCurrentLimitEnable = true;
-    current_limits.SupplyCurrentLimit =
-        constants::Values::kRollerSupplyCurrentLimit();
+    current_limits.SupplyCurrentLimit = units::current::ampere_t{
+        constants::Values::kRollerSupplyCurrentLimit()};
     current_limits.SupplyCurrentLimitEnable = true;
 
     ctre::phoenix6::configs::MotorOutputConfigs output_configs;
diff --git a/y2023_bot3/wpilib_interface.cc b/y2023_bot3/wpilib_interface.cc
index d8813be..0e43e10 100644
--- a/y2023_bot3/wpilib_interface.cc
+++ b/y2023_bot3/wpilib_interface.cc
@@ -147,11 +147,11 @@
     inverted_ = invert;
 
     ctre::phoenix6::configs::CurrentLimitsConfigs current_limits;
-    current_limits.StatorCurrentLimit =
-        constants::Values::kDrivetrainStatorCurrentLimit();
+    current_limits.StatorCurrentLimit = units::current::ampere_t{
+        constants::Values::kDrivetrainStatorCurrentLimit()};
     current_limits.StatorCurrentLimitEnable = true;
-    current_limits.SupplyCurrentLimit =
-        constants::Values::kDrivetrainSupplyCurrentLimit();
+    current_limits.SupplyCurrentLimit = units::current::ampere_t{
+        constants::Values::kDrivetrainSupplyCurrentLimit()};
     current_limits.SupplyCurrentLimitEnable = true;
 
     ctre::phoenix6::configs::MotorOutputConfigs output_configs;
diff --git a/y2024_swerve/y2024_swerve_roborio.json b/y2024_swerve/y2024_swerve_roborio.json
index 705370d..adf528d 100644
--- a/y2024_swerve/y2024_swerve_roborio.json
+++ b/y2024_swerve/y2024_swerve_roborio.json
@@ -305,8 +305,7 @@
             "name": "wpilib_interface",
             "executable_name": "wpilib_interface",
             "args": [
-                "--nodie_on_malloc",
-                "--ctre_diag_server"
+              "--nodie_on_malloc"
             ],
             "nodes": [
                 "roborio"