Fix agent tracking to use possible_agents instead of agents attribute

train_all_phases.py

@@ -3,12 +3,6 @@
 Full training pipeline: Phase 1 -> Phase 2 -> Phase 3
 TIL-26-AE Bomberman Agent Training
 
-Run with:
-  TOTAL_TIMESTEPS=500_000:500_000:1_000_000 \
-  HUB_MODEL_ID=E-Rong/til-26-ae-agent \
-  TRACKIO_PROJECT=til-26-ae \
-  python train_all_phases.py
-
 References:
 - Pommerman multi-agent RL: arxiv:2407.00662
 - MAPPO best practices: arxiv:2103.01955
@@ -101,6 +95,12 @@ class BombermanSingleAgentEnv(gym.Env):
             shape=(self._obs_size,), dtype=np.float32,
         )
 
+    def _get_agents(self):
+        """Get list of currently active agents from obs_dict."""
+        if self._last_obs_dict is not None:
+            return list(self._last_obs_dict.keys())
+        return self._parallel_env.possible_agents
+
     def reset(self, seed=None, options=None):
         if seed is not None:
             self._episode_seed = seed
@@ -109,18 +109,18 @@ class BombermanSingleAgentEnv(gym.Env):
             self._episode_count += 1
 
         obs_dict, info_dict = self._parallel_env.reset(seed=self._episode_seed, options=options)
-        self._store_action_mask(obs_dict[self.agent_id])
         self._last_obs_dict = obs_dict
+        self._store_action_mask(obs_dict[self.agent_id])
         return self._flatten_obs(obs_dict[self.agent_id]), {}
 
     def step(self, action):
         actions = {}
-        for agent_id in self._parallel_env.agents:
+        for agent_id in self._get_agents():
             if agent_id == self.agent_id:
                 actions[agent_id] = action
             else:
                 mask = (
-                    self._last_obs_dict[agent_id]["action_mask"]
+                    self._last_obs_dict[agent_id].get("action_mask")
                     if self._last_obs_dict and agent_id in self._last_obs_dict
                     else np.ones(6, dtype=np.int8)
                 )
@@ -374,6 +374,11 @@ class CurriculumEnv(gym.Env):
             shape=(self._obs_size,), dtype=np.float32,
         )
 
+    def _get_agents(self):
+        if self._last_obs_dict is not None:
+            return list(self._last_obs_dict.keys())
+        return self._parallel_env.possible_agents
+
     def _init_opponents(self):
         for i in range(1, self.cfg.env.num_teams):
             opp_id = f"agent_{i}"
@@ -422,14 +427,13 @@ class CurriculumEnv(gym.Env):
         obs_dict, info_dict = self._parallel_env.reset(
             seed=self._episode_seed, options=options
         )
-
-        self._store_action_mask(obs_dict[self.agent_id])
         self._last_obs_dict = obs_dict
+        self._store_action_mask(obs_dict[self.agent_id])
         return self._flatten_obs(obs_dict[self.agent_id]), {}
 
     def step(self, action):
         actions = {}
-        for agent_id in self._parallel_env.agents:
+        for agent_id in self._get_agents():
             if agent_id == self.agent_id:
                 actions[agent_id] = action
             else: