Merge pull request #6 from ProValarous/NehalBranch

PR #5

Author: ProValarous

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,39 @@
+name: "🕷️ Bug report"
+description: Report errors or unexpected behavior
+labels:
+- Issue-Bug
+- Needs-Triage
+
+body:
+- type: markdown
+  attributes:
+    value: Please make sure to [search for existing issues](https://github.com/ProValarous/Predator-Prey-Gridworld-Environment/issues) before filing a new one!
+
+- type: textarea
+  attributes:
+    label: Steps to reproduce
+    description: We highly suggest including screenshots and a bug report log (System tray > Report bug).
+    placeholder: Having detailed steps helps us reproduce the bug. 
+  validations:
+    required: true
+
+- type: textarea
+  attributes:
+    label: ✔️ Expected Behavior
+    placeholder: What were you expecting?
+  validations:
+    required: false
+
+- type: textarea
+  attributes:
+    label: ❌ Actual Behavior
+    placeholder: What happened instead?
+  validations:
+    required: false
+
+- id: additionalInfo
+  type: textarea
+  attributes:
+    label: Additional Information
+  validations:
+    required: false

src/baselines/CQL/cql_train.py

@@ -18,6 +18,7 @@
 python cql_train_cleaned.py --episodes 20000 --size 8 --alpha 0.25 --gamma 0.95
 
 """
+
 from __future__ import annotations
 
 import argparse
@@ -52,6 +53,7 @@
 
 # -------------------- helpers --------------------
 
+
 def setup_logging(level=logging.INFO):
     logging.basicConfig(
         level=level,
@@ -96,14 +98,21 @@ def make_agents(num_predators: int = 2, num_preys: int = 2) -> List[Agent]:
     for i in range(1, num_preys + 1):
         agents.append(Agent(agent_name=f"prey_{i}", agent_team=i, agent_type="prey"))
     for i in range(1, num_predators + 1):
-        agents.append(Agent(agent_name=f"predator_{i}", agent_team=i, agent_type="predator"))
+        agents.append(
+            Agent(agent_name=f"predator_{i}", agent_team=i, agent_type="predator")
+        )
     return agents
 
 
-def make_env_and_meta(agents: List[Agent], grid_size: int, seed: int) -> Tuple[GridWorldEnv, int, int]:
-    env = GridWorldEnv(agents=agents, render_mode=None, size=grid_size, perc_num_obstacle=10, seed=seed)
+def make_env_and_meta(
+    agents: List[Agent], grid_size: int, seed: int
+) -> Tuple[GridWorldEnv, int, int]:
+    env = GridWorldEnv(
+        agents=agents, render_mode=None, size=grid_size, perc_num_obstacle=10, seed=seed
+    )
     n_cells = grid_size * grid_size
-    # total joint states = n_cells ** n_agents
+
+    # total joint states = n_cells ** n_agents (may be very large)
     n_states = n_cells ** len(agents)
     n_actions = env.action_space.n
     return env, n_states, n_actions
@@ -113,11 +122,15 @@ def estimate_table_bytes(n_states: int, n_joint_actions: int, dtype=np.float32)
     return int(n_states) * int(n_joint_actions) * np.dtype(dtype).itemsize
 
 
-def init_joint_q_table(n_states: int, n_joint_actions: int, max_bytes: int | None = None) -> np.ndarray:
+def init_joint_q_table(
+    n_states: int, n_joint_actions: int, max_bytes: int | None = None
+) -> np.ndarray:
     """Create joint Q table; optionally check memory requirement first."""
     needed = estimate_table_bytes(n_states, n_joint_actions)
     if max_bytes is not None and needed > max_bytes:
-        raise MemoryError(f"Joint Q-table requires {needed/(1024**3):.2f} GiB > allowed {max_bytes/(1024**3):.2f} GiB")
+        raise MemoryError(
+            f"Joint Q-table requires {needed/(1024**3):.2f} GiB > allowed {max_bytes/(1024**3):.2f} GiB"
+        )
     return np.zeros((n_states, n_joint_actions), dtype=np.float32)
 
 
@@ -129,6 +142,7 @@ def save_q_table(path: str, Q: np.ndarray):
 
 # -------------------- training loop --------------------
 
+
 def train(
     episodes: int = 5000,
     max_steps: int = 200,
@@ -157,18 +171,26 @@ def train(
 
     env, n_states, n_actions = make_env_and_meta(agents, grid_size, seed)
 
+    n_agents = len(agent_names)
+
     # joint-action space size
-    n_joint_actions = n_actions ** n_agents
+    n_joint_actions = n_actions**n_agents
 
     # memory safety check (use a conservative default if not provided)
     if max_table_bytes is None:
         # set default max to 8 GiB for safety on typical dev machines
-        max_table_bytes = 16 * 1024 ** 3
+        max_table_bytes = 16 * 1024**3
 
-    LOGGER.info("Allocating joint Q-table: states=%d, joint_actions=%d", n_states, n_joint_actions)
+    LOGGER.info(
+        "Allocating joint Q-table: states=%d, joint_actions=%d",
+        n_states,
+        n_joint_actions,
+    )
     Q = init_joint_q_table(n_states, n_joint_actions, max_bytes=max_table_bytes)
 
-    save_path_Q = os.path.join(os.path.dirname(save_path) or ".", "central_cql_q_table.npz")
+    save_path_Q = os.path.join(
+        os.path.dirname(save_path) or ".", "central_cql_q_table.npz"
+    )
 
     eps = eps_start
 
@@ -200,9 +222,14 @@ def train(
             # select actions by marginalizing the joint-Q over others
             flat_row = Q[s]
             if flat_row.size != n_joint_actions:
-                raise ValueError("Unexpected joint-Q row length: %d != %d" % (flat_row.size, n_joint_actions))
+                raise ValueError(
+                    "Unexpected joint-Q row length: %d != %d"
+                    % (flat_row.size, n_joint_actions)
+                )
 
-            q_tensor = flat_row.reshape(action_shape)  # shape: (n_actions, n_actions, ...)
+            q_tensor = flat_row.reshape(
+                action_shape
+            )  # shape: (n_actions, n_actions, ...)
 
             # compute marginal per-agent action-values by averaging over other axes
             q_vals_per_agent = []
@@ -219,14 +246,18 @@ def train(
                 else:
                     row = np.asarray(q_vals_per_agent[i])
                     best = float(np.max(row))
-                    best_actions = np.flatnonzero(np.isclose(row, best)).astype(int).tolist()
+                    best_actions = (
+                        np.flatnonzero(np.isclose(row, best)).astype(int).tolist()
+                    )
                     a_i = int(rng.choice(best_actions))
                 chosen_actions.append(a_i)
 
             joint_idx = joint_actions_to_index(chosen_actions, n_actions)
 
             # build actions dict for env.step
-            actions = {agents[i].agent_name: int(chosen_actions[i]) for i in range(n_agents)}
+            actions = {
+                agents[i].agent_name: int(chosen_actions[i]) for i in range(n_agents)
+            }
 
             mgp = env.step(actions)
             next_obs, rewards = mgp["obs"], mgp["reward"]
@@ -260,7 +291,12 @@ def train(
             s2 = joint_state_index(next_positions, grid_size)
 
             # CQL update (centralized TD update)
-            td_target = central_r + (gamma * next_pot_sum) - current_pot_sum + gamma * np.max(Q[s2])
+            td_target = (
+                central_r
+                + (gamma * next_pot_sum)
+                - current_pot_sum
+                + gamma * np.max(Q[s2])
+            )
             td_error = td_target - Q[s, joint_idx]
             Q[s, joint_idx] += alpha * td_error
 
@@ -283,17 +319,32 @@ def train(
         writer.add_scalar("episode/captures", captures_this_episode, ep)
 
         for name in agent_names:
-            writer.add_scalar(f"episode/total_reward/{name}", float(total_reward_per_agent[name]), ep)
-            mean_reward_running = float(np.mean(rewards_per_ep[name][-window:])) if rewards_per_ep[name] else 0.0
+            writer.add_scalar(
+                f"episode/total_reward/{name}", float(total_reward_per_agent[name]), ep
+            )
+            mean_reward_running = (
+                float(np.mean(rewards_per_ep[name][-window:]))
+                if rewards_per_ep[name]
+                else 0.0
+            )
             writer.add_scalar(f"mean/{name}/reward", mean_reward_running, ep)
 
-        mean_captures_running = float(np.mean(captures_per_ep[-window:])) if captures_per_ep else 0.0
+        mean_captures_running = (
+            float(np.mean(captures_per_ep[-window:])) if captures_per_ep else 0.0
+        )
         writer.add_scalar("mean/captures", mean_captures_running, ep)
 
         # epsilon decay and logs
         if ep % 100 == 0:
             eps = max(eps_end, eps * eps_decay)
-            avg_per_agent = {name: np.mean(rewards_per_ep[name][-100:]) if rewards_per_ep[name] else 0.0 for name in agent_names}
+            avg_per_agent = {
+                name: (
+                    np.mean(rewards_per_ep[name][-100:])
+                    if rewards_per_ep[name]
+                    else 0.0
+                )
+                for name in agent_names
+            }
             LOGGER.info(
                 "Ep %d | eps=%.3f | averages(last100)=%s | mean captures(last100)=%.2f",
                 ep,
@@ -323,6 +374,7 @@ def train(
 
 # ---------------- CLI ----------------
 
+
 def parse_args():
     p = argparse.ArgumentParser("Train central CQL (tabular)")
     p.add_argument("--episodes", type=int, default=40000)
@@ -333,14 +385,19 @@ def parse_args():
     p.add_argument("--save-path", type=str, default="baselines/CQL/")
     p.add_argument("--predators", type=int, default=2)
     p.add_argument("--preys", type=int, default=2)
-    p.add_argument("--max-table-gb", type=float, default=16.0, help="Max allowed joint-Q memory in GiB before aborting")
+    p.add_argument(
+        "--max-table-gb",
+        type=float,
+        default=16.0,
+        help="Max allowed joint-Q memory in GiB before aborting",
+    )
     return p.parse_args()
 
 
 if __name__ == "__main__":
     setup_logging()
     args = parse_args()
-    max_table_bytes = int(args.max_table_gb * 1024 ** 3) if args.max_table_gb else None
+    max_table_bytes = int(args.max_table_gb * 1024**3) if args.max_table_gb else None
     try:
         train(
             episodes=args.episodes,

src/baselines/CQL/test_iql.py

@@ -40,9 +40,12 @@ def try_load_qs(file_path: str) -> Dict[str, np.ndarray]:
                 name = key[2:] if key.startswith("Q_") else key
                 qs[name] = arr.astype(np.float32)
         if not qs:
-            raise RuntimeError(f"No valid Q-table arrays found in '{file_path}'. Keys: {list(data.files)}")
+            raise RuntimeError(
+                f"No valid Q-table arrays found in '{file_path}'. Keys: {list(data.files)}"
+            )
         return qs
 
+
 def state_index_from_obs(obs: dict, predator: Agent, prey: Agent, size: int) -> int:
     """
     Build state index using predator + prey positions.
@@ -54,12 +57,7 @@ def state_index_from_obs(obs: dict, predator: Agent, prey: Agent, size: int) ->
     pred_x, pred_y = int(pos_pred[0]), int(pos_pred[1])
     prey_x, prey_y = int(pos_prey[0]), int(pos_prey[1])
 
-    return (
-        pred_x * size * size * size
-        + pred_y * size * size
-        + prey_x * size
-        + prey_y
-    )
+    return pred_x * size * size * size + pred_y * size * size + prey_x * size + prey_y
 
 
 def choose_action(agent: Agent, q_table: np.ndarray, s_idx: int) -> int:
@@ -83,7 +81,9 @@ def run_test(
 
     prey, predator = make_agents()
     agents = [prey, predator]
-    env = GridWorldEnv(agents=agents, render_mode="human", size=size, perc_num_obstacle=10)
+    env = GridWorldEnv(
+        agents=agents, render_mode="human", size=size, perc_num_obstacle=10
+    )
 
     try:
         for ep in range(1, episodes + 1):
@@ -112,7 +112,9 @@ def run_test(
                                 q_table = qs[k]
                                 break
                     if q_table is None:
-                        raise RuntimeError(f"No Q-table for agent '{ag.agent_name}'. Keys: {list(qs.keys())}")
+                        raise RuntimeError(
+                            f"No Q-table for agent '{ag.agent_name}'. Keys: {list(qs.keys())}"
+                        )
 
                     a = choose_action(ag, q_table, s_idx)
                     actions[ag.agent_name] = a
@@ -145,4 +147,10 @@ def parse_args() -> argparse.Namespace:
 if __name__ == "__main__":
     setup_logging()
     args = parse_args()
-    run_test(q_file=args.file, size=args.size, episodes=args.episodes, max_steps=args.max_steps, pause=args.pause)
+    run_test(
+        q_file=args.file,
+        size=args.size,
+        episodes=args.episodes,
+        max_steps=args.max_steps,
+        pause=args.pause,
+    )