""" SupportOps Arena - Core OpenEnv-compliant environment. Main environment class implementing the full OpenEnv interface. """ import uuid import logging from datetime import datetime from typing import Any from env.state import ( EnvState, EnvObservation, StepResult, HiddenState, TaskLevel, ActionRecord, NetworkStatus, VPNStatus, AuthStatus, ServiceHealth ) from env.actions import ActionType, ACTION_METADATA from env.scenarios import TicketGenerator from env.transitions import StateMachine from env.rewards import RewardCalculator # Setup logging logger = logging.getLogger(__name__) # ─── Configuration ──────────────────────────────────────────── MAX_STEPS = { TaskLevel.EASY: 10, TaskLevel.MEDIUM: 16, TaskLevel.HARD: 24, } class SupportOpsArena: """ OpenEnv-compliant environment for enterprise IT incident triage. An RL environment simulating IT support operations where agents must: - Gather evidence through sequential tool use - Reason about hidden root causes - Apply risk-weighted remediation actions - Escalate when autonomous resolution is unsafe Implements full OpenEnv interface: - async reset(task_level: str, seed: int | None) -> EnvObservation - async step(action: ActionType) -> StepResult - async state() -> EnvState """ def __init__(self, adversary: Any | None = None): """ Initialize the environment. Args: adversary: Optional AdaptiveAdversary instance for difficulty adaptation """ self.ticket_generator = TicketGenerator() self.state_machine = StateMachine() self.reward_calculator = RewardCalculator() self.adversary = adversary self._current_state: EnvState | None = None self._episode_active = False logger.info("SupportOpsArena environment initialized") async def reset( self, task_level: str = "easy", seed: int | None = None ) -> EnvObservation: """ Reset environment and start new episode. Args: task_level: Difficulty level - "easy", "medium", or "hard" seed: Random seed for reproducibility Returns: Initial observation for the episode Raises: ValueError: If task_level is not valid """ # Validate task level if task_level not in ["easy", "medium", "hard"]: raise ValueError(f"Invalid task_level: {task_level}. Must be 'easy', 'medium', or 'hard'") task_level_enum = TaskLevel(task_level) # Generate new episode ID episode_id = str(uuid.uuid4()) logger.info(f"Starting episode {episode_id} at {task_level} difficulty (seed={seed})") # Apply adversary sampling weights if available if self.adversary: weights = self.adversary.get_sampling_weights(task_level_enum) self.ticket_generator.set_sampling_weights(weights) # Generate ticket scenario ticket_data = self.ticket_generator.generate_ticket(task_level_enum, seed) root_cause_data = ticket_data["root_cause"] obs_template = ticket_data["observation_template"] misleading_idx = ticket_data.get("misleading_log_index") # Create hidden state hidden = HiddenState( root_cause=root_cause_data["description"], root_cause_category=root_cause_data["category"], correct_remediation=root_cause_data["correct_remediation"].value, correct_remediation_alts=[ alt.value if isinstance(alt, ActionType) else alt for alt in root_cause_data.get("correct_remediation_alts", []) ], misleading_log_index=misleading_idx, severity=root_cause_data["severity"], affected_users_count=root_cause_data["affected_users"], ticket_variant_id=root_cause_data["id"] ) # Create initial observation max_steps = MAX_STEPS[task_level_enum] observation = EnvObservation( ticket_id=f"INC-{episode_id[:8].upper()}", ticket_summary=obs_template["ticket_summary"], user_context=obs_template["user_context"], network_status=obs_template.get("network_status", NetworkStatus.UNKNOWN), vpn_status=obs_template.get("vpn_status", VPNStatus.NA), auth_status=obs_template.get("auth_status", AuthStatus.OK), service_health=obs_template.get("service_health", {}), system_logs=obs_template["system_logs"], action_history=[], step_count=0, escalation_allowed=False, task_level=task_level_enum, steps_remaining=max_steps ) # Create full state self._current_state = EnvState( episode_id=episode_id, task_level=task_level_enum, observation=observation, hidden=hidden, cumulative_reward=0.0, done=False, created_at=datetime.utcnow().isoformat(), action_log=[], diagnostic_steps_taken=0 ) self._episode_active = True logger.info(f"Episode {episode_id} initialized: {root_cause_data['id']}") return observation async def step(self, action: ActionType) -> StepResult: """ Execute one step in the environment. Args: action: Action to take (ActionType enum) Returns: StepResult containing (observation, reward, done, info) Raises: RuntimeError: If episode is not active or already terminated ValueError: If action is invalid """ if not self._episode_active or self._current_state is None: raise RuntimeError("Episode not active. Call reset() first.") if self._current_state.done: raise RuntimeError("Episode already terminated. Call reset() to start new episode.") # Validate action if not isinstance(action, ActionType): logger.error(f"Invalid action type: {action}") # Return penalty for invalid action info = { "episode_id": self._current_state.episode_id, "step": self._current_state.observation.step_count, "root_cause": None, "reward_breakdown": {"HARMFUL_ACTION": -0.35}, "correct": False } return StepResult( observation=self._current_state.observation, reward=-0.35, done=False, info=info ) logger.debug(f"Step {self._current_state.observation.step_count + 1}: {action.value}") # Apply state transition next_state, reward_events = await self.state_machine.transition( self._current_state, action ) # Check terminal conditions and add terminal rewards done, terminal_events = self._check_terminal_conditions(next_state, action) reward_events.extend(terminal_events) # Calculate reward reward, reward_breakdown = self.reward_calculator.calculate( reward_events, next_state, action ) # Update state next_state.cumulative_reward += reward next_state.observation.step_count += 1 next_state.observation.steps_remaining -= 1 next_state.done = done # Record action in history action_record = ActionRecord( step=next_state.observation.step_count, action_name=action.value, rationale=None, reward_received=reward, timestamp=datetime.utcnow().isoformat() ) next_state.observation.action_history.append(action_record) next_state.action_log.append(action_record) # Update adversary if available if self.adversary and done: await self.adversary.update_policy_tracker( next_state.hidden.ticket_variant_id, next_state.action_log, next_state.cumulative_reward ) # Build info dict info: dict[str, Any] = { "episode_id": next_state.episode_id, "step": next_state.observation.step_count, "root_cause": None, # Hidden until episode ends "reward_breakdown": reward_breakdown, "correct": False } # Reveal ground truth at episode end if done: info["root_cause"] = next_state.hidden.root_cause info["correct"] = self._check_correctness(next_state, action) info["episode_score"] = self.reward_calculator.calculate_episode_score( next_state.cumulative_reward, next_state.task_level ) logger.info( f"Episode {next_state.episode_id} completed: " f"score={info['episode_score']:.3f}, " f"correct={info['correct']}, " f"steps={next_state.observation.step_count}" ) # Update current state self._current_state = next_state if done: self._episode_active = False return StepResult( observation=next_state.observation, reward=reward, done=done, info=info ) async def state(self) -> EnvState: """ Return full internal state including hidden information. This method is for graders only and should not be called by agents during active episodes. Returns: Complete EnvState including hidden state Raises: RuntimeError: If no episode is active """ if self._current_state is None: raise RuntimeError("No active episode. Call reset() first.") return self._current_state def _check_terminal_conditions( self, state: EnvState, action: ActionType ) -> tuple[bool, list[str]]: """ Check if episode should terminate and return terminal reward events. Returns: Tuple of (done: bool, terminal_reward_events: list[str]) """ events = [] # Check if action is terminal action_meta = ACTION_METADATA.get(action) if action_meta and action_meta.is_terminal: if action == ActionType.ESCALATE_TICKET: # Check if escalation was premature (already handled in transition) if not state.observation.escalation_allowed: # Don't terminate if premature return False, events # Check if escalation is justified if state.diagnostic_steps_taken >= 3: events.append("JUSTIFIED_ESCALATION") else: events.append("UNJUSTIFIED_ESCALATION") # Escalation is always "correct" if allowed and justified if state.diagnostic_steps_taken >= 3: events.append("TERMINAL_CORRECT") return True, events elif action == ActionType.RESOLVE_TICKET: # Check if correct remediation correct_action = state.hidden.correct_remediation alt_actions = state.hidden.correct_remediation_alts # Check if we applied the correct fix earlier past_actions = [a.action_name for a in state.observation.action_history] if correct_action in past_actions: events.append("CORRECT_REMEDIATION") events.append("TERMINAL_CORRECT") # Check for optimal resolution if self.reward_calculator.is_optimal_resolution( state.observation.step_count, state.task_level ): events.append("OPTIMAL_RESOLUTION") else: # Check alternatives alt_found = any(alt in past_actions for alt in alt_actions) if alt_found: events.append("CORRECT_REMEDIATION") events.append("TERMINAL_CORRECT") else: events.append("INCORRECT_REMEDIATION") return True, events elif action == ActionType.CLOSE_WITHOUT_FIX: events.append("PREMATURE_CLOSURE") return True, events # Check step budget max_steps = MAX_STEPS[state.task_level] if state.observation.step_count >= max_steps: # Budget exceeded remaining_steps = max_steps - state.observation.step_count if remaining_steps < 0: for _ in range(abs(remaining_steps)): events.append("BUDGET_EXCEEDED_PER_STEP") return True, events return False, events def _check_correctness(self, state: EnvState, action: ActionType) -> bool: """Check if episode was resolved correctly.""" if action == ActionType.RESOLVE_TICKET: correct_action = state.hidden.correct_remediation alt_actions = state.hidden.correct_remediation_alts past_actions = [a.action_name for a in state.observation.action_history] return ( correct_action in past_actions or any(alt in past_actions for alt in alt_actions) ) elif action == ActionType.ESCALATE_TICKET: # Escalation with sufficient evidence is considered correct return state.diagnostic_steps_taken >= 3 else: return False