Change how the library of conceptual maping is called

Author: eldraco

agents/attackers/conceptual_q_learning/conceptual_q_agent.py

@@ -11,15 +11,14 @@
 import logging
 import subprocess
 import time
-import mlflow
 import wandb
 
 from os import path, makedirs
 # with the path fixed, we can import now
 from AIDojoCoordinator.game_components import Action, Observation, GameState, AgentStatus, ActionType
 from NetSecGameAgents.agents.base_agent import BaseAgent
 from NetSecGameAgents.agents.agent_utils import state_as_ordered_string, convert_ips_to_concepts, convert_concepts_to_actions, generate_valid_actions_concepts
-from concept_mapping_logger import ConceptMappingLogger
+from NetSecGameAgents.utils.concept_mapping_logger import ConceptMappingLogger
 
 class QAgent(BaseAgent):
 
@@ -268,9 +267,8 @@ def play_game(self, concept_observation, episode_num, testing=False):
     parser.add_argument("--models_dir", help="Folder to store models", default=path.join(path.dirname(path.abspath(__file__)), "models"))
     parser.add_argument("--previous_model", help="Load the previous model. If training, it will start from here. If testing, will use to test.", type=str)
     parser.add_argument("--testing", help="Test the agent. No train.", default=False, type=bool)
-    parser.add_argument("--experiment_id", help="Id of the experiment to record into Mlflow and/or Wandb.", default='', type=str)
+    parser.add_argument("--experiment_id", help="Id of the experiment to record into Wandb.", default='', type=str)
     # Logging platform selection arguments
-    parser.add_argument("--use_mlflow", help="Enable MLflow logging.", action='store_true')
     parser.add_argument("--disable_wandb", help="Disable Wandb logging (enabled by default).", action='store_true')
     # Wandb-specific arguments
     parser.add_argument("--wandb_project", help="Wandb project name.", default="netsec-conceptual-qlearning", type=str)
@@ -298,14 +296,9 @@ def play_game(self, concept_observation, episode_num, testing=False):
         agent.enable_enhanced_logging(verbose=True)
 
     # Set logging platform usage based on flags
-    # MLflow is disabled by default, wandb is enabled by default
+    # wandb is enabled by default
     args.use_wandb = not args.disable_wandb
 
-    # Validate logging platform selection
-    if not args.use_mlflow and not args.use_wandb:
-        agent._logger.warning("No logging platform selected. Enabling Wandb by default.")
-        args.use_wandb = True
-    
     # Early stop flag. Used to stop the training if the win rate goes over a threshold.
     early_stop = False
 
@@ -321,17 +314,13 @@ def play_game(self, concept_observation, episode_num, testing=False):
             print(message)
 
 
-    # Set mlflow for local tracking
+    # Set W&B for tracking
     if not args.testing:
         # Experiment name        
         experiment_name = "Training and Eval of Conceptual Q-learning Agent"
-        if args.use_mlflow:
-            mlflow.set_experiment(experiment_name)
     elif args.testing:
         # Experiment name        
         experiment_name = "Testing of Conceptual Q-learning Agent against defender agent"
-        if args.use_mlflow:
-            mlflow.set_experiment(experiment_name)
 
     # This code runs for both training and testing. 
     # How ti works:
@@ -340,7 +329,7 @@ def play_game(self, concept_observation, episode_num, testing=False):
     # - Test for --test_for episodes
     # - When each episode finishes you have: steps played, return, win/lose.
     # - For each episode, store all values and compute the avg and std of each of them
-    # - Every --test_for episodes and at the end of the testing, report results in log file, mlflow and console.
+    # - Every --test_for episodes and at the end of the testing, report results in log file, remote log and console.
 
     # Register the agent
     # Obsservation is in IPs
@@ -363,12 +352,6 @@ def play_game(self, concept_observation, episode_num, testing=False):
                 mode=args.wandb_mode
             )
 
-        # Start MLflow run if enabled
-        if args.use_mlflow:
-            mlflow_run = mlflow.start_run(run_name=experiment_name + f'. ID {args.experiment_id}')
-        else:
-            mlflow_run = None
-        
         try:
             # To keep statistics of each episode
             wins = 0
@@ -388,29 +371,6 @@ def play_game(self, concept_observation, episode_num, testing=False):
             agents_git_result = subprocess.run(agents_command, shell=True, capture_output=True, text=True).stdout
             agent._logger.info(f'Using commits. NetSecEnv: {netsecenv_git_result}. Agents: {agents_git_result}')
 
-            # Log configuration to MLflow if enabled
-            if args.use_mlflow:
-                mlflow.set_tag("experiment_name", experiment_name)
-                mlflow.set_tag("notes", "This is a training and evaluation of the conceptual Q-learning agent.")
-                if args.previous_model:
-                    mlflow.set_tag("Previous q-learning model loaded", str(args.previous_model))
-                mlflow.log_param("alpha", args.alpha)
-                mlflow.log_param("epsilon_start", args.epsilon_start)
-                mlflow.log_param("epsilon_end", args.epsilon_end)
-                mlflow.log_param("epsilon_max_episodes", args.epsilon_max_episodes)
-                mlflow.log_param("gamma", args.gamma)
-                mlflow.log_param("Episodes", args.episodes)
-                mlflow.log_param("Test each", str(args.test_each))
-                mlflow.log_param("Test for", str(args.test_for))
-                mlflow.log_param("Testing", str(args.testing))
-                mlflow.set_tag("NetSecEnv commit", netsecenv_git_result)
-                mlflow.set_tag("Agents commit", agents_git_result)
-                # Log the env conf
-                try:
-                    mlflow.log_artifact(args.env_conf)
-                except Exception as e:
-                    agent._logger.warning(f"Could not log env config file to MLflow: {e}")
-            
             # Log configuration to Wandb if enabled
             if args.use_wandb:
                 wandb.config.update({
@@ -520,23 +480,6 @@ def play_game(self, concept_observation, episode_num, testing=False):
                             '''
                         agent._logger.info(text)
 
-                        # Log evaluation metrics to MLflow if enabled
-                        if args.use_mlflow:
-                            mlflow.log_metric("eval_avg_win_rate", eval_win_rate, step=episode)
-                            mlflow.log_metric("eval_avg_detection_rate", eval_detection_rate, step=episode)
-                            mlflow.log_metric("eval_avg_returns", eval_average_returns, step=episode)
-                            mlflow.log_metric("eval_std_returns", eval_std_returns, step=episode)
-                            mlflow.log_metric("eval_avg_episode_steps", eval_average_episode_steps, step=episode)
-                            mlflow.log_metric("eval_std_episode_steps", eval_std_episode_steps, step=episode)
-                            mlflow.log_metric("eval_avg_win_steps", eval_average_win_steps, step=episode)
-                            mlflow.log_metric("eval_std_win_steps", eval_std_win_steps, step=episode)
-                            mlflow.log_metric("eval_avg_detected_steps", eval_average_detected_steps, step=episode)
-                            mlflow.log_metric("eval_std_detected_steps", eval_std_detected_steps, step=episode)
-                            mlflow.log_metric("eval_avg_max_steps_steps", eval_average_max_steps_steps, step=episode)
-                            mlflow.log_metric("eval_std_max_steps_steps", eval_std_max_steps_steps, step=episode)
-                            mlflow.log_metric("current_epsilon", agent.current_epsilon, step=episode)
-                            mlflow.log_metric("current_episode", episode, step=episode)
-                        
                         # Log evaluation metrics to Wandb if enabled
                         if args.use_wandb:
                             wandb.log({
@@ -640,23 +583,6 @@ def play_game(self, concept_observation, episode_num, testing=False):
                         agent._logger.info(text)
                         print(text)
 
-                        # Log test metrics to MLflow if enabled
-                        if args.use_mlflow:
-                            mlflow.log_metric("test_avg_win_rate", test_win_rate, step=episode)
-                            mlflow.log_metric("test_avg_detection_rate", test_detection_rate, step=episode)
-                            mlflow.log_metric("test_avg_returns", test_average_returns, step=episode)
-                            mlflow.log_metric("test_std_returns", test_std_returns, step=episode)
-                            mlflow.log_metric("test_avg_episode_steps", test_average_episode_steps, step=episode)
-                            mlflow.log_metric("test_std_episode_steps", test_std_episode_steps, step=episode)
-                            mlflow.log_metric("test_avg_win_steps", test_average_win_steps, step=episode)
-                            mlflow.log_metric("test_std_win_steps", test_std_win_steps, step=episode)
-                            mlflow.log_metric("test_avg_detected_steps", test_average_detected_steps, step=episode)
-                            mlflow.log_metric("test_std_detected_steps", test_std_detected_steps, step=episode)
-                            mlflow.log_metric("test_avg_max_steps_steps", test_average_max_steps_steps, step=episode)
-                            mlflow.log_metric("test_std_max_steps_steps", test_std_max_steps_steps, step=episode)
-                            mlflow.log_metric("current_epsilon", agent.current_epsilon, step=episode)
-                            mlflow.log_metric("current_episode", episode, step=episode)
-                        
                         # Log test metrics to Wandb if enabled
                         if args.use_wandb:
                             wandb.log({
@@ -704,8 +630,6 @@ def play_game(self, concept_observation, episode_num, testing=False):
 
         finally:
             # Clean up logging sessions
-            if args.use_mlflow and mlflow_run:
-                mlflow.end_run()
             if args.use_wandb:
                 wandb.finish()