Updates

MCAgent.py

@@ -7,7 +7,7 @@ from Shared import Shared
 
 class MCAgent(Shared):
     def __init__(self, /, **kwargs):
-        super().__init__(**kwargs)
+        super().__init__(run_name=self.__class__.__name__, **kwargs)
         self.reset()
 
     def reset(self):
@@ -79,6 +79,7 @@ class MCAgent(Shared):
         **kwargs,
     ):
         print(f"Training agent for {n_train_episodes} episodes...")
+        self.run_name = f"{self.run_name}_{update_type}"
 
         (
             train_running_success_rate,
@@ -140,7 +141,7 @@ class MCAgent(Shared):
             if log_wandb:
                 wandb.log(stats)
 
-            if test_running_success_rate > 0.999:
+            if test_running_success_rate > 0.99999:
                 if save_best:
                     if self.run_name is None:
                         print("WARNING: run_name is None, not saving best policy.")

Shared.py

@@ -12,19 +12,22 @@ class Shared:
         gamma=0.99,
         epsilon=0.1,
         run_name=None,
-        frozenlake_size=8,
+        seed=None,
         **kwargs,
     ):
         print("=" * 80)
         print(f"# Init Agent - {env}")
-        print(f"- epsilon: {epsilon}")
-        print(f"- gamma: {gamma}")
-        print(f"- run_name: {run_name}")
-        self.run_name = run_name
+
         self.env_name = env
-        self.epsilon, self.gamma = epsilon, gamma
+        self.epsilon, self.gamma = float(epsilon), float(gamma)
+        print(f"- epsilon: {self.epsilon}")
+        print(f"- gamma: {self.gamma}")
         self.epsilon_override = None
 
+        self.run_name = f"{run_name}_" if run_name is not None else ""
+        self.run_name += f"{env}_gamma:{gamma}_epsilon:{epsilon}"
+        print(f"- run_name: {run_name}")
+
         self.env_kwargs = {k: v for k, v in kwargs.items() if k in ["render_mode"]}
         if self.env_name == "FrozenLake-v1":
             # Can use defaults by defining map_name (4x4 or 8x8) or custom map by defining desc
@@ -39,7 +42,15 @@ class Shared:
             #     "FHFFHFHF",
             #     "FFFHFFFG",
             # ]
-            self.env_kwargs["desc"] = generate_random_map(size=frozenlake_size)
+            size = int(kwargs.get("size", 8))
+            print(f"- size: {size}")
+            self.run_name += f"_size:{size}"
+
+            seed = int(seed) if seed is not None else np.random.randint(0, 100000)
+            print(f"- seed: {seed}")
+            self.run_name += f"_seed:{seed}"
+
+            self.env_kwargs["desc"] = generate_random_map(size=size, seed=seed)
             self.env_kwargs["is_slippery"] = False
 
         self.env = gym.make(self.env_name, **self.env_kwargs)
@@ -150,13 +161,21 @@ class Shared:
             )
         return success_rate
 
-    def save_policy(self, fname="policy.npy", save_dir=None):
+    def save_policy(self, fname=None, save_dir=None):
+        if fname is None and self.run_name is None:
+            raise ValueError("Must provide a filename or a run name to save the policy")
+        elif fname is None:
+            fname = self.run_name
+
         if save_dir is not None:
             os.makedirs(save_dir, exist_ok=True)
             fname = os.path.join(save_dir, fname)
-        print(f"Saving policy to: {fname}")
+
+        print(f"Saving policy to: '{fname}'")
         np.save(fname, self.Pi)
 
     def load_policy(self, fname="policy.npy"):
-        print(f"Loading policy from: {fname}")
+        print(f"Loading policy from: '{fname}'")
+        if not fname.endswith(".npy"):
+            fname += ".npy"
         self.Pi = np.load(fname)

agents.py

@@ -2,7 +2,13 @@
 from MCAgent import MCAgent
 from DPAgent import DPAgent
 
-AGENTS_MAP = {
-    "MCAgent": MCAgent,
-    "DPAgent": DPAgent
-}
+AGENTS_MAP = {"MCAgent": MCAgent, "DPAgent": DPAgent}
+
+
+def load_agent(agent_name, **kwargs):
+    if agent_name not in AGENTS_MAP:
+        raise ValueError(
+            f"ERROR: Agent '{agent_name}' not valid. Must be one of: {AGENTS_MAP.keys()}"
+        )
+
+    return AGENTS_MAP[agent_name](**kwargs)

demo.py

@@ -1,5 +1,6 @@
 import os
 import time
+import warnings
 import numpy as np
 import gradio as gr
 
@@ -23,6 +24,7 @@ try:
     all_policies = [
         file for file in os.listdir(policies_folder) if file.endswith(".npy")
     ]
+    all_policies.sort()
 except FileNotFoundError:
     print("ERROR: No policies folder found!")
     all_policies = []
@@ -70,8 +72,10 @@ def reset(state, policy_fname):
     state.live_render_fps = default_render_fps
     state.live_epsilon = default_epsilon
     state.live_steps_forward = None
-    return gr.update(value=pause_val_map_inv[not state.live_paused]), gr.update(
-        interactive=state.live_paused
+    return (
+        state,
+        gr.update(value=pause_val_map_inv[not state.live_paused]),
+        gr.update(interactive=state.live_paused),
     )
 
 
@@ -135,15 +139,32 @@ def run(
     policy_path = os.path.join(policies_folder, policy_fname)
     props = policy_fname.split("_")
 
-    if len(props) < 2:
+    try:
+        agent_key, env_key = props[0], props[1]
+        agent_args = {}
+        for prop in props[2:]:
+            props_split = prop.split(":")
+            if len(props_split) == 2:
+                agent_args[props_split[0]] = props_split[1]
+            else:
+                warnings.warn(
+                    f"Skipping property {prop} as it does not have the format 'key:value'.",
+                    UserWarning,
+                )
+    except IndexError:
         yield localstate, None, None, None, None, None, None, None, None, None, None, "🚫 Please select a valid policy file."
         return
 
-    agent_type, env_name = props[0], props[1]
-
-    agent = AGENTS_MAP[agent_type](env=env_name, render_mode="rgb_array")
+    agent_args.update(
+        {
+            "env": env_key,
+            "render_mode": "rgb_array",
+        }
+    )
+    print("agent_args:", agent_args)
+    agent = AGENTS_MAP[agent_key](**agent_args)
     agent.load_policy(policy_path)
-    env_action_map = action_map.get(env_name)
+    env_action_map = action_map.get(env_key)
 
     solved, frame_env, frame_policy = None, None, None
     episode, step, state, action, reward, last_reward = (
@@ -255,7 +276,7 @@ def run(
                 f"Episode: {ep_str(episode + 1)} - step: {step_str(step)} - state: {state} - action: {action} - reward: {reward} (epsilon: {localstate.live_epsilon:.2f}) (frame time: {1 / localstate.live_render_fps:.2f}s)"
             )
 
-            yield localstate, agent_type, env_name, frame_env, frame_policy, ep_str(
+            yield localstate, agent_key, env_key, frame_env, frame_policy, ep_str(
                 episode + 1
             ), ep_str(episodes_solved), step_str(
                 step
@@ -272,7 +293,7 @@ def run(
                 time.sleep(1 / localstate.live_render_fps)
 
             while localstate.live_paused and localstate.live_steps_forward is None:
-                yield localstate, agent_type, env_name, frame_env, frame_policy, ep_str(
+                yield localstate, agent_key, env_key, frame_env, frame_policy, ep_str(
                     episode + 1
                 ), ep_str(episodes_solved), step_str(
                     step
@@ -285,8 +306,8 @@ def run(
                 localstate.should_reset = False
                 yield (
                     localstate,
-                    agent_type,
-                    env_name,
+                    agent_key,
+                    env_key,
                     np.ones((frame_env_h, frame_env_w, 3)),
                     np.ones((frame_policy_h, frame_policy_res)),
                     ep_str(episode + 1),
@@ -305,7 +326,7 @@ def run(
         time.sleep(0.25)
 
     localstate.current_policy = None
-    yield localstate, agent_type, env_name, frame_env, frame_policy, ep_str(
+    yield localstate, agent_key, env_key, frame_env, frame_policy, ep_str(
         episode + 1
     ), ep_str(episodes_solved), step_str(step), state, action, reward, "Done!"
 

requirements.txt

@@ -1,8 +1,7 @@
 gradio==3.27.0
 gymnasium>=0.28.1
-numpy==1.21.5
+numpy>=1.23
 opencv_python_headless==4.6.0.66
-pip==22.0.2
 scipy==1.8.0
 tabulate==0.9.0
 tqdm==4.64.1

run.py

@@ -1,7 +1,7 @@
 import argparse
 import wandb
 
-from agents import AGENTS_MAP
+from agents import AGENTS_MAP, load_agent
 
 
 def main():
@@ -96,6 +96,20 @@ def main():
         help="The Gymnasium environment to use. (default: CliffWalking-v0)",
     )
 
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="The seed to use when generating the FrozenLake environment. If not provided, a random seed is used. (default: None)",
+    )
+
+    parser.add_argument(
+        "--size",
+        type=int,
+        default=8,
+        help="The size to use when generating the FrozenLake environment. (default: 8)",
+    )
+
     parser.add_argument(
         "--render_mode",
         type=str,
@@ -123,13 +137,12 @@ def main():
 
     args = parser.parse_args()
     print(vars(args))
-    agent = AGENTS_MAP[args.agent](**dict(args._get_kwargs()))
 
-    run_name = f"{agent.__class__.__name__}_{args.env}_e{args.n_train_episodes}_s{args.max_steps}_g{args.gamma}_e{args.epsilon}_{args.update_type}"
-    if args.wandb_run_name_suffix is not None:
-        run_name += f"+{args.wandb_run_name_suffix}"
+    agent = load_agent(args.agent, **dict(args._get_kwargs()))
 
-    agent.run_name = run_name
+    agent.run_name += f"_e{args.n_train_episodes}_s{args.max_steps}"
+    if args.wandb_run_name_suffix is not None:
+        agent.run_name += f"+{args.wandb_run_name_suffix}"
 
     try:
         if args.train:
@@ -137,7 +150,7 @@ def main():
             if args.wandb_project is not None:
                 wandb.init(
                     project=args.wandb_project,
-                    name=run_name,
+                    name=agent.run_name,
                     group=args.agent,
                     job_type=args.wandb_job_type,
                     config=dict(args._get_kwargs()),
@@ -154,13 +167,8 @@ def main():
                 save_best_dir=args.save_dir,
             )
             if not args.no_save:
-                agent.save_policy(
-                    fname=f"{run_name}.npy",
-                    save_dir=args.save_dir,
-                )
+                agent.save_policy(save_dir=args.save_dir)
         elif args.test is not None:
-            if not args.test.endswith(".npy"):
-                args.test += ".npy"
             agent.load_policy(args.test)
             agent.test(
                 n_test_episodes=args.n_test_episodes,