Updates

MonteCarloAgent.py

@@ -5,50 +5,14 @@ from tqdm import tqdm
 import argparse
 from gymnasium.envs.toy_text.frozen_lake import generate_random_map
 import wandb
+from .Shared import Shared
 
-
-class MonteCarloAgent:
+class MonteCarloAgent(Shared):
     def __init__(
         self,
-        env_name="CliffWalking-v0",
-        gamma=0.99,
-        epsilon=0.1,
-        run_name=None,
         **kwargs,
     ):
-        print("=" * 80)
-        print(f"# MonteCarloAgent - {env_name}")
-        print(f"- epsilon: {epsilon}")
-        print(f"- gamma: {gamma}")
-        print(f"- run_name: {run_name}")
-        self.run_name = run_name
-        self.env_name = env_name
-        self.epsilon, self.gamma = epsilon, gamma
-
-        self.env_kwargs = kwargs
-        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["is_slippery"] = False
-
-        self.env = gym.make(self.env_name, **self.env_kwargs)
-
-        self.n_states, self.n_actions = (
-            self.env.observation_space.n,
-            self.env.action_space.n,
-        )
-        print(f"- n_states: {self.n_states}")
-        print(f"- n_actions: {self.n_actions}")
+        super().__init__(**kwargs)
         self.reset()
 
     def reset(self):
@@ -71,85 +35,6 @@ class MonteCarloAgent:
         print(self.Pi)
         print("=" * 80)
 
-    def choose_action(self, state, epsilon_override=None, greedy=False, **kwargs):
-        # Sample an action from the policy.
-        # The epsilon_override argument allows forcing the use of a new epsilon value than the one previously used during training.
-        # The ability to override was mostly added for testing purposes and for the demo.
-        greedy_action = np.argmax(self.Pi[state])
-
-        if greedy or epsilon_override == 0:
-            return greedy_action
-
-        if epsilon_override is None:
-            return np.random.choice(self.n_actions, p=self.Pi[state])
-
-        return np.random.choice(
-            [greedy_action, np.random.randint(self.n_actions)],
-            p=[1 - epsilon_override, epsilon_override],
-        )
-
-    def generate_episode(self, max_steps=500, render=False, **kwargs):
-        state, _ = self.env.reset()
-        episode_hist, solved, rgb_array = (
-            [],
-            False,
-            self.env.render() if render else None,
-        )
-
-        # Generate an episode following the current policy
-        for _ in range(max_steps):
-            # Sample an action from the policy
-            action = self.choose_action(state, **kwargs)
-            # Take the action and observe the reward and next state
-            next_state, reward, done, _, _ = self.env.step(action)
-
-            if self.env_name == "FrozenLake-v1":
-                if done:
-                    reward = 100 if reward == 1 else -10
-                else:
-                    reward = -1
-
-            # Keeping track of the trajectory
-            episode_hist.append((state, action, reward))
-            yield episode_hist, solved, rgb_array
-
-            # Rendering new frame if needed
-            rgb_array = self.env.render() if render else None
-
-            # For CliffWalking-v0 and Taxi-v3, the episode is solved when it terminates
-            if done and self.env_name in ["CliffWalking-v0", "Taxi-v3"]:
-                solved = True
-                break
-
-            # For FrozenLake-v1, the episode terminates when the agent moves into a hole or reaches the goal
-            # We consider the episode solved when the agent reaches the goal
-            if done and self.env_name == "FrozenLake-v1":
-                if next_state == self.env.nrow * self.env.ncol - 1:
-                    solved = True
-                    break
-                else:
-                    # Instead of terminating the episode when the agent moves into a hole, we reset the environment
-                    # This is to keep consistent with the other environments
-                    done = False
-                    next_state, _ = self.env.reset()
-
-            if solved or done:
-                break
-
-            state = next_state
-
-        rgb_array = self.env.render() if render else None
-        yield episode_hist, solved, rgb_array
-
-    def run_episode(self, max_steps=500, render=False, **kwargs):
-        # Run the generator until the end
-        episode_hist, solved, rgb_array = None, False, None
-        for episode_hist, solved, rgb_array in self.generate_episode(
-            max_steps, render, **kwargs
-        ):
-            pass
-        return episode_hist, solved, rgb_array
-
     def update_first_visit(self, episode_hist):
         G = 0
         # For each step of the episode, in reverse order
@@ -265,37 +150,6 @@ class MonteCarloAgent:
             if log_wandb:
                 wandb.log(stats)
 
-    def test(self, n_test_episodes=100, verbose=True, greedy=True, **kwargs):
-        if verbose:
-            print(f"Testing agent for {n_test_episodes} episodes...")
-        num_successes = 0
-        for e in range(n_test_episodes):
-            _, solved, _ = self.run_episode(greedy=greedy, **kwargs)
-            num_successes += solved
-            if verbose:
-                word = "reached" if solved else "did not reach"
-                emoji = "🏁" if solved else "🚫"
-                print(
-                    f"({e + 1:>{len(str(n_test_episodes))}}/{n_test_episodes}) - Agent {word} the goal {emoji}"
-                )
-
-        success_rate = num_successes / n_test_episodes
-        if verbose:
-            print(
-                f"Agent reached the goal in {num_successes}/{n_test_episodes} episodes ({success_rate * 100:.2f}%)"
-            )
-        return success_rate
-
-    def save_policy(self, fname="policy.npy", save_dir=None):
-        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}")
-        np.save(fname, self.Pi)
-
-    def load_policy(self, fname="policy.npy"):
-        print(f"Loading policy from: {fname}")
-        self.Pi = np.load(fname)
 
     def wandb_log_img(self, episode=None):
         caption_suffix = "Initial" if episode is None else f"After Episode {episode}"

Shared.py

@@ -0,0 +1,163 @@
+import os
+import numpy as np
+import gymnasium as gym
+from tqdm import tqdm
+import argparse
+import wandb
+
+class Shared:
+    
+    def __init__(
+        self,
+        env_name="CliffWalking-v0",
+        gamma=0.99,
+        epsilon=0.1,
+        run_name=None,
+        **kwargs,
+    ):
+        print("=" * 80)
+        print(f"# Init Agent - {env_name}")
+        print(f"- epsilon: {epsilon}")
+        print(f"- gamma: {gamma}")
+        print(f"- run_name: {run_name}")
+        self.run_name = run_name
+        self.env_name = env_name
+        self.epsilon, self.gamma = epsilon, gamma
+
+        self.env_kwargs = kwargs
+        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["is_slippery"] = False
+
+        self.env = gym.make(self.env_name, **self.env_kwargs)
+
+        self.n_states, self.n_actions = (
+            self.env.observation_space.n,
+            self.env.action_space.n,
+        )
+        print(f"- n_states: {self.n_states}")
+        print(f"- n_actions: {self.n_actions}")
+
+    def choose_action(self, state, epsilon_override=None, greedy=False, **kwargs):
+        # Sample an action from the policy.
+        # The epsilon_override argument allows forcing the use of a new epsilon value than the one previously used during training.
+        # The ability to override was mostly added for testing purposes and for the demo.
+        greedy_action = np.argmax(self.Pi[state])
+
+        if greedy or epsilon_override == 0:
+            return greedy_action
+
+        if epsilon_override is None:
+            return np.random.choice(self.n_actions, p=self.Pi[state])
+
+        return np.random.choice(
+            [greedy_action, np.random.randint(self.n_actions)],
+            p=[1 - epsilon_override, epsilon_override],
+        )
+        
+    def get_policy():
+        pass
+    
+    
+    def generate_episode(self, max_steps=500, render=False, **kwargs):
+        state, _ = self.env.reset()
+        episode_hist, solved, rgb_array = (
+            [],
+            False,
+            self.env.render() if render else None,
+        )
+
+        # Generate an episode following the current policy
+        for _ in range(max_steps):
+            # Sample an action from the policy
+            action = self.choose_action(state, **kwargs)
+            # Take the action and observe the reward and next state
+            next_state, reward, done, _, _ = self.env.step(action)
+
+            if self.env_name == "FrozenLake-v1":
+                if done:
+                    reward = 100 if reward == 1 else -10
+                else:
+                    reward = -1
+
+            # Keeping track of the trajectory
+            episode_hist.append((state, action, reward))
+            yield episode_hist, solved, rgb_array
+
+            # Rendering new frame if needed
+            rgb_array = self.env.render() if render else None
+
+            # For CliffWalking-v0 and Taxi-v3, the episode is solved when it terminates
+            if done and self.env_name in ["CliffWalking-v0", "Taxi-v3"]:
+                solved = True
+                break
+
+            # For FrozenLake-v1, the episode terminates when the agent moves into a hole or reaches the goal
+            # We consider the episode solved when the agent reaches the goal
+            if done and self.env_name == "FrozenLake-v1":
+                if next_state == self.env.nrow * self.env.ncol - 1:
+                    solved = True
+                    break
+                else:
+                    # Instead of terminating the episode when the agent moves into a hole, we reset the environment
+                    # This is to keep consistent with the other environments
+                    done = False
+                    next_state, _ = self.env.reset()
+
+            if solved or done:
+                break
+
+            state = next_state
+
+        rgb_array = self.env.render() if render else None
+        yield episode_hist, solved, rgb_array
+
+    def run_episode(self, max_steps=500, render=False, **kwargs):
+        # Run the generator until the end
+        episode_hist, solved, rgb_array = list(self.generate_episode(
+            max_steps, render, **kwargs
+        ))[-1]
+        return episode_hist, solved, rgb_array
+
+    def test(self, n_test_episodes=100, verbose=True, greedy=True, **kwargs):
+        if verbose:
+            print(f"Testing agent for {n_test_episodes} episodes...")
+        num_successes = 0
+        for e in range(n_test_episodes):
+            _, solved, _ = self.run_episode(greedy=greedy, **kwargs)
+            num_successes += solved
+            if verbose:
+                word = "reached" if solved else "did not reach"
+                emoji = "🏁" if solved else "🚫"
+                print(
+                    f"({e + 1:>{len(str(n_test_episodes))}}/{n_test_episodes}) - Agent {word} the goal {emoji}"
+                )
+
+        success_rate = num_successes / n_test_episodes
+        if verbose:
+            print(
+                f"Agent reached the goal in {num_successes}/{n_test_episodes} episodes ({success_rate * 100:.2f}%)"
+            )
+        return success_rate
+    
+    def save_policy(self, fname="policy.npy", save_dir=None):
+        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}")
+        np.save(fname, self.Pi)
+
+    def load_policy(self, fname="policy.npy"):
+        print(f"Loading policy from: {fname}")
+        self.Pi = np.load(fname)