Skip to content

Commit 7dcd242

Browse files
ondrej-lukasondrej-lukas
committed
remove mlflow support
1 parent 0a6bc74 commit 7dcd242

1 file changed

Lines changed: 77 additions & 102 deletions

File tree

‎agents/attackers/random/random_agent.py‎

Lines changed: 77 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import logging
44
import argparse
55
import numpy as np
6-
import mlflow
76
from os import path, makedirs
87
import random
98
from netsecgame import Action, Observation, BaseAgent, generate_valid_actions, AgentRole
@@ -49,7 +48,6 @@ def main():
4948
parser.add_argument("--episodes", help="Sets number of episodes to play", default=100, type=int)
5049
parser.add_argument("--seed", help="Sets random seed for agent's decisions", default=42, type=int)
5150
parser.add_argument("--logdir", help="Folder to store logs", default=path.join(path.dirname(path.abspath(__file__)), "logs"))
52-
parser.add_argument("--mlflow_url", help="URL for mlflow tracking server. If not provided, mlflow will store locally.", default=None)
5351
args = parser.parse_args()
5452

5553
if not path.exists(args.logdir):
@@ -58,118 +56,95 @@ def main():
5856

5957
# Create agent
6058
agent = RandomAttackerAgent(args.host, args.port, AgentRole.Attacker, seed=args.seed)
61-
62-
# Mlflow experiment name
63-
experiment_name = "Random Attacker Agent"
64-
if args.mlflow_url:
65-
mlflow.set_tracking_uri(args.mlflow_url)
66-
mlflow.set_experiment(experiment_name)
6759

6860
# Register in the game
6961
observation = agent.register()
62+
observation = agent.request_game_reset(randomize_topology=True, seed=0)
7063

71-
with mlflow.start_run(run_name=experiment_name) as run:
72-
# To keep statistics of each episode
73-
wins = 0
74-
detected = 0
75-
max_steps = 0
76-
num_win_steps = []
77-
num_detected_steps = []
78-
num_max_steps_steps = []
79-
num_detected_returns = []
80-
num_win_returns = []
81-
num_max_steps_returns = []
8264

83-
# Log more things in Mlflow
84-
mlflow.set_tag("experiment_name", experiment_name)
85-
mlflow.set_tag("episode_number", args.episodes)
65+
# To keep statistics of each episode
66+
wins = 0
67+
detected = 0
68+
max_steps = 0
69+
num_win_steps = []
70+
num_detected_steps = []
71+
num_max_steps_steps = []
72+
num_detected_returns = []
73+
num_win_returns = []
74+
num_max_steps_returns = []
8675

87-
for episode in range(1, args.episodes + 1):
88-
agent.logger.info(f'Starting episode {episode}')
89-
print(f'Starting episode {episode}')
9076

91-
# Play the game for one episode
92-
episodic_returns = []
93-
num_steps = 0
94-
95-
while observation and not observation.end:
96-
num_steps += 1
97-
agent.logger.debug(f'Observation received:{observation}')
98-
# Store returns in the episode
99-
episodic_returns.append(observation.reward)
100-
# Select the action randomly
101-
action = agent.select_action(observation)
102-
observation = agent.make_step(action)
103-
104-
agent.logger.debug(f'Observation received:{observation}')
105-
current_return = np.sum(episodic_returns)
106-
107-
agent.logger.info(f"Episode {episode} ended with return {current_return}.")
108-
109-
reward = current_return
110-
111-
if observation.info and observation.info.get('end_reason') == AgentStatus.Fail:
112-
detected +=1
113-
num_detected_steps.append(num_steps)
114-
num_detected_returns.append(reward)
115-
elif observation.info and observation.info.get('end_reason') == AgentStatus.Success:
116-
wins += 1
117-
num_win_steps.append(num_steps)
118-
num_win_returns.append(reward)
119-
elif observation.info and observation.info.get('end_reason') == AgentStatus.TimeoutReached:
120-
max_steps += 1
121-
num_max_steps_steps.append(num_steps)
122-
num_max_steps_returns.append(reward)
77+
for episode in range(1, args.episodes + 1):
78+
agent.logger.info(f'Starting episode {episode}')
79+
print(f'Starting episode {episode}')
12380

124-
# Reset the game - ONLY ONCE
125-
if episode < args.episodes:
126-
observation = agent.request_game_reset()
81+
# Play the game for one episode
82+
episodic_returns = []
83+
num_steps = 0
84+
85+
while observation and not observation.end:
86+
num_steps += 1
87+
agent.logger.debug(f'Observation received:{observation}')
88+
# Store returns in the episode
89+
episodic_returns.append(observation.reward)
90+
# Select the action randomly
91+
action = agent.select_action(observation)
92+
observation = agent.make_step(action)
93+
94+
agent.logger.debug(f'Observation received:{observation}')
95+
current_return = np.sum(episodic_returns)
96+
97+
agent.logger.info(f"Episode {episode} ended with return {current_return}.")
98+
99+
reward = current_return
100+
101+
if observation.info and observation.info.get('end_reason') == AgentStatus.Fail:
102+
detected +=1
103+
num_detected_steps.append(num_steps)
104+
num_detected_returns.append(reward)
105+
elif observation.info and observation.info.get('end_reason') == AgentStatus.Success:
106+
wins += 1
107+
num_win_steps.append(num_steps)
108+
num_win_returns.append(reward)
109+
elif observation.info and observation.info.get('end_reason') == AgentStatus.TimeoutReached:
110+
max_steps += 1
111+
num_max_steps_steps.append(num_steps)
112+
num_max_steps_returns.append(reward)
127113

128-
# Calculate stats for logging
129-
eval_win_rate = (wins/episode) * 100
130-
eval_detection_rate = (detected/episode) * 100
131-
132-
all_returns = num_detected_returns + num_win_returns + num_max_steps_returns
133-
eval_average_returns = np.mean(all_returns) if all_returns else 0
134-
eval_std_returns = np.std(all_returns) if all_returns else 0
135-
136-
all_steps = num_win_steps + num_detected_steps + num_max_steps_steps
137-
eval_average_episode_steps = np.mean(all_steps) if all_steps else 0
138-
eval_std_episode_steps = np.std(all_steps) if all_steps else 0
114+
# Reset the game - ONLY ONCE
115+
if episode < args.episodes:
116+
if episode % 10 == 0:
117+
observation = agent.request_game_reset(randomize_topology=True, seed=episode)
118+
else:
119+
observation = agent.request_game_reset(randomize_topology=False, seed=episode)
139120

140-
# Store in mlflow
141-
mlflow.log_metric("win_rate", eval_win_rate, step=episode)
142-
mlflow.log_metric("detection_rate", eval_detection_rate, step=episode)
143-
mlflow.log_metric("avg_returns", eval_average_returns, step=episode)
144-
mlflow.log_metric("std_returns", eval_std_returns, step=episode)
145-
mlflow.log_metric("avg_episode_steps", eval_average_episode_steps, step=episode)
121+
# Calculate stats for logging
122+
eval_win_rate = (wins/episode) * 100
123+
eval_detection_rate = (detected/episode) * 100
146124

147-
# Log the last final episode when it ends
148-
text = f'''Final results for {args.episodes} episodes:
149-
Wins={wins},
150-
Detections={detected},
151-
MaxSteps={max_steps},
152-
winrate={eval_win_rate:.3f}%,
153-
detection_rate={eval_detection_rate:.3f}%,
154-
average_returns={eval_average_returns:.3f} +- {eval_std_returns:.3f},
155-
average_episode_steps={eval_average_episode_steps:.3f} +- {eval_std_episode_steps:.3f}
156-
'''
157-
158-
agent.logger.info(text)
159-
print(text)
160-
agent._logger.info("Terminating interaction")
161-
agent.terminate_connection()
125+
all_returns = num_detected_returns + num_win_returns + num_max_steps_returns
126+
eval_average_returns = np.mean(all_returns) if all_returns else 0
127+
eval_std_returns = np.std(all_returns) if all_returns else 0
128+
129+
all_steps = num_win_steps + num_detected_steps + num_max_steps_steps
130+
eval_average_episode_steps = np.mean(all_steps) if all_steps else 0
131+
eval_std_episode_steps = np.std(all_steps) if all_steps else 0
132+
133+
# Log the last final episode when it ends
134+
text = f'''Final results for {args.episodes} episodes:
135+
Wins={wins},
136+
Detections={detected},
137+
MaxSteps={max_steps},
138+
winrate={eval_win_rate:.3f}%,
139+
detection_rate={eval_detection_rate:.3f}%,
140+
average_returns={eval_average_returns:.3f} +- {eval_std_returns:.3f},
141+
average_episode_steps={eval_average_episode_steps:.3f} +- {eval_std_episode_steps:.3f}
142+
'''
162143

163-
# Print and log the mlflow experiment ID, run ID, and storage location
164-
experiment_id = run.info.experiment_id
165-
run_id = run.info.run_id
166-
storage_location = "locally" if not args.mlflow_url else f"at {args.mlflow_url}"
167-
print(f"MLflow Experiment ID: {experiment_id}")
168-
print(f"MLflow Run ID: {run_id}")
169-
print(f"Experiment saved {storage_location}")
170-
agent._logger.info(f"MLflow Experiment ID: {experiment_id}")
171-
agent._logger.info(f"MLflow Run ID: {run_id}")
172-
agent._logger.info(f"Experiment saved {storage_location}")
144+
agent.logger.info(text)
145+
print(text)
146+
agent._logger.info("Terminating interaction")
147+
agent.terminate_connection()
173148

174149
if __name__ == '__main__':
175150
main()

0 commit comments

Comments
 (0)