""" NexusGrid-CyberPhysEnv — Core Environment Implementation. Full OpenEnv-compliant environment: step() / reset() / state(). Wires together: GridEngine + SpoofEngine + RewardCalculator + Graders. """ from __future__ import annotations from typing import Any, Dict, List, Optional from uuid import uuid4 try: from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import State except ImportError: class Environment: # type: ignore[override] """Lightweight fallback for local testing without openenv-core.""" class State: # type: ignore[override] def __init__(self, episode_id: str, step_count: int): self.episode_id = episode_id self.step_count = step_count try: from ..models import GridAction, GridObservation, ActionType except ImportError: from models import GridAction, GridObservation, ActionType from .grid_engine import GridEngine from .scenarios import build_scenario, MAX_TICKS, TASK_NAMES from .spoof_engine import SpoofEngine from .reward import RewardCalculator from .graders import grade_task from .rubric import evaluate_task_rubrics from .training_logger import TrainingLogger class NexusgridEnvironment(Environment): """ NexusGrid-CyberPhysEnv: National power grid defense under cyber-physical attack. A 20-node transmission network with DC power flow physics running alongside a deterministic SCADA sensor spoofing engine. The agent must distinguish real grid failures from adversarially fabricated telemetry. OpenEnv API: reset(seed) → GridObservation (initial state) step(action) → GridObservation (with reward, done, info embedded) state → GridState (full internal state for debugging) """ SUPPORTS_CONCURRENT_SESSIONS: bool = True def __init__(self): """Initialize with default state.""" self._episode_id = str(uuid4()) self._seed = 42 self._task_id = 0 self._tick = 0 self._done = False self._total_reward = 0.0 self._max_ticks = 3 # Core engines self._engine: Optional[GridEngine] = None self._spoof: Optional[SpoofEngine] = None self._reward_calc = RewardCalculator() # Action history for grading self._action_history: List[Dict[str, Any]] = [] # Frequency history for grading self._frequency_history: List[float] = [] # Spoofed telemetry cache self._current_spoofed_telemetry: Dict[str, Dict] = {} # Attack config from scenario self._attack_config: Dict[str, Any] = {} # State estimation tracking self._state_estimation_run = False self._last_estimation_result: Optional[Dict] = None # Black Start tracking self._hydro_stable_ticks = 0 self._max_island_count = 0 self._successful_mergers = 0 self._premature_mergers = 0 self._transformer_failures = 0 self._full_restoration_tick: Optional[int] = None # Track if packet logs have been "read" (agent called advance_tick to observe) self._logs_read_this_episode = False self._logs_read_before_estimation = False # Internal OpenEnv state self._state = State(episode_id=self._episode_id, step_count=0) def reset(self, seed: int = None, task_id: int = None, **kwargs) -> GridObservation: """ Reset the environment for a new episode. Args: seed: Episode seed for reproducibility. Defaults to 42. task_id: Which task to run (0-5). Can also be passed in kwargs. Returns: Initial GridObservation with all fields freshly computed. """ # Handle task_id from kwargs (OpenEnv may pass it there) if task_id is None: task_id = kwargs.get("task_id", self._task_id) if seed is None: seed = kwargs.get("seed", 42) self._seed = seed self._task_id = task_id self._tick = 0 self._done = False self._total_reward = 0.0 self._episode_id = str(uuid4()) self._action_history = [] self._frequency_history = [] self._state_estimation_run = False self._last_estimation_result = None self._logs_read_this_episode = False self._logs_read_before_estimation = False self._full_restoration_tick = None # Black Start tracking self._hydro_stable_ticks = 0 self._max_island_count = 0 self._successful_mergers = 0 self._premature_mergers = 0 self._transformer_failures = 0 # Build scenario scenario = build_scenario(task_id, seed) self._engine = scenario["engine"] self._max_ticks = scenario["max_ticks"] self._attack_config = scenario.get("attack_config", {}) # Initialize spoof engine self._spoof = SpoofEngine(seed) if self._attack_config: self._spoof.configure_attack(self._attack_config) # Initialize reward calculator self._reward_calc = RewardCalculator() # Record initial frequency self._frequency_history.append(self._engine.frequency_hz) # Generate initial telemetry (may be spoofed) ground_truth = self._engine.get_ground_truth_telemetry() self._current_spoofed_telemetry = self._spoof.apply_spoofs(ground_truth, self._tick) # Update OpenEnv state self._state = State(episode_id=self._episode_id, step_count=0) return self._build_observation() def step(self, action: GridAction) -> GridObservation: """ Execute an action in the environment. Returns GridObservation with reward, done, and info embedded. """ if self._done: return self._build_observation(error="Episode already finished") if self._engine is None: return self._build_observation(error="Environment not initialized. Call reset() first.") # Parse the action action_type = action.action_type.value if isinstance(action.action_type, ActionType) else str(action.action_type) action_record = { "action_type": action_type, "tick": self._tick, "node_id": action.node_id, "edge_id": action.edge_id, "mw": action.mw, "status": action.status, "subgraph": action.subgraph, "hz_offset": action.hz_offset, "duration": action.duration, } error_msg = None fault_isolated = False spoof_detected = False pre_islands = self._engine.get_stable_islands() if self._task_id == 5 else [] # Execute the action if action_type == "dispatch_generation": if not action.node_id or action.mw is None: error_msg = "dispatch_generation requires node_id and mw" else: result = self._engine.dispatch_generation(action.node_id, action.mw) if not result["success"]: error_msg = result.get("error", "dispatch failed") action_record["result"] = result elif action_type == "toggle_circuit_breaker": if not action.edge_id or not action.status: error_msg = "toggle_circuit_breaker requires edge_id and status" else: result = self._engine.toggle_circuit_breaker(action.edge_id, action.status) if not result["success"]: error_msg = result.get("error", "toggle failed") else: if action.status == "OPEN" and result.get("old_status") == "LIVE": fault_isolated = True elif action.status == "CLOSED": self._track_black_start_merger(action.edge_id, pre_islands) action_record["result"] = result elif action_type == "run_state_estimation": if not action.subgraph: error_msg = "run_state_estimation requires subgraph (list of node IDs)" else: result = self._engine.run_state_estimation( action.subgraph, self._current_spoofed_telemetry, ) self._state_estimation_run = True self._last_estimation_result = result action_record["result"] = result if not result.get("consistent", True): spoof_detected = True elif action_type == "quarantine_scada_node": if not action.node_id: error_msg = "quarantine_scada_node requires node_id" else: # Check anti-hallucination gate if not self._state_estimation_run: error_msg = "Must run state_estimation before quarantine (anti-hallucination penalty applies)" result = { "success": False, "error": error_msg, "skipped": True, } else: result = self._engine.quarantine_node(action.node_id) if result.get("success"): self._spoof.quarantine_node(action.node_id) if self._last_estimation_result and not self._last_estimation_result.get("consistent", True): if result.get("success") and not result.get("already_quarantined", False): spoof_detected = True action_record["result"] = result elif action_type == "inject_counter_signal": if not action.node_id or action.hz_offset is None or action.duration is None: error_msg = "inject_counter_signal requires node_id, hz_offset, duration" else: result = self._engine.inject_counter_signal( action.node_id, action.hz_offset, action.duration ) if not result["success"]: error_msg = result.get("error", "injection failed") action_record["result"] = result elif action_type == "advance_tick": # Advance simulation (weather, load, frequency) self._engine.advance_tick() self._spoof.advance_tick() self._logs_read_this_episode = True # Agent observes packet logs if not self._state_estimation_run: self._logs_read_before_estimation = True # Apply resonance effect if active if self._spoof.is_resonance_active(): resonance_effect = self._spoof.get_resonance_effect(self._tick) self._engine.frequency_hz += resonance_effect self._engine.frequency_hz = max(58.0, min(62.0, self._engine.frequency_hz)) else: error_msg = f"Unknown action type: {action_type}" if action_type != "advance_tick": self._engine._record_telemetry() # Record action self._action_history.append(action_record) # Advance tick counter self._tick += 1 # Update spoofed telemetry ground_truth = self._engine.get_ground_truth_telemetry() self._current_spoofed_telemetry = self._spoof.apply_spoofs(ground_truth, self._tick) # Generate packet logs all_node_ids = list(self._engine.nodes.keys()) self._spoof.generate_packet_logs(all_node_ids, self._tick) # Compute reward reward_breakdown = self._reward_calc.compute_tick_reward( action_type=action_type, action_params={ "subgraph": action.subgraph or [], "node_id": action.node_id, "mw": action.mw, }, frequency_hz=self._engine.frequency_hz, overloaded_edges=self._engine.get_overloaded_edges(), critical_nodes_shed=self._engine.get_critical_nodes_shed(), is_proactive=self._engine.is_dispatch_proactive(), spoof_detected=spoof_detected, fault_isolated=fault_isolated, has_read_logs_before_estimation=( self._logs_read_this_episode and action_type == "run_state_estimation" ), ) tick_reward = reward_breakdown["total"] self._total_reward += tick_reward # Record frequency self._frequency_history.append(self._engine.frequency_hz) # Update Black Start tracking self._update_black_start_tracking() self._update_full_restoration_tracking() # Check termination conditions if self._engine.frequency_hz < 59.0: self._done = True if self._tick >= self._max_ticks: self._done = True # Task 4 special: turbine destruction if self._task_id == 4: destruction_tick = self._attack_config.get("destruction_tick", 10) if self._tick >= destruction_tick and self._spoof.is_resonance_active(): self._done = True # Turbine destroyed # Update OpenEnv state self._state = State(episode_id=self._episode_id, step_count=self._tick) return self._build_observation( reward=tick_reward, error=error_msg, reward_breakdown=reward_breakdown, ) @property def state(self) -> State: """Get current OpenEnv state.""" return self._state # ------------------------------------------------------------------ # Grading # ------------------------------------------------------------------ def get_score(self) -> float: """Get the final grader score for the current episode.""" episode_state = self._build_episode_state() return grade_task(self._task_id, self._action_history, episode_state) # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ def _build_observation( self, reward: float = 0.0, error: Optional[str] = None, reward_breakdown: Optional[Dict] = None, ) -> GridObservation: """Build a GridObservation from current engine state.""" if self._engine is None: return GridObservation(done=True, last_action_error="Not initialized") # Build telemetry from spoofed readings for current tick current_telemetry = [] for node_id, readings in self._current_spoofed_telemetry.items(): current_telemetry.append(readings) # Telemetry history (includes spoofed values) telemetry_stream = self._engine.get_telemetry_history() # Add info dict fields info = { "task_id": self._task_id, "tick": self._tick, "grid_frequency_hz": self._engine.frequency_hz, "active_spoofs": self._spoof.get_active_spoofs() if self._spoof else [], "last_kirchhoff_result": self._last_estimation_result, "episode_seed": self._seed, "reward_breakdown": reward_breakdown or {}, "rubric_breakdown": self.get_rubric_breakdown(), "curriculum_ready": True, } if self._done: info["episode_summary"] = self.get_episode_summary() obs = GridObservation( topology_graph=self._engine.get_topology(), telemetry_stream=telemetry_stream, weather_forecast_matrix=self._engine.get_weather(), network_packet_logs=self._spoof.get_recent_packet_logs() if self._spoof else [], grid_frequency_hz=self._engine.frequency_hz, tick=self._tick, task_id=self._task_id, done=self._done, reward=reward, last_action_error=error, last_state_estimation=self._last_estimation_result, weather_summary=self._engine.get_weather_summary(), metadata=info, ) return obs def _build_episode_state(self) -> Dict[str, Any]: """Build episode state dict for grading.""" if self._engine is None: return {} # Compute load restoration fraction total_possible = self._engine.get_total_possible_mwh() total_served = self._engine.get_mwh_served() load_fraction = total_served / total_possible if total_possible > 0 else 0.0 # Check if all critical nodes are restored critical_nodes = [n for n in self._engine.nodes.values() if n["critical"]] critical_restored = all(n["energized"] for n in critical_nodes) recovered_above_59_5_in_3_ticks = _check_recovery(self._frequency_history, threshold=59.5, window=3) return { "frequency_history": self._frequency_history, "is_proactive_dispatch": self._engine.is_dispatch_proactive(), "critical_nodes_shed": self._engine.get_critical_nodes_shed(), "full_restoration_tick": self._get_full_restoration_tick(), "load_restored_fraction": load_fraction, "hydro_stable_ticks": self._hydro_stable_ticks, "energized_node_count": sum( 1 for n in self._engine.nodes.values() if n["energized"] ), "max_island_count": self._max_island_count, "successful_mergers": self._successful_mergers, "premature_mergers": self._premature_mergers, "transformer_failures": self._transformer_failures, "critical_nodes_restored": critical_restored, "logs_read_before_estimation": self._logs_read_before_estimation, "spoof_target": self._attack_config.get("target_node"), "recovered_above_59_5_in_3_ticks": recovered_above_59_5_in_3_ticks, } def _track_black_start_merger(self, edge_id: str, pre_islands: List[List[str]]) -> None: """Record Task 5 island merger outcomes when a breaker is closed.""" if self._task_id != 5 or self._engine is None: return edge = self._engine.edges.get(edge_id) if edge is None: return def find_island(node_id: str) -> Optional[List[str]]: for island in pre_islands: if node_id in island: return island return None source_island = find_island(edge["source"]) target_island = find_island(edge["target"]) if not source_island or not target_island or source_island == target_island: return if self._engine.check_phase_angle_compatible(source_island, target_island): self._successful_mergers += 1 else: self._premature_mergers += 1 self._transformer_failures += 1 def _update_black_start_tracking(self) -> None: """Track Black Start milestones for Task 5 grading.""" if self._task_id != 5 or self._engine is None: return # Track hydro stability hydro = self._engine.nodes.get("NODE_01") if hydro and hydro["energized"] and hydro["generation_mw"] > 0: self._hydro_stable_ticks += 1 else: self._hydro_stable_ticks = 0 # Track island count islands = self._engine.get_stable_islands() self._max_island_count = max(self._max_island_count, len(islands)) def _update_full_restoration_tracking(self) -> None: """Persist the first tick where load restoration crosses the grading threshold.""" if self._engine is None or self._full_restoration_tick is not None: return if self._task_id == 2: isolation_ticks = [ action.get("tick", 999) for action in self._action_history if action.get("action_type") == "toggle_circuit_breaker" and str(action.get("status", "")).upper() == "OPEN" ] fault_isolated = bool(isolation_ticks) first_isolation_tick = min(isolation_ticks) if isolation_ticks else None restorative_dispatches = sum( 1 for action in self._action_history if action.get("action_type") == "dispatch_generation" and (action.get("mw") or 0) > 0 and first_isolation_tick is not None and action.get("tick", 999) > first_isolation_tick ) frequency_stable = ( len(self._frequency_history) >= 2 and all(freq >= 59.7 for freq in self._frequency_history[-2:]) ) if ( fault_isolated and restorative_dispatches >= 2 and frequency_stable and not self._engine.get_overloaded_edges() and self._engine.get_critical_nodes_shed() == 0 ): self._full_restoration_tick = self._tick return total_possible = self._engine.get_total_possible_mwh() total_served = self._engine.get_mwh_served() if total_possible > 0 and total_served >= total_possible * 0.95: self._full_restoration_tick = self._tick def _get_full_restoration_tick(self) -> Optional[int]: """Get the tick at which full load was restored (for Task 2).""" return self._full_restoration_tick def get_rubric_breakdown(self) -> Dict[str, Any]: """Return the current rubric breakdown for the active episode.""" return evaluate_task_rubrics(self._task_id, self._action_history, self._build_episode_state()) def get_episode_summary(self) -> Dict[str, Any]: """Build a compact summary suitable for training logs and dashboards.""" rubric_eval = self.get_rubric_breakdown() frequency_history = self._frequency_history or [60.0] return { "task_id": self._task_id, "task_name": TASK_NAMES.get(self._task_id, f"task_{self._task_id}"), "seed": self._seed, "score": self.get_score(), "rubrics": rubric_eval["rubrics"], "weighted_rubric_score": rubric_eval["weighted_score"], "actions_taken": [action.get("action_type", "unknown") for action in self._action_history], "frequency_min": min(frequency_history), "ticks_used": self._tick, "done": self._done, } def emit_training_log(self, episode: int, logger: Optional[TrainingLogger] = None) -> str: """Write the current episode summary as one JSONL line.""" summary = self.get_episode_summary() training_logger = logger or TrainingLogger() record = training_logger.build_record( episode=episode, task_id=self._task_id, seed=self._seed, score=summary["score"], rubrics=summary["rubrics"], actions_taken=summary["actions_taken"], frequency_min=summary["frequency_min"], ticks_used=summary["ticks_used"], extra={"task_name": summary["task_name"], "done": summary["done"]}, ) return training_logger.write_episode(record) def _check_recovery(freq_history: List[float], threshold: float, window: int) -> bool: """Check if frequency recovered above threshold within `window` ticks after dipping.""" dip_started = False ticks_since_dip = 0 for freq in freq_history: if freq < threshold: dip_started = True ticks_since_dip = 0 elif dip_started: ticks_since_dip += 1 if freq >= threshold and ticks_since_dip <= window: return True return False