diff --git a/aos/input/drivetrain_input.cc b/aos/input/drivetrain_input.cc
index 631f97d..8c0fa6b 100644
--- a/aos/input/drivetrain_input.cc
+++ b/aos/input/drivetrain_input.cc
@@ -37,6 +37,14 @@
     robot_velocity_ = drivetrain_queue.status->robot_speed;
   }
 
+  // If we have a vision align function, and it is in control, don't run the
+  // normal driving code.
+  if (vision_align_fn_) {
+    if (vision_align_fn_(data)) {
+      return;
+    }
+  }
+
   bool is_control_loop_driving = false;
   bool is_line_following = false;