diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index 6a2822b..2e680a3 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -112,14 +112,14 @@
       break;
   }
 
-  kf_.set_controller_index(ControllerIndexFromGears());
+  kf_.set_index(ControllerIndexFromGears());
   {
     GearLogging gear_logging;
     gear_logging.left_state = static_cast<uint32_t>(left_gear_);
     gear_logging.right_state = static_cast<uint32_t>(right_gear_);
     gear_logging.left_loop_high = MaybeHigh(left_gear_);
     gear_logging.right_loop_high = MaybeHigh(right_gear_);
-    gear_logging.controller_index = kf_.controller_index();
+    gear_logging.controller_index = kf_.index();
     LOG_STRUCT(DEBUG, "state", gear_logging);
   }
   const bool is_latest_imu_values = ::frc971::imu_values.FetchLatest();
diff --git a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
index 41a4f3c..c66a4e3 100644
--- a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
@@ -166,15 +166,15 @@
 
     if (left_gear_high_) {
       if (right_gear_high_) {
-        drivetrain_plant_->set_plant_index(3);
+        drivetrain_plant_->set_index(3);
       } else {
-        drivetrain_plant_->set_plant_index(2);
+        drivetrain_plant_->set_index(2);
       }
     } else {
       if (right_gear_high_) {
-        drivetrain_plant_->set_plant_index(1);
+        drivetrain_plant_->set_index(1);
       } else {
-        drivetrain_plant_->set_plant_index(0);
+        drivetrain_plant_->set_index(0);
       }
     }
 
diff --git a/frc971/control_loops/drivetrain/polydrivetrain.cc b/frc971/control_loops/drivetrain/polydrivetrain.cc
index d3a220b..5fc4b17 100644
--- a/frc971/control_loops/drivetrain/polydrivetrain.cc
+++ b/frc971/control_loops/drivetrain/polydrivetrain.cc
@@ -1,16 +1,16 @@
 #include "frc971/control_loops/drivetrain/polydrivetrain.h"
 
-#include "aos/common/logging/logging.h"
-#include "aos/common/controls/polytope.h"
 #include "aos/common/commonmath.h"
-#include "aos/common/logging/queue_logging.h"
+#include "aos/common/controls/polytope.h"
+#include "aos/common/logging/logging.h"
 #include "aos/common/logging/matrix_logging.h"
+#include "aos/common/logging/queue_logging.h"
 
 #include "aos/common/messages/robot_state.q.h"
-#include "frc971/control_loops/state_feedback_loop.h"
 #include "frc971/control_loops/coerce_goal.h"
 #include "frc971/control_loops/drivetrain/drivetrain.q.h"
 #include "frc971/control_loops/drivetrain/drivetrain_config.h"
+#include "frc971/control_loops/state_feedback_loop.h"
 
 namespace frc971 {
 namespace control_loops {
@@ -22,13 +22,16 @@
       U_Poly_((Eigen::Matrix<double, 4, 2>() << /*[[*/ 1, 0 /*]*/,
                /*[*/ -1, 0 /*]*/,
                /*[*/ 0, 1 /*]*/,
-               /*[*/ 0, -1 /*]]*/).finished(),
+               /*[*/ 0, -1 /*]]*/)
+                  .finished(),
               (Eigen::Matrix<double, 4, 1>() << /*[[*/ 12 /*]*/,
                /*[*/ 12 /*]*/,
                /*[*/ 12 /*]*/,
-               /*[*/ 12 /*]]*/).finished(),
+               /*[*/ 12 /*]]*/)
+                  .finished(),
               (Eigen::Matrix<double, 2, 4>() << /*[[*/ 12, 12, -12, -12 /*]*/,
-               /*[*/ -12, 12, 12, -12 /*]*/).finished()),
+               /*[*/ -12, 12, 12, -12 /*]*/)
+                  .finished()),
       loop_(new StateFeedbackLoop<2, 2, 2>(dt_config.make_v_drivetrain_loop())),
       ttrust_(1.1),
       wheel_(0.0),
@@ -148,14 +151,14 @@
 
 double PolyDrivetrain::FilterVelocity(double throttle) const {
   const Eigen::Matrix<double, 2, 2> FF =
-      loop_->B().inverse() *
-      (Eigen::Matrix<double, 2, 2>::Identity() - loop_->A());
+      loop_->plant().B().inverse() *
+      (Eigen::Matrix<double, 2, 2>::Identity() - loop_->plant().A());
 
   constexpr int kHighGearController = 3;
   const Eigen::Matrix<double, 2, 2> FF_high =
-      loop_->controller(kHighGearController).plant.B.inverse() *
+      loop_->plant().coefficients(kHighGearController).B.inverse() *
       (Eigen::Matrix<double, 2, 2>::Identity() -
-       loop_->controller(kHighGearController).plant.A);
+       loop_->plant().coefficients(kHighGearController).A);
 
   ::Eigen::Matrix<double, 1, 2> FF_sum = FF.colwise().sum();
   int min_FF_sum_index;
@@ -173,14 +176,14 @@
 
 double PolyDrivetrain::MaxVelocity() {
   const Eigen::Matrix<double, 2, 2> FF =
-      loop_->B().inverse() *
-      (Eigen::Matrix<double, 2, 2>::Identity() - loop_->A());
+      loop_->plant().B().inverse() *
+      (Eigen::Matrix<double, 2, 2>::Identity() - loop_->plant().A());
 
   constexpr int kHighGearController = 3;
   const Eigen::Matrix<double, 2, 2> FF_high =
-      loop_->controller(kHighGearController).plant.B.inverse() *
+      loop_->plant().coefficients(kHighGearController).B.inverse() *
       (Eigen::Matrix<double, 2, 2>::Identity() -
-       loop_->controller(kHighGearController).plant.A);
+       loop_->plant().coefficients(kHighGearController).A);
 
   ::Eigen::Matrix<double, 1, 2> FF_sum = FF.colwise().sum();
   int min_FF_sum_index;
@@ -206,8 +209,8 @@
   if (IsInGear(left_gear_) && IsInGear(right_gear_)) {
     // FF * X = U (steady state)
     const Eigen::Matrix<double, 2, 2> FF =
-        loop_->B().inverse() *
-        (Eigen::Matrix<double, 2, 2>::Identity() - loop_->A());
+        loop_->plant().B().inverse() *
+        (Eigen::Matrix<double, 2, 2>::Identity() - loop_->plant().A());
 
     // Invert the plant to figure out how the velocity filter would have to
     // work
@@ -262,7 +265,7 @@
 
     if (dt_config_.loop_type == LoopType::OPEN_LOOP) {
       loop_->mutable_X_hat() =
-          loop_->A() * loop_->X_hat() + loop_->B() * loop_->U();
+          loop_->plant().A() * loop_->X_hat() + loop_->plant().B() * loop_->U();
     }
   } else {
     const double current_left_velocity =
diff --git a/frc971/control_loops/drivetrain/polydrivetrain.h b/frc971/control_loops/drivetrain/polydrivetrain.h
index 20f1a49..570b87f 100644
--- a/frc971/control_loops/drivetrain/polydrivetrain.h
+++ b/frc971/control_loops/drivetrain/polydrivetrain.h
@@ -17,7 +17,7 @@
   PolyDrivetrain(const DrivetrainConfig &dt_config,
                  StateFeedbackLoop<7, 2, 3> *kf);
 
-  int controller_index() const { return loop_->controller_index(); }
+  int controller_index() const { return loop_->index(); }
 
   // Computes the speed of the motor given the hall effect position and the
   // speed of the robot.