Tuned superstructure loop and added feed forwards.

Change-Id: Ia2e3a1746529a4c27395f2e9b6e875c5cddb7616
diff --git a/y2016/control_loops/python/BUILD b/y2016/control_loops/python/BUILD
index bd5c4bd..9396fea 100644
--- a/y2016/control_loops/python/BUILD
+++ b/y2016/control_loops/python/BUILD
@@ -56,11 +56,10 @@
     'intake.py',
   ],
   deps = [
-    ':polydrivetrain_lib',
     '//external:python-gflags',
     '//external:python-glog',
     '//frc971/control_loops/python:controls',
-  ]
+  ],
 )
 
 py_binary(
@@ -69,11 +68,11 @@
     'shoulder.py',
   ],
   deps = [
-    ':polydrivetrain_lib',
+    '//aos/common/util:py_trapezoid_profile',
     '//external:python-gflags',
     '//external:python-glog',
     '//frc971/control_loops/python:controls',
-  ]
+  ],
 )
 
 py_binary(
@@ -82,9 +81,64 @@
     'wrist.py',
   ],
   deps = [
-    ':polydrivetrain_lib',
+    '//aos/common/util:py_trapezoid_profile',
     '//external:python-gflags',
     '//external:python-glog',
     '//frc971/control_loops/python:controls',
-  ]
+  ],
+)
+
+py_library(
+  name = 'wrist_lib',
+  srcs = [
+    'wrist.py',
+  ],
+  deps = [
+    '//aos/common/util:py_trapezoid_profile',
+    '//external:python-gflags',
+    '//external:python-glog',
+    '//frc971/control_loops/python:controls',
+  ],
+)
+
+py_library(
+  name = 'shoulder_lib',
+  srcs = [
+    'shoulder.py',
+  ],
+  deps = [
+    '//aos/common/util:py_trapezoid_profile',
+    '//external:python-gflags',
+    '//external:python-glog',
+    '//frc971/control_loops/python:controls',
+  ],
+)
+
+py_library(
+  name = 'arm_lib',
+  srcs = [
+    'arm.py',
+  ],
+  deps = [
+    ':wrist_lib',
+    ':shoulder_lib',
+    '//external:python-gflags',
+    '//external:python-glog',
+    '//frc971/control_loops/python:controls',
+    '//aos/common/util:py_trapezoid_profile',
+  ],
+)
+py_binary(
+  name = 'arm',
+  srcs = [
+    'arm.py',
+  ],
+  deps = [
+    ':wrist_lib',
+    ':shoulder_lib',
+    '//external:python-gflags',
+    '//external:python-glog',
+    '//frc971/control_loops/python:controls',
+    '//aos/common/util:py_trapezoid_profile',
+  ],
 )
diff --git a/y2016/control_loops/python/arm.py b/y2016/control_loops/python/arm.py
new file mode 100644
index 0000000..28a704d
--- /dev/null
+++ b/y2016/control_loops/python/arm.py
@@ -0,0 +1,407 @@
+#!/usr/bin/python
+
+import numpy
+import sys
+import operator
+
+from frc971.control_loops.python import control_loop
+from frc971.control_loops.python import controls
+
+from y2016.control_loops.python.shoulder import Shoulder, IntegralShoulder
+from y2016.control_loops.python.wrist import Wrist, IntegralWrist
+from aos.common.util.trapezoid_profile import TrapizoidProfile
+
+from matplotlib import pylab
+import gflags
+import glog
+
+FLAGS = gflags.FLAGS
+
+try:
+  gflags.DEFINE_bool('plot', False, 'If true, plot the loop response.')
+except gflags.DuplicateFlagError:
+  pass
+
+
+class Arm(control_loop.ControlLoop):
+  def __init__(self, name="Arm"):
+    super(Arm, self).__init__(name=name)
+    self._shoulder = Shoulder(name=name)
+    self._shooter = Wrist(name=name)
+
+    # Do a coordinate transformation.
+    # X_shooter_grounded = X_shooter + X_shoulder
+    # dX_shooter_grounded/dt = A_shooter * X_shooter + A_shoulder * X_shoulder +
+    #                          B_shoulder * U_shoulder + B_shooter * U_shooter
+    # dX_shooter_grounded/dt = A_shooter * (X_shooter_grounded - X_shoulder) +
+    #                          A_shoulder * X_shoulder + B_shooter * U_shooter + B_shoulder * U_shoulder
+    # X = [X_shoulder; X_shooter + X_shoulder]
+    # dX/dt = [A_shoulder                       0] [X_shoulder        ] + [B_shoulder         0] [U_shoulder]
+    #         [(A_shoulder - A_shooter) A_shooter] [X_shooter_grounded] + [B_shoulder B_shooter] [ U_shooter]
+    # Y_shooter_grounded = Y_shooter + Y_shoulder
+
+    self.A_continuous = numpy.matrix(numpy.zeros((4, 4)))
+    self.A_continuous[0:2, 0:2] = self._shoulder.A_continuous
+    self.A_continuous[2:4, 0:2] = (self._shoulder.A_continuous -
+                                   self._shooter.A_continuous)
+    self.A_continuous[2:4, 2:4] = self._shooter.A_continuous
+
+    self.B_continuous = numpy.matrix(numpy.zeros((4, 2)))
+    self.B_continuous[0:2, 0:1] = self._shoulder.B_continuous
+    self.B_continuous[2:4, 1:2] = self._shooter.B_continuous
+    self.B_continuous[2:4, 0:1] = self._shoulder.B_continuous
+
+    self.C = numpy.matrix(numpy.zeros((2, 4)))
+    self.C[0:1, 0:2] = self._shoulder.C
+    self.C[1:2, 0:2] = -self._shoulder.C
+    self.C[1:2, 2:4] = self._shooter.C
+
+    # D is 0 for all our loops.
+    self.D = numpy.matrix(numpy.zeros((2, 2)))
+
+    self.dt = 0.005
+
+    self.A, self.B = self.ContinuousToDiscrete(
+        self.A_continuous, self.B_continuous, self.dt)
+
+    # Cost of error
+    self.Q = numpy.matrix(numpy.zeros((4, 4)))
+    q_pos_shoulder = 0.014
+    q_vel_shoulder = 4.00
+    q_pos_shooter = 0.014
+    q_vel_shooter = 4.00
+    self.Q[0, 0] = 1.0 / q_pos_shoulder ** 2.0
+    self.Q[1, 1] = 1.0 / q_vel_shoulder ** 2.0
+    self.Q[2, 2] = 1.0 / q_pos_shooter ** 2.0
+    self.Q[3, 3] = 1.0 / q_vel_shooter ** 2.0
+
+    # Cost of control effort
+    self.R = numpy.matrix(numpy.zeros((2, 2)))
+    r_voltage = 1.0 / 12.0
+    self.R[0, 0] = r_voltage ** 2.0
+    self.R[1, 1] = r_voltage ** 2.0
+
+    self.Kff = controls.TwoStateFeedForwards(self.B, self.Q)
+
+    glog.debug('Shoulder K')
+    glog.debug(self._shoulder.K)
+
+    # Compute controller gains.
+    # self.K = controls.dlqr(self.A, self.B, self.Q, self.R)
+    self.K = numpy.matrix(numpy.zeros((2, 4)))
+    self.K[0:1, 0:2] = self._shoulder.K
+    self.K[1:2, 0:2] = (
+        -self.Kff[1:2, 2:4] * self.B[2:4, 0:1] * self._shoulder.K
+        + self.Kff[1:2, 2:4] * self.A[2:4, 0:2])
+    self.K[1:2, 2:4] = self._shooter.K
+
+    glog.debug('Arm controller %s', repr(self.K))
+
+    # Cost of error
+    self.Q = numpy.matrix(numpy.zeros((4, 4)))
+    q_pos_shoulder = 0.05
+    q_vel_shoulder = 2.65
+    q_pos_shooter = 0.05
+    q_vel_shooter = 2.65
+    self.Q[0, 0] = q_pos_shoulder ** 2.0
+    self.Q[1, 1] = q_vel_shoulder ** 2.0
+    self.Q[2, 2] = q_pos_shooter ** 2.0
+    self.Q[3, 3] = q_vel_shooter ** 2.0
+
+    # Cost of control effort
+    self.R = numpy.matrix(numpy.zeros((2, 2)))
+    r_voltage = 0.025
+    self.R[0, 0] = r_voltage ** 2.0
+    self.R[1, 1] = r_voltage ** 2.0
+
+    self.KalmanGain, self.Q_steady = controls.kalman(
+        A=self.A, B=self.B, C=self.C, Q=self.Q, R=self.R)
+    self.L = self.A * self.KalmanGain
+
+    self.U_max = numpy.matrix([[12.0], [12.0]])
+    self.U_min = numpy.matrix([[-12.0], [-12.0]])
+
+    self.InitializeState()
+
+
+class IntegralArm(Arm):
+  def __init__(self, name="IntegralArm"):
+    super(IntegralArm, self).__init__(name=name)
+
+    self.A_continuous_unaugmented = self.A_continuous
+    self.B_continuous_unaugmented = self.B_continuous
+
+    self.A_continuous = numpy.matrix(numpy.zeros((6, 6)))
+    self.A_continuous[0:4, 0:4] = self.A_continuous_unaugmented
+    self.A_continuous[0:4, 4:6] = self.B_continuous_unaugmented
+
+    self.B_continuous = numpy.matrix(numpy.zeros((6, 2)))
+    self.B_continuous[0:4, 0:2] = self.B_continuous_unaugmented
+
+    self.C_unaugmented = self.C
+    self.C = numpy.matrix(numpy.zeros((2, 6)))
+    self.C[0:2, 0:4] = self.C_unaugmented
+
+    self.A, self.B = self.ContinuousToDiscrete(self.A_continuous, self.B_continuous, self.dt)
+
+    q_pos_shoulder = 0.08
+    q_vel_shoulder = 4.00
+    q_voltage_shoulder = 6.0
+    q_pos_shooter = 0.08
+    q_vel_shooter = 4.00
+    q_voltage_shooter = 6.0
+    self.Q = numpy.matrix(numpy.zeros((6, 6)))
+    self.Q[0, 0] = q_pos_shoulder ** 2.0
+    self.Q[1, 1] = q_vel_shoulder ** 2.0
+    self.Q[2, 2] = q_pos_shooter ** 2.0
+    self.Q[3, 3] = q_vel_shooter ** 2.0
+    self.Q[4, 4] = q_voltage_shoulder ** 2.0
+    self.Q[5, 5] = q_voltage_shooter ** 2.0
+
+    self.R = numpy.matrix(numpy.zeros((2, 2)))
+    r_pos = 0.05
+    self.R[0, 0] = r_pos ** 2.0
+    self.R[1, 1] = r_pos ** 2.0
+
+    self.KalmanGain, self.Q_steady = controls.kalman(
+        A=self.A, B=self.B, C=self.C, Q=self.Q, R=self.R)
+    self.L = self.A * self.KalmanGain
+
+    self.K_unaugmented = self.K
+    self.K = numpy.matrix(numpy.zeros((2, 6)))
+    self.K[0:2, 0:4] = self.K_unaugmented
+    self.K[0, 4] = 1
+    self.K[1, 5] = 1
+
+    self.Kff = numpy.concatenate((self.Kff, numpy.matrix(numpy.zeros((2, 2)))), axis=1)
+
+    self.InitializeState()
+
+
+class ScenarioPlotter(object):
+  def __init__(self):
+    # Various lists for graphing things.
+    self.t = []
+    self.x_shoulder = []
+    self.v_shoulder = []
+    self.a_shoulder = []
+    self.x_hat_shoulder = []
+    self.u_shoulder = []
+    self.offset_shoulder = []
+    self.x_shooter = []
+    self.v_shooter = []
+    self.a_shooter = []
+    self.x_hat_shooter = []
+    self.u_shooter = []
+    self.offset_shooter = []
+    self.goal_x_shoulder = []
+    self.goal_v_shoulder = []
+    self.goal_x_shooter = []
+    self.goal_v_shooter = []
+
+  def run_test(self, arm, end_goal,
+               iterations=200, controller=None, observer=None):
+    """Runs the plant with an initial condition and goal.
+
+      Args:
+        arm: Arm object to use.
+        end_goal: numpy.Matrix[6, 1], end goal state.
+        iterations: Number of timesteps to run the model for.
+        controller: Arm object to get K from, or None if we should
+            use arm.
+        observer: Arm object to use for the observer, or None if we should
+            use the actual state.
+    """
+
+    if controller is None:
+      controller = arm
+
+    vbat = 12.0
+
+    if self.t:
+      initial_t = self.t[-1] + arm.dt
+    else:
+      initial_t = 0
+
+    goal = numpy.concatenate((arm.X, numpy.matrix(numpy.zeros((2, 1)))), axis=0)
+    current_shoulder_goal = goal[0:2, 0].copy()
+    current_shooter_goal = goal[2:4, 0].copy()
+
+    shoulder_profile = TrapizoidProfile(arm.dt)
+    shoulder_profile.set_maximum_acceleration(50.0)
+    shoulder_profile.set_maximum_velocity(10.0)
+    shoulder_profile.SetGoal(current_shoulder_goal[0, 0])
+    shooter_profile = TrapizoidProfile(arm.dt)
+    shooter_profile.set_maximum_acceleration(50.0)
+    shooter_profile.set_maximum_velocity(10.0)
+    shooter_profile.SetGoal(current_shooter_goal[0, 0])
+
+    U_last = numpy.matrix(numpy.zeros((2, 1)))
+    for i in xrange(iterations):
+      X_hat = arm.X
+
+      if observer is not None:
+        observer.Y = arm.Y
+        observer.CorrectObserver(U_last)
+        self.offset_shoulder.append(observer.X_hat[4, 0])
+        self.offset_shooter.append(observer.X_hat[5, 0])
+
+      next_shoulder_goal = shoulder_profile.Update(end_goal[0, 0], end_goal[1, 0])
+      next_shooter_goal = shooter_profile.Update(end_goal[2, 0], end_goal[3, 0])
+
+      next_goal = numpy.concatenate((next_shoulder_goal, next_shooter_goal, numpy.matrix(numpy.zeros((2, 1)))), axis=0)
+      self.goal_x_shoulder.append(goal[0, 0])
+      self.goal_v_shoulder.append(goal[1, 0])
+      self.goal_x_shooter.append(goal[2, 0])
+      self.goal_v_shooter.append(goal[3, 0])
+
+      ff_U = controller.Kff * (next_goal - observer.A * goal)
+
+      if observer is not None:
+        X_hat = observer.X_hat
+        self.x_hat_shoulder.append(observer.X_hat[0, 0])
+        self.x_hat_shooter.append(observer.X_hat[2, 0])
+
+      U_uncapped = controller.K * (goal - X_hat) + ff_U
+      U = U_uncapped.copy()
+
+      U[0, 0] = numpy.clip(U[0, 0], -vbat, vbat)
+      U[1, 0] = numpy.clip(U[1, 0], -vbat, vbat)
+      self.x_shoulder.append(arm.X[0, 0])
+      self.x_shooter.append(arm.X[2, 0])
+
+      if self.v_shoulder:
+        last_v_shoulder = self.v_shoulder[-1]
+      else:
+        last_v_shoulder = 0
+      self.v_shoulder.append(arm.X[1, 0])
+      self.a_shoulder.append(
+          (self.v_shoulder[-1] - last_v_shoulder) / arm.dt)
+
+      if self.v_shooter:
+        last_v_shooter = self.v_shooter[-1]
+      else:
+        last_v_shooter = 0
+      self.v_shooter.append(arm.X[3, 0])
+      self.a_shooter.append(
+          (self.v_shooter[-1] - last_v_shooter) / arm.dt)
+
+      if i % 40 == 0:
+        # Test that if we move the shoulder, the shooter stays perfect.
+        #observer.X_hat[0, 0] += 0.20
+        #arm.X[0, 0] += 0.20
+        pass
+      U_error = numpy.matrix([[0.0], [0.0]])
+      # Kick it and see what happens.
+      #if (initial_t + i * arm.dt) % 0.4 > 0.2:
+        #U_error = numpy.matrix([[4.0], [0.0]])
+      #else:
+        #U_error = numpy.matrix([[-4.0], [0.0]])
+
+      arm.Update(U + U_error)
+
+      if observer is not None:
+        observer.PredictObserver(U)
+
+      self.t.append(initial_t + i * arm.dt)
+      self.u_shoulder.append(U[0, 0])
+      self.u_shooter.append(U[1, 0])
+
+      glog.debug('Time: %f', self.t[-1])
+
+      ff_U -= U_uncapped - U
+      goal = controller.A * goal + controller.B * ff_U
+
+      if U[0, 0] != U_uncapped[0, 0]:
+        glog.debug('Moving shoulder %s', repr(initial_t + i * arm.dt))
+        glog.debug('U error %s', repr(U_uncapped - U))
+        glog.debug('goal change is %s',
+                   repr(next_shoulder_goal -
+                        numpy.matrix([[goal[0, 0]], [goal[1, 0]]])))
+        shoulder_profile.MoveCurrentState(
+            numpy.matrix([[goal[0, 0]], [goal[1, 0]]]))
+      if U[1, 0] != U_uncapped[1, 0]:
+        glog.debug('Moving shooter %s', repr(initial_t + i * arm.dt))
+        glog.debug('U error %s', repr(U_uncapped - U))
+        shooter_profile.MoveCurrentState(
+            numpy.matrix([[goal[2, 0]], [goal[3, 0]]]))
+      U_last = U
+    glog.debug('End goal is %s', repr(end_goal))
+    glog.debug('last goal is %s', repr(goal))
+    glog.debug('End state is %s', repr(arm.X))
+
+
+  def Plot(self):
+    pylab.subplot(3, 1, 1)
+    pylab.plot(self.t, self.x_shoulder, label='x shoulder')
+    pylab.plot(self.t, self.goal_x_shoulder, label='goal x shoulder')
+    pylab.plot(self.t, self.x_hat_shoulder, label='x_hat shoulder')
+
+    pylab.plot(self.t, self.x_shooter, label='x shooter')
+    pylab.plot(self.t, self.x_hat_shooter, label='x_hat shooter')
+    pylab.plot(self.t, self.goal_x_shooter, label='goal x shooter')
+    pylab.plot(self.t, map(operator.add, self.x_shooter, self.x_shoulder),
+               label='x shooter ground')
+    pylab.plot(self.t, map(operator.add, self.x_hat_shooter, self.x_hat_shoulder),
+               label='x_hat shooter ground')
+    pylab.legend()
+
+    pylab.subplot(3, 1, 2)
+    pylab.plot(self.t, self.u_shoulder, label='u shoulder')
+    pylab.plot(self.t, self.offset_shoulder, label='voltage_offset shoulder')
+    pylab.plot(self.t, self.u_shooter, label='u shooter')
+    pylab.plot(self.t, self.offset_shooter, label='voltage_offset shooter')
+    pylab.legend()
+
+    pylab.subplot(3, 1, 3)
+    pylab.plot(self.t, self.a_shoulder, label='a_shoulder')
+    pylab.plot(self.t, self.a_shooter, label='a_shooter')
+    pylab.legend()
+
+    pylab.show()
+
+
+def main(argv):
+  argv = FLAGS(argv)
+  glog.init()
+
+  scenario_plotter = ScenarioPlotter()
+
+  arm = Arm()
+  arm_controller = IntegralArm()
+  arm_observer = IntegralArm()
+
+  # Test moving the shoulder with constant separation.
+  initial_X = numpy.matrix([[0.0], [0.0], [0.0], [0.0], [0.0], [0.0]])
+  R = numpy.matrix([[numpy.pi / 2.0],
+                    [0.0],
+                    [0.0], #[numpy.pi / 2.0],
+                    [0.0],
+                    [0.0],
+                    [0.0]])
+  arm.X = initial_X[0:4, 0]
+  arm_observer.X = initial_X
+
+  scenario_plotter.run_test(arm=arm,
+                            end_goal=R,
+                            iterations=300,
+                            controller=arm_controller,
+                            observer=arm_observer)
+
+  if len(argv) != 5:
+    glog.fatal('Expected .h file name and .cc file name for the wrist and integral wrist.')
+  else:
+    namespaces = ['y2016', 'control_loops', 'superstructure']
+    loop_writer = control_loop.ControlLoopWriter('Arm', [arm],
+                                                 namespaces=namespaces)
+    loop_writer.Write(argv[1], argv[2])
+
+    integral_loop_writer = control_loop.ControlLoopWriter(
+        'IntegralArm', [arm_controller], namespaces=namespaces)
+    integral_loop_writer.Write(argv[3], argv[4])
+
+  if FLAGS.plot:
+    scenario_plotter.Plot()
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv))
diff --git a/y2016/control_loops/python/intake.py b/y2016/control_loops/python/intake.py
index ae57730..2a623d7 100755
--- a/y2016/control_loops/python/intake.py
+++ b/y2016/control_loops/python/intake.py
@@ -2,8 +2,6 @@
 
 from frc971.control_loops.python import control_loop
 from frc971.control_loops.python import controls
-from frc971.control_loops.python import polytope
-from y2016.control_loops.python import polydrivetrain
 import numpy
 import sys
 import matplotlib
@@ -69,7 +67,7 @@
 
     controllability = controls.ctrb(self.A, self.B)
 
-    print "Free speed is", self.free_speed * numpy.pi * 2.0 / 60.0 / self.G
+    glog.debug("Free speed is %f", self.free_speed * numpy.pi * 2.0 / 60.0 / self.G)
 
     q_pos = 0.20
     q_vel = 5.5
@@ -79,15 +77,16 @@
     self.R = numpy.matrix([[(1.0 / (12.0 ** 2.0))]])
     self.K = controls.dlqr(self.A, self.B, self.Q, self.R)
 
-    print 'K', self.K
-    print 'Poles are', numpy.linalg.eig(self.A - self.B * self.K)[0]
+    glog.debug('K %s', repr(self.K))
+    glog.debug('Poles are %s',
+              repr(numpy.linalg.eig(self.A - self.B * self.K)[0]))
 
     self.rpl = 0.30
     self.ipl = 0.10
     self.PlaceObserverPoles([self.rpl + 1j * self.ipl,
                              self.rpl - 1j * self.ipl])
 
-    print 'L is', self.L
+    glog.debug('L is %s', repr(self.L))
 
     q_pos = 0.05
     q_vel = 2.65
@@ -100,9 +99,9 @@
     self.KalmanGain, self.Q_steady = controls.kalman(
         A=self.A, B=self.B, C=self.C, Q=self.Q, R=self.R)
 
-    print 'Kal', self.KalmanGain
+    glog.debug('Kal %s', repr(self.KalmanGain))
     self.L = self.A * self.KalmanGain
-    print 'KalL is', self.L
+    glog.debug('KalL is %s', repr(self.L))
 
     # The box formed by U_min and U_max must encompass all possible values,
     # or else Austin's code gets angry.
@@ -247,6 +246,7 @@
 
 def main(argv):
   argv = FLAGS(argv)
+  glog.init()
 
   scenario_plotter = ScenarioPlotter()
 
diff --git a/y2016/control_loops/python/shoulder.py b/y2016/control_loops/python/shoulder.py
index 554228e..9b70dbd 100755
--- a/y2016/control_loops/python/shoulder.py
+++ b/y2016/control_loops/python/shoulder.py
@@ -1,12 +1,10 @@
 #!/usr/bin/python
 
-from frc971.control_loops.python import control_loop
-from frc971.control_loops.python import controls
-from frc971.control_loops.python import polytope
-from y2016.control_loops.python import polydrivetrain
 import numpy
 import sys
-import matplotlib
+
+from frc971.control_loops.python import control_loop
+from frc971.control_loops.python import controls
 from matplotlib import pylab
 import gflags
 import glog
@@ -77,8 +75,10 @@
     self.R = numpy.matrix([[(1.0 / (12.0 ** 2.0))]])
     self.K = controls.dlqr(self.A, self.B, self.Q, self.R)
 
-    print 'K', self.K
-    print 'Poles are', numpy.linalg.eig(self.A - self.B * self.K)[0]
+    self.Kff = controls.TwoStateFeedForwards(self.B, self.Q)
+
+    glog.debug('Poles are %s for %s',
+               repr(numpy.linalg.eig(self.A - self.B * self.K)[0]), self._name)
 
     q_pos = 0.05
     q_vel = 2.65
@@ -91,9 +91,7 @@
     self.KalmanGain, self.Q_steady = controls.kalman(
         A=self.A, B=self.B, C=self.C, Q=self.Q, R=self.R)
 
-    print 'Kal', self.KalmanGain
     self.L = self.A * self.KalmanGain
-    print 'KalL is', self.L
 
     # The box formed by U_min and U_max must encompass all possible values,
     # or else Austin's code gets angry.
@@ -140,6 +138,9 @@
     self.K = numpy.matrix(numpy.zeros((1, 3)))
     self.K[0, 0:2] = self.K_unaugmented
     self.K[0, 2] = 1
+    self.Kff_unaugmented = self.Kff
+    self.Kff = numpy.matrix(numpy.zeros((1, 3)))
+    self.Kff[0, 0:2] = self.Kff_unaugmented
 
     self.InitializeState()
 
diff --git a/y2016/control_loops/python/wrist.py b/y2016/control_loops/python/wrist.py
index 79a115e..2ad1d07 100755
--- a/y2016/control_loops/python/wrist.py
+++ b/y2016/control_loops/python/wrist.py
@@ -2,8 +2,6 @@
 
 from frc971.control_loops.python import control_loop
 from frc971.control_loops.python import controls
-from frc971.control_loops.python import polytope
-from y2016.control_loops.python import polydrivetrain
 import numpy
 import sys
 import matplotlib
@@ -77,8 +75,8 @@
     self.R = numpy.matrix([[(1.0 / (12.0 ** 2.0))]])
     self.K = controls.dlqr(self.A, self.B, self.Q, self.R)
 
-    print 'K', self.K
-    print 'Poles are', numpy.linalg.eig(self.A - self.B * self.K)[0]
+    glog.debug('Poles are %s for %s',
+               repr(numpy.linalg.eig(self.A - self.B * self.K)[0]), self._name)
 
     q_pos = 0.05
     q_vel = 2.65
@@ -91,15 +89,15 @@
     self.KalmanGain, self.Q_steady = controls.kalman(
         A=self.A, B=self.B, C=self.C, Q=self.Q, R=self.R)
 
-    print 'Kal', self.KalmanGain
     self.L = self.A * self.KalmanGain
-    print 'KalL is', self.L
 
     # The box formed by U_min and U_max must encompass all possible values,
     # or else Austin's code gets angry.
     self.U_max = numpy.matrix([[12.0]])
     self.U_min = numpy.matrix([[-12.0]])
 
+    self.Kff = controls.TwoStateFeedForwards(self.B, self.Q)
+
     self.InitializeState()
 
 class IntegralWrist(Wrist):
@@ -140,6 +138,9 @@
     self.K = numpy.matrix(numpy.zeros((1, 3)))
     self.K[0, 0:2] = self.K_unaugmented
     self.K[0, 2] = 1
+    self.Kff_unaugmented = self.Kff
+    self.Kff = numpy.matrix(numpy.zeros((1, 3)))
+    self.Kff[0, 0:2] = self.Kff_unaugmented
 
     self.InitializeState()
 
@@ -260,13 +261,13 @@
   else:
     namespaces = ['y2016', 'control_loops', 'superstructure']
     wrist = Wrist("Wrist")
-    loop_writer = control_loop.ControlLoopWriter('Wrist', [wrist],
-                                                 namespaces=namespaces)
+    loop_writer = control_loop.ControlLoopWriter(
+        'Wrist', [wrist], namespaces=namespaces)
     loop_writer.Write(argv[1], argv[2])
 
-    integral_wrist = IntegralWrist("IntegralWrist")
-    integral_loop_writer = control_loop.ControlLoopWriter("IntegralWrist", [integral_wrist],
-                                                          namespaces=namespaces)
+    integral_wrist = IntegralWrist('IntegralWrist')
+    integral_loop_writer = control_loop.ControlLoopWriter(
+        'IntegralWrist', [integral_wrist], namespaces=namespaces)
     integral_loop_writer.Write(argv[3], argv[4])
 
 if __name__ == '__main__':