oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#1378 - Re-added post blue and snapshots

Browse Source
This commit is contained in:
Chris McCarthy
2023-05-25 12:37:42 +01:00
parent 95a0669e5c
commit 7bbdbd6997
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/primaite/environment/primaite_env.py
View File

@@ -346,11 +346,10 @@ class Primaite(Env):
# 1. Implement Blue Action # 1. Implement Blue Action
self.interpret_action_and_apply(action) self.interpret_action_and_apply(action)
# Take snapshots of nodes and links # Take snapshots of nodes and links
self.nodes_post_red = copy.deepcopy(self.nodes) self.nodes_post_blue = copy.deepcopy(self.nodes)
self.links_post_red = copy.deepcopy(self.links) self.links_post_blue = copy.deepcopy(self.links)
# 2. Perform any time-based activities (e.g. a component moving from # 2. Perform any time-based activities (e.g. a component moving from patching to good)
# patching to good)
self.apply_time_based_updates() self.apply_time_based_updates()
# 3. Apply PoL # 3. Apply PoL
@@ -389,6 +388,9 @@ class Primaite(Env):
apply_red_agent_node_pol( apply_red_agent_node_pol(
self.nodes, self.red_iers, self.red_node_pol, self.step_count self.nodes, self.red_iers, self.red_node_pol, self.step_count
) )
# Take snapshots of nodes and links
self.nodes_post_red = copy.deepcopy(self.nodes)
self.links_post_red = copy.deepcopy(self.links)
# 5. Calculate reward signal (for RL) # 5. Calculate reward signal (for RL)
reward = calculate_reward_function( reward = calculate_reward_function(
@@ -409,7 +411,7 @@ class Primaite(Env):
# step count is reached in order to prevent neverending episode # step count is reached in order to prevent neverending episode
done = True done = True
print("Average reward: " + str(self.average_reward)) print("Average reward: " + str(self.average_reward))
# Load the reward into the transaction # Load the reward into the transaction
transaction.set_reward(reward) transaction.set_reward(reward)
# 6. Output Verbose # 6. Output Verbose