diff --git a/y2018/control_loops/python/drivetrain.py b/y2018/control_loops/python/drivetrain.py
index 0e5ed43..ccc8a20 100644
--- a/y2018/control_loops/python/drivetrain.py
+++ b/y2018/control_loops/python/drivetrain.py
@@ -41,11 +41,11 @@
 
     if FLAGS.plot:
         drivetrain.PlotDrivetrainMotions(kDrivetrain)
-    elif len(argv) != 5:
-        print("Expected .h file name and .cc file name")
+    elif len(argv) != 7:
+        print("Expected .h, .cc, and .json filenames")
     else:
         # Write the generated constants out to a file.
-        drivetrain.WriteDrivetrain(argv[1:3], argv[3:5], 'y2018', kDrivetrain)
+        drivetrain.WriteDrivetrain(argv[1:4], argv[4:7], 'y2018', kDrivetrain)
 
 
 if __name__ == '__main__':
diff --git a/y2018/control_loops/python/polydrivetrain.py b/y2018/control_loops/python/polydrivetrain.py
index f406aa2..45ed5dc 100644
--- a/y2018/control_loops/python/polydrivetrain.py
+++ b/y2018/control_loops/python/polydrivetrain.py
@@ -20,10 +20,10 @@
 def main(argv):
     if FLAGS.plot:
         polydrivetrain.PlotPolyDrivetrainMotions(drivetrain.kDrivetrain)
-    elif len(argv) != 7:
-        glog.fatal('Expected .h file name and .cc file name')
+    elif len(argv) != 10:
+        glog.fatal('Expected .h, .cc, and .json filenames')
     else:
-        polydrivetrain.WritePolyDrivetrain(argv[1:3], argv[3:5], argv[5:7],
+        polydrivetrain.WritePolyDrivetrain(argv[1:4], argv[4:7], argv[7:10],
                                            'y2018', drivetrain.kDrivetrain)