diff --git a/marl_cyborg/core/state.py b/marl_cyborg/core/state.py index cc4f3b7..4eeb360 100644 --- a/marl_cyborg/core/state.py +++ b/marl_cyborg/core/state.py @@ -10,6 +10,7 @@ def __init__(self, ip: str, hostname: str, subnet_cidr: str): self.privilege = 'None' # "None", "User", "Root" self.decoy = 'inactive' # "inactive" or "active" self.compromised_by = 'None' # Tracks agent ID responsible for breach + self.edr_active = False # Track endpoint monitoring telemetry status def __repr__(self): return ( diff --git a/marl_cyborg/environment/parallel_env.py b/marl_cyborg/environment/parallel_env.py index 9c5e51e..978c691 100644 --- a/marl_cyborg/environment/parallel_env.py +++ b/marl_cyborg/environment/parallel_env.py @@ -224,8 +224,11 @@ def step( observations[agent] = obs_array - # Reward shaping applied here natively - rewards[agent] = self._calculate_reward(agent, self.global_state) + # Reward shaping applied here natively factoring in immediate action outcomes + agent_effect = resolved_effects.get(agent) + rewards[agent] = self._calculate_reward( + agent, self.global_state, agent_effect + ) self.agents = [ agent @@ -401,6 +404,8 @@ def _apply_state_deltas(self, effects: Dict[str, ActionEffect]): for delta_key, delta_val in effect.state_deltas.items(): self.global_state.apply_delta(delta_key, delta_val) - def _calculate_reward(self, agent_id: str, state) -> float: + def _calculate_reward( + self, agent_id: str, state, effect: ActionEffect = None + ) -> float: """Delegates reward logic directly to the localized Scenario module.""" - return self.scenario.calculate_reward(agent_id, state) + return self.scenario.calculate_reward(agent_id, state, effect) diff --git a/marl_cyborg/scenarios/apt_espionage.py b/marl_cyborg/scenarios/apt_espionage.py index 1153d52..da2dea1 100644 --- a/marl_cyborg/scenarios/apt_espionage.py +++ b/marl_cyborg/scenarios/apt_espionage.py @@ -3,6 +3,7 @@ if TYPE_CHECKING: from marl_cyborg.core.state import GlobalNetworkState + from marl_cyborg.core.action import ActionEffect class AptEspionageScenario(BaseScenario): @@ -17,10 +18,31 @@ def __init__(self, agents): self.agents = agents def calculate_reward( - self, agent_id: str, global_state: 'GlobalNetworkState' + self, + agent_id: str, + global_state: 'GlobalNetworkState', + effect: 'ActionEffect' = None, ) -> float: reward = 0.0 + if effect: + if getattr(effect, 'cost', 0) > 0: + reward -= ( + effect.cost * 0.2 + ) # Larger penalty for wasting energy in APT scope + + if 'red' in agent_id.lower() and effect.observation_data: + if 'shared' in effect.observation_data: + reward += 5.0 # Massive emphasis on coordinated reconnaissance + # Massive penalty for triggering Honeypots during stealth operation + if 'Failed against Decoy' in str(effect.observation_data.values()): + reward -= 15.0 + elif 'kernel panic' in str(effect.observation_data.values()): + reward -= 25.0 + # Direct penalty for being blocked natively by Firewalls + elif 'validation failed' in str(effect.observation_data.values()): + reward -= 1.0 + if 'red' in agent_id.lower(): for h in global_state.all_hosts.values(): # Reward persistence diff --git a/marl_cyborg/scenarios/base_scenario.py b/marl_cyborg/scenarios/base_scenario.py index 9f53be9..fd9ef8e 100644 --- a/marl_cyborg/scenarios/base_scenario.py +++ b/marl_cyborg/scenarios/base_scenario.py @@ -3,6 +3,7 @@ if TYPE_CHECKING: from marl_cyborg.core.state import GlobalNetworkState + from marl_cyborg.core.action import ActionEffect class BaseScenario(ABC): @@ -15,7 +16,10 @@ class BaseScenario(ABC): @abstractmethod def calculate_reward( - self, agent_id: str, global_state: 'GlobalNetworkState' + self, + agent_id: str, + global_state: 'GlobalNetworkState', + effect: 'ActionEffect' = None, ) -> float: """Dynamically calculates the reward for the specified agent.""" pass diff --git a/marl_cyborg/scenarios/ransomware.py b/marl_cyborg/scenarios/ransomware.py index 595f55e..fa7fe38 100644 --- a/marl_cyborg/scenarios/ransomware.py +++ b/marl_cyborg/scenarios/ransomware.py @@ -3,6 +3,7 @@ if TYPE_CHECKING: from marl_cyborg.core.state import GlobalNetworkState + from marl_cyborg.core.action import ActionEffect class RansomwareScenario(BaseScenario): @@ -18,10 +19,27 @@ def __init__(self, agents): self.agents = agents def calculate_reward( - self, agent_id: str, global_state: 'GlobalNetworkState' + self, + agent_id: str, + global_state: 'GlobalNetworkState', + effect: 'ActionEffect' = None, ) -> float: reward = 0.0 + if effect: + if getattr(effect, 'cost', 0) > 0: + reward -= effect.cost * 0.1 # Micro-penalty for expending energy + + if 'red' in agent_id.lower() and effect.observation_data: + # Reward successful intelligence sharing inherently + if 'shared' in effect.observation_data: + reward += 2.0 + # Penalize falling for Decoys or EDR telemetry traps + if 'Failed against Decoy' in str(effect.observation_data.values()): + reward -= 5.0 + elif 'kernel panic' in str(effect.observation_data.values()): + reward -= 10.0 + red_impact_count = sum( 1 for h in global_state.all_hosts.values()