Add acceptably tuned EKF for the arm.

I had to switch to a voltage error to get it to converge.  I don't know
if that's a residual problem from uninitialized stack or what, but it's
working much better now.

The disturbance estimator has a time constant of like 0.5 seconds right
now.  Faster would be nice, but I'll wait until we see it on a bot
before cranking it up much further.

Change-Id: I72d46aa308ce806a35cfed191ee3b15864e6905a
diff --git a/y2018/control_loops/superstructure/arm/BUILD b/y2018/control_loops/superstructure/arm/BUILD
index 6224df1..6961dd3 100644
--- a/y2018/control_loops/superstructure/arm/BUILD
+++ b/y2018/control_loops/superstructure/arm/BUILD
@@ -19,10 +19,10 @@
     srcs = [
         "trajectory_test.cc",
     ],
-    restricted_to = ["//tools:k8"],
     deps = [
         ":demo_path",
         ":dynamics",
+        ":ekf",
         ":trajectory",
         "//aos/testing:googletest",
         "//third_party/eigen",
@@ -71,8 +71,24 @@
     restricted_to = ["//tools:k8"],
     deps = [
         ":demo_path",
+        ":ekf",
         ":trajectory",
         "//third_party/eigen",
         "//third_party/matplotlib-cpp",
     ],
 )
+
+cc_library(
+    name = "ekf",
+    srcs = [
+        "ekf.cc",
+    ],
+    hdrs = [
+        "ekf.h",
+    ],
+    deps = [
+        ":dynamics",
+        "//frc971/control_loops:jacobian",
+        "//third_party/eigen",
+    ],
+)
diff --git a/y2018/control_loops/superstructure/arm/demo_path.cc b/y2018/control_loops/superstructure/arm/demo_path.cc
index 6be0a06..bcdda8c 100644
--- a/y2018/control_loops/superstructure/arm/demo_path.cc
+++ b/y2018/control_loops/superstructure/arm/demo_path.cc
@@ -1,5 +1,8 @@
 #include "y2018/control_loops/superstructure/arm/demo_path.h"
 
+#include <array>
+#include <initializer_list>
+
 namespace y2018 {
 namespace control_loops {
 namespace superstructure {
diff --git a/y2018/control_loops/superstructure/arm/dlqr.h b/y2018/control_loops/superstructure/arm/dlqr.h
deleted file mode 100644
index d33ee2b..0000000
--- a/y2018/control_loops/superstructure/arm/dlqr.h
+++ /dev/null
@@ -1,98 +0,0 @@
-#ifndef FRC971_CONTROL_LOOPS_DLQR_H_
-#define FRC971_CONTROL_LOOPS_DLQR_H_
-
-#include <Eigen/Dense>
-
-namespace frc971 {
-namespace controls {
-
-extern "C" {
-// Forward declaration for slicot fortran library.
-void sb02od_(char *DICO, char *JOBB, char *FACT, char *UPLO, char *JOBL,
-             char *SORT, int *N, int *M, int *P, double *A, int *LDA, double *B,
-             int *LDB, double *Q, int *LDQ, double *R, int *LDR, double *L,
-             int *LDL, double *RCOND, double *X, int *LDX, double *ALFAR,
-             double *ALFAI, double *BETA, double *S, int *LDS, double *T,
-             int *LDT, double *U, int *LDU, double *TOL, int *IWORK,
-             double *DWORK, int *LDWORK, int *BWORK, int *INFO);
-}
-
-template <int kN, int kM>
-void dlqr(::Eigen::Matrix<double, kN, kN> A, ::Eigen::Matrix<double, kN, kM> B,
-          ::Eigen::Matrix<double, kN, kN> Q, ::Eigen::Matrix<double, kM, kM> R,
-          ::Eigen::Matrix<double, kM, kN> *K,
-          ::Eigen::Matrix<double, kN, kN> *S) {
-  // Discrete (not continuous)
-  char DICO = 'D';
-  // B and R are provided instead of G.
-  char JOBB = 'B';
-  // Not factored, Q and R are provide.
-  char FACT = 'N';
-  // Store the upper triangle of Q and R.
-  char UPLO = 'U';
-  // L is zero.
-  char JOBL = 'Z';
-  // Stable eigenvalues first in the sort order
-  char SORT = 'S';
-
-  int N = 4;
-  int M = 2;
-  // Not needed since FACT = N
-  int P = 0;
-
-  int LDL = 1;
-
-  double RCOND = 0.0;
-  ::Eigen::Matrix<double, kN, kN> X = ::Eigen::Matrix<double, kN, kN>::Zero();
-
-  double ALFAR[kN * 2];
-  double ALFAI[kN * 2];
-  double BETA[kN * 2];
-  memset(ALFAR, 0, kN * 2);
-  memset(ALFAI, 0, kN * 2);
-  memset(BETA, 0, kN * 2);
-
-  int LDS = 2 * kN + kM;
-  ::Eigen::Matrix<double, 2 * kN + kM, 2 *kN + kM> S_schur =
-      ::Eigen::Matrix<double, 2 * kN + kM, 2 * kN + kM>::Zero();
-
-  ::Eigen::Matrix<double, 2 * kN + kM, 2 *kN> T =
-      ::Eigen::Matrix<double, 2 * kN + kM, 2 * kN>::Zero();
-  int LDT = 2 * kN + kM;
-
-  ::Eigen::Matrix<double, 2 * kN, 2 *kN> U =
-      ::Eigen::Matrix<double, 2 * kN, 2 * kN>::Zero();
-  int LDU = 2 * kN;
-
-  double TOL = 0.0;
-
-  int IWORK[2 * kN > kM ? 2 * kN : kM];
-  memset(IWORK, 0, 2 * kN > kM ? 2 * kN : kM);
-
-  int LDWORK = 16 * kN + 3 * kM + 16;
-
-  double DWORK[LDWORK];
-  memset(DWORK, 0, LDWORK);
-
-  int INFO = 0;
-
-  int BWORK[2 * kN];
-  memset(BWORK, 0, 2 * kN);
-
-  // TODO(austin): I can't tell if anything here is transposed...
-  sb02od_(&DICO, &JOBB, &FACT, &UPLO, &JOBL, &SORT, &N, &M, &P, A.data(), &N,
-          B.data(), &N, Q.data(), &N, R.data(), &M, nullptr, &LDL, &RCOND,
-          X.data(), &N, ALFAR, ALFAI, BETA, S_schur.data(), &LDS, T.data(),
-          &LDT, U.data(), &LDU, &TOL, IWORK, DWORK, &LDWORK, BWORK, &INFO);
-  //::std::cout << "slycot result: " << INFO << ::std::endl;
-
-  *K = (R + B.transpose() * X * B).inverse() * B.transpose() * X * A;
-  if (S != nullptr) {
-    *S = X;
-  }
-}
-
-}  // namespace controls
-}  // namespace frc971
-
-#endif  // FRC971_CONTROL_LOOPS_DLQR_H_
diff --git a/y2018/control_loops/superstructure/arm/dynamics.h b/y2018/control_loops/superstructure/arm/dynamics.h
index 832796a..d0194b0 100644
--- a/y2018/control_loops/superstructure/arm/dynamics.h
+++ b/y2018/control_loops/superstructure/arm/dynamics.h
@@ -104,6 +104,33 @@
         .finished();
   }
 
+  // Calculates the acceleration given the current augmented kalman filter state
+  // and control input.
+  static const ::Eigen::Matrix<double, 6, 1> EKFAcceleration(
+      const ::Eigen::Matrix<double, 6, 1> &X,
+      const ::Eigen::Matrix<double, 2, 1> &U) {
+    ::Eigen::Matrix<double, 2, 2> K1;
+    ::Eigen::Matrix<double, 2, 2> K2;
+
+    MatriciesForState(X.block<4, 1>(0, 0), &K1, &K2);
+
+    const ::Eigen::Matrix<double, 2, 1> velocity =
+        (::Eigen::Matrix<double, 2, 1>() << X(1, 0), X(3, 0)).finished();
+
+    const ::Eigen::Matrix<double, 2, 1> torque =
+        K3 *
+            (U +
+             (::Eigen::Matrix<double, 2, 1>() << X(4, 0), X(5, 0)).finished()) -
+        K4 * velocity;
+
+    const ::Eigen::Matrix<double, 2, 1> accel =
+        K1.inverse() * (torque - K2 * velocity);
+
+    return (::Eigen::Matrix<double, 6, 1>() << X(1, 0), accel(0, 0), X(3, 0),
+            accel(1, 0), 0.0, 0.0)
+        .finished();
+  }
+
   // Calculates the voltage required to follow the trajectory.  This requires
   // knowing the current state, desired angular velocity and acceleration.
   static const ::Eigen::Matrix<double, 2, 1> FF_U(
@@ -124,6 +151,13 @@
     return ::frc971::control_loops::RungeKutta(Dynamics::Acceleration, X, U,
                                                dt);
   }
+
+  static const ::Eigen::Matrix<double, 6, 1> UnboundedEKFDiscreteDynamics(
+      const ::Eigen::Matrix<double, 6, 1> &X,
+      const ::Eigen::Matrix<double, 2, 1> &U, double dt) {
+    return ::frc971::control_loops::RungeKutta(Dynamics::EKFAcceleration, X, U,
+                                               dt);
+  }
 };
 
 }  // namespace arm
diff --git a/y2018/control_loops/superstructure/arm/ekf.cc b/y2018/control_loops/superstructure/arm/ekf.cc
new file mode 100644
index 0000000..93c88d4
--- /dev/null
+++ b/y2018/control_loops/superstructure/arm/ekf.cc
@@ -0,0 +1,96 @@
+#include "y2018/control_loops/superstructure/arm/ekf.h"
+
+#include "Eigen/Dense"
+#include <iostream>
+
+#include "frc971/control_loops/jacobian.h"
+#include "y2018/control_loops/superstructure/arm/dynamics.h"
+
+namespace y2018 {
+namespace control_loops {
+namespace superstructure {
+namespace arm {
+
+namespace {
+// TODO(austin): When tuning this, make sure to verify that you are waiting
+// enough cycles to make sure it converges at startup. Otherwise you will have a
+// bad day.
+const ::Eigen::Matrix<double, 6, 6> Q_covariance(
+    (::Eigen::DiagonalMatrix<double, 6>().diagonal() << ::std::pow(0.1, 2),
+     ::std::pow(2.0, 2), ::std::pow(0.1, 2), ::std::pow(2.0, 2),
+     ::std::pow(0.80, 2), ::std::pow(0.70, 2))
+        .finished()
+        .asDiagonal());
+}  // namespace
+
+EKF::EKF() {
+  X_hat_.setZero();
+  P_ = Q_covariance;
+  //::std::cout << "Reset P: " << P_ << ::std::endl;
+  // TODO(austin): Running the EKF 2000 cycles works, but isn't super clever.
+  // We could just solve the DARE.
+
+  for (int i = 0; i < 1000; ++i) {
+    Predict(::Eigen::Matrix<double, 2, 1>::Zero(), 0.00505);
+    Correct(::Eigen::Matrix<double, 2, 1>::Zero(), 0.00505);
+  }
+  //::std::cout << "Stabilized P: " << P_ << ::std::endl;
+  for (int i = 0; i < 1000; ++i) {
+    Predict(::Eigen::Matrix<double, 2, 1>::Zero(), 0.00505);
+    Correct(::Eigen::Matrix<double, 2, 1>::Zero(), 0.00505);
+  }
+  //::std::cout << "Really stabilized P: " << P_ << ::std::endl;
+  P_reset_ = P_;
+
+  Reset(::Eigen::Matrix<double, 4, 1>::Zero());
+}
+
+void EKF::Reset(const ::Eigen::Matrix<double, 4, 1> &X) {
+  X_hat_.setZero();
+  P_ = P_reset_;
+  X_hat_.block<4, 1>(0, 0) = X;
+}
+
+void EKF::Predict(const ::Eigen::Matrix<double, 2, 1> &U, double dt) {
+  const ::Eigen::Matrix<double, 6, 6> A =
+      ::frc971::control_loops::NumericalJacobianX<6, 2>(
+          Dynamics::UnboundedEKFDiscreteDynamics, X_hat_, U, dt);
+
+  X_hat_ = Dynamics::UnboundedEKFDiscreteDynamics(X_hat_, U, dt);
+  P_ = A * P_ * A.transpose() + Q_covariance;
+}
+
+void EKF::Correct(const ::Eigen::Matrix<double, 2, 1> &Y, double /*dt*/) {
+  const ::Eigen::Matrix<double, 2, 2> R_covariance(
+      (::Eigen::DiagonalMatrix<double, 2>().diagonal() << ::std::pow(0.01, 2),
+       ::std::pow(0.01, 2))
+          .finished()
+          .asDiagonal());
+  // H is the jacobian of the h(x) measurement prediction function
+  const ::Eigen::Matrix<double, 2, 6> H_jacobian =
+      (::Eigen::Matrix<double, 2, 6>() << 1.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+      0.0, 0.0, 1.0, 0.0, 0.0, 0.0)
+          .finished();
+
+  // Update step Measurement residual error of proximal and distal joint
+  // angles.
+  const ::Eigen::Matrix<double, 2, 1> Y_hat =
+      Y - (::Eigen::Matrix<double, 2, 1>() << X_hat_(0), X_hat_(2)).finished();
+  // Residual covariance
+  const ::Eigen::Matrix<double, 2, 2> S =
+      H_jacobian * P_ * H_jacobian.transpose() + R_covariance;
+
+  // K is the Near-optimal Kalman gain
+  const ::Eigen::Matrix<double, 6, 2> kalman_gain =
+      P_ * H_jacobian.transpose() * S.inverse();
+  // Updated state estimate
+  X_hat_ = X_hat_ + kalman_gain * Y_hat;
+  // Updated covariance estimate
+  P_ = (::Eigen::Matrix<double, 6, 6>::Identity() - kalman_gain * H_jacobian) *
+       P_;
+}
+
+}  // namespace arm
+}  // namespace superstructure
+}  // namespace control_loops
+}  // namespace y2018
diff --git a/y2018/control_loops/superstructure/arm/ekf.h b/y2018/control_loops/superstructure/arm/ekf.h
new file mode 100644
index 0000000..9f45303
--- /dev/null
+++ b/y2018/control_loops/superstructure/arm/ekf.h
@@ -0,0 +1,44 @@
+#ifndef Y2018_CONTROL_LOOPS_SUPERSTRUCTURE_ARM_EKF_H_
+#define Y2018_CONTROL_LOOPS_SUPERSTRUCTURE_ARM_EKF_H_
+
+#include "Eigen/Dense"
+
+namespace y2018 {
+namespace control_loops {
+namespace superstructure {
+namespace arm {
+
+// An extended kalman filter for the Arm.
+// Our states are:
+//   [theta0, omega0, theta1, omega1, voltage error0, voltage error1]
+class EKF {
+ public:
+  EKF();
+
+  // Resets the internal state back to X.  Resets the torque disturbance to 0.
+  void Reset(const ::Eigen::Matrix<double, 4, 1> &X);
+  // TODO(austin): Offset the internal state when we calibrate.
+
+  // Corrects the state estimate with the provided sensor reading.
+  void Correct(const ::Eigen::Matrix<double, 2, 1> &Y, double dt);
+
+  // Predicts the next state and covariance given the control input.
+  void Predict(const ::Eigen::Matrix<double, 2, 1> &U, double dt);
+
+  // Returns the current state and covariance estimates.
+  const ::Eigen::Matrix<double, 6, 1> &X_hat() const { return X_hat_; }
+  double X_hat(int i) const { return X_hat_(i); }
+  const ::Eigen::Matrix<double, 6, 6> &P() const { return P_; }
+
+ private:
+  ::Eigen::Matrix<double, 6, 1> X_hat_;
+  ::Eigen::Matrix<double, 6, 6> P_;
+  ::Eigen::Matrix<double, 6, 6> P_reset_;
+};
+
+}  // namespace arm
+}  // namespace superstructure
+}  // namespace control_loops
+}  // namespace y2018
+
+#endif  // Y2018_CONTROL_LOOPS_SUPERSTRUCTURE_ARM_EKF_H_
diff --git a/y2018/control_loops/superstructure/arm/trajectory.cc b/y2018/control_loops/superstructure/arm/trajectory.cc
index 14ec413..80fca9c 100644
--- a/y2018/control_loops/superstructure/arm/trajectory.cc
+++ b/y2018/control_loops/superstructure/arm/trajectory.cc
@@ -314,8 +314,8 @@
   goal_acceleration_ = 0.0;
 }
 
-::Eigen::Matrix<double, 2, 4> TrajectoryFollower::K_at_state(
-    const ::Eigen::Matrix<double, 4, 1> &X,
+::Eigen::Matrix<double, 2, 6> TrajectoryFollower::K_at_state(
+    const ::Eigen::Matrix<double, 6, 1> &X,
     const ::Eigen::Matrix<double, 2, 1> &U) {
   constexpr double q_pos = 0.2;
   constexpr double q_vel = 4.0;
@@ -335,23 +335,29 @@
           .asDiagonal();
 
   const ::Eigen::Matrix<double, 4, 4> final_A =
-      ::frc971::control_loops::NumericalJacobianX(
-          Dynamics::UnboundedDiscreteDynamics, X, U, 0.00505);
+      ::frc971::control_loops::NumericalJacobianX<4, 2>(
+          Dynamics::UnboundedDiscreteDynamics, X.block<4, 1>(0, 0), U, 0.00505);
   const ::Eigen::Matrix<double, 4, 2> final_B =
-      ::frc971::control_loops::NumericalJacobianU(
-          Dynamics::UnboundedDiscreteDynamics, X, U, 0.00505);
+      ::frc971::control_loops::NumericalJacobianU<4, 2>(
+          Dynamics::UnboundedDiscreteDynamics, X.block<4, 1>(0, 0), U, 0.00505);
 
-  ::Eigen::Matrix<double, 2, 4> K;
   ::Eigen::Matrix<double, 4, 4> S;
-  ::frc971::controls::dlqr<4, 2>(final_A, final_B, Q, R, &K, &S);
+  ::Eigen::Matrix<double, 2, 4> sub_K;
+  ::frc971::controls::dlqr<4, 2>(final_A, final_B, Q, R, &sub_K, &S);
+
+  ::Eigen::Matrix<double, 2, 6> K;
+  K.setZero();
+  K.block<2, 4>(0, 0) = sub_K;
+  K(0, 4) = 1.0;
+  K(1, 5) = 1.0;
   return K;
 }
 
 void TrajectoryFollower::USaturationSearch(
     double goal_distance, double last_goal_distance, double goal_velocity,
     double last_goal_velocity, double fraction_along_path,
-    const ::Eigen::Matrix<double, 2, 4> &K,
-    const ::Eigen::Matrix<double, 4, 1> &X, const Trajectory &trajectory,
+    const ::Eigen::Matrix<double, 2, 6> &K,
+    const ::Eigen::Matrix<double, 6, 1> &X, const Trajectory &trajectory,
     ::Eigen::Matrix<double, 2, 1> *U, double *saturation_goal_velocity,
     double *saturation_goal_acceleration) {
   double saturation_goal_distance =
@@ -370,9 +376,9 @@
   const ::Eigen::Matrix<double, 2, 1> alpha_t =
       trajectory.AlphaT(saturation_goal_distance, *saturation_goal_velocity,
                         *saturation_goal_acceleration);
-  const ::Eigen::Matrix<double, 4, 1> R = trajectory.R(theta_t, omega_t);
+  const ::Eigen::Matrix<double, 6, 1> R = trajectory.R(theta_t, omega_t);
   const ::Eigen::Matrix<double, 2, 1> U_ff =
-      Dynamics::FF_U(R, omega_t, alpha_t);
+      Dynamics::FF_U(R.block<4, 1>(0, 0), omega_t, alpha_t);
 
   *U = U_ff + K * (R - X);
 }
@@ -405,7 +411,7 @@
   return next_goal;
 }
 
-void TrajectoryFollower::Update(const ::Eigen::Matrix<double, 4, 1> &X,
+void TrajectoryFollower::Update(const ::Eigen::Matrix<double, 6, 1> &X,
                                 double dt, double vmax) {
   // To avoid exposing the new goals before the outer code has a chance to
   // querry the internal state, move to the new goals here.
@@ -441,11 +447,11 @@
   const ::Eigen::Matrix<double, 2, 1> alpha_t =
       trajectory_->AlphaT(goal_(0), goal_(1), goal_acceleration_);
 
-  const ::Eigen::Matrix<double, 4, 1> R = trajectory_->R(theta_t, omega_t);
+  const ::Eigen::Matrix<double, 6, 1> R = trajectory_->R(theta_t, omega_t);
 
-  U_ff_ = Dynamics::FF_U(R, omega_t, alpha_t);
+  U_ff_ = Dynamics::FF_U(R.block<4, 1>(0, 0), omega_t, alpha_t);
 
-  const ::Eigen::Matrix<double, 2, 4> K = K_at_state(X, U_ff_);
+  const ::Eigen::Matrix<double, 2, 6> K = K_at_state(X, U_ff_);
   U_ = U_unsaturated_ = U_ff_ + K * (R - X);
 
   // Ok, now we know if we are staturated or not.  If we are, time to search
@@ -490,6 +496,8 @@
 
     goal_acceleration_ = trajectory_->InterpolateAcceleration(
         goal_(0), next_goal_(0), goal_(1), next_goal_(1));
+
+    U_ = U_.array().max(-12.0).min(12.0);
   }
 }
 
diff --git a/y2018/control_loops/superstructure/arm/trajectory.h b/y2018/control_loops/superstructure/arm/trajectory.h
index 5ecc783..b699fe0 100644
--- a/y2018/control_loops/superstructure/arm/trajectory.h
+++ b/y2018/control_loops/superstructure/arm/trajectory.h
@@ -1,6 +1,10 @@
 #ifndef Y2018_CONTROL_LOOPS_SUPERSTRUCTURE_ARM_TRAJECTORY_H_
 #define Y2018_CONTROL_LOOPS_SUPERSTRUCTURE_ARM_TRAJECTORY_H_
 
+#include <array>
+#include <initializer_list>
+#include <vector>
+
 #include "Eigen/Dense"
 
 namespace y2018 {
@@ -171,9 +175,10 @@
   // and that we are at point d between the two.
   static double InterpolateVelocity(double d, double d0, double d1, double v0,
                                     double v1) {
-    // We don't support negative velocities.
+    // We don't support negative velocities.  Report 0 velocity in that case.
+    // TODO(austin): Verify this doesn't show up in the real world.
     if (v0 < 0 || v1 < 0) {
-      abort();
+      return 0.0;
     }
     if (d <= d0) {
       return v0;
@@ -239,11 +244,11 @@
   }
 
   // Converts a theta and omega vector to a full state vector.
-  static ::Eigen::Matrix<double, 4, 1> R(
+  static ::Eigen::Matrix<double, 6, 1> R(
       ::Eigen::Matrix<double, 2, 1> theta_t,
       ::Eigen::Matrix<double, 2, 1> omega_t) {
-    return (::Eigen::Matrix<double, 4, 1>() << theta_t(0, 0), omega_t(0, 0),
-            theta_t(1, 0), omega_t(1, 0))
+    return (::Eigen::Matrix<double, 6, 1>() << theta_t(0, 0), omega_t(0, 0),
+            theta_t(1, 0), omega_t(1, 0), 0.0, 0.0)
         .finished();
   }
 
@@ -299,8 +304,8 @@
   void Reset();
 
   // Returns the controller gain at the provided state.
-  ::Eigen::Matrix<double, 2, 4> K_at_state(
-      const ::Eigen::Matrix<double, 4, 1> &X,
+  ::Eigen::Matrix<double, 2, 6> K_at_state(
+      const ::Eigen::Matrix<double, 6, 1> &X,
       const ::Eigen::Matrix<double, 2, 1> &U);
 
   // Returns the voltage, velocity and acceleration if we were to be partially
@@ -308,8 +313,8 @@
   void USaturationSearch(double goal_distance, double last_goal_distance,
                          double goal_velocity, double last_goal_velocity,
                          double fraction_along_path,
-                         const ::Eigen::Matrix<double, 2, 4> &K,
-                         const ::Eigen::Matrix<double, 4, 1> &X,
+                         const ::Eigen::Matrix<double, 2, 6> &K,
+                         const ::Eigen::Matrix<double, 6, 1> &X,
                          const Trajectory &trajectory,
                          ::Eigen::Matrix<double, 2, 1> *U,
                          double *saturation_goal_velocity,
@@ -321,7 +326,7 @@
       const ::Eigen::Matrix<double, 2, 1> &goal, double vmax, double dt);
 
   // Plans the next cycle and updates the internal state for consumption.
-  void Update(const ::Eigen::Matrix<double, 4, 1> &X, double dt, double vmax);
+  void Update(const ::Eigen::Matrix<double, 6, 1> &X, double dt, double vmax);
 
   // Returns the goal acceleration for this cycle.
   double goal_acceleration() const { return goal_acceleration_; }
diff --git a/y2018/control_loops/superstructure/arm/trajectory_plot.cc b/y2018/control_loops/superstructure/arm/trajectory_plot.cc
index 0b808f8..bfcf915 100644
--- a/y2018/control_loops/superstructure/arm/trajectory_plot.cc
+++ b/y2018/control_loops/superstructure/arm/trajectory_plot.cc
@@ -3,6 +3,7 @@
 #include "third_party/matplotlib-cpp/matplotlibcpp.h"
 #include "y2018/control_loops/superstructure/arm/demo_path.h"
 #include "y2018/control_loops/superstructure/arm/dynamics.h"
+#include "y2018/control_loops/superstructure/arm/ekf.h"
 
 namespace y2018 {
 namespace control_loops {
@@ -58,8 +59,9 @@
     const ::Eigen::Matrix<double, 2, 1> alpha_t =
         trajectory.AlphaT(distance, goal_velocity, goal_acceleration);
 
-    const ::Eigen::Matrix<double, 4, 1> R = trajectory.R(theta_t, omega_t);
-    const ::Eigen::Matrix<double, 2, 1> U = Dynamics::FF_U(R, omega_t, alpha_t);
+    const ::Eigen::Matrix<double, 6, 1> R = trajectory.R(theta_t, omega_t);
+    const ::Eigen::Matrix<double, 2, 1> U =
+        Dynamics::FF_U(R.block<4, 1>(0, 0), omega_t, alpha_t);
 
     Uff0_distance_array.push_back(U(0));
     Uff1_distance_array.push_back(U(1));
@@ -96,10 +98,22 @@
   ::std::vector<double> uff1_array;
   ::std::vector<double> u0_array;
   ::std::vector<double> u1_array;
+  ::std::vector<double> theta0_hat_t_array;
+  ::std::vector<double> omega0_hat_t_array;
+  ::std::vector<double> theta1_hat_t_array;
+  ::std::vector<double> omega1_hat_t_array;
+  ::std::vector<double> torque0_hat_t_array;
+  ::std::vector<double> torque1_hat_t_array;
 
-  while (t < 1.0) {
+  EKF arm_ekf;
+  arm_ekf.Reset(X);
+
+  while (t < 3.0) {
     t_array.push_back(t);
-    follower.Update(X, sim_dt, vmax);
+    arm_ekf.Correct((::Eigen::Matrix<double, 2, 1>() << X(0), X(2)).finished(),
+                    sim_dt);
+    // TODO(austin): Compensate for torque disturbance.
+    follower.Update(arm_ekf.X_hat(), sim_dt, vmax);
 
     const ::Eigen::Matrix<double, 2, 1> theta_t =
         trajectory.ThetaT(follower.goal()(0));
@@ -118,6 +132,12 @@
     omega0_t_array.push_back(X(1));
     theta1_t_array.push_back(X(2));
     omega1_t_array.push_back(X(3));
+    theta0_hat_t_array.push_back(arm_ekf.X_hat(0));
+    omega0_hat_t_array.push_back(arm_ekf.X_hat(1));
+    theta1_hat_t_array.push_back(arm_ekf.X_hat(2));
+    omega1_hat_t_array.push_back(arm_ekf.X_hat(3));
+    torque0_hat_t_array.push_back(arm_ekf.X_hat(4));
+    torque1_hat_t_array.push_back(arm_ekf.X_hat(5));
 
     distance_t_array.push_back(follower.goal()(0));
     velocity_t_array.push_back(follower.goal()(1));
@@ -126,10 +146,15 @@
     u0_unsaturated_array.push_back(follower.U_unsaturated()(0));
     u1_unsaturated_array.push_back(follower.U_unsaturated()(1));
 
-    const ::Eigen::Matrix<double, 4, 1> xdot =
-        Dynamics::Acceleration(X, follower.U());
+    ::Eigen::Matrix<double, 2, 1> actual_U = follower.U();
+    // Add in a disturbance force to see how well the arm learns it.
+    actual_U.array() += 1.0;
 
-    X = Dynamics::UnboundedDiscreteDynamics(X, follower.U(), sim_dt);
+    const ::Eigen::Matrix<double, 4, 1> xdot =
+        Dynamics::Acceleration(X, actual_U);
+
+    X = Dynamics::UnboundedDiscreteDynamics(X, actual_U, sim_dt);
+    arm_ekf.Predict(follower.U(), sim_dt);
 
     alpha0_t_array.push_back(xdot(1));
     alpha1_t_array.push_back(xdot(3));
@@ -185,9 +210,11 @@
   matplotlibcpp::plot(t_array, omega0_goal_t_array,
                       {{"label", "omega0_t_goal"}});
   matplotlibcpp::plot(t_array, omega0_t_array, {{"label", "omega0_t"}});
+  matplotlibcpp::plot(t_array, omega0_hat_t_array, {{"label", "omega0_hat_t"}});
   matplotlibcpp::plot(t_array, omega1_goal_t_array,
                       {{"label", "omega1_t_goal"}});
   matplotlibcpp::plot(t_array, omega1_t_array, {{"label", "omega1_t"}});
+  matplotlibcpp::plot(t_array, omega1_hat_t_array, {{"label", "omega1_hat_t"}});
   matplotlibcpp::legend();
 
   matplotlibcpp::figure();
@@ -198,6 +225,8 @@
   matplotlibcpp::plot(t_array, u1_unsaturated_array, {{"label", "u1_full"}});
   matplotlibcpp::plot(t_array, u1_array, {{"label", "u1"}});
   matplotlibcpp::plot(t_array, uff1_array, {{"label", "uff1"}});
+  matplotlibcpp::plot(t_array, torque0_hat_t_array, {{"label", "torque0_hat"}});
+  matplotlibcpp::plot(t_array, torque1_hat_t_array, {{"label", "torque1_hat"}});
   matplotlibcpp::legend();
 
   matplotlibcpp::figure();
@@ -205,9 +234,11 @@
   matplotlibcpp::plot(t_array, theta0_goal_t_array,
                       {{"label", "theta0_t_goal"}});
   matplotlibcpp::plot(t_array, theta0_t_array, {{"label", "theta0_t"}});
+  matplotlibcpp::plot(t_array, theta0_hat_t_array, {{"label", "theta0_hat_t"}});
   matplotlibcpp::plot(t_array, theta1_goal_t_array,
                       {{"label", "theta1_t_goal"}});
   matplotlibcpp::plot(t_array, theta1_t_array, {{"label", "theta1_t"}});
+  matplotlibcpp::plot(t_array, theta1_hat_t_array, {{"label", "theta1_hat_t"}});
   matplotlibcpp::legend();
 
 /*
diff --git a/y2018/control_loops/superstructure/arm/trajectory_test.cc b/y2018/control_loops/superstructure/arm/trajectory_test.cc
index 3c5d6a1..9687e40 100644
--- a/y2018/control_loops/superstructure/arm/trajectory_test.cc
+++ b/y2018/control_loops/superstructure/arm/trajectory_test.cc
@@ -3,6 +3,7 @@
 #include "gtest/gtest.h"
 #include "y2018/control_loops/superstructure/arm/demo_path.h"
 #include "y2018/control_loops/superstructure/arm/dynamics.h"
+#include "y2018/control_loops/superstructure/arm/ekf.h"
 
 namespace y2018 {
 namespace control_loops {
@@ -157,11 +158,17 @@
     X << theta_t(0), 0.0, theta_t(1), 0.0;
   }
 
+  EKF arm_ekf;
+  arm_ekf.Reset(X);
+
   TrajectoryFollower follower(&path, &trajectory, alpha_unitizer);
   constexpr double sim_dt = 0.00505;
   while (t < 1.0) {
-    follower.Update(X, sim_dt, vmax);
+    arm_ekf.Correct((::Eigen::Matrix<double, 2, 1>() << X(0), X(2)).finished(),
+                    sim_dt);
+    follower.Update(arm_ekf.X_hat(), sim_dt, vmax);
     X = Dynamics::UnboundedDiscreteDynamics(X, follower.U(), sim_dt);
+    arm_ekf.Predict(follower.U(), sim_dt);
     t += sim_dt;
   }