Add named command line arguments and optional arguments

main.py

@@ -1,12 +1,36 @@
-import sys
-import os
+#!/usr/bin/env python3
+
+import argparse
 import gym
+import os
+import sys
 import time
 import matplotlib.pyplot as plt
 from a3c.discrete_A3C import train, evaluate, evaluate_checkpoints
 from wordle_env.wordle import WordleEnvBase
 
 
+def training_mode(args, env, model_checkpoint_dir):
+    max_ep = args.games
+    start_time = time.time()
+    if args.model_name:
+        pretrained_model_path = os.path.join(
+            model_checkpoint_dir, args.model_name)
+        global_ep, win_ep, gnet, res = train(
+            env, max_ep, model_checkpoint_dir, pretrained_model_path)
+    else:
+        global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir)
+    print("--- %.0f seconds ---" % (time.time() - start_time))
+    print_results(global_ep, win_ep, res)
+    evaluate(gnet, env)
+
+
+def evaluation_mode(args, env, model_checkpoint_dir):
+    print("Evaluation mode")
+    results = evaluate_checkpoints(model_checkpoint_dir, env)
+    print(results)
+
+
 def print_results(global_ep, win_ep, res):
     print("Jugadas:", global_ep.value)
     print("Ganadas:", win_ep.value)
@@ -17,23 +41,29 @@ def print_results(global_ep, win_ep, res):
 
 
 if __name__ == "__main__":
-    max_ep = int(sys.argv[1]) if len(sys.argv) > 1 else 100000
-    env_id = sys.argv[2] if len(sys.argv) > 2 else 'WordleEnv100FullAction-v0'
-    evaluation = True if len(sys.argv) > 3 and sys.argv[3] == 'evaluation' else False
-    pretrained = True if len(sys.argv) > 3 and sys.argv[3] == 'pretrained' else False
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "enviroment", help="Enviroment (type of wordle game) used for training, example: WordleEnvFull-v0")
+    parser.add_argument(
+        "--models_dir", help="Directory where models are saved (default=checkpoints)", default='checkpoints')
+    subparsers = parser.add_subparsers(help='sub-command help')
+
+    parser_train = subparsers.add_parser(
+        'train', help='Train a model from scratch or train from pretrained model')
+    parser_train.add_argument(
+        "--games", "-g", help="Number of games to train", type=int, required=True)
+    parser_train.add_argument(
+        "--model_name", "-n", help="If want to train from a pretrained model, the name of the pretrained model file")
+    parser_train.add_argument(
+        "--gamma", help="Gamma hyperparameter value", type=float, default=0.)
+    parser_train.set_defaults(func=training_mode)
+
+    parser_eval = subparsers.add_parser(
+        'eval', help='Evaluate saved models for the enviroment')
+    parser_eval.set_defaults(func=evaluation_mode)
+
+    args = parser.parse_args()
+    env_id = args.enviroment
     env = gym.make(env_id)
-    model_checkpoint_dir = os.path.join('checkpoints', env.unwrapped.spec.id)
-    if not evaluation:
-        start_time = time.time()
-        if pretrained:
-            pretrained_model_path = os.path.join(model_checkpoint_dir, sys.argv[4]) if len(sys.argv) > 4 else ''
-            global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir, pretrained_model_path)
-        else:
-            global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir)
-        print("--- %.0f seconds ---" % (time.time() - start_time))
-        print_results(global_ep, win_ep, res)
-        evaluate(gnet, env)
-    else:
-        print("Evaluation mode")
-        results = evaluate_checkpoints(model_checkpoint_dir, env)
-        print(results)
+    model_checkpoint_dir = os.path.join(args.models_dir, env.unwrapped.spec.id)
+    args.func(args, env, model_checkpoint_dir)