# Cloud Arena Environment — Mathematical Model RL # Extracted from cloud_arena_final.py # This is the MATHEMATICAL model env, NOT the LLM model. import sys, math, random, copy from collections import deque from typing import Dict, List, Optional, Tuple import numpy as np import gymnasium as gym from gymnasium import spaces # ── Seeds ───────────────────────────────────────────────────────────────────── GLOBAL_SEED = 42 np.random.seed(GLOBAL_SEED) random.seed(GLOBAL_SEED) # ── Observation layout (must sum to OBS_DIM) ────────────────────────────────── MAX_RES_IN_OBS = 8 # fixed obs slots (pad unused with zeros) N_FEAT_PER_RES = 10 # features per resource slot in obs N_BLOCK_B = 8 # global security block N_BLOCK_C = 7 # global cost block N_BLOCK_D = 6 # environment state block N_BLOCK_E = 24 # history: 8 actions + 8 rewards + 8 progress OBS_DIM = MAX_RES_IN_OBS * N_FEAT_PER_RES + N_BLOCK_B + N_BLOCK_C + N_BLOCK_D + N_BLOCK_E # = 80 + 8 + 7 + 6 + 24 = 125 assert OBS_DIM == 125, f"OBS_DIM mismatch: {OBS_DIM}" # ── Action space ────────────────────────────────────────────────────────────── N_ACTION_TYPES = 15 MAX_RESOURCES = 10 N_ACTIONS = N_ACTION_TYPES * MAX_RESOURCES # 150 A_NOOP=0; A_ANALYZE=1; A_VERIFY_DEPS=2; A_RESIZE_DOWN=3; A_RESIZE_UP=4 A_STOP=5; A_RESTART=6; A_DELETE=7; A_PATCH=8; A_ENCRYPT=9 A_RESTRICT=10; A_ROTATE_CREDS=11; A_ENABLE_LOG=12; A_ARCHIVE=13; A_OPT_NET=14 # Action cost penalties (small friction — makes actions non-free) ACTION_COSTS = { A_NOOP: 0.0, A_ANALYZE: -0.01, A_VERIFY_DEPS: -0.01, A_RESIZE_DOWN: -0.02, A_RESIZE_UP: -0.02, A_STOP: -0.03, A_RESTART: -0.02, A_DELETE: -0.05, A_PATCH: -0.02, A_ENCRYPT: -0.02, A_RESTRICT: -0.02, A_ROTATE_CREDS: -0.02, A_ENABLE_LOG: -0.01, A_ARCHIVE: -0.03, A_OPT_NET: -0.02, } # ── Curriculum ──────────────────────────────────────────────────────────────── # n_resources active per phase N_RESOURCES_PHASE = {0: 4, 1: 5, 2: 6, 3: 7, 4: 8, 5: 10} # Phase feature flags PHASE_FOG = {0: False, 1: True, 2: True, 3: True, 4: True, 5: True} PHASE_EVENTS = {0: False, 1: False, 2: True, 3: True, 4: True, 5: True} PHASE_CHAOS = {0: False, 1: False, 2: False, 3: True, 4: True, 5: True} CHAOS_INIT_PROB = {0: 0.0, 1: 0.0, 2: 0.0, 3: 0.20, 4: 0.30, 5: 0.35} # Win thresholds: cost must drop to this fraction of initial AND security >= sec_thr WIN_COST_THR = {0: 0.55, 1: 0.60, 2: 0.60, 3: 0.65, 4: 0.65, 5: 0.70} WIN_SEC_THR = {0: 0.00, 1: 0.60, 2: 0.70, 3: 0.70, 4: 0.75, 5: 0.80} MAX_STEPS = 150 # ══════════════════════════════════════════════════════════════════════════════ # RESOURCE OBJECT # ══════════════════════════════════════════════════════════════════════════════ class ResourceObject: CRIT = {"LOW": 0.3, "MED": 0.6, "HIGH": 1.0} def __init__(self, idx: int, criticality: str = "MED", category: str = "compute", rng: random.Random = None): rng = rng or random.Random(idx) self.idx = idx self.criticality = self.CRIT[criticality] self.category = category # ── Cost state ────────────────────────────────────────────────────── self.allocated = rng.uniform(0.70, 1.00) # initially overprovisioned self.usage = rng.uniform(0.15, 0.50) # true usage (always < allocated) self.usage = min(self.usage, self.allocated - 0.10) self.cost_rate = self.allocated # cost ∝ allocated self.activity_status = 1.0 # 1=active, 0=idle # ── State flags ────────────────────────────────────────────────────── self.health = 1 self.is_stopped = False self.is_deleted = False self.alert_flag = 0 # ── Security state (hidden under fog) ──────────────────────────────── self.risk_score = rng.uniform(0.05, 0.20) self.vulnerability = False self.encryption = True self.over_permission = False self.logging_enabled = True self.credential_age = rng.uniform(0.0, 0.3) self.exposure = rng.uniform(0.0, 0.15) self.sensitivity = rng.uniform(0.3, 0.8) # ── Fog of war ─────────────────────────────────────────────────────── self.fog_active = True # True = attributes hidden until ANALYZE self.cost_known = False self.deps_known = False self.steps_since_analyze = 0 self.staleness = 0.0 self.STALE_STEPS = 15 # after this many steps, fog re-activates # ── Dependency ─────────────────────────────────────────────────────── self.dependency_children: List[int] = [] # indices of resources that depend on this self.dependency_parent: Optional[int] = None # ── Diagnostics ────────────────────────────────────────────────────── self.steps_broken = 0 self.time_broken = 0.0 # ── Derived properties ──────────────────────────────────────────────────── def overprovision_ratio(self) -> float: return max(0.0, (self.allocated - self.usage) / max(self.allocated, 1e-6)) def get_cost(self) -> float: if self.is_deleted: return 0.0 if self.is_stopped: return self.cost_rate * 0.05 # minimal maintenance cost return self.cost_rate # ── Observation vector (10 dims) ────────────────────────────────────────── def to_obs(self, fog: bool = False) -> np.ndarray: if fog and self.fog_active: risk_obs = 0.0 cost_obs = 0.5 # agent sees estimated cost when under fog exp_obs = 0.0 else: risk_obs = self.risk_score cost_obs = self.cost_rate exp_obs = self.exposure return np.array([ float(self.health), # 0 risk_obs, # 1 (hidden under fog) self.criticality, # 2 cost_obs, # 3 (hidden under fog) self.activity_status, # 4 exp_obs, # 5 (hidden under fog) self.sensitivity, # 6 self.staleness, # 7 (always visible) float(self.alert_flag), # 8 (always visible for critical) self.time_broken, # 9 ], dtype=np.float32) # ── Per-step tick ───────────────────────────────────────────────────────── def tick(self, rng: random.Random, phase: int, event_prob: float = 0.0): if self.is_deleted: return # Staleness self.steps_since_analyze += 1 self.staleness = min(self.steps_since_analyze / self.STALE_STEPS, 1.0) if self.steps_since_analyze >= self.STALE_STEPS: self.fog_active = True # knowledge expires # Usage drift (only when running) if not self.is_stopped and self.health: self.usage = float(np.clip( self.usage + rng.uniform(-0.03, 0.03), 0.10, self.allocated)) # Credential aging self.credential_age = min(self.credential_age + 0.01, 1.0) # Broken resource tracking if not self.health: self.steps_broken += 1 self.time_broken = min(self.steps_broken / MAX_STEPS, 1.0) self.risk_score = min(self.risk_score + 0.015, 1.0) if self.criticality >= 1.0: self.alert_flag = 1 # high-criticality broken = visible alert # Random security events (Phase 2+) if phase >= 2 and rng.random() < event_prob and self.health: ev = rng.choice(["vuln", "expose", "iam", "log_off"]) if ev == "vuln": self.vulnerability = True self.risk_score = min(self.risk_score + 0.25, 1.0) elif ev == "expose": self.exposure = min(self.exposure + 0.35, 1.0) self.risk_score = min(self.risk_score + 0.20, 1.0) elif ev == "iam": self.over_permission = True self.risk_score = min(self.risk_score + 0.15, 1.0) elif ev == "log_off": self.logging_enabled = False self.risk_score = min(self.risk_score + 0.05, 1.0) # ── Actions ─────────────────────────────────────────────────────────────── def do_analyze(self): self.fog_active = False self.cost_known = True self.steps_since_analyze = 0 self.staleness = 0.0 def do_verify_deps(self): self.deps_known = True def do_resize_down(self) -> float: """Returns cost delta (positive = saving).""" new_alloc = max(self.usage + 0.10, 0.25) if new_alloc < self.allocated - 0.02: saved = (self.allocated - new_alloc) self.allocated = new_alloc self.cost_rate = new_alloc return saved return 0.0 def do_resize_up(self): self.allocated = min(self.allocated + 0.20, 1.0) self.cost_rate = self.allocated def do_stop(self) -> float: if not self.is_stopped: self.is_stopped = True self.activity_status = 0.0 return self.cost_rate * 0.95 # 95% cost eliminated return 0.0 def do_restart(self): self.is_stopped = False self.activity_status = 1.0 self.health = 1 def do_delete(self) -> float: saved = self.cost_rate self.is_deleted = True self.health = 0 return saved def do_patch(self): self.vulnerability = False self.risk_score = max(self.risk_score - 0.30, 0.0) def do_encrypt(self): self.encryption = True self.risk_score = max(self.risk_score - 0.15, 0.0) def do_restrict(self): self.exposure = max(self.exposure - 0.40, 0.0) self.risk_score = max(self.risk_score - 0.20, 0.0) def do_rotate_creds(self): self.credential_age = 0.0 self.over_permission = False self.risk_score = max(self.risk_score - 0.10, 0.0) def do_enable_logging(self): self.logging_enabled = True self.risk_score = max(self.risk_score - 0.05, 0.0) def do_archive(self) -> float: if not self.is_stopped: self.is_stopped = True self.activity_status = 0.0 return self.cost_rate * 0.70 return 0.0 def do_opt_network(self): self.exposure = max(self.exposure - 0.15, 0.0) self.risk_score = max(self.risk_score - 0.08, 0.0) # ══════════════════════════════════════════════════════════════════════════════ # ENVIRONMENT # ══════════════════════════════════════════════════════════════════════════════ class CloudArenaEnv(gym.Env): """ Cloud Arena: multi-objective cloud operations RL environment. Observation: 125-dim flat float32. Action space: Discrete(150) = 15 types × 10 resource slots. """ metadata = {"render_modes": []} def __init__(self, curriculum_ref: List[int] = None, global_step_ref: List[int] = None): super().__init__() self._curriculum_ref = curriculum_ref or [0] self._global_step_ref = global_step_ref or [0] self.observation_space = spaces.Box( low=-np.inf, high=np.inf, shape=(OBS_DIM,), dtype=np.float32) self.action_space = spaces.Discrete(N_ACTIONS) # Episode state (set in reset) self.resources: List[ResourceObject] = [] self.n_active = 0 self.step_count = 0 self.chaos_active = False self.chaos_steps = 0 self.veto_count = 0 self.cascade_count = 0 self.initial_total_cost = 1.0 self.prev_total_cost = 1.0 self.prev_risk_agg = 0.0 self._action_hist = deque([0.0] * 8, maxlen=8) self._reward_hist = deque([0.0] * 8, maxlen=8) self._progress_hist= deque([0.0] * 8, maxlen=8) # ── Properties ──────────────────────────────────────────────────────────── @property def curriculum_level(self) -> int: return self._curriculum_ref[0] # ── Reset ───────────────────────────────────────────────────────────────── def reset(self, seed=None, options=None): super().reset(seed=seed) rng = random.Random(seed if seed is not None else GLOBAL_SEED + self.step_count) self.step_count = 0 self.chaos_active = False self.chaos_steps = 0 self.veto_count = 0 self.cascade_count = 0 phase = self.curriculum_level scenario = options.get("scenario", 0) if options else 0 if scenario > 0: self._setup_boss_scenario(scenario, rng) else: self._setup_normal_episode(phase, rng) self.initial_total_cost = max(sum(r.get_cost() for r in self.resources), 1e-6) self.prev_total_cost = self.initial_total_cost self.prev_risk_agg = self._risk_aggregate() self._action_hist = deque([0.0] * 8, maxlen=8) self._reward_hist = deque([0.0] * 8, maxlen=8) self._progress_hist = deque([0.0] * 8, maxlen=8) return self._build_obs(), {} def _setup_normal_episode(self, phase: int, rng: random.Random): """Standard episode with phase-appropriate resources.""" self.n_active = N_RESOURCES_PHASE[phase] n = self.n_active # Criticality distribution: ~20% HIGH, ~40% MED, ~40% LOW crits = [] for i in range(n): if i == 0: crits.append("HIGH") elif i < n // 2: crits.append("MED") else: crits.append("LOW") cats = ["compute", "compute", "storage", "database", "compute", "storage", "compute", "database", "compute", "storage"][:n] self.resources = [] for i in range(n): r = ResourceObject(i, crits[i], cats[i], rng) # Phase 0: full observability — reveal everything upfront if not PHASE_FOG[phase]: r.fog_active = False r.cost_known = True r.deps_known = True # Phase 0: no security issues to start (clean state) if phase == 0: r.risk_score = rng.uniform(0.02, 0.08) r.vulnerability = False r.encryption = True r.over_permission = False r.logging_enabled = True r.exposure = rng.uniform(0.0, 0.05) else: # 💥 ANTI-CHEAT FIX: Force the agent to actually do SecOps in Phase 1+! r.vulnerability = rng.random() < 0.40 r.encryption = rng.random() > 0.30 # 30% unencrypted r.over_permission = rng.random() < 0.30 r.logging_enabled = rng.random() > 0.20 r.exposure = rng.uniform(0.10, 0.40) r.risk_score = rng.uniform(0.30, 0.60) self.resources.append(r) # Set up simple dependency: resource 0 (HIGH) has children [1] # This means deleting resource 0 would cascade to resource 1 # Agent can't delete resource 0 anyway (HIGH criticality), so it's safe if n >= 2: self.resources[0].dependency_children = [1] self.resources[1].dependency_parent = 0 # Chaos initialization for Phase 3+ if PHASE_CHAOS[phase] and rng.random() < CHAOS_INIT_PROB[phase]: self.chaos_active = True # Break 1-2 non-critical resources victims = [r for r in self.resources if r.criticality < 1.0][:2] for v in victims: v.health = 0 v.risk_score = min(v.risk_score + 0.40, 1.0) v.alert_flag = 0 # hidden unless HIGH criticality def _setup_boss_scenario(self, scenario: int, rng: random.Random): """Boss fight: predefined stressful starting conditions.""" phase = max(self.curriculum_level, 3) # boss fights at phase 3+ difficulty self._setup_normal_episode(phase, rng) if scenario == 1: # Cost Crisis for r in self.resources: r.allocated = min(r.allocated + rng.uniform(0.10, 0.25), 1.0) r.cost_rate = r.allocated r.usage = max(r.usage - 0.10, 0.10) elif scenario == 2: # Security Breach for r in self.resources: r.fog_active = True # force fog — agent must analyze r.cost_known = False r.vulnerability = (rng.random() < 0.60) r.encryption = (rng.random() < 0.30) r.over_permission = (rng.random() < 0.50) r.logging_enabled = (rng.random() < 0.40) r.exposure = rng.uniform(0.30, 0.80) r.risk_score = rng.uniform(0.40, 0.90) elif scenario == 3: # Infrastructure Failure (NOOP Test) self.chaos_active = True for r in self.resources[:3]: r.health = 0 r.risk_score = min(r.risk_score + 0.50, 1.0) elif scenario == 4: # Traffic Surge (underprovisioned) for r in self.resources: r.usage = min(r.allocated - 0.05, rng.uniform(0.75, 0.95)) r.risk_score = min(r.risk_score + 0.10, 0.50) elif scenario == 5: # Final Boss: everything self.chaos_active = True for i, r in enumerate(self.resources): r.allocated = min(r.allocated + 0.15, 1.0) r.cost_rate = r.allocated r.vulnerability = (rng.random() < 0.50) r.encryption = (rng.random() < 0.40) r.exposure = rng.uniform(0.20, 0.70) r.risk_score = rng.uniform(0.30, 0.80) if i < 2: r.health = 0 # ── Step ────────────────────────────────────────────────────────────────── def step(self, action: int): action = int(action) self.step_count += 1 self._global_step_ref[0] += 1 atype = action // MAX_RESOURCES ridx = action % MAX_RESOURCES phase = self.curriculum_level # ── Tick all resources ──────────────────────────────────────────────── event_prob = 0.04 if PHASE_EVENTS[phase] else 0.0 rng = random.Random(self._global_step_ref[0]) for r in self.resources: r.tick(rng, phase, event_prob) # ── Chaos events (Phase 3+) ─────────────────────────────────────────── if PHASE_CHAOS[phase] and rng.random() < 0.03: healthy = [r for r in self.resources if r.health and not r.is_deleted and r.criticality < 1.0] if healthy: victim = rng.choice(healthy) victim.health = 0 victim.risk_score = min(victim.risk_score + 0.40, 1.0) self.chaos_active = True if self.chaos_active: self.chaos_steps += 1 if self.chaos_steps > 20: self.chaos_active = False # chaos resolves after ~20 steps # ── Snapshot pre-action state ───────────────────────────────────────── cost_before = sum(r.get_cost() for r in self.resources) risk_before = self._risk_aggregate() # ── Apply action ────────────────────────────────────────────────────── cost_delta, sec_delta, veto = self._apply_action(atype, ridx) if veto: self.veto_count += 1 # ── Post-action state ───────────────────────────────────────────────── cost_now = sum(r.get_cost() for r in self.resources) risk_now = self._risk_aggregate() # ── Compute reward ──────────────────────────────────────────────────── reward = self._compute_reward( atype, ridx, veto, cost_before, cost_now, risk_before, risk_now) # ── Check win/done ──────────────────────────────────────────────────── win = self._check_win(cost_now, risk_now, phase) terminated = win truncated = (self.step_count >= MAX_STEPS) if terminated or truncated: reward += self._terminal_reward(win, cost_now, risk_now, phase) reward = float(np.clip(reward, -30.0, 60.0)) else: reward = float(np.clip(reward, -2.0, 5.0)) # ── Update history ──────────────────────────────────────────────────── self._action_hist.append(atype / N_ACTION_TYPES) self._reward_hist.append(np.clip(reward / 5.0, -1.0, 1.0)) self._progress_hist.append(max(0.0, (self.initial_total_cost - cost_now) / max(self.initial_total_cost, 1e-6))) self.prev_total_cost = cost_now self.prev_risk_agg = risk_now info = { "win": int(win), "cost_score": float(np.clip(1.0 - cost_now / max(self.initial_total_cost, 1e-6), 0, 1)), "security_score": float(np.clip(1.0 - risk_now, 0, 1)), "reliability_score": self._reliability_score(), "savings_pct": float(np.clip( (self.initial_total_cost - cost_now) / max(self.initial_total_cost, 1e-6) * 100, 0, 100)), "veto_rate": self.veto_count / max(self.step_count, 1), "cascade_count": self.cascade_count, "risk": risk_now, "chaos_active": self.chaos_active, } return self._build_obs(), reward, terminated, truncated, info # ── Action application ──────────────────────────────────────────────────── def _apply_action(self, atype: int, ridx: int) -> Tuple[float, float, bool]: """Returns (cost_delta, security_delta, was_vetoed).""" if atype == A_NOOP: return 0.0, 0.0, False # NOOP is never a veto # Validate resource index if ridx >= len(self.resources): return 0.0, 0.0, True r = self.resources[ridx] if r.is_deleted: return 0.0, 0.0, True cost_before = r.get_cost() risk_before = r.risk_score veto = False if atype == A_ANALYZE: r.do_analyze() elif atype == A_VERIFY_DEPS: r.do_verify_deps() elif atype == A_RESIZE_DOWN: if r.overprovision_ratio() > 0.08 and not r.is_stopped: r.do_resize_down() else: veto = True elif atype == A_RESIZE_UP: if r.usage > r.allocated - 0.12: r.do_resize_up() else: veto = True elif atype == A_STOP: can_stop = (not r.is_stopped and (r.activity_status < 0.35 or r.criticality <= 0.3) and r.criticality < 1.0) if can_stop: r.do_stop() else: veto = True elif atype == A_RESTART: if r.is_stopped: r.do_restart() else: veto = True elif atype == A_DELETE: can_delete = (r.deps_known and r.criticality < 1.0 and not r.is_stopped) if can_delete: has_crit_child = any( (ci < len(self.resources) and not self.resources[ci].is_deleted and self.resources[ci].criticality >= 0.6) for ci in r.dependency_children) if has_crit_child: veto = True else: r.do_delete() for ci in r.dependency_children: if ci < len(self.resources) and not self.resources[ci].is_deleted: child = self.resources[ci] child.health = 0 child.risk_score = min(child.risk_score + 0.3, 1.0) self.cascade_count += 1 else: veto = True elif atype == A_PATCH: if r.vulnerability: r.do_patch() else: veto = True elif atype == A_ENCRYPT: if not r.encryption: r.do_encrypt() else: veto = True elif atype == A_RESTRICT: if r.exposure > 0.15: r.do_restrict() else: veto = True elif atype == A_ROTATE_CREDS: if r.credential_age > 0.40: r.do_rotate_creds() else: veto = True elif atype == A_ENABLE_LOG: if not r.logging_enabled: r.do_enable_logging() else: veto = True elif atype == A_ARCHIVE: if r.category == "storage" and r.activity_status < 0.35: r.do_archive() else: veto = True elif atype == A_OPT_NET: if r.exposure > 0.08: r.do_opt_network() else: veto = True cost_after = r.get_cost() if not r.is_deleted else 0.0 risk_after = r.risk_score if not r.is_deleted else 0.0 return (cost_before - cost_after), (risk_before - risk_after), veto # ── Reward ──────────────────────────────────────────────────────────────── def _compute_reward(self, atype, ridx, veto, cost_before, cost_now, risk_before, risk_now) -> float: phase = self.curriculum_level w_cost = 0.25 w_sec = 0.35 if phase >= 1 else 0.0 w_stab = 0.25 # ── 1. Dense cost channel ───────────────────────────────────────────── r_cost = -w_cost * (cost_now / max(self.initial_total_cost, 1e-6)) # ── 2. Dense security channel ───────────────────────────────────────── r_sec = -w_sec * risk_now # ── 3. Stability/reliability ────────────────────────────────────────── n_broken = sum(1 for r in self.resources if not r.health and not r.is_deleted) r_stab = -w_stab * (n_broken / max(len(self.resources), 1)) # ── 4. Delta reward (THE MOST IMPORTANT SIGNAL) ─────────────────────── # Positive when agent caused improvement, zero otherwise cost_improvement = (cost_before - cost_now) / max(self.initial_total_cost, 1e-6) risk_improvement = risk_before - risk_now r_delta = 3.0 * cost_improvement # strong signal for cost savings r_delta += 4.0 * risk_improvement # strong signal for security improvements r_delta = float(np.clip(r_delta, -1.0, 2.0)) # ── 5. NOOP shaping ─────────────────────────────────────────────────── if atype == A_NOOP: if self.chaos_active: r_noop = +0.10 # correct — don't touch things during chaos elif risk_now < 0.10 and cost_now < self.initial_total_cost * 0.60: r_noop = +0.05 # correct — system is genuinely healthy elif risk_now < 0.25: r_noop = +0.01 # acceptable elif risk_now < 0.50: r_noop = -0.05 # negligence else: r_noop = -0.15 # gross negligence else: r_noop = 0.0 # ── 6. Action cost penalty ──────────────────────────────────────────── r_action = ACTION_COSTS.get(atype, -0.02) # ── 7. Veto penalty ─────────────────────────────────────────────────── r_veto = -0.10 if veto else 0.0 # ── 8. Temporal neglect ─────────────────────────────────────────────── # Phase 1+: growing penalty for ignoring known high-risk resources r_neglect = 0.0 if phase >= 1: for r in self.resources: if (not r.fog_active and not r.is_deleted and r.risk_score > 0.60): neglect_scale = min(r.steps_broken / MAX_STEPS, 1.0) r_neglect -= 0.02 * (1.0 + neglect_scale) * r.criticality r_neglect = max(r_neglect, -0.20) total = r_cost + r_sec + r_stab + r_delta + r_noop + r_action + r_veto + r_neglect return float(total) def _terminal_reward(self, win: bool, cost_now: float, risk_now: float, phase: int) -> float: r = 0.0 if win: speed_bonus = 10.0 * (1.0 - self.step_count / MAX_STEPS) r += 15.0 + speed_bonus else: # Partial credit cost_reduction = (self.initial_total_cost - cost_now) / max(self.initial_total_cost, 1e-6) r += 3.0 * max(cost_reduction, 0.0) r -= 5.0 # timeout penalty r -= 10.0 * risk_now # end-state security penalty if self.cascade_count > 0: r -= 5.0 * min(self.cascade_count, 3) return r # ── Win condition ───────────────────────────────────────────────────────── def _check_win(self, cost_now: float, risk_now: float, phase: int) -> bool: cost_ratio = cost_now / max(self.initial_total_cost, 1e-6) cost_win = cost_ratio < WIN_COST_THR[phase] sec_score = 1.0 - risk_now sec_win = sec_score >= WIN_SEC_THR[phase] # No critical resources broken no_crit_broken = not any( r.criticality >= 1.0 and not r.health and not r.is_deleted for r in self.resources) return cost_win and sec_win and no_crit_broken # ── Observation ─────────────────────────────────────────────────────────── def _build_obs(self) -> np.ndarray: phase = self.curriculum_level fog = PHASE_FOG[phase] # Block A: resource observations (padded to MAX_RES_IN_OBS) block_a = np.zeros(MAX_RES_IN_OBS * N_FEAT_PER_RES, dtype=np.float32) for i, r in enumerate(self.resources[:MAX_RES_IN_OBS]): block_a[i * N_FEAT_PER_RES: (i + 1) * N_FEAT_PER_RES] = r.to_obs(fog) # Block B: global security (8 dims) active = [r for r in self.resources if not r.is_deleted] n_a = max(len(active), 1) risk_agg = self._risk_aggregate() n_vuln = sum(1 for r in active if r.vulnerability) n_exposed = sum(1 for r in active if r.exposure > 0.3) n_unenc = sum(1 for r in active if not r.encryption) n_no_log = sum(1 for r in active if not r.logging_enabled) n_overperm = sum(1 for r in active if r.over_permission) block_b = np.array([ risk_agg, n_vuln / n_a, n_exposed / n_a, n_unenc / n_a, n_no_log / n_a, n_overperm / n_a, min(sum(r.credential_age for r in active) / n_a, 1.0), float(self.chaos_active), ], dtype=np.float32) # Block C: global cost (7 dims) total_cost = sum(r.get_cost() for r in self.resources) n_idle = sum(1 for r in active if r.activity_status < 0.3) n_overprov = sum(1 for r in active if r.overprovision_ratio() > 0.2) n_stopped = sum(1 for r in self.resources if r.is_stopped) n_deleted = sum(1 for r in self.resources if r.is_deleted) block_c = np.array([ total_cost / max(self.initial_total_cost, 1e-6), n_idle / n_a, n_overprov / n_a, n_stopped / max(len(self.resources), 1), n_deleted / max(len(self.resources), 1), (self.initial_total_cost - total_cost) / max(self.initial_total_cost, 1e-6), float(self._check_win(total_cost, risk_agg, self.curriculum_level)), ], dtype=np.float32) # Block D: environment state (6 dims) n_broken = sum(1 for r in active if not r.health) block_d = np.array([ self.step_count / MAX_STEPS, self.curriculum_level / 5.0, float(self.chaos_active), n_broken / n_a, self.veto_count / max(self.step_count, 1), self.cascade_count / max(n_a, 1), ], dtype=np.float32) # Block E: history (24 dims) block_e = np.array( list(self._action_hist) + list(self._reward_hist) + list(self._progress_hist), dtype=np.float32) obs = np.concatenate([block_a, block_b, block_c, block_d, block_e]) assert obs.shape == (OBS_DIM,), f"Obs shape {obs.shape} != {OBS_DIM}" return obs # ── Action masks ────────────────────────────────────────────────────────── def action_masks(self) -> np.ndarray: mask = np.zeros(N_ACTIONS, dtype=bool) # NOOP (action 0) — always valid mask[A_NOOP * MAX_RESOURCES] = True for ridx in range(MAX_RESOURCES): # Resources beyond active set are always invalid if ridx >= len(self.resources): # Only NOOP is already set; skip rest continue r = self.resources[ridx] if r.is_deleted: continue aid = lambda atype: atype * MAX_RESOURCES + ridx # noqa # ANALYZE — always valid (costs a small amount) mask[aid(A_ANALYZE)] = True # VERIFY_DEPS — always valid mask[aid(A_VERIFY_DEPS)] = True # 💥 ANTI-CHEAT FIX: If fog is active, the agent CANNOT execute these actions! if r.fog_active: continue # Skips evaluating the rest, keeping them False (Masked) # --- ONLY EVALUATED IF FOG IS LIFTED --- # RESIZE_DOWN — valid if overprovisioned and running mask[aid(A_RESIZE_DOWN)] = (r.overprovision_ratio() > 0.08 and not r.is_stopped) # RESIZE_UP — valid if near capacity mask[aid(A_RESIZE_UP)] = (r.usage > r.allocated - 0.12 and not r.is_stopped) # STOP — valid if idle or LOW criticality and currently running mask[aid(A_STOP)] = (not r.is_stopped and r.criticality < 1.0 and (r.activity_status < 0.35 or r.criticality <= 0.3)) # RESTART — valid if stopped mask[aid(A_RESTART)] = r.is_stopped # DELETE — valid if deps known, not critical, no critical children has_crit_child = any( (ci < len(self.resources) and not self.resources[ci].is_deleted and self.resources[ci].criticality >= 0.6) for ci in r.dependency_children) mask[aid(A_DELETE)] = (r.deps_known and r.criticality < 1.0 and not has_crit_child) # Security fixes (Phase 1+) mask[aid(A_PATCH)] = r.vulnerability mask[aid(A_ENCRYPT)] = not r.encryption mask[aid(A_RESTRICT)] = r.exposure > 0.15 mask[aid(A_ROTATE_CREDS)] = r.credential_age > 0.40 mask[aid(A_ENABLE_LOG)] = not r.logging_enabled mask[aid(A_ARCHIVE)] = (r.category == "storage" and r.activity_status < 0.35) mask[aid(A_OPT_NET)] = r.exposure > 0.08 # Collapse guard: always at least 3 valid actions if mask.sum() < 3: mask[A_NOOP * MAX_RESOURCES] = True if len(self.resources) > 0: mask[A_ANALYZE * MAX_RESOURCES] = True if len(self.resources) > 1: mask[A_ANALYZE * MAX_RESOURCES + 1] = True return mask # ── Helpers ─────────────────────────────────────────────────────────────── def _risk_aggregate(self) -> float: active = [r for r in self.resources if not r.is_deleted] if not active: return 0.0 weighted = sum(r.risk_score * r.criticality for r in active) total_w = sum(r.criticality for r in active) return weighted / max(total_w, 1e-6) def _reliability_score(self) -> float: active = [r for r in self.resources if not r.is_deleted] if not active: return 0.0 broken_w = sum(r.criticality for r in active if not r.health) total_w = sum(r.criticality for r in active) return max(0.0, 1.0 - broken_w / max(total_w, 1e-6)) def render(self): pass # ── Gymnasium wrapper ───────────────────────────────────────────────────────── from sb3_contrib.common.wrappers import ActionMasker def get_action_masks(env) -> np.ndarray: """Extract mask through ActionMasker wrapper.""" inner = env while hasattr(inner, "env"): inner = inner.env return inner.action_masks()