Go back to having agent configs in dicts, backwards compatible

Author: mrkickling

malsim/agents/defenders/heuristic_agent.py

@@ -34,7 +34,7 @@ def get_next_action(
 
         """Return an action that disables a compromised node"""
 
-        self.compromised_nodes |= agent_state.step_all_compromised_nodes
+        self.compromised_nodes |= agent_state.step_compromised_nodes
 
         selected_node_cost = math.inf
         selected_node = None
@@ -92,7 +92,7 @@ def get_next_action(
 
         """Return an action that disables a compromised node"""
 
-        self.compromised_nodes |= agent_state.step_all_compromised_nodes
+        self.compromised_nodes |= agent_state.step_compromised_nodes
 
         selected_node_cost = math.inf
         selected_node = None

malsim/envs/gym_envs.py

@@ -13,7 +13,7 @@
 import numpy as np
 
 from ..scenario import load_scenario
-from ..mal_simulator import MalSimulator, AgentConfig, AttackerAgentConfig
+from ..mal_simulator import MalSimulator, AgentType
 from ..envs import MalSimVectorizedObsEnv
 from ..agents import DecisionAgent
 
@@ -36,7 +36,7 @@ def __init__(self, scenario_file: str, **kwargs: Any) -> None:
 
         attacker_agents = [
             agent for agent in scenario.agents
-            if isinstance(agent, AttackerAgentConfig)
+            if agent['type'] == AgentType.ATTACKER
         ]
 
         assert len(attacker_agents) == 1, (
@@ -45,11 +45,11 @@ def __init__(self, scenario_file: str, **kwargs: Any) -> None:
         )
 
         attacker_agent = attacker_agents[0]
-        self.attacker_agent_name = attacker_agent.name
+        self.attacker_agent_name = attacker_agent['name']
 
         self.sim.register_attacker(
             self.attacker_agent_name,
-            attacker_agent.entry_points
+            attacker_agent['entry_points']
         )
         self.sim.reset()
 
@@ -127,26 +127,26 @@ def __init__(self, scenario_file: str, **kwargs: Any) -> None:
         self.action_space = \
             self.sim.action_space(self.defender_agent_name)
 
-    def _register_attacker_agents(self, agents: list[AgentConfig]) -> None:
+    def _register_attacker_agents(self, agents: list[dict[str, Any]]) -> None:
         """Register attackers in simulator"""
         for agent_config in agents:
-            if isinstance(agent_config, AttackerAgentConfig):
+            if agent_config['type'] == AgentType.ATTACKER:
                 self.sim.register_attacker(
-                    agent_config.name, agent_config.entry_points
+                    agent_config['name'], agent_config['entry_points']
                 )
 
     def _create_attacker_decision_agents(
-            self, agents: list[AgentConfig], seed: Optional[int] = None
+            self, agents: list[dict[str, Any]], seed: Optional[int] = None
         ) -> dict[str, DecisionAgent]:
         """Create decision agents for each attacker"""
 
         attacker_agents = {}
         for agent_config in agents:
-            if isinstance(agent_config, AttackerAgentConfig):
-                agent_name = agent_config.name
-                if agent_config.agent_class:
+            if agent_config['type'] == AgentType.ATTACKER:
+                agent_name = agent_config['name']
+                if agent_config['agent_class']:
                     attacker_agents[agent_name] = (
-                        agent_config.agent_class(
+                        agent_config['agent_class'](
                             {'seed': seed, 'randomize': self.randomize}
                         )
                     )

malsim/mal_simulator.py

@@ -96,7 +96,7 @@ class MalSimDefenderState(MalSimAgentState):
     @property
     def step_all_compromised_nodes(self) -> frozenset[AttackGraphNode]:
         print(
-            "Deprecated in mal-simulator 1.1.0, "
+            "'step_all_compromised_nodes' deprecated in mal-simulator 1.1.0, "
             "please use 'step_compromised_nodes'"
         )
         return self.step_compromised_nodes
@@ -1038,7 +1038,7 @@ def _defender_step_reward(
         - reward_mode: which way to calculate reward
         """
         step_enabled_defenses = defender_state.step_performed_nodes
-        step_compromised_nodes = defender_state.step_all_compromised_nodes
+        step_compromised_nodes = defender_state.step_compromised_nodes
 
         # Defender is penalized for compromised steps and enabled defenses
         step_reward = - sum(
@@ -1112,7 +1112,7 @@ def step(
         self.recording[self.cur_iter] = {}
 
         # Populate these from the results for all agents' actions.
-        step_all_compromised_nodes: set[AttackGraphNode] = set()
+        step_compromised_nodes: set[AttackGraphNode] = set()
         step_enabled_defenses: set[AttackGraphNode] = set()
         step_nodes_made_unviable: set[AttackGraphNode] = set()
 
@@ -1130,7 +1130,7 @@ def step(
             agent_compromised, agent_attempted = self._attacker_step(
                 attacker_state, actions.get(attacker_state.name, [])
             )
-            step_all_compromised_nodes |= agent_compromised
+            step_compromised_nodes |= agent_compromised
             self.recording[self.cur_iter][attacker_state.name] = (
                 list(agent_compromised)
             )
@@ -1160,7 +1160,7 @@ def step(
                 # Update defender state
                 updated_defender_state = self._update_defender_state(
                     agent_state,
-                    step_all_compromised_nodes,
+                    step_compromised_nodes,
                     step_enabled_defenses,
                     step_nodes_made_unviable
                 )
@@ -1189,14 +1189,14 @@ def render(self) -> None:
 
 
 def run_simulation(
-        sim: MalSimulator, agents: list[AgentConfig]
+        sim: MalSimulator, agents: list[dict[str, Any]]
     ) -> dict[str, list[AttackGraphNode]]:
     """Run a simulation with agents
 
     Return selected actions by each agent in each step
     """
     agent_actions: dict[str, list[AttackGraphNode]] = {}
-    total_rewards = {agent_config.name: 0.0 for agent_config in agents}
+    total_rewards = {agent_config['name']: 0.0 for agent_config in agents}
 
     logger.info("Starting CLI env simulator.")
     states = sim.reset()
@@ -1207,8 +1207,8 @@ def run_simulation(
 
         # Select actions for each agent
         for agent_config in agents:
-            decision_agent: Optional[DecisionAgent] = agent_config.agent
-            agent_name = agent_config.name
+            decision_agent: Optional[DecisionAgent] = agent_config['agent']
+            agent_name = agent_config['name']
             if decision_agent is None:
                 print(
                     f'Agent "{agent_name}" has no decision agent class '
@@ -1232,14 +1232,15 @@ def run_simulation(
         # Perform next step of simulation
         states = sim.step(actions)
         for agent_config in agents:
-            total_rewards[agent_config.name] += sim.agent_reward(agent_config.name)
+            total_rewards[agent_config['name']] += sim.agent_reward(agent_config['name'])
 
         print("---")
 
     print(f"Simulation over after {sim.cur_iter} steps.")
 
     # Print total rewards
     for agent_config in agents:
-        print(f'Total reward "{agent_config.name}"', total_rewards[agent_config.name])
+        agent_name = agent_config['name']
+        print(f'Total reward "{agent_name}"', total_rewards[agent_config['name']])
 
     return agent_actions

malsim/scenario.py

@@ -12,7 +12,7 @@
 """
 from __future__ import annotations
 import os
-from dataclasses import dataclass, asdict
+from dataclasses import dataclass
 from typing import Any, Optional, TextIO
 from enum import Enum
 import logging
@@ -83,6 +83,7 @@ class AgentType(Enum):
 
 @dataclass
 class AgentConfig:
+    # Will be used for agents in the future instead of dicts
     name: str
     agent_class: Any
     agent: Any
@@ -151,20 +152,18 @@ def __init__(
             self.attack_graph, false_negative_rates or {}
         )
         self.is_observable = apply_scenario_node_property(
-            self.attack_graph, is_observable or {}
+            self.attack_graph, is_observable or {}, default_value=False
         )
         self.is_actionable = apply_scenario_node_property(
-            self.attack_graph, is_actionable or {}
+            self.attack_graph, is_actionable or {}, default_value=False
         )
 
     def to_dict(self) -> dict[str, Any]:
         assert self._lang_file, "Can not save scenario to file if lang file was not given"
         scenario_dict = {
             # 'version': ?
             'lang_file': self._lang_file,
-            'agents': {
-                a.name: asdict(a) for a in self.agents
-            },
+            'agents': self.agents,
             'rewards': {},
             'false_positive_rates': {},
             'false_negative_rates': {},
@@ -437,9 +436,8 @@ def get_entry_point_nodes(
 
 
 def load_simulator_agents(
-        attack_graph: AttackGraph,
-        scenario_agents: dict[str, Any],
-    ) -> list[AgentConfig]:
+        attack_graph: AttackGraph, scenario_agents: dict[str, Any]
+    ) -> list[dict[str, Any]]:
     """Load agents to be registered in MALSimulator
 
     Create the agents from the specified classes,
@@ -449,7 +447,7 @@ def load_simulator_agents(
     - attack_graph: the attack graph
     - scenario: the scenario in question as a dict
     Return:
-    - agents: a list of agent configurations
+    - agents: a list of agent configurations (dicts)
     """
 
     # Create list of agents dicts
@@ -470,29 +468,29 @@ def load_simulator_agents(
             )
 
         if agent_type == AgentType.ATTACKER:
-            agent_config = AttackerAgentConfig(
-                name=agent_name,
-                agent_class=agent_class,
-                agent=agent,
-                policy=policy,
-                config=agent_config,
-                entry_points=get_entry_point_nodes(
+            agent_config = {
+                'name': agent_name,
+                'agent_class': agent_class,
+                'agent': agent,
+                'policy': policy,
+                'config': agent_config,
+                'entry_points': get_entry_point_nodes(
                     attack_graph, agent_info['entry_points'] # Required
                 ),
-                goals=get_entry_point_nodes(
+                'goals': get_entry_point_nodes(
                     attack_graph, agent_info.get('goals', []) # Optional
                 ),
-                type=AgentType.ATTACKER
-            )
+                'type': AgentType.ATTACKER
+            }
         elif agent_type == AgentType.DEFENDER:
-            agent_config = DefenderAgentConfig(
-                name=agent_name,
-                agent_class=agent_class,
-                agent=agent,
-                policy=policy,
-                config=agent_config,
-                type=AgentType.DEFENDER
-            )
+            agent_config = {
+                'name': agent_name,
+                'agent_class': agent_class,
+                'agent': agent,
+                'policy': policy,
+                'config': agent_config,
+                'type':AgentType.DEFENDER
+            }
 
         agents.append(agent_config)
 

tests/envs/test_example_scenarios.py

@@ -41,12 +41,12 @@ def test_bfs_vs_bfs_state_and_reward() -> None:
     attacker_agent_name = "attacker1"
 
     attacker_agent = next(
-        agent.agent for agent in scenario.agents
-        if agent.name == attacker_agent_name
+        agent['agent'] for agent in scenario.agents
+        if agent['name'] == attacker_agent_name
     )
     defender_agent = next(
-        agent.agent for agent in scenario.agents
-        if agent.name == defender_agent_name
+        agent['agent'] for agent in scenario.agents
+        if agent['name'] == defender_agent_name
     )
 
     total_reward_defender = 0.0
@@ -78,7 +78,7 @@ def test_bfs_vs_bfs_state_and_reward() -> None:
         # If actions were performed, add them to respective list
         if attacker_node and attacker_node in attacker_state.step_performed_nodes:
             attacker_actions.append(attacker_node.full_name)
-            assert attacker_node in defender_state.step_all_compromised_nodes
+            assert attacker_node in defender_state.step_compromised_nodes
 
         if defender_node and defender_node in defender_state.step_performed_nodes:
             defender_actions.append(defender_node.full_name)
@@ -194,12 +194,12 @@ def test_bfs_vs_bfs_state_and_reward_per_step_ttc() -> None:
     attacker_agent_name = "attacker1"
 
     attacker_agent = next(
-        agent.agent for agent in scenario.agents
-        if agent.name == attacker_agent_name
+        agent['agent'] for agent in scenario.agents
+        if agent['name'] == attacker_agent_name
     )
     defender_agent = next(
-        agent.agent for agent in scenario.agents
-        if agent.name == defender_agent_name
+        agent['agent'] for agent in scenario.agents
+        if agent['name'] == defender_agent_name
     )
 
     total_reward_defender = 0.0
@@ -230,7 +230,7 @@ def test_bfs_vs_bfs_state_and_reward_per_step_ttc() -> None:
         # If actions were performed, add them to respective list
         if attacker_node and attacker_node in attacker_state.step_performed_nodes:
             attacker_actions.append(attacker_node.full_name)
-            assert attacker_node in defender_state.step_all_compromised_nodes
+            assert attacker_node in defender_state.step_compromised_nodes
 
         if defender_node and defender_node in \
                 states['defender1'].step_performed_nodes:
@@ -325,12 +325,12 @@ def test_bfs_vs_bfs_state_and_reward_per_step_effort_based() -> None:
     attacker_agent_name = "attacker1"
 
     attacker_agent = next(
-        agent_info.agent for agent_info in scenario.agents
-        if agent_info.name == attacker_agent_name
+        agent['agent'] for agent in scenario.agents
+        if agent['name'] == attacker_agent_name
     )
     defender_agent = next(
-        agent_info.agent for agent_info in scenario.agents
-        if agent_info.name == defender_agent_name
+        agent['agent'] for agent in scenario.agents
+        if agent['name'] == defender_agent_name
     )
 
     total_reward_defender = 0.0
@@ -361,7 +361,7 @@ def test_bfs_vs_bfs_state_and_reward_per_step_effort_based() -> None:
         # If actions were performed, add them to respective list
         if attacker_node and attacker_node in attacker_state.step_performed_nodes:
             attacker_actions.append(attacker_node.full_name)
-            assert attacker_node in defender_state.step_all_compromised_nodes
+            assert attacker_node in defender_state.step_compromised_nodes
 
         if defender_node and defender_node in defender_state.step_performed_nodes:
             defender_actions.append(defender_node.full_name)
@@ -432,12 +432,12 @@ def test_bfs_vs_bfs_state_and_reward_expected_value_ttc() -> None:
     attacker_agent_name = "attacker1"
 
     attacker_agent = next(
-        agent_info.agent for agent_info in scenario.agents
-        if agent_info.name == attacker_agent_name
+        agent['agent'] for agent in scenario.agents
+        if agent['name'] == attacker_agent_name
     )
     defender_agent = next(
-        agent_info.agent for agent_info in scenario.agents
-        if agent_info.name == defender_agent_name
+        agent['agent'] for agent in scenario.agents
+        if agent['name'] == defender_agent_name
     )
 
     total_reward_defender = 0.0
@@ -468,7 +468,7 @@ def test_bfs_vs_bfs_state_and_reward_expected_value_ttc() -> None:
         # If actions were performed, add them to respective list
         if attacker_node and attacker_node in attacker_state.step_performed_nodes:
             attacker_actions.append(attacker_node.full_name)
-            assert attacker_node in defender_state.step_all_compromised_nodes
+            assert attacker_node in defender_state.step_compromised_nodes
 
         if defender_node and defender_node in defender_state.step_performed_nodes:
             defender_actions.append(defender_node.full_name)