AINetworkDefense_Competitive/trainer.py

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()