Update AINetworkDefense_Competitive/env.py

AINetworkDefense_Competitive/env.py

@@ -5,14 +5,13 @@ from config import Config
 
 class CyberPettingZooEnv(ParallelEnv):
     metadata = {'render_modes': ['human'], "name": "cyber_network_defense_v2_gnn"}
-
+    
     def __init__(self):
         super().__init__()
         
         self.topology = Config.NETWORK_TOPOLOGY
         self.total_nodes = Config.TOTAL_HOSTS()  
-        self.features_per_node = 2                  
-        
+        self.features_per_node = 2                                  
         self.possible_agents = ["red", "blue"]
         self.agents = self.possible_agents[:]
         
@@ -24,8 +23,8 @@ class CyberPettingZooEnv(ParallelEnv):
         }
         
         self.action_spaces = {
-            'red': gym.spaces.Discrete(self.total_nodes * 4),               
-            'blue': gym.spaces.Discrete(self.total_nodes * 2)           
+            'red': gym.spaces.Discrete(self.total_nodes * 4),                           
+            'blue': gym.spaces.Discrete(self.total_nodes * 2)                   
         }
         
         self.node_id_to_specs = {}
@@ -43,7 +42,7 @@ class CyberPettingZooEnv(ParallelEnv):
                 node_id += 1
 
     def _generate_static_adjacency_matrix(self):
-        """🚀 GNN CORE: Programmatically wires lateral communication pathways between subnets."""
+        """ GNN CORE: Programmatically wires lateral communication pathways between subnets."""
         A = np.zeros((self.total_nodes, self.total_nodes), dtype=np.float32)
         
         # Define cross-subnet routing permissions
@@ -63,7 +62,6 @@ class CyberPettingZooEnv(ParallelEnv):
                 if sub_j in allowed_connections[sub_i]:
                     A[i, j] = 1.0
                     
-        # Add Self-Loops (A + I)
         A_tilde = A + np.eye(self.total_nodes, dtype=np.float32)
         
         # Compute Symmetric Degree Normalization:
@@ -78,6 +76,9 @@ class CyberPettingZooEnv(ParallelEnv):
         self.agents = self.possible_agents[:]
         self.state_matrix = np.zeros((self.total_nodes, self.features_per_node), dtype=np.float32)
         
+        # Initialize step tracker
+        self.current_step = 0
+        
         obs_dict = {
             'node_features': self.state_matrix.copy(),
             'adjacency_matrix': self.normalized_adj.copy()
@@ -96,54 +97,77 @@ class CyberPettingZooEnv(ParallelEnv):
         if not actions:
             self.agents = []
             return {}, {}, {}, {}, {}
-
+            
+        self.current_step += 1
+        
         red_action = actions.get('red', 0)
         blue_action = actions.get('blue', 0)
-
+        
         red_target_node = red_action // 4
         red_action_type = red_action % 4
         blue_target_node = blue_action // 2
         blue_action_type = blue_action % 2
-
+        
         red_target_node = min(max(0, red_target_node), self.total_nodes - 1)
         blue_target_node = min(max(0, blue_target_node), self.total_nodes - 1)
-
+        
         red_target_specs = self.node_id_to_specs[red_target_node]
         red_target_subnet = red_target_specs["subnet"]
-
-        # 1. Defense Execution
+        
+        # Defense Execution
         if blue_action_type == 1:
             self.state_matrix[blue_target_node, 1] = 1.0
-
-        # 2. Attack Execution with Infrastructure Co-Dependencies
+            
+        # Attack Execution with Infrastructure Co-Dependencies
         if red_action_type == 1:
             if self.state_matrix[red_target_node, 1] == 0.0:
                 exploit_probability = red_target_specs["base_vulnerability"]
                 
                 if red_target_subnet == "secure_core":
                     if not self._is_subnet_compromised("active_directory"):
-                        exploit_probability = 0.0  
+                        exploit_probability = 0.0
                 elif red_target_subnet == "active_directory":
                     if self._is_subnet_compromised("dns_services"):
                         exploit_probability = min(1.0, exploit_probability * 2.0)
-                
+                        
                 if np.random.rand() < exploit_probability:
                     self.state_matrix[red_target_node, 0] = 1.0
-
-        # 3. Telemetry Preparation
+                    
+        # Telemetry Preparation
         total_compromised = np.sum(self.state_matrix[:, 0])
         total_patched = np.sum(self.state_matrix[:, 1])
         score_diff = float(total_compromised - total_patched)
         progress = float(total_compromised / self.total_nodes) if self.total_nodes > 0 else 0.0
-
+        
         obs_dict = {
             'node_features': self.state_matrix.copy(),
             'adjacency_matrix': self.normalized_adj.copy()
         }
         observations = {'red': obs_dict, 'blue': obs_dict}
         rewards = {'red': 0.0, 'blue': 0.0}
-        terminations = {'red': False, 'blue': False}
-        truncations = {'red': False, 'blue': False}
+        
+        # DYNAMIC EPISODE TERMINATIONS AND TRUNCATIONS
+        reached_max_steps = self.current_step >= Config.MAX_STEPS_PER_EPISODE
+        
+        if str(Config.MODE).lower() == "cooperative":
+            # Cooperative: Runs the full duration of the step limit
+            is_terminated = False
+            is_truncated = reached_max_steps
+        else:
+            # Competitive: Allow early wins (e.g.Red conquers all, or Blue erases threat)
+            red_won = (total_compromised == self.total_nodes)
+            blue_won = (total_compromised == 0 and self.current_step > 5)
+            
+            is_terminated = red_won or blue_won
+            is_truncated = reached_max_steps
+            
+        terminations = {agent: is_terminated for agent in self.possible_agents}
+        truncations = {agent: is_truncated for agent in self.possible_agents}
+    
         infos = {'score_difference': score_diff, 'team_progress': progress}
-
+        
+        # Cleanup agent lists if active sequence ends
+        if is_terminated or is_truncated:
+            self.agents = []
+            
         return observations, rewards, terminations, truncations, infos