Get the db admin green agent working

tests/assets/configs/bad_primaite_session.yaml

@@ -21,7 +21,7 @@ game:
 agents:
   - ref: client_2_green_user
     team: GREEN
-    type: GreenWebBrowsingAgent
+    type: probabilistic_agent
     observation_space:
       type: UC2GreenObservation
     action_space:

tests/assets/configs/basic_switched_network.yaml

@@ -33,7 +33,7 @@ game:
 agents:
   - ref: client_2_green_user
     team: GREEN
-    type: GreenWebBrowsingAgent
+    type: probabilistic_agent
     observation_space:
       type: UC2GreenObservation
     action_space:

tests/assets/configs/eval_only_primaite_session.yaml

@@ -25,7 +25,7 @@ game:
 agents:
   - ref: client_2_green_user
     team: GREEN
-    type: GreenWebBrowsingAgent
+    type: probabilistic_agent
     observation_space:
       type: UC2GreenObservation
     action_space:

tests/assets/configs/multi_agent_session.yaml

@@ -31,7 +31,7 @@ game:
 agents:
   - ref: client_2_green_user
     team: GREEN
-    type: GreenWebBrowsingAgent
+    type: probabilistic_agent
     observation_space:
       type: UC2GreenObservation
     action_space:

tests/assets/configs/test_primaite_session.yaml

@@ -29,7 +29,7 @@ game:
 agents:
   - ref: client_2_green_user
     team: GREEN
-    type: GreenWebBrowsingAgent
+    type: probabilistic_agent
     observation_space:
       type: UC2GreenObservation
     action_space:

tests/assets/configs/train_only_primaite_session.yaml

@@ -25,7 +25,7 @@ game:
 agents:
   - ref: client_2_green_user
     team: GREEN
-    type: GreenWebBrowsingAgent
+    type: probabilistic_agent
     observation_space:
       type: UC2GreenObservation
     action_space:

tests/conftest.py

@@ -1,6 +1,6 @@
 # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
 from pathlib import Path
-from typing import Any, Dict, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import pytest
 import yaml
@@ -309,7 +309,7 @@ class ControlledAgent(AbstractAgent):
         )
         self.most_recent_action: Tuple[str, Dict]
 
-    def get_action(self, obs: None, reward: float = 0.0) -> Tuple[str, Dict]:
+    def get_action(self, obs: None, reward: float = 0.0, timestep: Optional[int] = None) -> Tuple[str, Dict]:
         """Return the agent's most recent action, formatted in CAOS format."""
         return self.most_recent_action
 
@@ -478,7 +478,6 @@ def game_and_agent():
     ]
 
     action_space = ActionManager(
-        game=game,
         actions=actions,  # ALL POSSIBLE ACTIONS
         nodes=[
             {

tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py

@@ -0,0 +1,84 @@
+from primaite.game.agent.actions import ActionManager
+from primaite.game.agent.observations import ICSObservation, ObservationManager
+from primaite.game.agent.rewards import RewardFunction
+from primaite.game.agent.scripted_agents import ProbabilisticAgent
+
+
+def test_probabilistic_agent():
+    """
+    Check that the probabilistic agent selects actions with approximately the right probabilities.
+
+    Using a binomial probability calculator (https://www.wolframalpha.com/input/?i=binomial+distribution+calculator),
+    we can generate some lower and upper bounds of how many times we expect the agent to take each action. These values
+    were chosen to guarantee a less than 1 in a million chance of the test failing due to unlucky random number
+    generation.
+    """
+    N_TRIALS = 10_000
+    P_DO_NOTHING = 0.1
+    P_NODE_APPLICATION_EXECUTE = 0.3
+    P_NODE_FILE_DELETE = 0.6
+    MIN_DO_NOTHING = 850
+    MAX_DO_NOTHING = 1150
+    MIN_NODE_APPLICATION_EXECUTE = 2800
+    MAX_NODE_APPLICATION_EXECUTE = 3200
+    MIN_NODE_FILE_DELETE = 5750
+    MAX_NODE_FILE_DELETE = 6250
+
+    action_space = ActionManager(
+        actions=[
+            {"type": "DONOTHING"},
+            {"type": "NODE_APPLICATION_EXECUTE"},
+            {"type": "NODE_FILE_DELETE"},
+        ],
+        nodes=[
+            {
+                "node_name": "client_1",
+                "applications": [{"application_name": "WebBrowser"}],
+                "folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}],
+            },
+        ],
+        max_folders_per_node=2,
+        max_files_per_folder=2,
+        max_services_per_node=2,
+        max_applications_per_node=2,
+        max_nics_per_node=2,
+        max_acl_rules=10,
+        protocols=["TCP", "UDP", "ICMP"],
+        ports=["HTTP", "DNS", "ARP"],
+        act_map={
+            0: {"action": "DONOTHING", "options": {}},
+            1: {"action": "NODE_APPLICATION_EXECUTE", "options": {"node_id": 0, "application_id": 0}},
+            2: {"action": "NODE_FILE_DELETE", "options": {"node_id": 0, "folder_id": 0, "file_id": 0}},
+        },
+    )
+    observation_space = ObservationManager(ICSObservation())
+    reward_function = RewardFunction()
+
+    pa = ProbabilisticAgent(
+        agent_name="test_agent",
+        action_space=action_space,
+        observation_space=observation_space,
+        reward_function=reward_function,
+        settings={
+            "action_probabilities": {0: P_DO_NOTHING, 1: P_NODE_APPLICATION_EXECUTE, 2: P_NODE_FILE_DELETE},
+            "random_seed": 120,
+        },
+    )
+
+    do_nothing_count = 0
+    node_application_execute_count = 0
+    node_file_delete_count = 0
+    for _ in range(N_TRIALS):
+        a = pa.get_action(0, timestep=0)
+        if a == ("DONOTHING", {}):
+            do_nothing_count += 1
+        elif a == ("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}):
+            node_application_execute_count += 1
+        elif a == ("NODE_FILE_DELETE", {"node_id": 0, "folder_id": 0, "file_id": 0}):
+            node_file_delete_count += 1
+        else:
+            raise AssertionError("Probabilistic agent produced an unexpected action.")
+
+    assert MIN_DO_NOTHING < do_nothing_count < MAX_DO_NOTHING
+    assert MIN_NODE_APPLICATION_EXECUTE < node_application_execute_count < MAX_NODE_APPLICATION_EXECUTE
+    assert MIN_NODE_FILE_DELETE < node_file_delete_count < MAX_NODE_FILE_DELETE