Updates

MCAgent.py

@@ -4,8 +4,8 @@ from Shared import Shared
 import wandb
 from Shared import Shared
 
+
 class MCAgent(Shared):
-    
     def __init__(self, /, **kwargs):
         super().__init__(**kwargs)
         self.reset()
@@ -126,12 +126,6 @@ class MCAgent(Shared):
             # Test the agent every test_every episodes with the greedy policy (by default)
             if e % test_every == 0:
                 test_success_rate = self.test(verbose=False, **kwargs)
-                if save_best and test_success_rate > 0.9:
-                    if self.run_name is None:
-                        print(f"Warning: run_name is None, not saving best policy")
-                    else:
-                        self.save_policy(self.run_name, save_best_dir)
-
                 if log_wandb:
                     self.wandb_log_img(episode=e)
 
@@ -145,6 +139,16 @@ class MCAgent(Shared):
             if log_wandb:
                 wandb.log(stats)
 
+            if test_running_success_rate > 0.999:
+                print(
+                    f"CONVERGED: test success rate running avg reached 100% after {e} episodes."
+                )
+                if save_best:
+                    if self.run_name is None:
+                        print("WARNING: run_name is None, not saving best policy.")
+                    else:
+                        self.save_policy(self.run_name, save_best_dir)
+                break
 
     def wandb_log_img(self, episode=None):
         caption_suffix = "Initial" if episode is None else f"After Episode {episode}"
@@ -160,4 +164,3 @@ class MCAgent(Shared):
                 ),
             }
         )
-

README.md

@@ -75,10 +75,10 @@ Parameter testing results:
 
 ```bash
 # Training: Policy will be saved as a `.npy` file.
-python3 MonteCarloAgent.py --train
+python3 run.py --agent "MCAgent" --train
 
 # Testing: Use the `--test` flag with the path to the policy file.
-python3 MonteCarloAgent.py --test policy_mc_CliffWalking-v0_e2000_s500_g0.99_e0.1.npy --render_mode human
+python3 run.py --agent "MCAgent" --test "policies/MCAgent_CliffWalking-v0_e2500_s200_g1.0_e0.4_first_visit.npy" --render_mode human
 ```
 
 **MC Usage**

Shared.py

@@ -1,6 +1,7 @@
 import os
 import numpy as np
 import gymnasium as gym
+from gymnasium.envs.toy_text.frozen_lake import generate_random_map
 
 
 class Shared:
@@ -11,6 +12,7 @@ class Shared:
         gamma=0.99,
         epsilon=0.1,
         run_name=None,
+        frozenlake_size=8,
         **kwargs,
     ):
         print("=" * 80)
@@ -27,16 +29,17 @@ class Shared:
         if self.env_name == "FrozenLake-v1":
             # Can use defaults by defining map_name (4x4 or 8x8) or custom map by defining desc
             # self.env_kwargs["map_name"] = "8x8"
-            self.env_kwargs["desc"] = [
-                "SFFFFFFF",
-                "FFFFFFFH",
-                "FFFHFFFF",
-                "FFFFFHFF",
-                "FFFHFFFF",
-                "FHHFFFHF",
-                "FHFFHFHF",
-                "FFFHFFFG",
-            ]
+            # self.env_kwargs["desc"] = [
+            #     "SFFFFFFF",
+            #     "FFFFFFFH",
+            #     "FFFHFFFF",
+            #     "FFFFFHFF",
+            #     "FFFHFFFF",
+            #     "FHHFFFHF",
+            #     "FHFFHFHF",
+            #     "FFFHFFFG",
+            # ]
+            self.env_kwargs["desc"] = generate_random_map(size=frozenlake_size)
             self.env_kwargs["is_slippery"] = False
 
         self.env = gym.make(self.env_name, **self.env_kwargs)

run.py

@@ -3,6 +3,7 @@ import wandb
 
 from agents import AGENTS_MAP
 
+
 def main():
     parser = argparse.ArgumentParser()
     ### Train/Test parameters
@@ -72,7 +73,7 @@ def main():
         choices=AGENTS_MAP.keys(),
         help=f"The agent to use. One of: {AGENTS_MAP.keys()}",
     )
-        
+
     parser.add_argument(
         "--gamma",
         type=float,
@@ -94,7 +95,7 @@ def main():
         choices=["CliffWalking-v0", "FrozenLake-v1", "Taxi-v3"],
         help="The Gymnasium environment to use. (default: CliffWalking-v0)",
     )
-    
+
     parser.add_argument(
         "--render_mode",
         type=str,