Merge pull request #97 from stratosphereips/sebas-add-w&b-to-qlearning-attacker

Sebas add w&b to qlearning attacker

Author: eldraco

.gitignore

@@ -153,3 +153,4 @@ agents/mlruns*
 agents/*/*/mlruns/
 agents/*/*/logs
 aim/*
+wandb/

agents/attackers/q_learning/q_agent.py

@@ -7,7 +7,7 @@
 import pickle
 import argparse
 import logging
-import mlflow
+import wandb
 import subprocess
 import time
 
@@ -185,7 +185,7 @@ def play_game(self, observation, episode_num, testing=False):
     parser.add_argument("--logdir", help="Folder to store logs", default=path.join(path.dirname(path.abspath(__file__)), "logs"))
     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.", default='', type=str)
+    parser.add_argument("--experiment_id", help="Id of the experiment to record into Weights & Biases.", default='', type=str)
     parser.add_argument("--store_actions", help="Store actions in the log file q_agents_actions.log.", default=False, type=bool)
     parser.add_argument("--store_models_every", help="Store a model to disk every these number of episodes.", default=2000, type=int)
     parser.add_argument("--env_conf", help="Configuration file of the env. Only for logging purposes.", required=False, default='./env/netsecenv_conf.yaml', type=str)
@@ -225,71 +225,89 @@ def play_game(self, observation, episode_num, testing=False):
 
 
     if not args.testing:
-        # Mlflow experiment name        
+        # Wandb experiment name
         experiment_name = "Training and Eval of Q-learning Agent"
-        mlflow.set_experiment(experiment_name)
     elif args.testing:
         # Evaluate the agent performance
 
-        # Mlflow experiment name        
-        experiment_name = "Testing of Q-learning Agent against defender agent"
-        mlflow.set_experiment(experiment_name)
+        # Wandb experiment name
+        experiment_name = "Testing of Q-learning Agent"
 
 
     # This code runs for both training and testing. The difference is in the args.testing variable that is passed along
     # How it works:
     # - Evaluate for several 'episodes' (parameter)
     # - Each episode finishes with: steps played, return, win/lose. Store all
     # - Each episode compute the avg and std of all.
-    # - Every X episodes (parameter), report in log and mlflow
-    # - At the end, report in log and mlflow and console
+    # - Every X episodes (parameter), report in log and wandb
+    # - At the end, report in log and wandb and console
 
     # Register the agent
     observation = agent.register()
 
     try:
-        with mlflow.start_run(run_name=experiment_name + f'. ID {args.experiment_id}') as run:
-            # To keep statistics of each episode
-            wins = 0
-            detected = 0
-            max_steps = 0
-            num_win_steps = []
-            num_detected_steps = []
-            num_max_steps_steps = []
-            num_detected_returns = []
-            num_win_returns = []
-            num_max_steps_returns = []
-
-            # Log more things in Mlflow
-            mlflow.set_tag("experiment_name", experiment_name)
-            # Log notes or additional information
-            mlflow.set_tag("notes", "This is an evaluation")
-            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))
-            # Use subprocess.run to get the commit hash
-            netsecenv_command = "git rev-parse HEAD"
-            netsecenv_git_result = subprocess.run(netsecenv_command, shell=True, capture_output=True, text=True).stdout
-            agents_command = "cd NetSecGameAgents; git rev-parse HEAD"
-            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}')
-            mlflow.set_tag("NetSecEnv commit", netsecenv_git_result)
-            mlflow.set_tag("Agents commit", agents_git_result)
-            # Log the env conf
-            mlflow.log_artifact(args.env_conf)
-            agent._logger.info(f'Epsilon Start: {agent.epsilon_start}')
-            agent._logger.info(f'Epsilon End: {agent.epsilon_end}')
-            agent._logger.info(f'Epsilon Max Episodes: {agent.epsilon_max_episodes}')
-
-            for episode in range(1, args.episodes + 1):
+        # Initialize wandb
+        wandb.init(
+            entity='Stratosphere',
+            project='UTEP-Collaboration',
+            group='sebas-qlearning',
+            name=experiment_name + f'. ID {args.experiment_id}'
+        )
+
+        # To keep statistics of each episode
+        wins = 0
+        detected = 0
+        max_steps = 0
+        num_win_steps = []
+        num_detected_steps = []
+        num_max_steps_steps = []
+        num_detected_returns = []
+        num_win_returns = []
+        num_max_steps_returns = []
+
+        # Configure wandb with parameters and tags
+        wandb.config.update({
+            "alpha": args.alpha,
+            "epsilon_start": args.epsilon_start,
+            "epsilon_end": args.epsilon_end,
+            "epsilon_max_episodes": args.epsilon_max_episodes,
+            "gamma": args.gamma,
+            "episodes": args.episodes,
+            "test_each": args.test_each,
+            "test_for": args.test_for,
+            "testing": args.testing,
+            "experiment_name": experiment_name,
+            "notes": "This is an evaluation"
+        })
+
+        if args.previous_model:
+            wandb.config.update({"previous_model_loaded": str(args.previous_model)})
+
+        # Use subprocess.run to get the commit hash
+        netsecenv_command = "git rev-parse HEAD"
+        netsecenv_git_result = subprocess.run(netsecenv_command, shell=True, capture_output=True, text=True).stdout
+        agents_command = "cd NetSecGameAgents; git rev-parse HEAD"
+        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}')
+        wandb.config.update({
+            "netsecenv_commit": netsecenv_git_result.strip(),
+            "agents_commit": agents_git_result.strip()
+        })
+        # Log the env conf
+        try:
+            if path.exists(args.env_conf):
+                wandb.save(args.env_conf, base_path=path.dirname(path.abspath(args.env_conf)))
+            else:
+                agent._logger.warning(f"Environment config file not found: {args.env_conf}")
+                wandb.config.update({"env_conf_path": args.env_conf})
+        except Exception as e:
+            agent._logger.warning(f"Could not save env config file: {e}")
+            wandb.config.update({"env_conf_path": args.env_conf})
+        agent._logger.info(f'Epsilon Start: {agent.epsilon_start}')
+        agent._logger.info(f'Epsilon End: {agent.epsilon_end}')
+        agent._logger.info(f'Epsilon Max Episodes: {agent.epsilon_max_episodes}')
+
+        for episode in range(1, args.episodes + 1):
                 if not early_stop:
                     # Play 1 episode
                     observation, num_steps = agent.play_game(observation, testing=args.testing, episode_num=episode)       
@@ -333,6 +351,24 @@ def play_game(self, observation, episode_num, testing=False):
                     eval_average_max_steps_steps = np.mean(num_max_steps_steps)
                     eval_std_max_steps_steps = np.std(num_max_steps_steps)
 
+                    # Log results for testing mode every episode
+                    if args.testing:
+                        wandb.log({
+                            "test_avg_win_rate": eval_win_rate,
+                            "test_avg_detection_rate": eval_detection_rate,
+                            "test_avg_returns": eval_average_returns,
+                            "test_std_returns": eval_std_returns,
+                            "test_avg_episode_steps": eval_average_episode_steps,
+                            "test_std_episode_steps": eval_std_episode_steps,
+                            "test_avg_win_steps": eval_average_win_steps,
+                            "test_std_win_steps": eval_std_win_steps,
+                            "test_avg_detected_steps": eval_average_detected_steps,
+                            "test_std_detected_steps": eval_std_detected_steps,
+                            "test_avg_max_steps_steps": eval_average_max_steps_steps,
+                            "test_std_max_steps_steps": eval_std_max_steps_steps,
+                            "current_episode": episode
+                        }, step=episode)
+
                     # Now Test, log and report. This happens every X training episodes
                     if episode % args.test_each == 0 and episode != 0:
                         # If we are training, every these number of episodes, we need to test for some episodes.
@@ -354,20 +390,22 @@ def play_game(self, observation, episode_num, testing=False):
                                 epsilon={agent.current_epsilon}
                                 '''
                             agent._logger.info(text)
-                            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)
+                            wandb.log({
+                                "eval_avg_win_rate": eval_win_rate,
+                                "eval_avg_detection_rate": eval_detection_rate,
+                                "eval_avg_returns": eval_average_returns,
+                                "eval_std_returns": eval_std_returns,
+                                "eval_avg_episode_steps": eval_average_episode_steps,
+                                "eval_std_episode_steps": eval_std_episode_steps,
+                                "eval_avg_win_steps": eval_average_win_steps,
+                                "eval_std_win_steps": eval_std_win_steps,
+                                "eval_avg_detected_steps": eval_average_detected_steps,
+                                "eval_std_detected_steps": eval_std_detected_steps,
+                                "eval_avg_max_steps_steps": eval_average_max_steps_steps,
+                                "eval_std_max_steps_steps": eval_std_max_steps_steps,
+                                "current_epsilon": agent.current_epsilon,
+                                "current_episode": episode
+                            }, step=episode)
 
                             # To keep statistics of testing each episode
                             test_wins = 0
@@ -441,45 +479,50 @@ def play_game(self, observation, episode_num, testing=False):
                                 '''
                             agent._logger.info(text)
                             print(text)
-                            # Store in 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)
+                            # Store in wandb
+                            wandb.log({
+                                "test_avg_win_rate": test_win_rate,
+                                "test_avg_detection_rate": test_detection_rate,
+                                "test_avg_returns": test_average_returns,
+                                "test_std_returns": test_std_returns,
+                                "test_avg_episode_steps": test_average_episode_steps,
+                                "test_std_episode_steps": test_std_episode_steps,
+                                "test_avg_win_steps": test_average_win_steps,
+                                "test_std_win_steps": test_std_win_steps,
+                                "test_avg_detected_steps": test_average_detected_steps,
+                                "test_std_detected_steps": test_std_detected_steps,
+                                "test_avg_max_steps_steps": test_average_max_steps_steps,
+                                "test_std_max_steps_steps": test_std_max_steps_steps,
+                                "test_current_epsilon": agent.current_epsilon,
+                                "test_current_episode": episode
+                            }, step=episode)
 
                             if test_win_rate >= args.early_stop_threshold:
                                 agent.logger.info(f'Early stopping. Test win rate: {test_win_rate}. Threshold: {args.early_stop_threshold}')
                                 early_stop = True
 
-            
-            # Log the last final episode when it ends
-            text = f'''Final model performance after {episode} episodes.
-                Wins={wins},
-                Detections={detected},
-                winrate={eval_win_rate:.3f}%,
-                detection_rate={eval_detection_rate:.3f}%,
-                average_returns={eval_average_returns:.3f} +- {eval_std_returns:.3f},
-                average_episode_steps={eval_average_episode_steps:.3f} +- {eval_std_episode_steps:.3f},
-                average_win_steps={eval_average_win_steps:.3f} +- {eval_std_win_steps:.3f},
-                average_detected_steps={eval_average_detected_steps:.3f} +- {eval_std_detected_steps:.3f}
-                average_max_steps_steps={eval_std_max_steps_steps:.3f} +- {eval_std_max_steps_steps:.3f},
-                epsilon={agent.current_epsilon}
-                '''
-
-            agent._logger.info(text)
-            print(text)
-            agent._logger.error("Terminating interaction")
-            agent.terminate_connection()
+
+        # Log the last final episode when it ends
+        text = f'''Final model performance after {episode} episodes.
+            Wins={wins},
+            Detections={detected},
+            winrate={eval_win_rate:.3f}%,
+            detection_rate={eval_detection_rate:.3f}%,
+            average_returns={eval_average_returns:.3f} +- {eval_std_returns:.3f},
+            average_episode_steps={eval_average_episode_steps:.3f} +- {eval_std_episode_steps:.3f},
+            average_win_steps={eval_average_win_steps:.3f} +- {eval_std_win_steps:.3f},
+            average_detected_steps={eval_average_detected_steps:.3f} +- {eval_std_detected_steps:.3f}
+            average_max_steps_steps={eval_std_max_steps_steps:.3f} +- {eval_std_max_steps_steps:.3f},
+            epsilon={agent.current_epsilon}
+            '''
+
+        agent._logger.info(text)
+        print(text)
+        agent._logger.error("Terminating interaction")
+        agent.terminate_connection()
+
+        # Finish wandb run
+        wandb.finish()
 
     except KeyboardInterrupt:
         # Store the q-table