Moved Drivetrain from y2017 python to frc971

Change-Id: If931cf988d2615acc286d288fc0e5c9e7e3a5b90
diff --git a/frc971/control_loops/python/BUILD b/frc971/control_loops/python/BUILD
index b1897e0..4b25bb4 100644
--- a/frc971/control_loops/python/BUILD
+++ b/frc971/control_loops/python/BUILD
@@ -40,3 +40,27 @@
   ],
   restricted_to = ['//tools:k8'],
 )
+
+py_library(
+  name = 'drivetrain',
+  srcs = [
+    'drivetrain.py',
+  ],
+  deps = [
+    ':controls',
+  ],
+  restricted_to = ['//tools:k8'],
+)
+
+py_library(
+  name = 'polydrivetrain',
+  srcs = [
+    'polydrivetrain.py',
+  ],
+  deps = [
+    ':controls',
+    ':drivetrain',
+    '//external:python-glog',
+  ],
+  restricted_to = ['//tools:k8'],
+)
diff --git a/frc971/control_loops/python/control_loop.py b/frc971/control_loops/python/control_loop.py
index 805e079..900c06f 100644
--- a/frc971/control_loops/python/control_loop.py
+++ b/frc971/control_loops/python/control_loop.py
@@ -475,3 +475,21 @@
       self.P = self.Q_steady_state
     else:
       self.P = numpy.matrix(numpy.zeros((self.A.shape[0], self.A.shape[0])))
+
+
+class CIM(object):
+  def __init__(self):
+    # Stall Torque in N m
+    self.stall_torque = 2.42
+    # Stall Current in Amps
+    self.stall_current = 133.0
+    # Free Speed in rad/s
+    self.free_speed = 5500.0 / 60.0 * 2.0 * numpy.pi
+    # Free Current in Amps
+    self.free_current = 4.7
+    # Resistance of the motor
+    self.resistance = 12.0 / self.stall_current
+    # Motor velocity constant
+    self.Kv = (self.free_speed / (12.0 - self.resistance * self.free_current))
+    # Torque constant
+    self.Kt = self.stall_torque / self.stall_current
diff --git a/frc971/control_loops/python/drivetrain.py b/frc971/control_loops/python/drivetrain.py
new file mode 100644
index 0000000..5665fd9
--- /dev/null
+++ b/frc971/control_loops/python/drivetrain.py
@@ -0,0 +1,416 @@
+#!/usr/bin/python
+
+from frc971.control_loops.python import control_loop
+from frc971.control_loops.python import controls
+import numpy
+import sys
+from matplotlib import pylab
+import glog
+
+class DrivetrainParams(object):
+  def __init__(self, J, mass, robot_radius, wheel_radius, G_high, G_low,
+               q_pos_low, q_pos_high, q_vel_low, q_vel_high,
+               motor_type = control_loop.CIM(), num_motors = 2, dt = 0.00505,
+               controller_poles=[0.90, 0.90], observer_poles=[0.02, 0.02]):
+    """Defines all constants of a drivetrain.
+
+    Args:
+      J: float, Moment of inertia of drivetrain in kg m^2
+      mass: float, Mass of the robot in kg.
+      robot_radius: float, Radius of the robot, in meters (requires tuning by
+        hand).
+      wheel_radius: float, Radius of the wheels, in meters.
+      G_high: float, Gear ratio for high gear.
+      G_low: float, Gear ratio for low gear.
+      dt: float, Control loop time step.
+      q_pos_low: float, q position low gear.
+      q_pos_high: float, q position high gear.
+      q_vel_low: float, q velocity low gear.
+      q_vel_high: float, q velocity high gear.
+      motor_type: object, class of values defining the motor in drivetrain.
+      num_motors: int, number of motors on one side of drivetrain.
+      controller_poles: array, An array of poles. (See control_loop.py)
+      observer_poles: array, An array of poles. (See control_loop.py)
+    """
+
+    self.J = J
+    self.mass = mass
+    self.robot_radius = robot_radius
+    self.wheel_radius = wheel_radius
+    self.G_high = G_high
+    self.G_low = G_low
+    self.dt = dt
+    self.q_pos_low = q_pos_low
+    self.q_pos_high = q_pos_high
+    self.q_vel_low = q_vel_low
+    self.q_vel_high = q_vel_high
+    self.motor_type = motor_type
+    self.num_motors = num_motors
+    self.controller_poles = controller_poles
+    self.observer_poles = observer_poles
+
+class Drivetrain(control_loop.ControlLoop):
+  def __init__(self, drivetrain_params, name="Drivetrain", left_low=True,
+               right_low=True):
+    """Defines a base drivetrain for a robot.
+
+    Args:
+      drivetrain_params: DrivetrainParams, class of values defining the drivetrain.
+      name: string, Name of this drivetrain.
+      left_low: bool, Whether the left is in high gear.
+      right_low: bool, Whether the right is in high gear.
+    """
+    super(Drivetrain, self).__init__(name)
+
+    # Moment of inertia of the drivetrain in kg m^2
+    self.J = drivetrain_params.J
+    # Mass of the robot, in kg.
+    self.mass = drivetrain_params.mass
+    # Radius of the robot, in meters (requires tuning by hand)
+    self.robot_radius = drivetrain_params.robot_radius
+    # Radius of the wheels, in meters.
+    self.r = drivetrain_params.wheel_radius
+
+    # Gear ratios
+    self.G_low = drivetrain_params.G_low
+    self.G_high = drivetrain_params.G_high
+    if left_low:
+      self.Gl = self.G_low
+    else:
+      self.Gl = self.G_high
+    if right_low:
+      self.Gr = self.G_low
+    else:
+      self.Gr = self.G_high
+
+    # Control loop time step
+    self.dt = drivetrain_params.dt
+
+    self.BuildDrivetrain(drivetrain_params.motor_type, drivetrain_params.num_motors);
+
+    if left_low or right_low:
+      q_pos = drivetrain_params.q_pos_low
+      q_vel = drivetrain_params.q_vel_low
+    else:
+      q_pos = drivetrain_params.q_pos_high
+      q_vel = drivetrain_params.q_vel_high
+
+    self.BuildDrivetrainController(q_pos, q_vel)
+
+    self.InitializeState()
+
+  def BuildDrivetrain(self, motor, num_motors_per_side):
+    self.motor = motor
+    # Number of motors per side
+    self.num_motors = num_motors_per_side
+    # Stall Torque in N m
+    self.stall_torque = motor.stall_torque * self.num_motors * 0.60
+    # Stall Current in Amps
+    self.stall_current = motor.stall_current * self.num_motors
+    # Free Speed in rad/s
+    self.free_speed = motor.free_speed
+    # Free Current in Amps
+    self.free_current = motor.free_current * self.num_motors
+
+    # Effective motor resistance in ohms.
+    self.resistance = 12.0 / self.stall_current
+
+    # Resistance of the motor, divided by the number of motors.
+    # Motor velocity constant
+    self.Kv = (self.free_speed / (12.0 - self.resistance * self.free_current))
+    # Torque constant
+    self.Kt = self.stall_torque / self.stall_current
+
+    # These describe the way that a given side of a robot will be influenced
+    # by the other side. Units of 1 / kg.
+    self.msp = 1.0 / self.mass + self.robot_radius * self.robot_radius / self.J
+    self.msn = 1.0 / self.mass - self.robot_radius * self.robot_radius / self.J
+    # The calculations which we will need for A and B.
+    self.tcl = self.Kt / self.Kv / (self.Gl * self.Gl * self.resistance * self.r * self.r)
+    self.tcr = self.Kt / self.Kv / (self.Gr * self.Gr * self.resistance * self.r * self.r)
+    self.mpl = self.Kt / (self.Gl * self.resistance * self.r)
+    self.mpr = self.Kt / (self.Gr * self.resistance * self.r)
+
+    # State feedback matrices
+    # X will be of the format
+    # [[positionl], [velocityl], [positionr], velocityr]]
+    self.A_continuous = numpy.matrix(
+        [[0, 1, 0, 0],
+         [0, -self.msp * self.tcl, 0, -self.msn * self.tcr],
+         [0, 0, 0, 1],
+         [0, -self.msn * self.tcl, 0, -self.msp * self.tcr]])
+    self.B_continuous = numpy.matrix(
+        [[0, 0],
+         [self.msp * self.mpl, self.msn * self.mpr],
+         [0, 0],
+         [self.msn * self.mpl, self.msp * self.mpr]])
+    self.C = numpy.matrix([[1, 0, 0, 0],
+                           [0, 0, 1, 0]])
+    self.D = numpy.matrix([[0, 0],
+                           [0, 0]])
+
+    self.A, self.B = self.ContinuousToDiscrete(
+        self.A_continuous, self.B_continuous, self.dt)
+
+  def BuildDrivetrainController(self, q_pos, q_vel):
+    # Tune the LQR controller
+    self.Q = numpy.matrix([[(1.0 / (q_pos ** 2.0)), 0.0, 0.0, 0.0],
+                           [0.0, (1.0 / (q_vel ** 2.0)), 0.0, 0.0],
+                           [0.0, 0.0, (1.0 / (q_pos ** 2.0)), 0.0],
+                           [0.0, 0.0, 0.0, (1.0 / (q_vel ** 2.0))]])
+
+    self.R = numpy.matrix([[(1.0 / (12.0 ** 2.0)), 0.0],
+                           [0.0, (1.0 / (12.0 ** 2.0))]])
+    self.K = controls.dlqr(self.A, self.B, self.Q, self.R)
+
+    glog.debug('DT q_pos %f q_vel %s %s', q_pos, q_vel, self._name)
+    glog.debug(str(numpy.linalg.eig(self.A - self.B * self.K)[0]))
+    glog.debug('K %s', repr(self.K))
+
+    self.hlp = 0.3
+    self.llp = 0.4
+    self.PlaceObserverPoles([self.hlp, self.hlp, self.llp, self.llp])
+
+    self.U_max = numpy.matrix([[12.0], [12.0]])
+    self.U_min = numpy.matrix([[-12.0], [-12.0]])
+
+class KFDrivetrain(Drivetrain):
+  def __init__(self, drivetrain_params, name="KFDrivetrain",
+               left_low=True, right_low=True):
+    """Kalman filter values of a drivetrain.
+
+    Args:
+      drivetrain_params: DrivetrainParams, class of values defining the drivetrain.
+      name: string, Name of this drivetrain.
+      left_low: bool, Whether the left is in high gear.
+      right_low: bool, Whether the right is in high gear.
+    """
+    super(KFDrivetrain, self).__init__(drivetrain_params, name, left_low, right_low)
+
+    self.unaugmented_A_continuous = self.A_continuous
+    self.unaugmented_B_continuous = self.B_continuous
+
+    # The practical voltage applied to the wheels is
+    #   V_left = U_left + left_voltage_error
+    #
+    # The states are
+    # [left position, left velocity, right position, right velocity,
+    #  left voltage error, right voltage error, angular_error]
+    #
+    # The left and right positions are filtered encoder positions and are not
+    # adjusted for heading error.
+    # The turn velocity as computed by the left and right velocities is
+    # adjusted by the gyro velocity.
+    # The angular_error is the angular velocity error between the wheel speed
+    # and the gyro speed.
+    self.A_continuous = numpy.matrix(numpy.zeros((7, 7)))
+    self.B_continuous = numpy.matrix(numpy.zeros((7, 2)))
+    self.A_continuous[0:4,0:4] = self.unaugmented_A_continuous
+    self.A_continuous[0:4,4:6] = self.unaugmented_B_continuous
+    self.B_continuous[0:4,0:2] = self.unaugmented_B_continuous
+    self.A_continuous[0,6] = 1
+    self.A_continuous[2,6] = -1
+
+    self.A, self.B = self.ContinuousToDiscrete(
+        self.A_continuous, self.B_continuous, self.dt)
+
+    self.C = numpy.matrix([[1, 0, 0, 0, 0, 0, 0],
+                           [0, 0, 1, 0, 0, 0, 0],
+                           [0, -0.5 / drivetrain_params.robot_radius, 0, 0.5 / drivetrain_params.robot_radius, 0, 0, 0]])
+
+    self.D = numpy.matrix([[0, 0],
+                           [0, 0],
+                           [0, 0]])
+
+    q_pos = 0.05
+    q_vel = 1.00
+    q_voltage = 10.0
+    q_encoder_uncertainty = 2.00
+
+    self.Q = numpy.matrix([[(q_pos ** 2.0), 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                           [0.0, (q_vel ** 2.0), 0.0, 0.0, 0.0, 0.0, 0.0],
+                           [0.0, 0.0, (q_pos ** 2.0), 0.0, 0.0, 0.0, 0.0],
+                           [0.0, 0.0, 0.0, (q_vel ** 2.0), 0.0, 0.0, 0.0],
+                           [0.0, 0.0, 0.0, 0.0, (q_voltage ** 2.0), 0.0, 0.0],
+                           [0.0, 0.0, 0.0, 0.0, 0.0, (q_voltage ** 2.0), 0.0],
+                           [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, (q_encoder_uncertainty ** 2.0)]])
+
+    r_pos =  0.0001
+    r_gyro = 0.000001
+    self.R = numpy.matrix([[(r_pos ** 2.0), 0.0, 0.0],
+                           [0.0, (r_pos ** 2.0), 0.0],
+                           [0.0, 0.0, (r_gyro ** 2.0)]])
+
+    # Solving for kf gains.
+    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
+
+    unaug_K = self.K
+
+    # Implement a nice closed loop controller for use by the closed loop
+    # controller.
+    self.K = numpy.matrix(numpy.zeros((self.B.shape[1], self.A.shape[0])))
+    self.K[0:2, 0:4] = unaug_K
+    self.K[0, 4] = 1.0
+    self.K[1, 5] = 1.0
+
+    self.Qff = numpy.matrix(numpy.zeros((4, 4)))
+    qff_pos = 0.005
+    qff_vel = 1.00
+    self.Qff[0, 0] = 1.0 / qff_pos ** 2.0
+    self.Qff[1, 1] = 1.0 / qff_vel ** 2.0
+    self.Qff[2, 2] = 1.0 / qff_pos ** 2.0
+    self.Qff[3, 3] = 1.0 / qff_vel ** 2.0
+    self.Kff = numpy.matrix(numpy.zeros((2, 7)))
+    self.Kff[0:2, 0:4] = controls.TwoStateFeedForwards(self.B[0:4,:], self.Qff)
+
+    self.InitializeState()
+
+
+def WriteDrivetrain(drivetrain_files, kf_drivetrain_files, year_namespace,
+                    drivetrain_params):
+
+  # Write the generated constants out to a file.
+  drivetrain_low_low = Drivetrain(name="DrivetrainLowLow",
+          left_low=True, right_low=True, drivetrain_params=drivetrain_params)
+  drivetrain_low_high = Drivetrain(name="DrivetrainLowHigh",
+          left_low=True, right_low=False, drivetrain_params=drivetrain_params)
+  drivetrain_high_low = Drivetrain(name="DrivetrainHighLow",
+          left_low=False, right_low=True, drivetrain_params=drivetrain_params)
+  drivetrain_high_high = Drivetrain(name="DrivetrainHighHigh",
+          left_low=False, right_low=False, drivetrain_params=drivetrain_params)
+
+  kf_drivetrain_low_low = KFDrivetrain(name="KFDrivetrainLowLow",
+          left_low=True, right_low=True, drivetrain_params=drivetrain_params)
+  kf_drivetrain_low_high = KFDrivetrain(name="KFDrivetrainLowHigh",
+          left_low=True, right_low=False, drivetrain_params=drivetrain_params)
+  kf_drivetrain_high_low = KFDrivetrain(name="KFDrivetrainHighLow",
+          left_low=False, right_low=True, drivetrain_params=drivetrain_params)
+  kf_drivetrain_high_high = KFDrivetrain(name="KFDrivetrainHighHigh",
+          left_low=False, right_low=False, drivetrain_params=drivetrain_params)
+
+  namespaces = [year_namespace, 'control_loops', 'drivetrain']
+  dog_loop_writer = control_loop.ControlLoopWriter(
+      "Drivetrain", [drivetrain_low_low, drivetrain_low_high,
+                     drivetrain_high_low, drivetrain_high_high],
+      namespaces = namespaces)
+  dog_loop_writer.AddConstant(control_loop.Constant("kDt", "%f",
+        drivetrain_low_low.dt))
+  dog_loop_writer.AddConstant(control_loop.Constant("kStallTorque", "%f",
+        drivetrain_low_low.stall_torque))
+  dog_loop_writer.AddConstant(control_loop.Constant("kStallCurrent", "%f",
+        drivetrain_low_low.stall_current))
+  dog_loop_writer.AddConstant(control_loop.Constant("kFreeSpeed", "%f",
+        drivetrain_low_low.free_speed))
+  dog_loop_writer.AddConstant(control_loop.Constant("kFreeCurrent", "%f",
+        drivetrain_low_low.free_current))
+  dog_loop_writer.AddConstant(control_loop.Constant("kJ", "%f",
+        drivetrain_low_low.J))
+  dog_loop_writer.AddConstant(control_loop.Constant("kMass", "%f",
+        drivetrain_low_low.mass))
+  dog_loop_writer.AddConstant(control_loop.Constant("kRobotRadius", "%f",
+        drivetrain_low_low.robot_radius))
+  dog_loop_writer.AddConstant(control_loop.Constant("kWheelRadius", "%f",
+        drivetrain_low_low.r))
+  dog_loop_writer.AddConstant(control_loop.Constant("kR", "%f",
+        drivetrain_low_low.resistance))
+  dog_loop_writer.AddConstant(control_loop.Constant("kV", "%f",
+        drivetrain_low_low.Kv))
+  dog_loop_writer.AddConstant(control_loop.Constant("kT", "%f",
+        drivetrain_low_low.Kt))
+  dog_loop_writer.AddConstant(control_loop.Constant("kLowGearRatio", "%f",
+        drivetrain_low_low.G_low))
+  dog_loop_writer.AddConstant(control_loop.Constant("kHighGearRatio", "%f",
+        drivetrain_high_high.G_high))
+  dog_loop_writer.AddConstant(control_loop.Constant("kHighOutputRatio", "%f",
+        drivetrain_high_high.G_high * drivetrain_high_high.r))
+
+  dog_loop_writer.Write(drivetrain_files[0], drivetrain_files[1])
+
+  kf_loop_writer = control_loop.ControlLoopWriter(
+      "KFDrivetrain", [kf_drivetrain_low_low, kf_drivetrain_low_high,
+                       kf_drivetrain_high_low, kf_drivetrain_high_high],
+      namespaces = namespaces)
+  kf_loop_writer.Write(kf_drivetrain_files[0], kf_drivetrain_files[1])
+
+def PlotDrivetrainMotions(drivetrain_params):
+  # Simulate the response of the system to a step input.
+  drivetrain = Drivetrain(left_low=False, right_low=False, drivetrain_params=drivetrain_params)
+  simulated_left = []
+  simulated_right = []
+  for _ in xrange(100):
+    drivetrain.Update(numpy.matrix([[12.0], [12.0]]))
+    simulated_left.append(drivetrain.X[0, 0])
+    simulated_right.append(drivetrain.X[2, 0])
+
+  pylab.rc('lines', linewidth=4)
+  pylab.plot(range(100), simulated_left, label='left position')
+  pylab.plot(range(100), simulated_right, 'r--', label='right position')
+  pylab.suptitle('Acceleration Test\n12 Volt Step Input')
+  pylab.legend(loc='lower right')
+  pylab.show()
+
+  # Simulate forwards motion.
+  drivetrain = Drivetrain(left_low=False, right_low=False, drivetrain_params=drivetrain_params)
+  close_loop_left = []
+  close_loop_right = []
+  left_power = []
+  right_power = []
+  R = numpy.matrix([[1.0], [0.0], [1.0], [0.0]])
+  for _ in xrange(300):
+    U = numpy.clip(drivetrain.K * (R - drivetrain.X_hat),
+                   drivetrain.U_min, drivetrain.U_max)
+    drivetrain.UpdateObserver(U)
+    drivetrain.Update(U)
+    close_loop_left.append(drivetrain.X[0, 0])
+    close_loop_right.append(drivetrain.X[2, 0])
+    left_power.append(U[0, 0])
+    right_power.append(U[1, 0])
+
+  pylab.plot(range(300), close_loop_left, label='left position')
+  pylab.plot(range(300), close_loop_right, 'm--', label='right position')
+  pylab.plot(range(300), left_power, label='left power')
+  pylab.plot(range(300), right_power, '--', label='right power')
+  pylab.suptitle('Linear Move\nLeft and Right Position going to 1')
+  pylab.legend()
+  pylab.show()
+
+  # Try turning in place
+  drivetrain = Drivetrain(drivetrain_params=drivetrain_params)
+  close_loop_left = []
+  close_loop_right = []
+  R = numpy.matrix([[-1.0], [0.0], [1.0], [0.0]])
+  for _ in xrange(200):
+    U = numpy.clip(drivetrain.K * (R - drivetrain.X_hat),
+                   drivetrain.U_min, drivetrain.U_max)
+    drivetrain.UpdateObserver(U)
+    drivetrain.Update(U)
+    close_loop_left.append(drivetrain.X[0, 0])
+    close_loop_right.append(drivetrain.X[2, 0])
+
+  pylab.plot(range(200), close_loop_left, label='left position')
+  pylab.plot(range(200), close_loop_right, label='right position')
+  pylab.suptitle('Angular Move\nLeft position going to -1 and right position going to 1')
+  pylab.legend(loc='center right')
+  pylab.show()
+
+  # Try turning just one side.
+  drivetrain = Drivetrain(drivetrain_params=drivetrain_params)
+  close_loop_left = []
+  close_loop_right = []
+  R = numpy.matrix([[0.0], [0.0], [1.0], [0.0]])
+  for _ in xrange(300):
+    U = numpy.clip(drivetrain.K * (R - drivetrain.X_hat),
+                   drivetrain.U_min, drivetrain.U_max)
+    drivetrain.UpdateObserver(U)
+    drivetrain.Update(U)
+    close_loop_left.append(drivetrain.X[0, 0])
+    close_loop_right.append(drivetrain.X[2, 0])
+
+  pylab.plot(range(300), close_loop_left, label='left position')
+  pylab.plot(range(300), close_loop_right, label='right position')
+  pylab.suptitle('Pivot\nLeft position not changing and right position going to 1')
+  pylab.legend(loc='center right')
+  pylab.show()
diff --git a/frc971/control_loops/python/polydrivetrain.py b/frc971/control_loops/python/polydrivetrain.py
new file mode 100644
index 0000000..c9c9efe
--- /dev/null
+++ b/frc971/control_loops/python/polydrivetrain.py
@@ -0,0 +1,488 @@
+#!/usr/bin/python
+
+import numpy
+from frc971.control_loops.python import polytope
+import frc971.control_loops.python.drivetrain
+from frc971.control_loops.python import control_loop
+from frc971.control_loops.python import controls
+from frc971.control_loops.python.cim import CIM
+from matplotlib import pylab
+
+import glog
+
+def CoerceGoal(region, K, w, R):
+  """Intersects a line with a region, and finds the closest point to R.
+
+  Finds a point that is closest to R inside the region, and on the line
+  defined by K X = w.  If it is not possible to find a point on the line,
+  finds a point that is inside the region and closest to the line.  This
+  function assumes that
+
+  Args:
+    region: HPolytope, the valid goal region.
+    K: numpy.matrix (2 x 1), the matrix for the equation [K1, K2] [x1; x2] = w
+    w: float, the offset in the equation above.
+    R: numpy.matrix (2 x 1), the point to be closest to.
+
+  Returns:
+    numpy.matrix (2 x 1), the point.
+  """
+  return DoCoerceGoal(region, K, w, R)[0]
+
+def DoCoerceGoal(region, K, w, R):
+  if region.IsInside(R):
+    return (R, True)
+
+  perpendicular_vector = K.T / numpy.linalg.norm(K)
+  parallel_vector = numpy.matrix([[perpendicular_vector[1, 0]],
+                                  [-perpendicular_vector[0, 0]]])
+
+  # We want to impose the constraint K * X = w on the polytope H * X <= k.
+  # We do this by breaking X up into parallel and perpendicular components to
+  # the half plane.  This gives us the following equation.
+  #
+  #  parallel * (parallel.T \dot X) + perpendicular * (perpendicular \dot X)) = X
+  #
+  # Then, substitute this into the polytope.
+  #
+  #  H * (parallel * (parallel.T \dot X) + perpendicular * (perpendicular \dot X)) <= k
+  #
+  # Substitute K * X = w
+  #
+  # H * parallel * (parallel.T \dot X) + H * perpendicular * w <= k
+  #
+  # Move all the knowns to the right side.
+  #
+  # H * parallel * ([parallel1 parallel2] * X) <= k - H * perpendicular * w
+  #
+  # Let t = parallel.T \dot X, the component parallel to the surface.
+  #
+  # H * parallel * t <= k - H * perpendicular * w
+  #
+  # This is a polytope which we can solve, and use to figure out the range of X
+  # that we care about!
+
+  t_poly = polytope.HPolytope(
+      region.H * parallel_vector,
+      region.k - region.H * perpendicular_vector * w)
+
+  vertices = t_poly.Vertices()
+
+  if vertices.shape[0]:
+    # The region exists!
+    # Find the closest vertex
+    min_distance = numpy.infty
+    closest_point = None
+    for vertex in vertices:
+      point = parallel_vector * vertex + perpendicular_vector * w
+      length = numpy.linalg.norm(R - point)
+      if length < min_distance:
+        min_distance = length
+        closest_point = point
+
+    return (closest_point, True)
+  else:
+    # Find the vertex of the space that is closest to the line.
+    region_vertices = region.Vertices()
+    min_distance = numpy.infty
+    closest_point = None
+    for vertex in region_vertices:
+      point = vertex.T
+      length = numpy.abs((perpendicular_vector.T * point)[0, 0])
+      if length < min_distance:
+        min_distance = length
+        closest_point = point
+
+    return (closest_point, False)
+
+class VelocityDrivetrainModel(control_loop.ControlLoop):
+  def __init__(self, drivetrain_params, left_low=True, right_low=True,
+                   name="VelocityDrivetrainModel"):
+    super(VelocityDrivetrainModel, self).__init__(name)
+    self._drivetrain = frc971.control_loops.python.drivetrain.Drivetrain(
+                           left_low=left_low, right_low=right_low,
+                           drivetrain_params=drivetrain_params)
+    self.dt = drivetrain_params.dt
+    self.A_continuous = numpy.matrix(
+        [[self._drivetrain.A_continuous[1, 1], self._drivetrain.A_continuous[1, 3]],
+         [self._drivetrain.A_continuous[3, 1], self._drivetrain.A_continuous[3, 3]]])
+
+    self.B_continuous = numpy.matrix(
+        [[self._drivetrain.B_continuous[1, 0], self._drivetrain.B_continuous[1, 1]],
+         [self._drivetrain.B_continuous[3, 0], self._drivetrain.B_continuous[3, 1]]])
+    self.C = numpy.matrix(numpy.eye(2))
+    self.D = numpy.matrix(numpy.zeros((2, 2)))
+
+    self.A, self.B = self.ContinuousToDiscrete(self.A_continuous,
+                                               self.B_continuous, self.dt)
+
+    # FF * X = U (steady state)
+    self.FF = self.B.I * (numpy.eye(2) - self.A)
+
+    self.PlaceControllerPoles(drivetrain_params.controller_poles)
+    self.PlaceObserverPoles(drivetrain_params.observer_poles)
+
+    self.G_high = self._drivetrain.G_high
+    self.G_low = self._drivetrain.G_low
+    self.resistance = self._drivetrain.resistance
+    self.r = self._drivetrain.r
+    self.Kv = self._drivetrain.Kv
+    self.Kt = self._drivetrain.Kt
+
+    self.U_max = self._drivetrain.U_max
+    self.U_min = self._drivetrain.U_min
+
+
+class VelocityDrivetrain(object):
+  HIGH = 'high'
+  LOW = 'low'
+  SHIFTING_UP = 'up'
+  SHIFTING_DOWN = 'down'
+
+  def __init__(self, drivetrain_params):
+    self.drivetrain_low_low = VelocityDrivetrainModel(
+        left_low=True, right_low=True, name='VelocityDrivetrainLowLow',
+        drivetrain_params=drivetrain_params)
+    self.drivetrain_low_high = VelocityDrivetrainModel(
+        left_low=True, right_low=False, name='VelocityDrivetrainLowHigh',
+        drivetrain_params=drivetrain_params)
+    self.drivetrain_high_low = VelocityDrivetrainModel(
+        left_low=False, right_low=True, name = 'VelocityDrivetrainHighLow',
+        drivetrain_params=drivetrain_params)
+    self.drivetrain_high_high = VelocityDrivetrainModel(
+        left_low=False, right_low=False, name = 'VelocityDrivetrainHighHigh',
+        drivetrain_params=drivetrain_params)
+
+    # X is [lvel, rvel]
+    self.X = numpy.matrix(
+        [[0.0],
+         [0.0]])
+
+    self.U_poly = polytope.HPolytope(
+        numpy.matrix([[1, 0],
+                      [-1, 0],
+                      [0, 1],
+                      [0, -1]]),
+        numpy.matrix([[12],
+                      [12],
+                      [12],
+                      [12]]))
+
+    self.U_max = numpy.matrix(
+        [[12.0],
+         [12.0]])
+    self.U_min = numpy.matrix(
+        [[-12.0000000000],
+         [-12.0000000000]])
+
+    self.dt = 0.00505
+
+    self.R = numpy.matrix(
+        [[0.0],
+         [0.0]])
+
+    self.U_ideal = numpy.matrix(
+        [[0.0],
+         [0.0]])
+
+    # ttrust is the comprimise between having full throttle negative inertia,
+    # and having no throttle negative inertia.  A value of 0 is full throttle
+    # inertia.  A value of 1 is no throttle negative inertia.
+    self.ttrust = 1.0
+
+    self.left_gear = VelocityDrivetrain.LOW
+    self.right_gear = VelocityDrivetrain.LOW
+    self.left_shifter_position = 0.0
+    self.right_shifter_position = 0.0
+    self.left_cim = CIM()
+    self.right_cim = CIM()
+
+  def IsInGear(self, gear):
+    return gear is VelocityDrivetrain.HIGH or gear is VelocityDrivetrain.LOW
+
+  def MotorRPM(self, shifter_position, velocity):
+    if shifter_position > 0.5:
+      return (velocity / self.CurrentDrivetrain().G_high /
+              self.CurrentDrivetrain().r)
+    else:
+      return (velocity / self.CurrentDrivetrain().G_low /
+              self.CurrentDrivetrain().r)
+
+  def CurrentDrivetrain(self):
+    if self.left_shifter_position > 0.5:
+      if self.right_shifter_position > 0.5:
+        return self.drivetrain_high_high
+      else:
+        return self.drivetrain_high_low
+    else:
+      if self.right_shifter_position > 0.5:
+        return self.drivetrain_low_high
+      else:
+        return self.drivetrain_low_low
+
+  def SimShifter(self, gear, shifter_position):
+    if gear is VelocityDrivetrain.HIGH or gear is VelocityDrivetrain.SHIFTING_UP:
+      shifter_position = min(shifter_position + 0.5, 1.0)
+    else:
+      shifter_position = max(shifter_position - 0.5, 0.0)
+
+    if shifter_position == 1.0:
+      gear = VelocityDrivetrain.HIGH
+    elif shifter_position == 0.0:
+      gear = VelocityDrivetrain.LOW
+
+    return gear, shifter_position
+
+  def ComputeGear(self, wheel_velocity, should_print=False, current_gear=False, gear_name=None):
+    high_omega = (wheel_velocity / self.CurrentDrivetrain().G_high /
+                  self.CurrentDrivetrain().r)
+    low_omega = (wheel_velocity / self.CurrentDrivetrain().G_low /
+                 self.CurrentDrivetrain().r)
+    #print gear_name, "Motor Energy Difference.", 0.5 * 0.000140032647 * (low_omega * low_omega - high_omega * high_omega), "joules"
+    high_torque = ((12.0 - high_omega / self.CurrentDrivetrain().Kv) *
+                   self.CurrentDrivetrain().Kt / self.CurrentDrivetrain().resistance)
+    low_torque = ((12.0 - low_omega / self.CurrentDrivetrain().Kv) *
+                  self.CurrentDrivetrain().Kt / self.CurrentDrivetrain().resistance)
+    high_power = high_torque * high_omega
+    low_power = low_torque * low_omega
+    #if should_print:
+    #  print gear_name, "High omega", high_omega, "Low omega", low_omega
+    #  print gear_name, "High torque", high_torque, "Low torque", low_torque
+    #  print gear_name, "High power", high_power, "Low power", low_power
+
+    # Shift algorithm improvements.
+    # TODO(aschuh):
+    # It takes time to shift.  Shifting down for 1 cycle doesn't make sense
+    # because you will end up slower than without shifting.  Figure out how
+    # to include that info.
+    # If the driver is still in high gear, but isn't asking for the extra power
+    # from low gear, don't shift until he asks for it.
+    goal_gear_is_high = high_power > low_power
+    #goal_gear_is_high = True
+
+    if not self.IsInGear(current_gear):
+      glog.debug('%s Not in gear.', gear_name)
+      return current_gear
+    else:
+      is_high = current_gear is VelocityDrivetrain.HIGH
+      if is_high != goal_gear_is_high:
+        if goal_gear_is_high:
+          glog.debug('%s Shifting up.', gear_name)
+          return VelocityDrivetrain.SHIFTING_UP
+        else:
+          glog.debug('%s Shifting down.', gear_name)
+          return VelocityDrivetrain.SHIFTING_DOWN
+      else:
+        return current_gear
+
+  def FilterVelocity(self, throttle):
+    # Invert the plant to figure out how the velocity filter would have to work
+    # out in order to filter out the forwards negative inertia.
+    # This math assumes that the left and right power and velocity are equal.
+
+    # The throttle filter should filter such that the motor in the highest gear
+    # should be controlling the time constant.
+    # Do this by finding the index of FF that has the lowest value, and computing
+    # the sums using that index.
+    FF_sum = self.CurrentDrivetrain().FF.sum(axis=1)
+    min_FF_sum_index = numpy.argmin(FF_sum)
+    min_FF_sum = FF_sum[min_FF_sum_index, 0]
+    min_K_sum = self.CurrentDrivetrain().K[min_FF_sum_index, :].sum()
+    # Compute the FF sum for high gear.
+    high_min_FF_sum = self.drivetrain_high_high.FF[0, :].sum()
+
+    # U = self.K[0, :].sum() * (R - x_avg) + self.FF[0, :].sum() * R
+    # throttle * 12.0 = (self.K[0, :].sum() + self.FF[0, :].sum()) * R
+    #                   - self.K[0, :].sum() * x_avg
+
+    # R = (throttle * 12.0 + self.K[0, :].sum() * x_avg) /
+    #     (self.K[0, :].sum() + self.FF[0, :].sum())
+
+    # U = (K + FF) * R - K * X
+    # (K + FF) ^-1 * (U + K * X) = R
+
+    # Scale throttle by min_FF_sum / high_min_FF_sum.  This will make low gear
+    # have the same velocity goal as high gear, and so that the robot will hold
+    # the same speed for the same throttle for all gears.
+    adjusted_ff_voltage = numpy.clip(throttle * 12.0 * min_FF_sum / high_min_FF_sum, -12.0, 12.0)
+    return ((adjusted_ff_voltage + self.ttrust * min_K_sum * (self.X[0, 0] + self.X[1, 0]) / 2.0)
+            / (self.ttrust * min_K_sum + min_FF_sum))
+
+  def Update(self, throttle, steering):
+    # Shift into the gear which sends the most power to the floor.
+    # This is the same as sending the most torque down to the floor at the
+    # wheel.
+
+    self.left_gear = self.right_gear = True
+    if True:
+      self.left_gear = self.ComputeGear(self.X[0, 0], should_print=True,
+                                        current_gear=self.left_gear,
+                                        gear_name="left")
+      self.right_gear = self.ComputeGear(self.X[1, 0], should_print=True,
+                                         current_gear=self.right_gear,
+                                         gear_name="right")
+      if self.IsInGear(self.left_gear):
+        self.left_cim.X[0, 0] = self.MotorRPM(self.left_shifter_position, self.X[0, 0])
+
+      if self.IsInGear(self.right_gear):
+        self.right_cim.X[0, 0] = self.MotorRPM(self.right_shifter_position, self.X[0, 0])
+
+    if self.IsInGear(self.left_gear) and self.IsInGear(self.right_gear):
+      # Filter the throttle to provide a nicer response.
+      fvel = self.FilterVelocity(throttle)
+
+      # Constant radius means that angualar_velocity / linear_velocity = constant.
+      # Compute the left and right velocities.
+      steering_velocity = numpy.abs(fvel) * steering
+      left_velocity = fvel - steering_velocity
+      right_velocity = fvel + steering_velocity
+
+      # Write this constraint in the form of K * R = w
+      # angular velocity / linear velocity = constant
+      # (left - right) / (left + right) = constant
+      # left - right = constant * left + constant * right
+
+      # (fvel - steering * numpy.abs(fvel) - fvel - steering * numpy.abs(fvel)) /
+      #  (fvel - steering * numpy.abs(fvel) + fvel + steering * numpy.abs(fvel)) =
+      #       constant
+      # (- 2 * steering * numpy.abs(fvel)) / (2 * fvel) = constant
+      # (-steering * sign(fvel)) = constant
+      # (-steering * sign(fvel)) * (left + right) = left - right
+      # (steering * sign(fvel) + 1) * left + (steering * sign(fvel) - 1) * right = 0
+
+      equality_k = numpy.matrix(
+          [[1 + steering * numpy.sign(fvel), -(1 - steering * numpy.sign(fvel))]])
+      equality_w = 0.0
+
+      self.R[0, 0] = left_velocity
+      self.R[1, 0] = right_velocity
+
+      # Construct a constraint on R by manipulating the constraint on U
+      # Start out with H * U <= k
+      # U = FF * R + K * (R - X)
+      # H * (FF * R + K * R - K * X) <= k
+      # H * (FF + K) * R <= k + H * K * X
+      R_poly = polytope.HPolytope(
+          self.U_poly.H * (self.CurrentDrivetrain().K + self.CurrentDrivetrain().FF),
+          self.U_poly.k + self.U_poly.H * self.CurrentDrivetrain().K * self.X)
+
+      # Limit R back inside the box.
+      self.boxed_R = CoerceGoal(R_poly, equality_k, equality_w, self.R)
+
+      FF_volts = self.CurrentDrivetrain().FF * self.boxed_R
+      self.U_ideal = self.CurrentDrivetrain().K * (self.boxed_R - self.X) + FF_volts
+    else:
+      glog.debug('Not all in gear')
+      if not self.IsInGear(self.left_gear) and not self.IsInGear(self.right_gear):
+        # TODO(austin): Use battery volts here.
+        R_left = self.MotorRPM(self.left_shifter_position, self.X[0, 0])
+        self.U_ideal[0, 0] = numpy.clip(
+            self.left_cim.K * (R_left - self.left_cim.X) + R_left / self.left_cim.Kv,
+            self.left_cim.U_min, self.left_cim.U_max)
+        self.left_cim.Update(self.U_ideal[0, 0])
+
+        R_right = self.MotorRPM(self.right_shifter_position, self.X[1, 0])
+        self.U_ideal[1, 0] = numpy.clip(
+            self.right_cim.K * (R_right - self.right_cim.X) + R_right / self.right_cim.Kv,
+            self.right_cim.U_min, self.right_cim.U_max)
+        self.right_cim.Update(self.U_ideal[1, 0])
+      else:
+        assert False
+
+    self.U = numpy.clip(self.U_ideal, self.U_min, self.U_max)
+
+    # TODO(austin): Model the robot as not accelerating when you shift...
+    # This hack only works when you shift at the same time.
+    if self.IsInGear(self.left_gear) and self.IsInGear(self.right_gear):
+      self.X = self.CurrentDrivetrain().A * self.X + self.CurrentDrivetrain().B * self.U
+
+    self.left_gear, self.left_shifter_position = self.SimShifter(
+        self.left_gear, self.left_shifter_position)
+    self.right_gear, self.right_shifter_position = self.SimShifter(
+        self.right_gear, self.right_shifter_position)
+
+    glog.debug('U is %s %s', str(self.U[0, 0]), str(self.U[1, 0]))
+    glog.debug('Left shifter %s %d Right shifter %s %d',
+               self.left_gear, self.left_shifter_position,
+               self.right_gear, self.right_shifter_position)
+
+def WritePolyDrivetrain(drivetrain_files, motor_files, year_namespace,
+                        drivetrain_params):
+  vdrivetrain = VelocityDrivetrain(drivetrain_params)
+  namespaces = [year_namespace, 'control_loops', 'drivetrain']
+  dog_loop_writer = control_loop.ControlLoopWriter(
+      "VelocityDrivetrain", [vdrivetrain.drivetrain_low_low,
+                     vdrivetrain.drivetrain_low_high,
+                     vdrivetrain.drivetrain_high_low,
+                     vdrivetrain.drivetrain_high_high],
+                     namespaces=namespaces)
+
+  dog_loop_writer.Write(drivetrain_files[0], drivetrain_files[1])
+
+  cim_writer = control_loop.ControlLoopWriter("CIM", [CIM()])
+
+  cim_writer.Write(motor_files[0], motor_files[1])
+
+def PlotPolyDrivetrainMotions(drivetrain_params):
+  vdrivetrain = VelocityDrivetrain(drivetrain_params)
+  vl_plot = []
+  vr_plot = []
+  ul_plot = []
+  ur_plot = []
+  radius_plot = []
+  t_plot = []
+  left_gear_plot = []
+  right_gear_plot = []
+  vdrivetrain.left_shifter_position = 0.0
+  vdrivetrain.right_shifter_position = 0.0
+  vdrivetrain.left_gear = VelocityDrivetrain.LOW
+  vdrivetrain.right_gear = VelocityDrivetrain.LOW
+
+  glog.debug('K is %s', str(vdrivetrain.CurrentDrivetrain().K))
+
+  if vdrivetrain.left_gear is VelocityDrivetrain.HIGH:
+    glog.debug('Left is high')
+  else:
+    glog.debug('Left is low')
+  if vdrivetrain.right_gear is VelocityDrivetrain.HIGH:
+    glog.debug('Right is high')
+  else:
+    glog.debug('Right is low')
+
+  for t in numpy.arange(0, 1.7, vdrivetrain.dt):
+    if t < 0.5:
+      vdrivetrain.Update(throttle=0.00, steering=1.0)
+    elif t < 1.2:
+      vdrivetrain.Update(throttle=0.5, steering=1.0)
+    else:
+      vdrivetrain.Update(throttle=0.00, steering=1.0)
+    t_plot.append(t)
+    vl_plot.append(vdrivetrain.X[0, 0])
+    vr_plot.append(vdrivetrain.X[1, 0])
+    ul_plot.append(vdrivetrain.U[0, 0])
+    ur_plot.append(vdrivetrain.U[1, 0])
+    left_gear_plot.append((vdrivetrain.left_gear is VelocityDrivetrain.HIGH) * 2.0 - 10.0)
+    right_gear_plot.append((vdrivetrain.right_gear is VelocityDrivetrain.HIGH) * 2.0 - 10.0)
+
+    fwd_velocity = (vdrivetrain.X[1, 0] + vdrivetrain.X[0, 0]) / 2
+    turn_velocity = (vdrivetrain.X[1, 0] - vdrivetrain.X[0, 0])
+    if abs(fwd_velocity) < 0.0000001:
+      radius_plot.append(turn_velocity)
+    else:
+      radius_plot.append(turn_velocity / fwd_velocity)
+
+  # TODO(austin):
+  # Shifting compensation.
+
+  # Tighten the turn.
+  # Closed loop drive.
+
+  pylab.plot(t_plot, vl_plot, label='left velocity')
+  pylab.plot(t_plot, vr_plot, label='right velocity')
+  pylab.plot(t_plot, ul_plot, label='left voltage')
+  pylab.plot(t_plot, ur_plot, label='right voltage')
+  pylab.plot(t_plot, radius_plot, label='radius')
+  pylab.plot(t_plot, left_gear_plot, label='left gear high')
+  pylab.plot(t_plot, right_gear_plot, label='right gear high')
+  pylab.legend()
+  pylab.show()
diff --git a/y2017/control_loops/python/BUILD b/y2017/control_loops/python/BUILD
index 7365124..83473ae 100644
--- a/y2017/control_loops/python/BUILD
+++ b/y2017/control_loops/python/BUILD
@@ -8,7 +8,7 @@
   deps = [
     '//external:python-gflags',
     '//external:python-glog',
-    '//frc971/control_loops/python:controls',
+    '//frc971/control_loops/python:drivetrain',
   ],
   restricted_to = ['//tools:k8'],
 )
@@ -22,7 +22,7 @@
   deps = [
     '//external:python-gflags',
     '//external:python-glog',
-    '//frc971/control_loops/python:controls',
+    '//frc971/control_loops/python:polydrivetrain',
   ],
   restricted_to = ['//tools:k8'],
 )
@@ -37,6 +37,7 @@
     '//external:python-gflags',
     '//external:python-glog',
     '//frc971/control_loops/python:controls',
+    '//frc971/control_loops/python:drivetrain',
   ],
   restricted_to = ['//tools:k8'],
 )
diff --git a/y2017/control_loops/python/drivetrain.py b/y2017/control_loops/python/drivetrain.py
index 541d158..b4635ee 100755
--- a/y2017/control_loops/python/drivetrain.py
+++ b/y2017/control_loops/python/drivetrain.py
@@ -1,10 +1,7 @@
 #!/usr/bin/python
 
-from frc971.control_loops.python import control_loop
-from frc971.control_loops.python import controls
-import numpy
+from frc971.control_loops.python import drivetrain
 import sys
-from matplotlib import pylab
 
 import gflags
 import glog
@@ -13,351 +10,28 @@
 
 gflags.DEFINE_bool('plot', False, 'If true, plot the loop response.')
 
-
-class Drivetrain(control_loop.ControlLoop):
-  def __init__(self, name="Drivetrain", left_low=True, right_low=True):
-    super(Drivetrain, self).__init__(name)
-    # Number of motors per side
-    self.num_motors = 2
-    # Stall Torque in N m
-    self.stall_torque = 2.42 * self.num_motors * 0.60
-    # Stall Current in Amps
-    self.stall_current = 133.0 * self.num_motors
-    self.free_speed_rpm = 5500.0
-    # Free Speed in rotations/second.
-    self.free_speed = self.free_speed_rpm / 60
-    # Free Current in Amps
-    self.free_current = 4.7 * self.num_motors
-    # Moment of inertia of the drivetrain in kg m^2
-    self.J = 6.0
-    # Mass of the robot, in kg.
-    self.m = 52
-    # Radius of the robot, in meters (requires tuning by hand)
-    self.rb = 0.59055 / 2.0
-    # Radius of the wheels, in meters.
-    self.r = 0.08255 / 2.0
-    # Resistance of the motor, divided by the number of motors.
-    self.resistance = 12.0 / self.stall_current
-    # Motor velocity constant
-    self.Kv = ((self.free_speed * 2.0 * numpy.pi) /
-               (12.0 - self.resistance * self.free_current))
-    # Torque constant
-    self.Kt = self.stall_torque / self.stall_current
-    # Gear ratios
-    self.G_low = 11.0 / 60.0
-    self.G_high = 11.0 / 60.0
-    if left_low:
-      self.Gl = self.G_low
-    else:
-      self.Gl = self.G_high
-    if right_low:
-      self.Gr = self.G_low
-    else:
-      self.Gr = self.G_high
-
-    # Control loop time step
-    self.dt = 0.00505
-
-    # These describe the way that a given side of a robot will be influenced
-    # by the other side. Units of 1 / kg.
-    self.msp = 1.0 / self.m + self.rb * self.rb / self.J
-    self.msn = 1.0 / self.m - self.rb * self.rb / self.J
-    # The calculations which we will need for A and B.
-    self.tcl = -self.Kt / self.Kv / (self.Gl * self.Gl * self.resistance * self.r * self.r)
-    self.tcr = -self.Kt / self.Kv / (self.Gr * self.Gr * self.resistance * self.r * self.r)
-    self.mpl = self.Kt / (self.Gl * self.resistance * self.r)
-    self.mpr = self.Kt / (self.Gr * self.resistance * self.r)
-
-    # State feedback matrices
-    # X will be of the format
-    # [[positionl], [velocityl], [positionr], velocityr]]
-    self.A_continuous = numpy.matrix(
-        [[0, 1, 0, 0],
-         [0, self.msp * self.tcl, 0, self.msn * self.tcr],
-         [0, 0, 0, 1],
-         [0, self.msn * self.tcl, 0, self.msp * self.tcr]])
-    self.B_continuous = numpy.matrix(
-        [[0, 0],
-         [self.msp * self.mpl, self.msn * self.mpr],
-         [0, 0],
-         [self.msn * self.mpl, self.msp * self.mpr]])
-    self.C = numpy.matrix([[1, 0, 0, 0],
-                           [0, 0, 1, 0]])
-    self.D = numpy.matrix([[0, 0],
-                           [0, 0]])
-
-    self.A, self.B = self.ContinuousToDiscrete(
-        self.A_continuous, self.B_continuous, self.dt)
-
-    if left_low or right_low:
-      q_pos = 0.12
-      q_vel = 1.0
-    else:
-      q_pos = 0.14
-      q_vel = 0.95
-
-    # Tune the LQR controller
-    self.Q = numpy.matrix([[(1.0 / (q_pos ** 2.0)), 0.0, 0.0, 0.0],
-                           [0.0, (1.0 / (q_vel ** 2.0)), 0.0, 0.0],
-                           [0.0, 0.0, (1.0 / (q_pos ** 2.0)), 0.0],
-                           [0.0, 0.0, 0.0, (1.0 / (q_vel ** 2.0))]])
-
-    self.R = numpy.matrix([[(1.0 / (12.0 ** 2.0)), 0.0],
-                           [0.0, (1.0 / (12.0 ** 2.0))]])
-    self.K = controls.dlqr(self.A, self.B, self.Q, self.R)
-
-    glog.debug('DT q_pos %f q_vel %s %s', q_pos, q_vel, name)
-    glog.debug(str(numpy.linalg.eig(self.A - self.B * self.K)[0]))
-    glog.debug('K %s', repr(self.K))
-
-    self.hlp = 0.3
-    self.llp = 0.4
-    self.PlaceObserverPoles([self.hlp, self.hlp, self.llp, self.llp])
-
-    self.U_max = numpy.matrix([[12.0], [12.0]])
-    self.U_min = numpy.matrix([[-12.0], [-12.0]])
-
-    self.InitializeState()
-
-
-class KFDrivetrain(Drivetrain):
-  def __init__(self, name="KFDrivetrain", left_low=True, right_low=True):
-    super(KFDrivetrain, self).__init__(name, left_low, right_low)
-
-    self.unaugmented_A_continuous = self.A_continuous
-    self.unaugmented_B_continuous = self.B_continuous
-
-    # The practical voltage applied to the wheels is
-    #   V_left = U_left + left_voltage_error
-    #
-    # The states are
-    # [left position, left velocity, right position, right velocity,
-    #  left voltage error, right voltage error, angular_error]
-    #
-    # The left and right positions are filtered encoder positions and are not
-    # adjusted for heading error.
-    # The turn velocity as computed by the left and right velocities is
-    # adjusted by the gyro velocity.
-    # The angular_error is the angular velocity error between the wheel speed
-    # and the gyro speed.
-    self.A_continuous = numpy.matrix(numpy.zeros((7, 7)))
-    self.B_continuous = numpy.matrix(numpy.zeros((7, 2)))
-    self.A_continuous[0:4,0:4] = self.unaugmented_A_continuous
-    self.A_continuous[0:4,4:6] = self.unaugmented_B_continuous
-    self.B_continuous[0:4,0:2] = self.unaugmented_B_continuous
-    self.A_continuous[0,6] = 1
-    self.A_continuous[2,6] = -1
-
-    self.A, self.B = self.ContinuousToDiscrete(
-        self.A_continuous, self.B_continuous, self.dt)
-
-    self.C = numpy.matrix([[1, 0, 0, 0, 0, 0, 0],
-                           [0, 0, 1, 0, 0, 0, 0],
-                           [0, -0.5 / self.rb, 0, 0.5 / self.rb, 0, 0, 0]])
-
-    self.D = numpy.matrix([[0, 0],
-                           [0, 0],
-                           [0, 0]])
-
-    q_pos = 0.05
-    q_vel = 1.00
-    q_voltage = 10.0
-    q_encoder_uncertainty = 2.00
-
-    self.Q = numpy.matrix([[(q_pos ** 2.0), 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
-                           [0.0, (q_vel ** 2.0), 0.0, 0.0, 0.0, 0.0, 0.0],
-                           [0.0, 0.0, (q_pos ** 2.0), 0.0, 0.0, 0.0, 0.0],
-                           [0.0, 0.0, 0.0, (q_vel ** 2.0), 0.0, 0.0, 0.0],
-                           [0.0, 0.0, 0.0, 0.0, (q_voltage ** 2.0), 0.0, 0.0],
-                           [0.0, 0.0, 0.0, 0.0, 0.0, (q_voltage ** 2.0), 0.0],
-                           [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, (q_encoder_uncertainty ** 2.0)]])
-
-    r_pos =  0.0001
-    r_gyro = 0.000001
-    self.R = numpy.matrix([[(r_pos ** 2.0), 0.0, 0.0],
-                           [0.0, (r_pos ** 2.0), 0.0],
-                           [0.0, 0.0, (r_gyro ** 2.0)]])
-
-    # Solving for kf gains.
-    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
-
-    unaug_K = self.K
-
-    # Implement a nice closed loop controller for use by the closed loop
-    # controller.
-    self.K = numpy.matrix(numpy.zeros((self.B.shape[1], self.A.shape[0])))
-    self.K[0:2, 0:4] = unaug_K
-    self.K[0, 4] = 1.0
-    self.K[1, 5] = 1.0
-
-    self.Qff = numpy.matrix(numpy.zeros((4, 4)))
-    qff_pos = 0.005
-    qff_vel = 1.00
-    self.Qff[0, 0] = 1.0 / qff_pos ** 2.0
-    self.Qff[1, 1] = 1.0 / qff_vel ** 2.0
-    self.Qff[2, 2] = 1.0 / qff_pos ** 2.0
-    self.Qff[3, 3] = 1.0 / qff_vel ** 2.0
-    self.Kff = numpy.matrix(numpy.zeros((2, 7)))
-    self.Kff[0:2, 0:4] = controls.TwoStateFeedForwards(self.B[0:4,:], self.Qff)
-
-    self.InitializeState()
-
+kDrivetrain = drivetrain.DrivetrainParams(J = 6.0,
+                                          mass = 52,
+                                          robot_radius = 0.59055 / 2.0,
+                                          wheel_radius = 0.08255 / 2.0,
+                                          G_high = 11.0 / 60.0,
+                                          G_low = 11.0 / 60.0,
+                                          q_pos_low = 0.12,
+                                          q_pos_high = 0.14,
+                                          q_vel_low = 1.0,
+                                          q_vel_high = 0.95)
 
 def main(argv):
   argv = FLAGS(argv)
   glog.init()
 
-  # Simulate the response of the system to a step input.
-  drivetrain = Drivetrain(left_low=False, right_low=False)
-  simulated_left = []
-  simulated_right = []
-  for _ in xrange(100):
-    drivetrain.Update(numpy.matrix([[12.0], [12.0]]))
-    simulated_left.append(drivetrain.X[0, 0])
-    simulated_right.append(drivetrain.X[2, 0])
-
   if FLAGS.plot:
-    pylab.rc('lines', linewidth=4)
-    pylab.plot(range(100), simulated_left, label='left position')
-    pylab.plot(range(100), simulated_right, 'r--', label='right position')
-    pylab.suptitle('Acceleration Test\n12 Volt Step Input')
-    pylab.legend(loc='lower right')
-    pylab.show()
-
-  # Simulate forwards motion.
-  drivetrain = Drivetrain(left_low=False, right_low=False)
-  close_loop_left = []
-  close_loop_right = []
-  left_power = []
-  right_power = []
-  R = numpy.matrix([[1.0], [0.0], [1.0], [0.0]])
-  for _ in xrange(300):
-    U = numpy.clip(drivetrain.K * (R - drivetrain.X_hat),
-                   drivetrain.U_min, drivetrain.U_max)
-    drivetrain.UpdateObserver(U)
-    drivetrain.Update(U)
-    close_loop_left.append(drivetrain.X[0, 0])
-    close_loop_right.append(drivetrain.X[2, 0])
-    left_power.append(U[0, 0])
-    right_power.append(U[1, 0])
-
-  if FLAGS.plot:
-    pylab.plot(range(300), close_loop_left, label='left position')
-    pylab.plot(range(300), close_loop_right, 'm--', label='right position')
-    pylab.plot(range(300), left_power, label='left power')
-    pylab.plot(range(300), right_power, '--', label='right power')
-    pylab.suptitle('Linear Move\nLeft and Right Position going to 1')
-    pylab.legend()
-    pylab.show()
-
-  # Try turning in place
-  drivetrain = Drivetrain()
-  close_loop_left = []
-  close_loop_right = []
-  R = numpy.matrix([[-1.0], [0.0], [1.0], [0.0]])
-  for _ in xrange(200):
-    U = numpy.clip(drivetrain.K * (R - drivetrain.X_hat),
-                   drivetrain.U_min, drivetrain.U_max)
-    drivetrain.UpdateObserver(U)
-    drivetrain.Update(U)
-    close_loop_left.append(drivetrain.X[0, 0])
-    close_loop_right.append(drivetrain.X[2, 0])
-
-  if FLAGS.plot:
-    pylab.plot(range(200), close_loop_left, label='left position')
-    pylab.plot(range(200), close_loop_right, label='right position')
-    pylab.suptitle('Angular Move\nLeft position going to -1 and right position going to 1')
-    pylab.legend(loc='center right')
-    pylab.show()
-
-  # Try turning just one side.
-  drivetrain = Drivetrain()
-  close_loop_left = []
-  close_loop_right = []
-  R = numpy.matrix([[0.0], [0.0], [1.0], [0.0]])
-  for _ in xrange(300):
-    U = numpy.clip(drivetrain.K * (R - drivetrain.X_hat),
-                   drivetrain.U_min, drivetrain.U_max)
-    drivetrain.UpdateObserver(U)
-    drivetrain.Update(U)
-    close_loop_left.append(drivetrain.X[0, 0])
-    close_loop_right.append(drivetrain.X[2, 0])
-
-  if FLAGS.plot:
-    pylab.plot(range(300), close_loop_left, label='left position')
-    pylab.plot(range(300), close_loop_right, label='right position')
-    pylab.suptitle('Pivot\nLeft position not changing and right position going to 1')
-    pylab.legend(loc='center right')
-    pylab.show()
-
-  # Write the generated constants out to a file.
-  drivetrain_low_low = Drivetrain(
-      name="DrivetrainLowLow", left_low=True, right_low=True)
-  drivetrain_low_high = Drivetrain(
-      name="DrivetrainLowHigh", left_low=True, right_low=False)
-  drivetrain_high_low = Drivetrain(
-      name="DrivetrainHighLow", left_low=False, right_low=True)
-  drivetrain_high_high = Drivetrain(
-      name="DrivetrainHighHigh", left_low=False, right_low=False)
-
-  kf_drivetrain_low_low = KFDrivetrain(
-      name="KFDrivetrainLowLow", left_low=True, right_low=True)
-  kf_drivetrain_low_high = KFDrivetrain(
-      name="KFDrivetrainLowHigh", left_low=True, right_low=False)
-  kf_drivetrain_high_low = KFDrivetrain(
-      name="KFDrivetrainHighLow", left_low=False, right_low=True)
-  kf_drivetrain_high_high = KFDrivetrain(
-      name="KFDrivetrainHighHigh", left_low=False, right_low=False)
-
-  if len(argv) != 5:
+    drivetrain.PlotDrivetrainMotions(kDrivetrain)
+  elif len(argv) != 5:
     print "Expected .h file name and .cc file name"
   else:
-    namespaces = ['y2017', 'control_loops', 'drivetrain']
-    dog_loop_writer = control_loop.ControlLoopWriter(
-        "Drivetrain", [drivetrain_low_low, drivetrain_low_high,
-                       drivetrain_high_low, drivetrain_high_high],
-        namespaces = namespaces)
-    dog_loop_writer.AddConstant(control_loop.Constant("kDt", "%f",
-          drivetrain_low_low.dt))
-    dog_loop_writer.AddConstant(control_loop.Constant("kStallTorque", "%f",
-          drivetrain_low_low.stall_torque))
-    dog_loop_writer.AddConstant(control_loop.Constant("kStallCurrent", "%f",
-          drivetrain_low_low.stall_current))
-    dog_loop_writer.AddConstant(control_loop.Constant("kFreeSpeed", "%f",
-          drivetrain_low_low.free_speed))
-    dog_loop_writer.AddConstant(control_loop.Constant("kFreeCurrent", "%f",
-          drivetrain_low_low.free_current))
-    dog_loop_writer.AddConstant(control_loop.Constant("kJ", "%f",
-          drivetrain_low_low.J))
-    dog_loop_writer.AddConstant(control_loop.Constant("kMass", "%f",
-          drivetrain_low_low.m))
-    dog_loop_writer.AddConstant(control_loop.Constant("kRobotRadius", "%f",
-          drivetrain_low_low.rb))
-    dog_loop_writer.AddConstant(control_loop.Constant("kWheelRadius", "%f",
-          drivetrain_low_low.r))
-    dog_loop_writer.AddConstant(control_loop.Constant("kR", "%f",
-          drivetrain_low_low.resistance))
-    dog_loop_writer.AddConstant(control_loop.Constant("kV", "%f",
-          drivetrain_low_low.Kv))
-    dog_loop_writer.AddConstant(control_loop.Constant("kT", "%f",
-          drivetrain_low_low.Kt))
-    dog_loop_writer.AddConstant(control_loop.Constant("kLowGearRatio", "%f",
-          drivetrain_low_low.G_low))
-    dog_loop_writer.AddConstant(control_loop.Constant("kHighGearRatio", "%f",
-          drivetrain_high_high.G_high))
-    dog_loop_writer.AddConstant(control_loop.Constant("kHighOutputRatio", "%f",
-          drivetrain_high_high.G_high * drivetrain_high_high.r))
-
-    dog_loop_writer.Write(argv[1], argv[2])
-
-    kf_loop_writer = control_loop.ControlLoopWriter(
-        "KFDrivetrain", [kf_drivetrain_low_low, kf_drivetrain_low_high,
-                         kf_drivetrain_high_low, kf_drivetrain_high_high],
-        namespaces = namespaces)
-    kf_loop_writer.Write(argv[3], argv[4])
+    # Write the generated constants out to a file.
+    drivetrain.WriteDrivetrain(argv[1:3], argv[3:5], 'y2017', kDrivetrain)
 
 if __name__ == '__main__':
   sys.exit(main(sys.argv))
diff --git a/y2017/control_loops/python/polydrivetrain.py b/y2017/control_loops/python/polydrivetrain.py
index d3a5683..701308e 100755
--- a/y2017/control_loops/python/polydrivetrain.py
+++ b/y2017/control_loops/python/polydrivetrain.py
@@ -1,13 +1,8 @@
 #!/usr/bin/python
 
-import numpy
 import sys
-from frc971.control_loops.python import polytope
 from y2017.control_loops.python import drivetrain
-from frc971.control_loops.python import control_loop
-from frc971.control_loops.python import controls
-from frc971.control_loops.python.cim import CIM
-from matplotlib import pylab
+from frc971.control_loops.python import polydrivetrain
 
 import gflags
 import glog
@@ -21,479 +16,14 @@
 except gflags.DuplicateFlagError:
   pass
 
-def CoerceGoal(region, K, w, R):
-  """Intersects a line with a region, and finds the closest point to R.
-
-  Finds a point that is closest to R inside the region, and on the line
-  defined by K X = w.  If it is not possible to find a point on the line,
-  finds a point that is inside the region and closest to the line.  This
-  function assumes that
-
-  Args:
-    region: HPolytope, the valid goal region.
-    K: numpy.matrix (2 x 1), the matrix for the equation [K1, K2] [x1; x2] = w
-    w: float, the offset in the equation above.
-    R: numpy.matrix (2 x 1), the point to be closest to.
-
-  Returns:
-    numpy.matrix (2 x 1), the point.
-  """
-  return DoCoerceGoal(region, K, w, R)[0]
-
-def DoCoerceGoal(region, K, w, R):
-  if region.IsInside(R):
-    return (R, True)
-
-  perpendicular_vector = K.T / numpy.linalg.norm(K)
-  parallel_vector = numpy.matrix([[perpendicular_vector[1, 0]],
-                                  [-perpendicular_vector[0, 0]]])
-
-  # We want to impose the constraint K * X = w on the polytope H * X <= k.
-  # We do this by breaking X up into parallel and perpendicular components to
-  # the half plane.  This gives us the following equation.
-  #
-  #  parallel * (parallel.T \dot X) + perpendicular * (perpendicular \dot X)) = X
-  #
-  # Then, substitute this into the polytope.
-  #
-  #  H * (parallel * (parallel.T \dot X) + perpendicular * (perpendicular \dot X)) <= k
-  #
-  # Substitute K * X = w
-  #
-  # H * parallel * (parallel.T \dot X) + H * perpendicular * w <= k
-  #
-  # Move all the knowns to the right side.
-  #
-  # H * parallel * ([parallel1 parallel2] * X) <= k - H * perpendicular * w
-  #
-  # Let t = parallel.T \dot X, the component parallel to the surface.
-  #
-  # H * parallel * t <= k - H * perpendicular * w
-  #
-  # This is a polytope which we can solve, and use to figure out the range of X
-  # that we care about!
-
-  t_poly = polytope.HPolytope(
-      region.H * parallel_vector,
-      region.k - region.H * perpendicular_vector * w)
-
-  vertices = t_poly.Vertices()
-
-  if vertices.shape[0]:
-    # The region exists!
-    # Find the closest vertex
-    min_distance = numpy.infty
-    closest_point = None
-    for vertex in vertices:
-      point = parallel_vector * vertex + perpendicular_vector * w
-      length = numpy.linalg.norm(R - point)
-      if length < min_distance:
-        min_distance = length
-        closest_point = point
-
-    return (closest_point, True)
-  else:
-    # Find the vertex of the space that is closest to the line.
-    region_vertices = region.Vertices()
-    min_distance = numpy.infty
-    closest_point = None
-    for vertex in region_vertices:
-      point = vertex.T
-      length = numpy.abs((perpendicular_vector.T * point)[0, 0])
-      if length < min_distance:
-        min_distance = length
-        closest_point = point
-
-    return (closest_point, False)
-
-
-class VelocityDrivetrainModel(control_loop.ControlLoop):
-  def __init__(self, left_low=True, right_low=True, name="VelocityDrivetrainModel"):
-    super(VelocityDrivetrainModel, self).__init__(name)
-    self._drivetrain = drivetrain.Drivetrain(left_low=left_low,
-                                             right_low=right_low)
-    self.dt = 0.00505
-    self.A_continuous = numpy.matrix(
-        [[self._drivetrain.A_continuous[1, 1], self._drivetrain.A_continuous[1, 3]],
-         [self._drivetrain.A_continuous[3, 1], self._drivetrain.A_continuous[3, 3]]])
-
-    self.B_continuous = numpy.matrix(
-        [[self._drivetrain.B_continuous[1, 0], self._drivetrain.B_continuous[1, 1]],
-         [self._drivetrain.B_continuous[3, 0], self._drivetrain.B_continuous[3, 1]]])
-    self.C = numpy.matrix(numpy.eye(2))
-    self.D = numpy.matrix(numpy.zeros((2, 2)))
-
-    self.A, self.B = self.ContinuousToDiscrete(self.A_continuous,
-                                               self.B_continuous, self.dt)
-
-    # FF * X = U (steady state)
-    self.FF = self.B.I * (numpy.eye(2) - self.A)
-
-    self.PlaceControllerPoles([0.90, 0.90])
-    self.PlaceObserverPoles([0.02, 0.02])
-
-    self.G_high = self._drivetrain.G_high
-    self.G_low = self._drivetrain.G_low
-    self.resistance = self._drivetrain.resistance
-    self.r = self._drivetrain.r
-    self.Kv = self._drivetrain.Kv
-    self.Kt = self._drivetrain.Kt
-
-    self.U_max = self._drivetrain.U_max
-    self.U_min = self._drivetrain.U_min
-
-
-class VelocityDrivetrain(object):
-  HIGH = 'high'
-  LOW = 'low'
-  SHIFTING_UP = 'up'
-  SHIFTING_DOWN = 'down'
-
-  def __init__(self):
-    self.drivetrain_low_low = VelocityDrivetrainModel(
-        left_low=True, right_low=True, name='VelocityDrivetrainLowLow')
-    self.drivetrain_low_high = VelocityDrivetrainModel(left_low=True, right_low=False, name='VelocityDrivetrainLowHigh')
-    self.drivetrain_high_low = VelocityDrivetrainModel(left_low=False, right_low=True, name = 'VelocityDrivetrainHighLow')
-    self.drivetrain_high_high = VelocityDrivetrainModel(left_low=False, right_low=False, name = 'VelocityDrivetrainHighHigh')
-
-    # X is [lvel, rvel]
-    self.X = numpy.matrix(
-        [[0.0],
-         [0.0]])
-
-    self.U_poly = polytope.HPolytope(
-        numpy.matrix([[1, 0],
-                      [-1, 0],
-                      [0, 1],
-                      [0, -1]]),
-        numpy.matrix([[12],
-                      [12],
-                      [12],
-                      [12]]))
-
-    self.U_max = numpy.matrix(
-        [[12.0],
-         [12.0]])
-    self.U_min = numpy.matrix(
-        [[-12.0000000000],
-         [-12.0000000000]])
-
-    self.dt = 0.00505
-
-    self.R = numpy.matrix(
-        [[0.0],
-         [0.0]])
-
-    self.U_ideal = numpy.matrix(
-        [[0.0],
-         [0.0]])
-
-    # ttrust is the comprimise between having full throttle negative inertia,
-    # and having no throttle negative inertia.  A value of 0 is full throttle
-    # inertia.  A value of 1 is no throttle negative inertia.
-    self.ttrust = 1.0
-
-    self.left_gear = VelocityDrivetrain.LOW
-    self.right_gear = VelocityDrivetrain.LOW
-    self.left_shifter_position = 0.0
-    self.right_shifter_position = 0.0
-    self.left_cim = CIM()
-    self.right_cim = CIM()
-
-  def IsInGear(self, gear):
-    return gear is VelocityDrivetrain.HIGH or gear is VelocityDrivetrain.LOW
-
-  def MotorRPM(self, shifter_position, velocity):
-    if shifter_position > 0.5:
-      return (velocity / self.CurrentDrivetrain().G_high /
-              self.CurrentDrivetrain().r)
-    else:
-      return (velocity / self.CurrentDrivetrain().G_low /
-              self.CurrentDrivetrain().r)
-
-  def CurrentDrivetrain(self):
-    if self.left_shifter_position > 0.5:
-      if self.right_shifter_position > 0.5:
-        return self.drivetrain_high_high
-      else:
-        return self.drivetrain_high_low
-    else:
-      if self.right_shifter_position > 0.5:
-        return self.drivetrain_low_high
-      else:
-        return self.drivetrain_low_low
-
-  def SimShifter(self, gear, shifter_position):
-    if gear is VelocityDrivetrain.HIGH or gear is VelocityDrivetrain.SHIFTING_UP:
-      shifter_position = min(shifter_position + 0.5, 1.0)
-    else:
-      shifter_position = max(shifter_position - 0.5, 0.0)
-
-    if shifter_position == 1.0:
-      gear = VelocityDrivetrain.HIGH
-    elif shifter_position == 0.0:
-      gear = VelocityDrivetrain.LOW
-
-    return gear, shifter_position
-
-  def ComputeGear(self, wheel_velocity, should_print=False, current_gear=False, gear_name=None):
-    high_omega = (wheel_velocity / self.CurrentDrivetrain().G_high /
-                  self.CurrentDrivetrain().r)
-    low_omega = (wheel_velocity / self.CurrentDrivetrain().G_low /
-                 self.CurrentDrivetrain().r)
-    #print gear_name, "Motor Energy Difference.", 0.5 * 0.000140032647 * (low_omega * low_omega - high_omega * high_omega), "joules"
-    high_torque = ((12.0 - high_omega / self.CurrentDrivetrain().Kv) *
-                   self.CurrentDrivetrain().Kt / self.CurrentDrivetrain().resistance)
-    low_torque = ((12.0 - low_omega / self.CurrentDrivetrain().Kv) *
-                  self.CurrentDrivetrain().Kt / self.CurrentDrivetrain().resistance)
-    high_power = high_torque * high_omega
-    low_power = low_torque * low_omega
-    #if should_print:
-    #  print gear_name, "High omega", high_omega, "Low omega", low_omega
-    #  print gear_name, "High torque", high_torque, "Low torque", low_torque
-    #  print gear_name, "High power", high_power, "Low power", low_power
-
-    # Shift algorithm improvements.
-    # TODO(aschuh):
-    # It takes time to shift.  Shifting down for 1 cycle doesn't make sense
-    # because you will end up slower than without shifting.  Figure out how
-    # to include that info.
-    # If the driver is still in high gear, but isn't asking for the extra power
-    # from low gear, don't shift until he asks for it.
-    goal_gear_is_high = high_power > low_power
-    #goal_gear_is_high = True
-
-    if not self.IsInGear(current_gear):
-      glog.debug('%s Not in gear.', gear_name)
-      return current_gear
-    else:
-      is_high = current_gear is VelocityDrivetrain.HIGH
-      if is_high != goal_gear_is_high:
-        if goal_gear_is_high:
-          glog.debug('%s Shifting up.', gear_name)
-          return VelocityDrivetrain.SHIFTING_UP
-        else:
-          glog.debug('%s Shifting down.', gear_name)
-          return VelocityDrivetrain.SHIFTING_DOWN
-      else:
-        return current_gear
-
-  def FilterVelocity(self, throttle):
-    # Invert the plant to figure out how the velocity filter would have to work
-    # out in order to filter out the forwards negative inertia.
-    # This math assumes that the left and right power and velocity are equal.
-
-    # The throttle filter should filter such that the motor in the highest gear
-    # should be controlling the time constant.
-    # Do this by finding the index of FF that has the lowest value, and computing
-    # the sums using that index.
-    FF_sum = self.CurrentDrivetrain().FF.sum(axis=1)
-    min_FF_sum_index = numpy.argmin(FF_sum)
-    min_FF_sum = FF_sum[min_FF_sum_index, 0]
-    min_K_sum = self.CurrentDrivetrain().K[min_FF_sum_index, :].sum()
-    # Compute the FF sum for high gear.
-    high_min_FF_sum = self.drivetrain_high_high.FF[0, :].sum()
-
-    # U = self.K[0, :].sum() * (R - x_avg) + self.FF[0, :].sum() * R
-    # throttle * 12.0 = (self.K[0, :].sum() + self.FF[0, :].sum()) * R
-    #                   - self.K[0, :].sum() * x_avg
-
-    # R = (throttle * 12.0 + self.K[0, :].sum() * x_avg) /
-    #     (self.K[0, :].sum() + self.FF[0, :].sum())
-
-    # U = (K + FF) * R - K * X
-    # (K + FF) ^-1 * (U + K * X) = R
-
-    # Scale throttle by min_FF_sum / high_min_FF_sum.  This will make low gear
-    # have the same velocity goal as high gear, and so that the robot will hold
-    # the same speed for the same throttle for all gears.
-    adjusted_ff_voltage = numpy.clip(throttle * 12.0 * min_FF_sum / high_min_FF_sum, -12.0, 12.0)
-    return ((adjusted_ff_voltage + self.ttrust * min_K_sum * (self.X[0, 0] + self.X[1, 0]) / 2.0)
-            / (self.ttrust * min_K_sum + min_FF_sum))
-
-  def Update(self, throttle, steering):
-    # Shift into the gear which sends the most power to the floor.
-    # This is the same as sending the most torque down to the floor at the
-    # wheel.
-
-    self.left_gear = self.right_gear = True
-    if True:
-      self.left_gear = self.ComputeGear(self.X[0, 0], should_print=True,
-                                        current_gear=self.left_gear,
-                                        gear_name="left")
-      self.right_gear = self.ComputeGear(self.X[1, 0], should_print=True,
-                                         current_gear=self.right_gear,
-                                         gear_name="right")
-      if self.IsInGear(self.left_gear):
-        self.left_cim.X[0, 0] = self.MotorRPM(self.left_shifter_position, self.X[0, 0])
-
-      if self.IsInGear(self.right_gear):
-        self.right_cim.X[0, 0] = self.MotorRPM(self.right_shifter_position, self.X[0, 0])
-
-    if self.IsInGear(self.left_gear) and self.IsInGear(self.right_gear):
-      # Filter the throttle to provide a nicer response.
-      fvel = self.FilterVelocity(throttle)
-
-      # Constant radius means that angualar_velocity / linear_velocity = constant.
-      # Compute the left and right velocities.
-      steering_velocity = numpy.abs(fvel) * steering
-      left_velocity = fvel - steering_velocity
-      right_velocity = fvel + steering_velocity
-
-      # Write this constraint in the form of K * R = w
-      # angular velocity / linear velocity = constant
-      # (left - right) / (left + right) = constant
-      # left - right = constant * left + constant * right
-
-      # (fvel - steering * numpy.abs(fvel) - fvel - steering * numpy.abs(fvel)) /
-      #  (fvel - steering * numpy.abs(fvel) + fvel + steering * numpy.abs(fvel)) =
-      #       constant
-      # (- 2 * steering * numpy.abs(fvel)) / (2 * fvel) = constant
-      # (-steering * sign(fvel)) = constant
-      # (-steering * sign(fvel)) * (left + right) = left - right
-      # (steering * sign(fvel) + 1) * left + (steering * sign(fvel) - 1) * right = 0
-
-      equality_k = numpy.matrix(
-          [[1 + steering * numpy.sign(fvel), -(1 - steering * numpy.sign(fvel))]])
-      equality_w = 0.0
-
-      self.R[0, 0] = left_velocity
-      self.R[1, 0] = right_velocity
-
-      # Construct a constraint on R by manipulating the constraint on U
-      # Start out with H * U <= k
-      # U = FF * R + K * (R - X)
-      # H * (FF * R + K * R - K * X) <= k
-      # H * (FF + K) * R <= k + H * K * X
-      R_poly = polytope.HPolytope(
-          self.U_poly.H * (self.CurrentDrivetrain().K + self.CurrentDrivetrain().FF),
-          self.U_poly.k + self.U_poly.H * self.CurrentDrivetrain().K * self.X)
-
-      # Limit R back inside the box.
-      self.boxed_R = CoerceGoal(R_poly, equality_k, equality_w, self.R)
-
-      FF_volts = self.CurrentDrivetrain().FF * self.boxed_R
-      self.U_ideal = self.CurrentDrivetrain().K * (self.boxed_R - self.X) + FF_volts
-    else:
-      glog.debug('Not all in gear')
-      if not self.IsInGear(self.left_gear) and not self.IsInGear(self.right_gear):
-        # TODO(austin): Use battery volts here.
-        R_left = self.MotorRPM(self.left_shifter_position, self.X[0, 0])
-        self.U_ideal[0, 0] = numpy.clip(
-            self.left_cim.K * (R_left - self.left_cim.X) + R_left / self.left_cim.Kv,
-            self.left_cim.U_min, self.left_cim.U_max)
-        self.left_cim.Update(self.U_ideal[0, 0])
-
-        R_right = self.MotorRPM(self.right_shifter_position, self.X[1, 0])
-        self.U_ideal[1, 0] = numpy.clip(
-            self.right_cim.K * (R_right - self.right_cim.X) + R_right / self.right_cim.Kv,
-            self.right_cim.U_min, self.right_cim.U_max)
-        self.right_cim.Update(self.U_ideal[1, 0])
-      else:
-        assert False
-
-    self.U = numpy.clip(self.U_ideal, self.U_min, self.U_max)
-
-    # TODO(austin): Model the robot as not accelerating when you shift...
-    # This hack only works when you shift at the same time.
-    if self.IsInGear(self.left_gear) and self.IsInGear(self.right_gear):
-      self.X = self.CurrentDrivetrain().A * self.X + self.CurrentDrivetrain().B * self.U
-
-    self.left_gear, self.left_shifter_position = self.SimShifter(
-        self.left_gear, self.left_shifter_position)
-    self.right_gear, self.right_shifter_position = self.SimShifter(
-        self.right_gear, self.right_shifter_position)
-
-    glog.debug('U is %s %s', str(self.U[0, 0]), str(self.U[1, 0]))
-    glog.debug('Left shifter %s %d Right shifter %s %d',
-               self.left_gear, self.left_shifter_position,
-               self.right_gear, self.right_shifter_position)
-
-
 def main(argv):
-  vdrivetrain = VelocityDrivetrain()
-
-  if not FLAGS.plot:
-    if len(argv) != 5:
-      glog.fatal('Expected .h file name and .cc file name')
-    else:
-      namespaces = ['y2017', 'control_loops', 'drivetrain']
-      dog_loop_writer = control_loop.ControlLoopWriter(
-          "VelocityDrivetrain", [vdrivetrain.drivetrain_low_low,
-                         vdrivetrain.drivetrain_low_high,
-                         vdrivetrain.drivetrain_high_low,
-                         vdrivetrain.drivetrain_high_high],
-                         namespaces=namespaces)
-
-      dog_loop_writer.Write(argv[1], argv[2])
-
-      cim_writer = control_loop.ControlLoopWriter("CIM", [CIM()])
-
-      cim_writer.Write(argv[3], argv[4])
-      return
-
-  vl_plot = []
-  vr_plot = []
-  ul_plot = []
-  ur_plot = []
-  radius_plot = []
-  t_plot = []
-  left_gear_plot = []
-  right_gear_plot = []
-  vdrivetrain.left_shifter_position = 0.0
-  vdrivetrain.right_shifter_position = 0.0
-  vdrivetrain.left_gear = VelocityDrivetrain.LOW
-  vdrivetrain.right_gear = VelocityDrivetrain.LOW
-
-  glog.debug('K is %s', str(vdrivetrain.CurrentDrivetrain().K))
-
-  if vdrivetrain.left_gear is VelocityDrivetrain.HIGH:
-    glog.debug('Left is high')
+  if FLAGS.plot:
+    polydrivetrain.PlotPolyDrivetrainMotions(drivetrain.kDrivetrain)
+  elif len(argv) != 5:
+    glog.fatal('Expected .h file name and .cc file name')
   else:
-    glog.debug('Left is low')
-  if vdrivetrain.right_gear is VelocityDrivetrain.HIGH:
-    glog.debug('Right is high')
-  else:
-    glog.debug('Right is low')
-
-  for t in numpy.arange(0, 1.7, vdrivetrain.dt):
-    if t < 0.5:
-      vdrivetrain.Update(throttle=0.00, steering=1.0)
-    elif t < 1.2:
-      vdrivetrain.Update(throttle=0.5, steering=1.0)
-    else:
-      vdrivetrain.Update(throttle=0.00, steering=1.0)
-    t_plot.append(t)
-    vl_plot.append(vdrivetrain.X[0, 0])
-    vr_plot.append(vdrivetrain.X[1, 0])
-    ul_plot.append(vdrivetrain.U[0, 0])
-    ur_plot.append(vdrivetrain.U[1, 0])
-    left_gear_plot.append((vdrivetrain.left_gear is VelocityDrivetrain.HIGH) * 2.0 - 10.0)
-    right_gear_plot.append((vdrivetrain.right_gear is VelocityDrivetrain.HIGH) * 2.0 - 10.0)
-
-    fwd_velocity = (vdrivetrain.X[1, 0] + vdrivetrain.X[0, 0]) / 2
-    turn_velocity = (vdrivetrain.X[1, 0] - vdrivetrain.X[0, 0])
-    if abs(fwd_velocity) < 0.0000001:
-      radius_plot.append(turn_velocity)
-    else:
-      radius_plot.append(turn_velocity / fwd_velocity)
-
-  # TODO(austin):
-  # Shifting compensation.
-
-  # Tighten the turn.
-  # Closed loop drive.
-
-  pylab.plot(t_plot, vl_plot, label='left velocity')
-  pylab.plot(t_plot, vr_plot, label='right velocity')
-  pylab.plot(t_plot, ul_plot, label='left voltage')
-  pylab.plot(t_plot, ur_plot, label='right voltage')
-  pylab.plot(t_plot, radius_plot, label='radius')
-  pylab.plot(t_plot, left_gear_plot, label='left gear high')
-  pylab.plot(t_plot, right_gear_plot, label='right gear high')
-  pylab.legend()
-  pylab.show()
-  return 0
+    polydrivetrain.WritePolyDrivetrain(argv[1:3], argv[3:5], 'y2017',
+                                       drivetrain.kDrivetrain)
 
 if __name__ == '__main__':
   argv = FLAGS(sys.argv)