Merge "Add an experience buffer and tests." into main
diff --git a/frc971/control_loops/swerve/casadi_velocity_mpc.py b/frc971/control_loops/swerve/casadi_velocity_mpc.py
index 239f49a..517de5a 100644
--- a/frc971/control_loops/swerve/casadi_velocity_mpc.py
+++ b/frc971/control_loops/swerve/casadi_velocity_mpc.py
@@ -47,16 +47,30 @@
         # TODO(austin): Do we need a disturbance torque per module to account for friction?
         # Do it only in the observer/post?
 
-        self.next_X = self.make_physics()
-        self.cost = self.make_cost()
         self.force = [
             dynamics.F(i, casadi.SX.sym("X", 25, 1), casadi.SX.sym("U", 8, 1))
             for i in range(4)
         ]
+        self.force_vel = [
+            dynamics.F_vel(i,
+                           casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
+                           casadi.SX.sym("U", 8, 1)) for i in range(4)
+        ]
         self.slip_angle = [
             dynamics.slip_angle(i, casadi.SX.sym("X", 25, 1),
                                 casadi.SX.sym("U", 8, 1)) for i in range(4)
         ]
+        self.mounting_location = [
+            dynamics.mounting_location(
+                i, casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
+                casadi.SX.sym("U", 8, 1)) for i in range(4)
+        ]
+        self.torque_cross = self.torque_cross_func(casadi.SX.sym("r", 2, 1),
+                                                   casadi.SX.sym("F", 2, 1))
+        self.force_cross = self.force_cross_func(casadi.SX.sym("Tau", 1, 1),
+                                                 casadi.SX.sym("r", 2, 1))
+        self.next_X = self.make_physics()
+        self.cost = self.make_cost()
 
         self.N = 200
 
@@ -209,6 +223,8 @@
 
         # TODO(austin): Do we want to do something more special for 0?
 
+        # cost velocity a lot more in the perpendicular direction to allow tire to spin up
+        # we also only want to get moving in the correct direction as fast as possible
         J += 75 * ((R[0] - X[dynamics.VELOCITY_STATE_VX]) * vnormx +
                    (R[1] - X[dynamics.VELOCITY_STATE_VY]) * vnormy)**2.0
 
@@ -231,11 +247,29 @@
         J += kSteerPositionGain * (X[dynamics.VELOCITY_STATE_THETAS3])**2.0
         J += kSteerVelocityGain * (X[dynamics.VELOCITY_STATE_OMEGAS3])**2.0
 
+        # cost variance of the force by a tire and the expected average force and torque on it
+        total_force = self.force_vel[0](X, U)
+        total_torque = self.torque_cross(self.mounting_location[0](X, U),
+                                         self.force_vel[0](X, U))
+        for i in range(3):
+            total_force += self.force_vel[i + 1](X, U)
+            total_torque += self.torque_cross(
+                self.mounting_location[i + 1](X, U), self.force_vel[i + 1](X,
+                                                                           U))
+
+        total_force /= 4
+        total_torque /= 4
+        for i in range(4):
+            f_diff = (total_force +
+                      self.force_cross(total_torque, self.mounting_location[i]
+                                       (X, U))) - self.force_vel[i](X, U)
+            J += 0.01 * (f_diff[0, 0]**2.0 + f_diff[1, 0]**2.0)
+
         # TODO(austin): Don't penalize torque steering current.
         for i in range(4):
             Is = U[2 * i + 0]
             Id = U[2 * i + 1]
-            # Steer
+            # Steer, cost it a lot less than drive to be more agressive in steering
             J += ((Is + dynamics.STEER_CURRENT_COUPLING_FACTOR * Id)**
                   2.0) / 100000.0
             # Drive
@@ -243,6 +277,17 @@
 
         return casadi.Function("Jn", [X, U, R], [J])
 
+    def torque_cross_func(self, r, F):
+        result = casadi.SX.sym('Tau', 1, 1)
+        result[0, 0] = r[0, 0] * F[1, 0] - r[1, 0] * F[0, 0]
+        return casadi.Function('Tau', [r, F], [result])
+
+    def force_cross_func(self, Tau, r):
+        result = casadi.SX.sym('F', 2, 1)
+        result[0, 0] = -r[1, 0] * Tau[0, 0] / casadi.norm_2(r)**2.0
+        result[1, 0] = r[0, 0] * Tau[0, 0] / casadi.norm_2(r)**2.0
+        return casadi.Function('F', [Tau, r], [result])
+
     def solve(self, p, seed=None):
         if seed is None:
             seed = []
@@ -334,6 +379,9 @@
                                  self.X_plot[dynamics.STATE_VX, 0:i + 1],
                                  label="vx")
                     axs0[0].plot(self.t,
+                                 self.X_plot[dynamics.STATE_OMEGA, 0:i + 1],
+                                 label="omega")
+                    axs0[0].plot(self.t,
                                  self.X_plot[dynamics.STATE_VY, 0:i + 1],
                                  label="vy")
                     axs0[0].legend()
diff --git a/frc971/control_loops/swerve/generate_physics.cc b/frc971/control_loops/swerve/generate_physics.cc
index 748e3cb..77dfbb0 100644
--- a/frc971/control_loops/swerve/generate_physics.cc
+++ b/frc971/control_loops/swerve/generate_physics.cc
@@ -785,7 +785,16 @@
         },
         &result_py);
 
-    DefineVector2dFunction(
+    DefineVector2dVelocityFunction(
+        "F_vel",
+        "Returns the force on the wheel in absolute coordinates based on the "
+        "velocity controller",
+        [](const Module &m, int dimension) {
+          return ccode(*m.direct.F.get(dimension, 0));
+        },
+        &result_py);
+
+    DefineVector2dVelocityFunction(
         "mounting_location",
         "Returns the mounting location of wheel in robot coordinates",
         [](const Module &m, int dimension) {
@@ -852,6 +861,34 @@
         "    return casadi.Function('$0', [X, U], [result])", name));
   }
 
+  void DefineVector2dVelocityFunction(
+      std::string_view name, std::string_view documentation,
+      std::function<std::string(const Module &, int)> scalar_fn,
+      std::vector<std::string> *result_py) {
+    result_py->emplace_back("");
+    result_py->emplace_back(absl::Substitute("# $0.", documentation));
+    result_py->emplace_back(absl::Substitute("def $0(i, X, U):", name));
+    WriteCasadiVelocityVariables(result_py);
+    result_py->emplace_back(
+        absl::Substitute("    result = casadi.SX.sym('$0', 2, 1)", name));
+    for (size_t m = 0; m < kNumModules; ++m) {
+      if (m == 0) {
+        result_py->emplace_back("    if i == 0:");
+      } else {
+        result_py->emplace_back(absl::Substitute("    elif i == $0:", m));
+      }
+      for (int j = 0; j < 2; ++j) {
+        result_py->emplace_back(absl::Substitute("        result[$0, 0] = $1",
+                                                 j, scalar_fn(modules_[m], j)));
+      }
+    }
+    result_py->emplace_back("    else:");
+    result_py->emplace_back(
+        "        raise ValueError(\"Invalid module number\")");
+    result_py->emplace_back(absl::Substitute(
+        "    return casadi.Function('$0', [X, U], [result])", name));
+  }
+
  private:
   static constexpr uint8_t kNumModules = 4;