import os import tensorflow as tf from tf_agents.utils import common from tf_agents.trajectories import trajectory # Custom Modular Imports from config import Config from rewards import compute_rewards from utilities import get_system_telemetry class MultiAgentTrainer: def __init__(self, config, tf_env, agent_1, agent_2): self.config = config self.tf_env = tf_env self.agent_1 = agent_1 self.agent_2 = agent_2 log_dir = self.config.LOG_DIR checkpoint_dir = self.config.CHECKPOINT_DIR self.global_step = tf.Variable(0, name="global_step", trainable=False, dtype=tf.int64) self.summary_writer = tf.summary.create_file_writer(log_dir) self.checkpoint = tf.train.Checkpoint( step=self.global_step, agent_1_actor=self.agent_1._actor_net, agent_1_critic=self.agent_1._value_net, agent_2_actor=self.agent_2._actor_net, agent_2_critic=self.agent_2._value_net ) self.checkpoint_manager = tf.train.CheckpointManager( self.checkpoint, directory=checkpoint_dir, max_to_keep=1000 ) if self.checkpoint_manager.latest_checkpoint: self.checkpoint.restore(self.checkpoint_manager.latest_checkpoint) print(f"πŸ”„ Checkpoint restored successfully from: {self.checkpoint_manager.latest_checkpoint}") else: print("πŸ†• No checkpoints found. Initializing brand new weights.") self.era_red_rewards = [] self.era_blue_rewards = [] def _evaluate_and_print_era(self, epoch_idx, current_step): """Analyzes historical trends over the last 1000 episodes, saves to file, and displays report.""" if not self.era_red_rewards or not self.era_blue_rewards: return avg_red = sum(self.era_red_rewards) / len(self.era_red_rewards) avg_blue = sum(self.era_blue_rewards) / len(self.era_blue_rewards) # 1. Terminal Output Presentation print("\n" + "="*60) print(f"πŸ“Š πŸ—ΊοΈ ERA TRACKING REPORT | EPISODES {epoch_idx - 999} - {epoch_idx}") print("="*60) print(f"πŸ”΄ Red Agent Avg Reward: {avg_red:+.2f}") print(f"πŸ”΅ Blue Agent Avg Reward: {avg_blue:+.2f}") print("-"*60) if avg_red > 200: dominance = "πŸ”΄ RED DOMINANT (Attackers are breaching defenses with high margin)" elif avg_blue > 200: dominance = "πŸ”΅ BLUE DOMINANT (Defenders are shutting down and patching threats efficiently)" else: dominance = "βš–οΈ NASH EQUILIBRIUM / STALEMATE (Highly contested stalemate block)" print(f"πŸ‘‘ Tactical Dominance Status: {dominance}") if avg_blue >= 500: security_status = "πŸ”’ VERIFIED SECURE (Exploits mitigated, network completely hardened)" elif avg_blue >= -200: security_status = "⚠️ CONDITIONALLY STABLE (Minor compromises occurring, but controlled)" else: security_status = "🚨 CRITICAL VULNERABILITY BREACH (Active Directory/Core compromised consistently)" print(f"πŸ›‘οΈ Network Security State: {security_status}") print("="*60 + "\n") # 2. Append snapshot data securely to structured flat text asset txt_path = os.path.join(self.config.LOG_DIR, "era_history.txt") os.makedirs(self.config.LOG_DIR, exist_ok=True) with open(txt_path, "a") as f: f.write(f"EPISODE:{epoch_idx}|RED_REWARD:{avg_red:.2f}|BLUE_REWARD:{avg_blue:.2f}|STEP:{current_step}\n") # Clear data structures for next epoch block tracking self.era_red_rewards.clear() self.era_blue_rewards.clear() def train_epoch(self, epoch_idx): """Runs a complete training cycle: collects trajectories and applies gradient updates.""" time_step = self.tf_env.reset() episode_reward_1 = 0.0 episode_reward_2 = 0.0 steps = 0 red_trajectory_buffer = [] blue_trajectory_buffer = [] with self.summary_writer.as_default(): while steps < self.config.MAX_STEPS_PER_EPISODE: steps += 1 # Split nested structures for isolated calculation steps red_time_step = time_step._replace( observation=time_step.observation['red'], reward=time_step.reward['red'] ) blue_time_step = time_step._replace( observation=time_step.observation['blue'], reward=time_step.reward['blue'] ) # Collect structural policy logit info action_step_1 = self.agent_1.collect_policy.action(red_time_step) action_step_2 = self.agent_2.collect_policy.action(blue_time_step) combined_actions = { 'red': action_step_1.action, 'blue': action_step_2.action } next_time_step = self.tf_env.step(combined_actions) #print(f"DEBUG | Red Action: {combined_actions['red']} | Info: {self.tf_env.pyenv.envs[0].get_current_info()}") #print(f"DEBUG | Blue Action: {combined_actions['blue']} | Info: {self.tf_env.pyenv.envs[0].get_current_info()}") # Process Environment Reward Signals env_info = self.tf_env.pyenv.envs[0].get_current_info() r1, r2 = compute_rewards(self.config.MODE, env_info) #print(f"DEBUG | Red Reward: {r1} | Blue Reward: {r2}") if hasattr(r1, 'numpy'): r1 = r1.numpy() if hasattr(r2, 'numpy'): r2 = r2.numpy() episode_reward_1 += float(r1) episode_reward_2 += float(r2) # Stretch shape arrays matching batch environments configuration batch_size = tf.shape(time_step.step_type)[0] broadcast_r1 = tf.fill([batch_size], tf.cast(r1, tf.float32)) broadcast_r2 = tf.fill([batch_size], tf.cast(r2, tf.float32)) next_red_time_step = next_time_step._replace( observation=next_time_step.observation['red'], reward=broadcast_r1 ) next_blue_time_step = next_time_step._replace( observation=next_time_step.observation['blue'], reward=broadcast_r2 ) # Construct Trajectories preserving distribution signatures red_traj = trajectory.Trajectory( step_type=red_time_step.step_type, observation=red_time_step.observation, action=action_step_1.action, policy_info=action_step_1.info, next_step_type=next_red_time_step.step_type, reward=next_red_time_step.reward, discount=next_red_time_step.discount ) blue_traj = trajectory.Trajectory( step_type=blue_time_step.step_type, observation=blue_time_step.observation, action=action_step_2.action, policy_info=action_step_2.info, next_step_type=next_blue_time_step.step_type, reward=next_blue_time_step.reward, discount=next_blue_time_step.discount ) red_trajectory_buffer.append(red_traj) blue_trajectory_buffer.append(blue_traj) if steps % 100 == 0: print(f"⏳ Episode {epoch_idx:04d} | Simulating step: {steps}/{self.config.MAX_STEPS_PER_EPISODE}...", end="\r") time_step = next_time_step if time_step.is_last(): break # Backpropagation Backpass Phase red_experience = tf.nest.map_structure(lambda *args: tf.stack(args, axis=1), *red_trajectory_buffer) blue_experience = tf.nest.map_structure(lambda *args: tf.stack(args, axis=1), *blue_trajectory_buffer) red_loss_info = self.agent_1.train(experience=red_experience) blue_loss_info = self.agent_2.train(experience=blue_experience) self.global_step.assign_add(1) current_step = self.global_step.numpy() # Record metrics into rolling lists self.era_red_rewards.append(episode_reward_1) self.era_blue_rewards.append(episode_reward_2) # TensorBoard metrics execution tf.summary.scalar("Rewards/Red_Agent", float(episode_reward_1), step=current_step) tf.summary.scalar("Rewards/Blue_Agent", float(episode_reward_2), step=current_step) tf.summary.scalar("Losses/Red_Total_Loss", red_loss_info.loss, step=current_step) tf.summary.scalar("Losses/Blue_Total_Loss", blue_loss_info.loss, step=current_step) telemetry = get_system_telemetry() for metric_name, metric_value in telemetry.items(): tf.summary.scalar(metric_name, metric_value, step=current_step) # Trigger Evaluation Checkpoints and Write logs if epoch_idx % 1000 == 0: self._evaluate_and_print_era(epoch_idx, current_step) if epoch_idx % self.config.CHECKPOINT_INTERVAL == 0: save_path = self.checkpoint_manager.save(checkpoint_number=epoch_idx) print(f"πŸ’Ύ Checkpoint safely saved for episode {epoch_idx} at: {save_path}") self.summary_writer.flush()