From 2bdf3c788fabca5ba71f0bdb401138549850fe21 Mon Sep 17 00:00:00 2001
From: Quentin18 <quentindeschamps18@gmail.com>
Date: Tue, 24 Oct 2023 20:57:33 +0200
Subject: [PATCH] fix tdl-small in enjoy

---
 scripts/enjoy.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/scripts/enjoy.py b/scripts/enjoy.py
index b4d56ba..9f7d95a 100644
--- a/scripts/enjoy.py
+++ b/scripts/enjoy.py
@@ -9,6 +9,7 @@
     NTupleNetworkBasePolicy,
     NTupleNetworkQLearningPolicy,
     NTupleNetworkTDPolicy,
+    NTupleNetworkTDPolicySmall,
 )
 
 
@@ -21,7 +22,7 @@ def parse_args() -> argparse.Namespace:
         "--algo",
         default="tdl",
         help="RL Algorithm",
-        choices=["ql", "tdl"],
+        choices=["ql", "tdl", "tdl-small"],
     )
     parser.add_argument(
         "--env",
@@ -62,6 +63,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)