fix load tdl-small policy

Quentin18 · Oct 24, 2023 · d09c78c · d09c78c
1 parent 9eb4cf0
commit d09c78c
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/gymnasium_2048/agents/ntuple/policy.py b/src/gymnasium_2048/agents/ntuple/policy.py
@@ -249,3 +249,9 @@ def _get_tuples(state: np.ndarray) -> Sequence[Sequence[int]]:
             *get_all_straight_3_tuples(state=state),
             *get_all_corners_3_tuples(state=state),
         ]
+
+    @classmethod
+    def load(cls, path: str | bytes | os.PathLike) -> NTupleNetworkBasePolicy:
+        policy = NTupleNetworkTDPolicySmall()
+        policy.net = NTupleNetwork.load(path=path)
+        return policy