"""OpenEnv parliamentary environment for the Nation Simulator.""" from __future__ import annotations from typing import Any, Optional from openenv.core.env_server import Environment from core.config import GameConfig from core.game import NationGame from schemas.phases import Phase, valid_action_types_for_phase from server.models import ( ParliamentaryAction, ParliamentaryObservation, NationAction, NationObservation, NationState, EventModel, ProposalModel, VoteModel, OwnDepartmentModel, ) import numpy as np # Maximum retry attempts for rejected proposals before fallback MAX_PROPOSAL_RETRIES = 2 DEFAULT_ENV_NAME = "nation_optimizer_rl" class NationEnvironment(Environment): """ OpenEnv wrapper exposing the full 9-phase parliamentary cycle. Each call to step() processes ONE agent action in the current phase. System-only phases (5-9) are auto-advanced after all agent actions in a phase are complete. Retry loop: after voting, if any proposals were rejected, the environment loops back to Phase 3 for rejected departments (up to MAX_PROPOSAL_RETRIES times). After retries are exhausted, rejected departments receive baseline demand as a fallback. """ def __init__(self, seed: int | None = None): super().__init__() self.game = NationGame(seed=seed) self.departments = list(self.game.config.SECTOR_ORDER) self._retry_count = 0 self._round_reward = 0.0 @property def state(self) -> NationState: """Returns the current internal state of the environment.""" return NationState( step_count=self.game.round, raw_game_state=self.game.state(), ) def reset( self, seed: Optional[int] = None, **kwargs ) -> tuple[ParliamentaryObservation, dict]: """ Reset the environment. Phase 1 (event revelation) runs automatically. Returns observation ready for Phase 2 (debate). """ super().reset(seed=seed) if seed is not None: self.game = NationGame(seed=seed) else: self.game.reset() self._retry_count = 0 self._round_reward = 0.0 # Phase 1 already ran inside game.reset() → _start_round() # Engine starts in Phase 1. We need to advance to Phase 2 for debate. self.game.force_advance_phase() obs = self._build_observation(agent_id=self.departments[0]) return obs, {} def step( self, action: ParliamentaryAction ) -> tuple[ParliamentaryObservation, float, bool, bool, dict]: """ Process one agent action in the current phase. """ # Special case: FINISH_DEBATE or empty DEBATE message can be used to signal "pass/finish" if action.type == "FINISH_DEBATE" or (action.type == "DEBATE" and not (action.message or "").strip()): self.game.force_advance_phase() next_agent = self._determine_next_agent() obs = self._build_observation(agent_id=next_agent) return obs, 0.0, False, False, {"action": "debate_finish"} # Convert to engine dict and step action_dict = action.to_engine_dict() result = self.game.step(action_dict) # Check if the episode ended during this step if self.game.done: self._round_reward = result.reward.total obs = self._build_observation(agent_id=action.agent_id) return obs, self._round_reward, True, False, self._build_info(result) # Handle phase transitions if self.game.phase == Phase.BUDGET_EXECUTION: rejected = self._get_rejected_departments() if rejected and self._retry_count < MAX_PROPOSAL_RETRIES: self._retry_count += 1 self.game.reopen_proposal_phase(rejected) obs = self._build_observation( agent_id=rejected[0], rejected_departments=rejected, ) return ( obs, 0.0, False, False, { "retry": True, "retry_count": self._retry_count, "rejected_departments": rejected, }, ) elif rejected and self._retry_count >= MAX_PROPOSAL_RETRIES: self.game.apply_fallback_allocations(rejected) return self._run_system_phases(action.agent_id) # Still in an agent phase — return observation for next action next_agent = self._determine_next_agent() obs = self._build_observation(agent_id=next_agent) return obs, 0.0, False, False, self._build_info(result) def _run_system_phases( self, last_agent_id: str ) -> tuple[ParliamentaryObservation, float, bool, bool, dict]: """ Auto-advance through system phases 5-9. Returns the final observation + reward for this round. """ # Step through system phases until we reach next round's agent phase or done result = None while not self.game.done: phase = self.game.phase # If in VOTING phase but no proposals, advance if phase == Phase.VOTING: pending = [p for p in self.game.proposals if p.status == "pending"] if not pending: result = self.game.step(None) continue if phase in ( Phase.BUDGET_EXECUTION, Phase.CONSUMPTION_AND_EVENT_IMPACT, Phase.REVENUE_CALCULATION, Phase.SURPLUS_ROLLOVER, Phase.TERMINATION_CHECK, ): result = self.game.step(None) continue # Not a system phase or done break # Capture the round reward if available if result: self._round_reward = result.reward.total if self.game.done: obs = self._build_observation(agent_id=last_agent_id) return obs, self._round_reward, True, False, self._build_info(result) # New round started — Phase 1 runs automatically in game.step # Advance to Phase 2 for agents if self.game.phase == Phase.EVENT_REVELATION: self.game.phase = Phase.DEBATE self._retry_count = 0 next_agent = self._determine_next_agent() obs = self._build_observation(agent_id=next_agent) return obs, self._round_reward, False, False, self._build_info(result) def _get_rejected_departments(self) -> list[str]: """Return departments whose proposals were rejected in voting.""" rejected = [] proposed_depts = set() for p in self.game.proposals: proposed_depts.add(p.department) if p.status == "rejected" or p.status == "rejected_invalid": rejected.append(p.department) # Also include departments that never proposed for dept in self.departments: if dept not in proposed_depts: # Department never proposed — they need a proposal if dept not in rejected: rejected.append(dept) return rejected def _determine_next_agent(self) -> str: """Determine which agent should act next based on current phase.""" phase = self.game.phase if phase == Phase.DEBATE: # Cycle through all departments during debate, but cap at 20 messages n = len(self.departments) msg_count = len(self.game.debate_messages) if msg_count >= 18: # Force transition to Proposal phase self.game.force_advance_phase() return self._get_proposal_order()[0] next_idx = msg_count % n return self.departments[next_idx] if phase == Phase.PROPOSAL: # Return first department that hasn't submitted yet for dept in self._get_proposal_order(): if dept not in self.game._submitted_departments: return dept return self.departments[0] if phase == Phase.VOTING: # Return first agent who has a pending vote on the current proposal pending = [p for p in self.game.proposals if p.status == "pending"] for target in pending: # Everyone except the proposer must vote required_votes = len(self.departments) - 1 if len(target.votes) < required_votes: # Find someone who hasn't voted yet and is NOT the proposer proposer = target.agent_id for dept in self.departments: if dept == proposer: continue if dept not in target.votes: return dept # If we reach here, either no proposals are pending or all are fully voted # Fallback to the first department to avoid returning None return self.departments[0] return self.departments[0] def _get_proposal_order(self) -> list[str]: """Get the rotating proposal order for the current round.""" n = len(self.departments) start_idx = (self.game.round - 1) % n return self.departments[start_idx:] + self.departments[:start_idx] def _build_observation( self, agent_id: str, rejected_departments: list[str] | None = None, ) -> ParliamentaryObservation: """Build a spec-compliant observation for a specific agent.""" gs = self.game.state() # Current events current_events = [ EventModel( name=e.get("name", ""), severity=e.get("severity", 0), category=e.get("category", ""), narrative=e.get("narrative", ""), affected_departments=_affected_departments(e), round=e.get("round"), cost=e.get("cost"), ) for e in gs.get("current_events", []) ] # Proposals proposals = [ ProposalModel( proposal_id=p.get("proposal_id", ""), agent_id=p.get("agent_id", ""), department=p.get("department", ""), amount=p.get("amount", 0.0), justification=p.get("justification", ""), status=p.get("status", "pending"), votes=p.get("votes", {}), rejection_reason=p.get("rejection_reason"), ) for p in gs.get("proposals", []) ] # Votes votes = [ VoteModel( proposal_id=v.get("proposal_id", ""), agent_id=v.get("agent_id", ""), vote=v.get("vote", ""), ) for v in gs.get("votes", []) ] # Own department private info own_dept = None sectors = gs.get("sectors", {}) if agent_id in sectors: s = sectors[agent_id] own_dept = OwnDepartmentModel( name=agent_id, allocated_budget=s.get("allocation"), consumption=s.get("consumption"), surplus=s.get("surplus"), efficiency_rating=s.get("revenue_factor"), treasury_surplus_returned_this_round=s.get("surplus"), baseline=s.get("baseline"), ) # Valid actions for current phase phase_int = gs.get("phase") if phase_int is not None: valid = list(valid_action_types_for_phase(phase_int)) else: valid = [] # Determine target_proposal_id for voting phase target_proposal_id = None if phase_int is not None and Phase(phase_int) == Phase.VOTING: # We must focus on the same proposal that _determine_next_agent uses pending = [p for p in proposals if p.status == "pending"] for target in pending: required_votes = len(self.departments) - 1 if len(target.votes) < required_votes: target_proposal_id = target.proposal_id break return ParliamentaryObservation( round=gs.get("round", 0), phase=int(gs.get("phase", 1)), phase_name=gs.get("phase_name", ""), year=gs.get("year", 1), quarter=gs.get("quarter", 1), treasury=gs.get("treasury", 0.0), population=gs.get("population", 0), productivity=gs.get("productivity", 1.0), event_ledger=gs.get("event_ledger", []), current_events=current_events, proposals=proposals, votes=votes, debate_messages=gs.get("debate_messages", []), own_department=own_dept, valid_actions=valid, target_proposal_id=target_proposal_id, termination=gs.get("termination", {}), current_agent=agent_id, retry_count=self._retry_count, rejected_departments=rejected_departments or [], # OpenEnv base fields reward=0.0, done=self.game.done, ) def _build_info(self, result: Any) -> dict[str, Any]: """Build the info dict from a StepResult.""" return { "round": result.round_num, "termination_reason": result.termination_reason, "retry_count": self._retry_count, } class NationOpenEnv(Environment): """Thin OpenEnv wrapper around NationGame for whole-game smoke clients.""" SUPPORTS_CONCURRENT_SESSIONS = True def __init__(self, config: GameConfig | None = None, seed: int | None = None) -> None: super().__init__() self._config = config or GameConfig.from_json() self._game = NationGame(config=self._config, seed=seed) self._step_count = 0 def reset( self, seed: int | None = None, episode_id: str | None = None, **kwargs: Any, ) -> NationObservation: del episode_id, kwargs self._game = NationGame(config=self._config, seed=seed) self._step_count = 0 return self._observation(info={"reset": True}) def step( self, action: NationAction, timeout_s: float | None = None, **kwargs: Any, ) -> NationObservation: del timeout_s, kwargs self._step_count += 1 result = self._game.step(action.to_core_action()) return self._observation( state=result.observation, reward=float(result.reward.total), done=result.done, info=result.info, ) @property def state(self) -> NationState: state = self._game.state() return NationState( step_count=self._step_count, raw_game_state=state, core_state=state, ) def _observation( self, *, state: dict[str, Any] | None = None, reward: float | None = None, done: bool | None = None, info: dict[str, Any] | None = None, ) -> NationObservation: state = state or self._game.state() return NationObservation( done=self._game.done if done is None else done, reward=float(self._game.last_reward.total) if reward is None else reward, state=state, info=info or {}, metadata={ "env_name": DEFAULT_ENV_NAME, "step_count": self._step_count, }, ) def _affected_departments(event: dict[str, Any]) -> list[str]: affected = event.get("affected_departments") or event.get("affected_sectors") or [] if isinstance(affected, dict): return list(affected) return list(affected)