add tdl-small policy in evaluate

Quentin18 · Oct 24, 2023 · 9eb4cf0 · 9eb4cf0
1 parent 2b21b1a
commit 9eb4cf0
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/scripts/evaluate.py b/scripts/evaluate.py
@@ -9,6 +9,7 @@
     NTupleNetworkBasePolicy,
     NTupleNetworkQLearningPolicy,
     NTupleNetworkTDPolicy,
+    NTupleNetworkTDPolicySmall,
 )
 
 
@@ -20,7 +21,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--algo",
         help="RL Algorithm",
-        choices=["ql", "tdl"],
+        choices=["ql", "tdl", "tdl-small"],
     )
     parser.add_argument(
         "--env",
@@ -59,6 +60,7 @@ def make_policy(algo: str, trained_agent: str) -> NTupleNetworkBasePolicy:
     algo_policy_map = {
         "ql": NTupleNetworkQLearningPolicy,
         "tdl": NTupleNetworkTDPolicy,
+        "tdl-small": NTupleNetworkTDPolicySmall,
     }
     policy = algo_policy_map[algo]
     return policy.load(trained_agent)