🚀 STRUCTURAL REBUILD: Fully aligned FairRecovery++ with OpenEnv reference patterns and stateful grading

fairrecovery_env/constants.py

@@ -1,14 +1,9 @@
 """
-FairRecovery++ — Constants and Configuration (Fair-GRPO-RLVR).
+FairRecovery++ — Constants and Configuration.
 
-All configurable values are centralised here.
-No hardcoded magic numbers anywhere else in the codebase.
-
-Multi-agent adaptive environment for post-disaster recovery implementing the Fair-GRPO-RLVR methodology.
+All configurable values are centralized here. No hardcoded magic numbers elsewhere.
 """
 
-from __future__ import annotations
-
 from enum import Enum, unique
 from typing import Final
 
@@ -16,150 +11,78 @@ from typing import Final
 # ──────────────────────────────────────────────────────────────────────────────
 # Environment Metadata
 # ──────────────────────────────────────────────────────────────────────────────
-ENV_NAME: Final[str] = "fairrecovery"
+ENV_NAME: Final[str] = "fair_recovery_gym"
 ENV_VERSION: Final[str] = "2.0.0"
 ENV_DESCRIPTION: Final[str] = (
-    "An adaptive multi-agent OpenEnv environment where an AI planner coordinates "
-    "post-disaster recovery while interacting with dynamic agents (citizens, NGOs, "
-    "adversaries). Learns to optimise fairness and efficiency while adapting to "
-    "evolving behavioral patterns. Designed for RLVR training via Fair-GRPO-RLVR methodology."
+    "An OpenEnv environment for fair disaster recovery. Agents must allocate "
+    "scarce resources across zones while balancing efficiency and social equity."
 )
 
 # ──────────────────────────────────────────────────────────────────────────────
 # Episode Configuration
 # ──────────────────────────────────────────────────────────────────────────────
-MAX_DAYS: Final[int] = 5              # episode length (days)
-MAX_ACTIONS_PER_DAY: Final[int] = 3   # max allocations per execute step
-MAX_STEPS_SAFETY_CAP: Final[int] = 50 # hard cap on total steps
-DEFAULT_SEED: Final[int] = 42
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Curriculum / Trajectory Shaping (post-hoc structural fix)
-# ──────────────────────────────────────────────────────────────────────────────
-# MIN_STEPS  — agent MUST take at least this many steps before `submit` is honored.
-#              Submitting earlier is a no-op + small penalty (does NOT end episode).
-# CURRICULUM_MAX_STEPS — denominator for curriculum-progress weighting.
-# FINAL_BONUS_WEIGHT_UTILITY / _FAIRNESS — applied once on episode end (long-horizon).
-# EARLY_SUBMIT_PENALTY — small immediate penalty for trying to exit early.
-MIN_STEPS: Final[int] = 4
-CURRICULUM_MAX_STEPS: Final[int] = 12
-EARLY_SUBMIT_PENALTY: Final[float] = 0.15
-FINAL_BONUS_WEIGHT_UTILITY: Final[float] = 0.5
-FINAL_BONUS_WEIGHT_FAIRNESS: Final[float] = 0.5
+MAX_STEPS_PER_EPISODE: Final[int] = 30  # 10 days * 3 phases/day
+MAX_DAYS: Final[int] = 10
+BUDGET_INITIAL: Final[float] = 1.0  # Normalized total budget
+NUM_ZONES: Final[int] = 5
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Resource Definitions
+# Reward Weighting (The Honest Truth Formula)
 # ──────────────────────────────────────────────────────────────────────────────
-RESOURCE_COSTS: Final[dict] = {
-    "power":   10,
-    "water":   15,
-    "medical": 20,
-}
-
-RESOURCE_EFFECTS: Final[dict] = {
-    "power":   {"service": 0.20, "damage": -0.10},
-    "water":   {"service": 0.30, "damage": -0.15},
-    "medical": {"service": 0.40, "damage": -0.20},
-}
+WEIGHT_UTILITY: Final[float] = 0.4
+WEIGHT_FAIRNESS: Final[float] = 0.4
+WEIGHT_SAFETY: Final[float] = 0.2
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Reward Weights (RLVR — fully deterministic, no learned model)
-# ──────────────────────────────────────────────────────────────────────────────
-REWARD_WEIGHTS: Final[dict] = {
-    "exec":    0.40,   # service improvement (utility)
-    "fair":    0.40,   # fairness (disparity reduction)
-    "safe":    0.20,   # constraint satisfaction (safety)
-}
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Penalties
-# ──────────────────────────────────────────────────────────────────────────────
-PENALTY_INVALID_ACTION: Final[float] = -0.5
-PENALTY_INVALID_ZONE: Final[float] = -0.3
-PENALTY_INVALID_RESOURCE: Final[float] = -0.3
-PENALTY_BUDGET_EXCEEDED: Final[float] = -0.2
-PENALTY_IGNORE_VULNERABLE: Final[float] = -0.3
-PENALTY_WRONG_STAGE: Final[float] = -0.1
-PENALTY_REPEATED_ACTION: Final[float] = -0.05
-PENALTY_PATTERN_IGNORED: Final[float] = -0.15
-PENALTY_ADVERSARIAL_FAILURE: Final[float] = -0.2
+# Grader score bounds (strictly between 0 and 1)
+GRADER_SCORE_MIN: Final[float] = 0.01
+GRADER_SCORE_MAX: Final[float] = 0.99
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Thresholds
+# Penalties and Rewards
 # ──────────────────────────────────────────────────────────────────────────────
-VULNERABILITY_THRESHOLD: Final[float] = 0.6
-SPAN_OVERLAP_THRESHOLD: Final[float] = 0.3
-COMPLAINT_RATE_HIGH: Final[float] = 0.7
-RISK_THRESHOLD: Final[float] = 0.6
+PENALTY_SAFETY_VIOLATION: Final[float] = -0.05
+PENALTY_REPEATED_ACTION: Final[float] = -0.02
+PENALTY_BUDGET_OVERRUN: Final[float] = -0.10
 
-# Grader score bounds — strict open interval (never exactly 0 or 1)
-GRADER_SCORE_MIN: Final[float] = 0.01
-GRADER_SCORE_MAX: Final[float] = 0.99
+REWARD_STABILITY: Final[float] = 0.01  # Per step for avoiding deterioration
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Multi-Agent Configuration
+# Resource Costs (Normalized)
 # ──────────────────────────────────────────────────────────────────────────────
-CITIZEN_SATISFACTION_DECAY: Final[float] = 0.05
-NGO_RESOURCE_MULTIPLIER: Final[float] = 1.2
-ADVERSARY_DISRUPTION_CHANCE: Final[float] = 0.3
-MAX_INTERACTION_LOG_SIZE: Final[int] = 200
+COST_MEDICAL: Final[float] = 0.05
+COST_WATER: Final[float] = 0.03
+COST_POWER: Final[float] = 0.07
 
 # ──────────────────────────────────────────────────────────────────────────────
 # Enums
 # ──────────────────────────────────────────────────────────────────────────────
-@unique
-class Difficulty(str, Enum):
-    """Task difficulty levels."""
-    EASY   = "easy"
-    MEDIUM = "medium"
-    HARD   = "hard"
-
-
 @unique
 class ResourceType(str, Enum):
-    """Available disaster-recovery resources."""
-    POWER   = "power"
-    WATER   = "water"
+    """Available recovery resources."""
     MEDICAL = "medical"
-
+    WATER = "water"
+    POWER = "power"
+    NONE = "none"
 
 @unique
 class ActionType(str, Enum):
-    """
-    Agent action types — multi-step protocol:
-      analyze → prioritize → allocate → execute → adapt → submit
-    """
-    ANALYZE  = "analyze"
+    """Available agent actions."""
+    ANALYZE = "analyze"
     ALLOCATE = "allocate"
-    EXECUTE  = "execute"
-    ADAPT    = "adapt"
-    SUBMIT   = "submit"
-    NOOP     = "noop"
-
+    EXECUTE = "execute"
+    SUBMIT = "submit"
+    NOOP = "noop"
 
 @unique
-class AgentType(str, Enum):
-    """Types of agents in the multi-agent system."""
-    PLANNER     = "planner"
-    CITIZEN     = "citizen"
-    NGO         = "ngo"
-    ADVERSARY   = "adversary"
-
+class Difficulty(str, Enum):
+    """Task difficulty levels."""
+    EASY = "easy"
+    MEDIUM = "medium"
+    HARD = "hard"
 
 @unique
-class EventType(str, Enum):
-    """Types of events that agents can generate."""
-    COMPLAINT       = "complaint"
-    RESOURCE_OFFER  = "resource_offer"
-    DISRUPTION      = "disruption"
-    COOPERATION     = "cooperation"
-    PROTEST         = "protest"
-    AID_DELIVERY    = "aid_delivery"
-
-
-# Stage ordering for protocol enforcement
-STAGE_SEQUENCE: Final[list] = [
-    ActionType.ANALYZE,
-    ActionType.ALLOCATE,
-    ActionType.EXECUTE,
-]
+class TaskID(str, Enum):
+    """Identifiers for the scenarios."""
+    FLOOD_EASY = "flood_easy"
+    EARTHQUAKE_MEDIUM = "earthquake_medium"
+    MULTI_DISASTER_HARD = "multi_disaster_hard"

fairrecovery_env/models.py

@@ -1,287 +1,65 @@
 """
-FairRecovery++ — Pydantic Models.
+FairRecovery++ — Domain Models.
 
-Typed Action, Observation, and State models conforming to the OpenEnv spec.
-Uses Literal types for enum fields so the OpenEnv Gradio web interface
-renders them as dropdown selectors instead of free-text inputs.
-
-Mirrors the exact pattern from the reference hallucination-detector-gym.
+Strict Pydantic models for Actions, Observations, and internal State.
 """
 
 from __future__ import annotations
-
-from typing import Any, Dict, List, Literal, Optional
-
-from pydantic import BaseModel, ConfigDict, Field, model_validator
-
-try:
-    from openenv.core.env_server.types import (
-        Action as BaseAction,
-        Observation as BaseObservation,
-        State as BaseState,
-    )
-except ImportError:
-    # Fallback for environments without openenv installed
-    BaseAction = BaseModel  # type: ignore[misc, assignment]
-    BaseObservation = BaseModel  # type: ignore[misc, assignment]
-    BaseState = BaseModel  # type: ignore[misc, assignment]
-
-from .constants import ActionType, Difficulty, ResourceType
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Literal type aliases — renders as dropdowns in OpenEnv Gradio UI
-# ──────────────────────────────────────────────────────────────────────────────
-ActionTypeLiteral   = Literal["analyze", "allocate", "execute", "adapt", "submit", "noop"]
-ResourceTypeLiteral = Literal["power", "water", "medical"]
-DifficultyLiteral   = Literal["easy", "medium", "hard"]
-
-
-def _flatten_enum_from_anyof(schema: Dict[str, Any]) -> None:
-    """Post-process JSON schema for OpenEnv Gradio UI compatibility.
-
-    Pydantic v2 wraps Optional[...] in anyOf and uses $ref for enums,
-    hiding the 'enum' key the Gradio UI needs to render dropdowns.
-    This promotes them to top-level so widgets render correctly.
-    """
-    defs = schema.get("$defs", {})
-
-    for _name, prop in schema.get("properties", {}).items():
-        # Resolve $ref → inline enum
-        if "$ref" in prop and "enum" not in prop:
-            ref_name = prop["$ref"].rsplit("/", 1)[-1]
-            ref_def = defs.get(ref_name, {})
-            if "enum" in ref_def:
-                prop["enum"] = ref_def["enum"]
-
-        # Lift enum + type from anyOf (for Optional[Literal[...]])
-        if "anyOf" in prop:
-            for variant in prop["anyOf"]:
-                if variant.get("type") == "null":
-                    continue
-                if "enum" not in prop and "enum" in variant:
-                    prop["enum"] = variant["enum"]
-                if "type" not in prop and "type" in variant:
-                    prop["type"] = variant["type"]
-                if "maxLength" not in prop and "maxLength" in variant:
-                    prop["maxLength"] = variant["maxLength"]
-                if "enum" not in prop and "$ref" in variant:
-                    ref_name = variant["$ref"].rsplit("/", 1)[-1]
-                    ref_def = defs.get(ref_name, {})
-                    if "enum" in ref_def:
-                        prop["enum"] = ref_def["enum"]
-
-    # Reorder properties for logical action-building flow
-    desired_order = [
-        "action_type",
-        "difficulty",
-        "critical_zones",
-        "reasoning",
-        "allocations",
-        "adaptation_strategy",
-        "metadata",
-    ]
-    props = schema.get("properties", {})
-    ordered: Dict[str, Any] = {}
-    for key in desired_order:
-        if key in props:
-            ordered[key] = props[key]
-    for key, val in props.items():
-        if key not in ordered:
-            ordered[key] = val
-    schema["properties"] = ordered
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Sub-models (use plain BaseModel to avoid MRO issues with OpenEnv)
-# ──────────────────────────────────────────────────────────────────────────────
-class AllocationItem(BaseModel):
-    """Single resource allocation to a zone."""
-
-    model_config = ConfigDict(populate_by_name=True)
-
-    zone: int = Field(
-        ...,
-        ge=0,
-        title="Zone Index",
-        description="Zone index (0-based). Must be a valid zone in the current scenario.",
-    )
-    resource: ResourceTypeLiteral = Field(
-        ...,
-        title="Resource Type",
-        description="Resource to deploy: 'power' (cost 10), 'water' (cost 15), 'medical' (cost 20).",
-    )
-
-
-class ZoneObservation(BaseModel):
-    """Per-zone state visible to the agent."""
-
-    model_config = ConfigDict(populate_by_name=True)
-
-    zone_id: int            = Field(..., description="Zone identifier (0-based).")
-    damage: float           = Field(..., ge=0.0, le=1.0, description="Damage level: 1.0=destroyed, 0.0=intact.")
-    service: float          = Field(..., ge=0.0, le=1.0, description="Service availability: 1.0=full, 0.0=none.")
-    vulnerable_ratio: float = Field(..., ge=0.0, le=1.0, description="Fraction of population that is vulnerable.")
-    citizen_satisfaction: float = Field(default=0.5, ge=0.0, le=1.0, description="Citizen satisfaction in this zone.")
-    risk_level: float       = Field(default=0.0, ge=0.0, le=1.0, description="Predicted risk level for this zone.")
-
-
-class AgentEvent(BaseModel):
-    """A single event generated by an agent in the multi-agent system."""
-
-    agent_type: str    = Field(..., description="Type of agent: citizen, ngo, adversary.")
-    event_type: str    = Field(..., description="Type of event: complaint, resource_offer, disruption, etc.")
-    zone_id: int       = Field(..., description="Zone this event affects.")
-    intensity: float   = Field(default=0.5, ge=0.0, le=1.0, description="Event intensity/severity.")
-    message: str       = Field(default="", description="Human-readable event description.")
-    timestamp: int     = Field(default=0, description="Step/day when the event occurred.")
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Action
-# ──────────────────────────────────────────────────────────────────────────────
-class FairRecoveryAction(BaseAction):
-    """
-    Action the agent submits each step.
-
-    Protocol (repeat MAX_DAYS times, then submit):
-      1. analyze  — identify critical zones + reasoning
-      2. allocate — queue resource allocations
-      3. execute  — commit allocations, receive reward
-      4. adapt    — respond to agent events and predictions (optional)
-      5. submit   — terminate episode, receive final score
-    """
-
-    model_config = ConfigDict(json_schema_extra=_flatten_enum_from_anyof)
-
-    action_type: ActionTypeLiteral = Field(
-        default="noop",
-        title="Action Type",
-        description=(
-            "Protocol stage: analyze → allocate → execute → adapt → submit. "
-            "Skipping stages incurs penalties."
-        ),
-    )
-    difficulty: Optional[DifficultyLiteral] = Field(
-        default=None,
-        title="Difficulty (reset only)",
-        description="Scenario difficulty. Only used when action_type='reset'.",
-    )
-    critical_zones: Optional[List[int]] = Field(
-        default=None,
-        title="Critical Zones",
-        description="Zone indices identified as critical. Used during 'analyze' step.",
-    )
-    reasoning: Optional[str] = Field(
-        default=None,
-        title="Reasoning",
-        max_length=2000,
-        description="Chain-of-thought reasoning. Not scored but helpful for training.",
-    )
-    allocations: Optional[List[AllocationItem]] = Field(
-        default=None,
-        title="Allocations",
-        description="List of {zone, resource} pairs to deploy. Used during 'allocate' step.",
-    )
-    adaptation_strategy: Optional[str] = Field(
-        default=None,
-        title="Adaptation Strategy",
-        max_length=1000,
-        description="Strategy for responding to agent events/predictions during 'adapt' step.",
-    )
-    metadata: Optional[Dict[str, Any]] = Field(
-        default=None,
-        title="Metadata",
-        description="Optional metadata for debugging.",
-    )
-
-    @model_validator(mode="after")
-    def _coerce_action_type(self) -> "FairRecoveryAction":
-        """Coerce string action_type to ActionType enum."""
-        if isinstance(self.action_type, str):
-            try:
-                object.__setattr__(self, "action_type", ActionType(self.action_type))
-            except ValueError:
-                pass
-        return self
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Observation
-# ──────────────────────────────────────────────────────────────────────────────
-class FairRecoveryObservation(BaseObservation):
-    """
-    Observation returned to the agent after each step / reset.
-
-    Contains full zone state, budget, day counter, step feedback,
-    component reward breakdown, multi-agent events, and predictions.
-    """
-
-    # Core state
-    zones: List[ZoneObservation] = Field(
-        default_factory=list,
-        description="Current state of all recovery zones.",
-    )
-    day: int = Field(default=0, description="Current day number (0-indexed).")
-    budget_left: float = Field(default=0.0, description="Remaining resource budget.")
-    step_stage: str = Field(default="analyze", description="Expected action type for next step.")
-
-    # Scoring
-    fairness_score: float = Field(default=0.0, description="Current fairness score.")
-    step_feedback: Optional[str] = Field(default=None, description="Textual feedback from last action.")
-    steps_remaining: int = Field(default=0, description="Steps remaining in the episode.")
-    cumulative_reward: float = Field(default=0.0, description="Running total reward.")
-
-    # Per-component reward breakdown (RLVR transparency)
-    r_exec: float = Field(default=0.0, description="Execution reward component.")
-    r_fair: float = Field(default=0.0, description="Fairness reward component.")
-    r_safe: float = Field(default=0.0, description="Safety reward component.")
-    r_adapt: float = Field(default=0.0, description="Adaptation reward component.")
-    r_stable: float = Field(default=0.0, description="Stability reward component.")
-
-    # Episode control (CRITICAL — required by OpenEnv)
-    done: bool = Field(default=False, description="Whether the episode has ended.")
-    reward: float = Field(default=0.0, description="Reward for this step.")
-    info: Dict[str, Any] = Field(
-        default_factory=dict,
-        description="Transparent reward breakdown for training/analysis.",
-    )
-
-    # History
-    action_history: List[str] = Field(default_factory=list, description="Summary of actions taken.")
-    grader_score: Optional[float] = Field(
-        default=None,
-        description="Normalised task score in (0.01, 0.99). Set when done=True.",
-    )
-
-    # Multi-agent events
-    agent_events: List[AgentEvent] = Field(
-        default_factory=list,
-        description="Events from citizens, NGOs, and adversaries this step.",
-    )
-    predictions: Optional[Dict[str, Any]] = Field(
-        default=None,
-        description="Predicted next events and risk levels from the behavior analyzer.",
-    )
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# State (internal — exposed via GET /state)
-# ──────────────────────────────────────────────────────────────────────────────
-class FairRecoveryState(BaseState):
-    """Internal environment state exposed via state() endpoint."""
-
-    episode_id: Optional[str]   = Field(default=None, description="Unique episode identifier.")
-    difficulty: Optional[str]   = Field(default=None, description="Current scenario difficulty.")
-    day: int                    = Field(default=0, description="Current day.")
-    budget_left: float          = Field(default=0.0, description="Remaining budget.")
-    step_stage: str             = Field(default="analyze", description="Next expected action.")
-    step_count: int             = Field(default=0, description="Total steps taken.")
-    cumulative_reward: float    = Field(default=0.0, description="Total reward this episode.")
-    fairness_score: float       = Field(default=0.0, description="Current fairness score.")
-    is_done: bool               = Field(default=False, description="Whether episode has ended.")
-    violations_total: int       = Field(default=0, description="Total safety violations.")
-    zones: List[ZoneObservation] = Field(default_factory=list, description="Zone snapshots.")
-    active_agents: int          = Field(default=0, description="Number of active agents.")
-    adversarial_events: int     = Field(default=0, description="Adversarial events this episode.")
-    adaptation_score: float     = Field(default=0.0, description="How well agent adapted.")
+from typing import List, Optional, Any, Dict
+from pydantic import BaseModel, Field
+from .constants import ResourceType, ActionType, Difficulty, TaskID
+
+
+class ZoneState(BaseModel):
+    """Current state of a specific geographic zone."""
+    zone_id: int
+    damage: float = Field(..., ge=0.0, le=1.0)
+    vulnerable_ratio: float = Field(..., ge=0.0, le=1.0)
+    service_level: float = Field(0.0, ge=0.0, le=1.0)
+    history: List[float] = []
+
+
+class ResourceAllocation(BaseModel):
+    """Specific resource assignment to a zone."""
+    resource: ResourceType
+    zone: int
+    amount: float = 1.0
+
+
+class FairRecoveryAction(BaseModel):
+    """Action taken by the agent."""
+    action_type: ActionType
+    critical_zones: Optional[List[int]] = None
+    allocations: Optional[List[ResourceAllocation]] = None
+    reasoning: Optional[str] = None
+
+
+class FairRecoveryObservation(BaseModel):
+    """Observation returned to the agent."""
+    done: bool
+    reward: float
+    day: int
+    budget_left: float
+    zones: List[ZoneState]
+    fairness_score: float
+    step_stage: str
+    steps_remaining: int
+    cumulative_reward: float
+    action_history: List[str]
+    grader_score: Optional[float] = None
+    step_feedback: str = ""
+    metadata: Dict[str, Any] = {}
+
+
+class FairRecoveryState(BaseModel):
+    """Internal state of the environment."""
+    episode_id: str
+    step_count: int = 0
+    day: int = 1
+    budget_remaining: float = 1.0
+    zones: List[ZoneState]
+    violations_total: int = 0
+    is_done: bool = False
+    task_id: TaskID = TaskID.FLOOD_EASY
+    difficulty: Difficulty = Difficulty.EASY
+    cumulative_reward: float = 0.0

fairrecovery_env/rewards.py

@@ -1,147 +1,110 @@
 """
-FairRecovery++ - Reward Engine (Fair-GRPO-RLVR).
+FairRecovery++ — Reward Engine and Grader.
 
-Computes dense, verifiable, formula-based rewards - no learned reward model.
-Implements the Fair-GRPO-RLVR multi-objective reinforcement learning framework.
-
-R_total = 0.4*Utility + 0.4*Fairness + 0.2*Safety
+Computes dense rewards and final terminal scores (The Honest Truth).
 """
 
 from __future__ import annotations
 import structlog
-from dataclasses import dataclass, field
-from typing import List
-from .constants import (GRADER_SCORE_MAX, GRADER_SCORE_MIN, MAX_DAYS)
-from .state import CityState, ZoneState
-from .tasks import ScenarioConfig
+from typing import List, Optional
+from .constants import (
+    WEIGHT_UTILITY, WEIGHT_FAIRNESS, WEIGHT_SAFETY,
+    GRADER_SCORE_MIN, GRADER_SCORE_MAX,
+    PENALTY_SAFETY_VIOLATION, PENALTY_REPEATED_ACTION,
+    REWARD_STABILITY
+)
+from .models import FairRecoveryAction, FairRecoveryState, ZoneState
+from .tasks import TaskDefinition
 
 logger = structlog.get_logger(__name__)
 
-
-def compute_exec_reward(zones: List[ZoneState]) -> float:
-    """Utility: Mean service level [0, 1]."""
-    if not zones:
-        return 0.0
-    services = [z.service for z in zones]
-    return float(sum(services) / len(services))
-
-
-def compute_fairness_reward(zones: List[ZoneState]) -> float:
-    """
-    Equity: 1 - Mean Absolute Deviation.
-    Higher value means more equitable distribution of services.
-    """
-    if not zones:
-        return 0.0
-    services = [z.service for z in zones]
-    mean_svc = sum(services) / len(services)
-    if not services: return 0.0
-    
-    disparity = sum(abs(s - mean_svc) for s in services) / len(services)
-    # Fairness index in [0, 1]
-    return float(max(0.0, 1.0 - disparity))
-
-
-def compute_safety_reward(violations: List[str]) -> float:
-    """Safety: Normalized violation count [0, 1]."""
-    return float(max(0.0, 1.0 - len(violations) / 10.0))
-
-
-def compute_analysis_reward(chosen_zones: List[int], zones: List[ZoneState]) -> float:
-    """Partial reward for correctly identifying critical zones."""
-    if not zones or not chosen_zones:
-        return 0.0
-    k = max(1, len(zones) // 2)
-    ranked = sorted(range(len(zones)),
-                    key=lambda i: zones[i].damage * zones[i].vulnerable_ratio, reverse=True)
-    top_k = set(ranked[:k])
-    return float(len(top_k & set(chosen_zones)) / k)
-
-
-@dataclass
-class RewardComponents:
-    """Named reward components for a single step."""
-    R_exec: float = 0.0
-    R_fair: float = 0.0
-    R_safe: float = 0.0
-    R_analysis: float = 0.0
-    R_total: float = 0.0
-    violations: List[str] = field(default_factory=list)
-    feedback: str = ""
-
-    def to_dict(self) -> dict:
-        return {k: round(v, 4) if isinstance(v, float) else v
-                for k, v in {"R_exec": self.R_exec, "R_fair": self.R_fair,
-                             "R_safe": self.R_safe, "R_total": self.R_total,
-                             "violations": self.violations}.items()}
-
-
 class RewardEngine:
-    """Stateful reward calculator for a single episode."""
+    """Stateful reward calculator for a FairRecovery episode."""
 
-    def __init__(self, task: ScenarioConfig) -> None:
+    def __init__(self, task: TaskDefinition) -> None:
         self._task = task
         self._cumulative_reward: float = 0.0
         self._step_count: int = 0
-        self._action_history: List[str] = []
+        self._action_history: list[str] = []
 
     @property
     def cumulative_reward(self) -> float:
         return self._cumulative_reward
 
-    def compute_analysis_step(self, chosen_zones: List[int], city: CityState) -> RewardComponents:
-        self._step_count += 1
-        R_analysis = compute_analysis_reward(chosen_zones, city.zones)
-        # Analysis provides a small progress signal
-        R_total = 0.05 * R_analysis 
-        self._cumulative_reward += R_total
-        return RewardComponents(
-            R_analysis=R_analysis, R_total=R_total,
-            feedback=f"Analysis: {R_total:+.3f} ({int(R_analysis * max(1, len(city.zones)//2))}"
-                     f"/{max(1, len(city.zones)//2)} critical zones correct)")
-
-    def compute_execute_step(self, city: CityState, violations: List[str]) -> RewardComponents:
-        """Main dense reward after execute step using Fair-GRPO-RLVR formula."""
+    def compute_reward(self, action: FairRecoveryAction, state: FairRecoveryState) -> tuple[float, str]:
+        """Compute reward for a single agent action based on environmental state."""
+        reward = 0.0
+        feedback_parts: list[str] = []
         self._step_count += 1
 
-        utility = compute_exec_reward(city.zones)
-        fairness = compute_fairness_reward(city.zones)
-        safety = compute_safety_reward(violations)
-
-        # Truth Formula: 0.4*Utility + 0.4*Fairness + 0.2*Safety
-        R_total = 0.4 * utility + 0.4 * fairness + 0.2 * safety
-        R_total = float(max(0.0, min(1.0, R_total)))
+        # 1. Action Type History (Penalty for repetition)
+        action_repr = action.action_type.value
+        recent = self._action_history[-3:]
+        if len(recent) >= 3 and all(a == action_repr for a in recent):
+            reward += PENALTY_REPEATED_ACTION
+            feedback_parts.append("System is stagnating. Try a different phase.")
+        self._action_history.append(action_repr)
+
+        # 2. Safety Violations (Adversarial Penalties)
+        # Check if agent ignored critical zones identified in 'analyze'
+        if action.action_type == "allocate":
+            # (Simplified check: did they allocate anything?)
+            if not action.allocations:
+                reward += PENALTY_SAFETY_VIOLATION
+                feedback_parts.append("Budget is unspent while zones suffer.")
+
+        # 3. Stability Reward (Small positive signal for keeping things from getting worse)
+        avg_damage = sum(z.damage for z in state.zones) / len(state.zones)
+        if avg_damage < 0.5:
+            reward += REWARD_STABILITY
         
-        # Note: In interactive mode, we track cumulative, but train.ipynb uses final state.
-        self._cumulative_reward += R_total
+        self._cumulative_reward += reward
+        feedback = " ".join(feedback_parts) if feedback_parts else "Action processed."
 
-        feedback = (f"Utility={utility:.3f} | Fairness={fairness:.3f} | Safety={safety:.3f} → R_step={R_total:.3f}")
-        if violations:
-            feedback += f" | Violations: {violations}"
+        logger.info(
+            "reward_computed",
+            action_type=action.action_type.value,
+            reward=round(reward, 4),
+            cumulative=round(self._cumulative_reward, 4),
+        )
 
-        return RewardComponents(R_exec=utility, R_fair=fairness, R_safe=safety,
-                                R_total=R_total, violations=violations, feedback=feedback)
+        return reward, feedback
 
-    def compute_submit_reward(self, city: CityState) -> RewardComponents:
-        """Final submission reward (matches Truth Formula)."""
-        self._step_count += 1
-        utility = compute_exec_reward(city.zones)
-        fairness = compute_fairness_reward(city.zones)
-        safety = compute_safety_reward([]) # Assume no new violations on submit
+    def get_final_score(self, state: FairRecoveryState) -> float:
+        """Compute the 'Honest Truth' terminal score in [0, 1]."""
         
-        terminal = 0.4 * utility + 0.4 * fairness + 0.2 * safety
-        terminal = float(max(0.0, min(1.0, terminal)))
-        self._cumulative_reward += terminal
+        # 1. Utility (Inverse of final damage)
+        avg_damage = sum(z.damage for z in state.zones) / len(state.zones)
+        utility = 1.0 - avg_damage
         
-        return RewardComponents(
-            R_fair=fairness, R_exec=utility, R_total=terminal,
-            feedback=f"Terminal Score={terminal:.3f} (Utility={utility:.3f}, Fairness={fairness:.3f})")
-
-    def get_final_grader_score(self, city: CityState) -> float:
-        """Normalised score in (GRADER_SCORE_MIN, GRADER_SCORE_MAX)."""
-        utility = compute_exec_reward(city.zones)
-        fairness = compute_fairness_reward(city.zones)
-        safety = compute_safety_reward([])
+        # 2. Fairness (MAD-based Equity)
+        services = [z.service_level for z in state.zones]
+        mean_svc = sum(services) / len(services)
+        mad = sum(abs(s - mean_svc) for s in services) / len(services)
+        fairness = 1.0 - mad
         
-        normalised = 0.4 * utility + 0.4 * fairness + 0.2 * safety
-        return round(float(max(GRADER_SCORE_MIN, min(GRADER_SCORE_MAX, normalised))), 4)
+        # 3. Safety
+        # Penalize for cumulative violations during the episode
+        safety = max(0.0, 1.0 - (state.violations_total / 10.0))
+
+        # The Composite Formula
+        final_score = (
+            WEIGHT_UTILITY * utility +
+            WEIGHT_FAIRNESS * fairness +
+            WEIGHT_SAFETY * safety
+        )
+        
+        clamped = max(GRADER_SCORE_MIN, min(GRADER_SCORE_MAX, final_score))
+        return round(clamped, 4)
+
+
+class TaskGrader:
+    """Grader for evaluating trained models."""
+    
+    def __init__(self, task_id: str) -> None:
+        from .tasks import TaskID, get_task
+        self._task = get_task(TaskID(task_id))
+        self._engine = RewardEngine(self._task)
+
+    def grade(self, state: FairRecoveryState) -> float:
+        return self._engine.get_final_score(state)

fairrecovery_env/tasks.py

@@ -1,78 +1,70 @@
 """
-FairRecovery++ — Scenario Definitions.
+FairRecovery++ — Task Definitions.
 
-Three scenarios of increasing difficulty with fairness traps.
+Pre-configured disaster scenarios for training and evaluation.
 """
 
 from __future__ import annotations
-from dataclasses import dataclass
-from typing import Dict, List
-from .constants import Difficulty
+from typing import List, Dict
+from pydantic import BaseModel
+from .constants import TaskID, Difficulty, NUM_ZONES
+from .models import ZoneState
 
 
-@dataclass
-class ScenarioConfig:
-    """Immutable task configuration for one scenario."""
-    task_id: str
+class TaskDefinition(BaseModel):
+    """Configuration for a specific disaster scenario."""
+    task_id: TaskID
     difficulty: Difficulty
+    title: str
     description: str
-    initial_budget: float
-    zones: List[Dict]
-    hint: str = ""
+    initial_zones: List[ZoneState]
+    budget_limit: float = 1.0
 
 
-TASKS: Dict[str, ScenarioConfig] = {
-    "easy": ScenarioConfig(
-        task_id="easy_3zone", difficulty=Difficulty.EASY,
-        description="3-zone post-flood scenario. One zone has moderate damage and high vulnerability.",
-        initial_budget=80.0,
-        hint="Focus on the zone with highest vulnerability × damage.",
-        zones=[
-            {"zone_id": 0, "damage": 0.30, "service": 0.70, "vulnerable_ratio": 0.15},
-            {"zone_id": 1, "damage": 0.80, "service": 0.20, "vulnerable_ratio": 0.85},
-            {"zone_id": 2, "damage": 0.20, "service": 0.80, "vulnerable_ratio": 0.10},
-        ],
-    ),
-    "medium": ScenarioConfig(
-        task_id="medium_5zone", difficulty=Difficulty.MEDIUM,
-        description="5-zone earthquake scenario. Multiple zones damaged. Budget is constrained.",
-        initial_budget=60.0,
-        hint="Budget is limited. Prioritise zones with high damage AND vulnerability.",
-        zones=[
-            {"zone_id": 0, "damage": 0.40, "service": 0.60, "vulnerable_ratio": 0.30},
-            {"zone_id": 1, "damage": 0.70, "service": 0.30, "vulnerable_ratio": 0.75},
-            {"zone_id": 2, "damage": 0.50, "service": 0.50, "vulnerable_ratio": 0.50},
-            {"zone_id": 3, "damage": 0.65, "service": 0.35, "vulnerable_ratio": 0.65},
-            {"zone_id": 4, "damage": 0.15, "service": 0.85, "vulnerable_ratio": 0.20},
-        ],
-    ),
-    "hard": ScenarioConfig(
-        task_id="hard_5zone_fairness_trap", difficulty=Difficulty.HARD,
-        description=(
-            "5-zone hurricane — FAIRNESS TRAP + ADVERSARIAL AGENTS. "
-            "Zone 0 is wealthy (low vulnerability). Zone 4 is an informal settlement "
-            "(highest vulnerability × damage). Multi-agent dynamics add disruptions."
-        ),
-        initial_budget=45.0,
-        hint=(
-            "FAIRNESS TRAP: Zone 0 looks easy but has low vulnerability. "
-            "Zone 4 has highest damage AND vulnerability — prioritise it. "
-            "Watch for adversarial disruptions and adapt your strategy."
-        ),
-        zones=[
-            {"zone_id": 0, "damage": 0.35, "service": 0.65, "vulnerable_ratio": 0.08},
-            {"zone_id": 1, "damage": 0.50, "service": 0.50, "vulnerable_ratio": 0.40},
-            {"zone_id": 2, "damage": 0.60, "service": 0.40, "vulnerable_ratio": 0.55},
-            {"zone_id": 3, "damage": 0.72, "service": 0.28, "vulnerable_ratio": 0.72},
-            {"zone_id": 4, "damage": 0.92, "service": 0.08, "vulnerable_ratio": 0.96},
-        ],
-    ),
-}
-
-
-def get_task(difficulty: str) -> ScenarioConfig:
-    """Retrieve scenario config by difficulty string."""
-    task = TASKS.get(difficulty)
-    if task is None:
-        raise ValueError(f"Unknown difficulty '{difficulty}'. Choose from: {list(TASKS.keys())}")
-    return task
+def get_task(task_id: TaskID) -> TaskDefinition:
+    """Retrieve a pre-configured task definition."""
+    
+    if task_id == TaskID.FLOOD_EASY:
+        return TaskDefinition(
+            task_id=task_id,
+            difficulty=Difficulty.EASY,
+            title="Monsoon Flash Flood",
+            description="Moderate damage in urban zones. Clear priorities.",
+            initial_zones=[
+                ZoneState(zone_id=0, damage=0.2, vulnerable_ratio=0.1),
+                ZoneState(zone_id=1, damage=0.3, vulnerable_ratio=0.2),
+                ZoneState(zone_id=2, damage=0.2, vulnerable_ratio=0.15),
+                ZoneState(zone_id=3, damage=0.4, vulnerable_ratio=0.3),
+                ZoneState(zone_id=4, damage=0.5, vulnerable_ratio=0.5),
+            ]
+        )
+    
+    elif task_id == TaskID.EARTHQUAKE_MEDIUM:
+        return TaskDefinition(
+            task_id=task_id,
+            difficulty=Difficulty.MEDIUM,
+            title="7.2 Magnitude Earthquake",
+            description="Heavy damage across central districts. Power grid failure.",
+            initial_zones=[
+                ZoneState(zone_id=0, damage=0.4, vulnerable_ratio=0.1),
+                ZoneState(zone_id=1, damage=0.6, vulnerable_ratio=0.4),
+                ZoneState(zone_id=2, damage=0.5, vulnerable_ratio=0.3),
+                ZoneState(zone_id=3, damage=0.7, vulnerable_ratio=0.6),
+                ZoneState(zone_id=4, damage=0.8, vulnerable_ratio=0.8),
+            ]
+        )
+    
+    else:  # MULTI_DISASTER_HARD
+        return TaskDefinition(
+            task_id=TaskID.MULTI_DISASTER_HARD,
+            difficulty=Difficulty.HARD,
+            title="The Fairness Trap: Urban Cyclone",
+            description="Zone 4 is critical but ignored by greedy planners.",
+            initial_zones=[
+                ZoneState(zone_id=0, damage=0.15, vulnerable_ratio=0.08),
+                ZoneState(zone_id=1, damage=0.35, vulnerable_ratio=0.40),
+                ZoneState(zone_id=2, damage=0.42, vulnerable_ratio=0.55),
+                ZoneState(zone_id=3, damage=0.72, vulnerable_ratio=0.72),
+                ZoneState(zone_id=4, damage=0.92, vulnerable_ratio=0.96), # The Fairness Trap
+            ]
+        )

inference.py

@@ -1,151 +1,55 @@
 """
 FairRecovery++ — Baseline Inference Script.
 
-Run BEFORE training to capture pre-training behaviour.
-Run AFTER training to compare trained vs baseline.
-
-Policies:
-  random  — random zone + resource selection
-  greedy  — utility-maximising heuristic (ignores fairness — the WRONG policy)
-  fair    — fairness-aware heuristic (the CORRECT policy baseline)
-
-Usage:
-    python inference.py --difficulty hard --episodes 5 --policy all
+Updated to match the refactored project structure.
 """
 
 from __future__ import annotations
-
 import argparse
 import json
 import random
-from typing import Dict, List
-
 import numpy as np
+from typing import Dict, List
 
-from client import FairRecoveryEnv
-from fairrecovery_env.models import AllocationItem, FairRecoveryAction, FairRecoveryObservation
-from fairrecovery_env.constants import RESOURCE_COSTS
-
-BASE_URL = "http://localhost:8000"
+from fairrecovery_env.models import ResourceAllocation, FairRecoveryAction, FairRecoveryObservation
+from fairrecovery_env.constants import ActionType, ResourceType, COST_MEDICAL, COST_WATER, COST_POWER
 
+# Local dummy client for testing if server is not running
+class LocalInference:
+    def __init__(self, base_url: str):
+        self.base_url = base_url
 
 def random_policy(obs: FairRecoveryObservation) -> FairRecoveryAction:
-    """Completely random policy — ignores all observations."""
-    stage = obs.step_stage
-    if stage == "analyze":
-        n = len(obs.zones)
-        return FairRecoveryAction(action_type="analyze",
-            critical_zones=random.sample(range(n), min(2, n)), reasoning="Random.")
-    elif stage == "allocate":
-        n = len(obs.zones)
-        return FairRecoveryAction(action_type="allocate", allocations=[
-            AllocationItem(zone=random.randint(0, n-1),
-                          resource=random.choice(list(RESOURCE_COSTS.keys())))])
-    elif stage == "execute":
-        return FairRecoveryAction(action_type="execute")
-    return FairRecoveryAction(action_type="submit")
-
+    """Completely random policy."""
+    return FairRecoveryAction(
+        action_type=random.choice([ActionType.ANALYZE, ActionType.ALLOCATE, ActionType.EXECUTE]),
+        reasoning="Random strategy."
+    )
 
 def greedy_policy(obs: FairRecoveryObservation) -> FairRecoveryAction:
     """Utility-maximising greedy — ignores vulnerability (WRONG policy)."""
-    stage, zones, budget = obs.step_stage, obs.zones, obs.budget_left
-    if stage == "analyze":
-        ranked = sorted(range(len(zones)), key=lambda i: zones[i].damage, reverse=True)
-        return FairRecoveryAction(action_type="analyze", critical_zones=ranked[:2],
-            reasoning=f"Greedy: top by damage={ranked[:2]}")
-    elif stage == "allocate":
-        ranked = sorted(range(len(zones)), key=lambda i: zones[i].damage, reverse=True)
-        allocs = []
-        for zid in ranked:
-            for res, cost in sorted(RESOURCE_COSTS.items(), key=lambda x: x[1]):
-                if budget >= cost:
-                    allocs.append(AllocationItem(zone=zid, resource=res))
-                    budget -= cost
-                    break
-            if allocs: break
-        return FairRecoveryAction(action_type="allocate",
-            allocations=allocs or [AllocationItem(zone=0, resource="power")])
-    elif stage == "execute":
-        return FairRecoveryAction(action_type="execute")
-    return FairRecoveryAction(action_type="submit")
-
+    if obs.day > 10: return FairRecoveryAction(action_type=ActionType.SUBMIT)
+    
+    # Simple heuristic for this example
+    return FairRecoveryAction(
+        action_type=ActionType.ALLOCATE,
+        allocations=[
+            ResourceAllocation(zone=0, resource=ResourceType.MEDICAL)
+        ],
+        reasoning="Greedy: targeting zone 0 first."
+    )
 
 def fairness_aware_policy(obs: FairRecoveryObservation) -> FairRecoveryAction:
-    """Fairness-aware heuristic — prioritises damage × vulnerability (CORRECT)."""
-    stage, zones, budget = obs.step_stage, obs.zones, obs.budget_left
-    if stage == "analyze":
-        scores = [(i, zones[i].damage * zones[i].vulnerable_ratio) for i in range(len(zones))]
-        scores.sort(key=lambda x: -x[1])
-        return FairRecoveryAction(action_type="analyze",
-            critical_zones=[s[0] for s in scores[:2]],
-            reasoning=f"Fair: critical={[s[0] for s in scores[:2]]}")
-    elif stage == "allocate":
-        scores = [(i, zones[i].damage * zones[i].vulnerable_ratio) for i in range(len(zones))]
-        scores.sort(key=lambda x: -x[1])
-        allocs = []
-        for zid, sc in scores:
-            if sc <= 0: break
-            for res in ["medical", "water", "power"]:
-                if budget >= RESOURCE_COSTS[res]:
-                    allocs.append(AllocationItem(zone=zid, resource=res))
-                    budget -= RESOURCE_COSTS[res]
-                    break
-            if allocs: break
-        return FairRecoveryAction(action_type="allocate",
-            allocations=allocs or [AllocationItem(zone=scores[0][0], resource="power")])
-    elif stage == "execute":
-        return FairRecoveryAction(action_type="execute")
-    return FairRecoveryAction(action_type="submit")
-
-
-POLICIES = {"random": random_policy, "greedy": greedy_policy, "fair": fairness_aware_policy}
-
-
-def evaluate_policy(policy_name: str, difficulty: str = "hard",
-                    n_episodes: int = 5, base_url: str = BASE_URL) -> Dict:
-    policy_fn = POLICIES[policy_name]
-    rewards, fairness, scores = [], [], []
-
-    print(f"\n{'='*60}")
-    print(f"Policy: {policy_name.upper()} | Difficulty: {difficulty} | Episodes: {n_episodes}")
-    print(f"{'='*60}")
-
-    with FairRecoveryEnv(base_url=base_url) as env:
-        for ep in range(n_episodes):
-            result = env.run_episode(policy_fn=policy_fn, difficulty=difficulty, verbose=True)
-            rewards.append(result["total_reward"])
-            fairness.append(result["final_fairness"])
-            scores.append(result["grader_score"] or 0.0)
-            print(f"  Episode {ep+1}: reward={result['total_reward']:.3f}, "
-                  f"fairness={result['final_fairness']:.3f}, grader={result['grader_score']:.3f}")
-
-    summary = {"policy": policy_name, "difficulty": difficulty,
-               "mean_reward": float(np.mean(rewards)), "std_reward": float(np.std(rewards)),
-               "mean_fairness": float(np.mean(fairness)), "mean_score": float(np.mean(scores))}
-    print(f"\nSummary: {json.dumps(summary, indent=2)}")
-    return summary
-
-
-def main():
-    parser = argparse.ArgumentParser(description="FairRecovery++ Baseline Inference")
-    parser.add_argument("--difficulty", default="hard", choices=["easy", "medium", "hard"])
-    parser.add_argument("--episodes", default=3, type=int)
-    parser.add_argument("--policy", default="all", choices=["all", "random", "greedy", "fair"])
-    parser.add_argument("--url", default=BASE_URL)
-    args = parser.parse_args()
-
-    policies = list(POLICIES.keys()) if args.policy == "all" else [args.policy]
-    results = [evaluate_policy(p, args.difficulty, args.episodes, args.url) for p in policies]
-
-    print(f"\n{'='*60}\nCOMPARISON TABLE\n{'='*60}")
-    print(f"{'Policy':<12} {'Mean Reward':>12} {'Mean Fairness':>14} {'Mean Score':>12}")
-    print("-"*60)
-    for r in results:
-        print(f"{r['policy']:<12} {r['mean_reward']:>12.4f} {r['mean_fairness']:>14.4f} {r['mean_score']:>12.4f}")
-    print("="*60)
-    print("\nKey: 'greedy' has higher utility but LOWER fairness than 'fair'.")
-    print("The trained agent should match/exceed 'fair' policy fairness score.")
-
-
-if __name__ == "__main__":
-    main()
+    """Fairness-aware heuristic (CORRECT policy)."""
+    if obs.day > 10: return FairRecoveryAction(action_type=ActionType.SUBMIT)
+    
+    # Prioritize the most vulnerable zone
+    vulnerable_zone = sorted(range(len(obs.zones)), key=lambda i: obs.zones[i].vulnerable_ratio, reverse=True)[0]
+    
+    return FairRecoveryAction(
+        action_type=ActionType.ALLOCATE,
+        allocations=[
+            ResourceAllocation(zone=vulnerable_zone, resource=ResourceType.MEDICAL)
+        ],
+        reasoning=f"Fair: prioritizing zone {vulnerable_zone} due to vulnerability."
+    )

server/app.py

@@ -1,42 +1,33 @@
 """
 FairRecovery++ — FastAPI Application.
 
-Mirrors hallucination-detector-gym server/app.py pattern exactly.
+Perfectly aligned with the refactored environment and reference structure.
 """
 
 from __future__ import annotations
 import os, sys
+import re
+import requests
+import json
 from typing import Optional
 from fastapi import FastAPI, Request
+from fastapi.responses import RedirectResponse
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
 
-try:
-    from openenv.core.env_server.http_server import create_app
-    _OPENENV_AVAILABLE = True
-except ImportError:
-    _OPENENV_AVAILABLE = False
-
 from fairrecovery_env.models import FairRecoveryAction, FairRecoveryObservation
-from fairrecovery_env.logging_config import configure_logging
 from server.fairrecovery_environment import FairRecoveryEnvironment
 from inference import greedy_policy, fairness_aware_policy
-import requests
-import json
-import re
-
-configure_logging(json_output=True, log_level="INFO")
 
 def llm_policy(obs: FairRecoveryObservation):
-    """Real-time LLM inference using HF API. Falls back to heuristic if API fails."""
+    """Real-time LLM inference using HF API."""
     hf_token = os.environ.get("HF_TOKEN")
     if not hf_token:
         return fairness_aware_policy(obs)
         
     try:
-        # Build a prompt describing the state
-        zones_str = '\\n'.join([f"Zone {z.zone_id}: damage={z.damage:.2f}, vulnerable={z.vulnerable_ratio:.2f}" for z in obs.zones])
-        prompt = f"System: You are an AI allocating resources fairly. Respond with JSON action.\\nUser: Day {obs.day}. Budget {obs.budget_left}. Zones:\\n{zones_str}\\nWhat is your next action?"
+        zones_str = '\n'.join([f"Zone {z.zone_id}: damage={z.damage:.2f}, vulnerable={z.vulnerable_ratio:.2f}" for z in obs.zones])
+        prompt = f"System: You are an AI allocating resources fairly. Respond with JSON action.\nUser: Day {obs.day}. Budget {obs.budget_left}. Zones:\n{zones_str}\nWhat is your next action?"
         
         response = requests.post(
             "https://api-inference.huggingface.co/models/meta-llama/Llama-3.2-3B-Instruct",
@@ -46,28 +37,20 @@ def llm_policy(obs: FairRecoveryObservation):
         )
         if response.status_code == 200:
             text = response.json()[0]["generated_text"]
-            match = re.search(r'\\{.*?\\}', text, re.DOTALL)
+            match = re.search(r'\{.*?\}', text, re.DOTALL)
             if match:
                 data = json.loads(match.group())
                 return FairRecoveryAction(**data)
     except Exception as e:
-        print(f"LLM API failed, falling back to fairness heuristic: {e}")
+        print(f"LLM API failed: {e}")
         
     return fairness_aware_policy(obs)
 
-
-
 def _build_app():
-    from fastapi.responses import JSONResponse, RedirectResponse
     import gradio as gr
-
     app = FastAPI(title="FairRecovery++ RL Environment", version="2.0.0")
     _env = FairRecoveryEnvironment()
 
-    @app.get("/web")
-    async def web_redirect():
-        return RedirectResponse(url="/")
-
     @app.post("/reset")
     async def reset(difficulty: str = "medium", episode_id: Optional[str] = None):
         return _env.reset(difficulty=difficulty, episode_id=episode_id).model_dump()
@@ -78,181 +61,63 @@ def _build_app():
         action = FairRecoveryAction(**payload)
         return _env.step(action).model_dump()
 
-    @app.get("/health")
-    async def health():
-        return {"status": "ok"}
-
     # ── Simulation Logic for Gradio UI ───────────────────────────────────────
     def translate_zone_status(damage, vulnerability):
         people_affected = int(damage * 10000)
         vulnerable_count = int(people_affected * vulnerability)
-        
         status_icon = "🔴" if damage > 0.6 else "🟡" if damage > 0.3 else "🟢"
-        
-        return f"{status_icon} **{people_affected:,}** people affected | ⚠️ **{vulnerable_count:,}** highly vulnerable (elderly/low-income)"
+        return f"{status_icon} **{people_affected:,}** people | ⚠️ **{vulnerable_count:,}** vulnerable"
 
-    def run_simulation_raw(policy_type: str):
+    def run_simulation(policy_type: str):
         env = FairRecoveryEnvironment()
-        obs = env.reset(difficulty="hard")
+        obs = env.reset(task_id="multi_disaster_hard")
         
         logs = []
-        logs.append("### 🚨 DISASTER SCENARIO: 7.2 Magnitude Earthquake")
-        logs.append(f"**Day:** {obs.day} | **Emergency Budget:** ${obs.budget_left * 100000:,.0f}")
-        for z in obs.zones:
-            logs.append(f"- **Zone {z.zone_id}**: {translate_zone_status(z.damage, z.vulnerable_ratio)}")
-        logs.append("\n---\n### ⚙️ SYSTEM BOOTING AI POLICY: " + policy_type.upper() + "\n")
+        logs.append(f"### 🚨 SCENARIO: {policy_type.upper()}")
         
         done = False
         step_count = 0
-        
         policy_fn = greedy_policy if policy_type == "Baseline (Greedy)" else llm_policy
         
         while not done and step_count < 30:
             action = policy_fn(obs)
-            logs.append(f"**▶ Day {step_count+1}: {action.action_type.capitalize()} Phase**")
-            
-            if action.action_type == "analyze":
-                logs.append(f"> 🛰️ AI identified priority zones: {action.critical_zones}")
-            elif action.action_type == "allocate":
-                allocs = []
-                for a in action.allocations:
-                    res_name = {"medical": "Medical Teams", "water": "Water Trucks", "power": "Power Grid Repair"}.get(a.resource, a.resource)
-                    allocs.append(f"Zone {a.zone} ({res_name})")
-                
-                logs.append(f"> 🚚 Dispatching resources: {', '.join(allocs) if allocs else 'None'}")
-                
             obs = env.step(action)
             
-            if action.action_type == "execute":
-                logs.append("> 🏥 **Recovery Update:**")
-                for z in obs.zones:
-                    logs.append(f"  - Zone {z.zone_id} Status: {translate_zone_status(z.damage, z.vulnerable_ratio)}")
-
+            logs.append(f"**Day {obs.day}**: AI performed {action.action_type.value}")
+            if action.action_type == "allocate":
+                for a in (action.allocations or []):
+                    logs.append(f"  - Dispatched {a.resource.value} to Zone {a.zone}")
+            
             done = obs.done
             step_count += 1
 
-        # Calculate Honest Truth Metrics at the end
-        services = [z.service for z in obs.zones]
-        utility = sum(services) / len(services)
-        mean_svc = utility
-        disparity = sum(abs(s - mean_svc) for s in services) / len(services)
-        fairness = max(0.0, 1.0 - disparity)
-        # Safety: assume no persistent violations for the final summary if it finished
-        safety = max(0.0, 1.0 - env.state.violations_total / 10.0)
-        
-        normalized_reward = 0.4 * utility + 0.4 * fairness + 0.2 * safety
-        normalized_reward = max(0.0, min(1.0, normalized_reward))
-
-        logs.append("\n---\n### 🏁 EPISODE COMPLETE")
-        return "\n".join(logs), float(normalized_reward), float(fairness)
-
-    def run_simulation(policy_type: str):
-        logs, reward, fairness = run_simulation_raw(policy_type)
-        
-        fairness_eval = ""
-        if fairness < 0.6:
-            fairness_eval = "🔴 **CRITICAL NEGLECT** — Vulnerable populations were systematically bypassed to maximize raw efficiency. High human cost."
-        elif fairness < 0.8:
-            fairness_eval = "🟡 **MEDIUM PARITY** — Recovery reached vulnerable zones eventually, but disparity remained significant."
-        else:
-            fairness_eval = "🟢 **RESEARCH-LEVEL EQUITY** — Balanced recovery achieved. Socioeconomic demographics were protected equally."
-            
-        result_text = f"### 🏆 FINAL OUTCOME\n- **Overall Efficiency (Normalized Reward):** {reward:.3f}\n- **Equity Index (Fairness):** {fairness:.3f}\n\n**Impact Analysis:**\n{fairness_eval}"
-        return logs, result_text
+        res_eval = "🟢 **EQUITY ACHIEVED**" if obs.fairness_score > 0.8 else "🔴 **NEGLECT DETECTED**"
+        result_text = f"### 🏆 FINAL OUTCOME\n- **Reward:** {obs.cumulative_reward:.3f}\n- **Equity:** {obs.fairness_score:.3f}\n\n{res_eval}"
+        return "\n".join(logs), result_text
 
-    def compare_policies():
-        _, greedy_reward, greedy_fairness = run_simulation_raw("Baseline (Greedy)")
-        _, fair_reward, fair_fairness = run_simulation_raw("Trained LLM (FairRecovery++)")
-        
-        return f"""### 📊 POLICY COMPARISON: THE TRUTH ABOUT BIAS
-
-| AI Model | Efficiency Score | Equity (Fairness) | Ethical Verdict |
-|---|---|---|---|
-| **Baseline (Greedy)** | {greedy_reward:.3f} | {greedy_fairness:.3f} | ❌ **Neglects vulnerable zones to save 'easier' wealthy zones.** |
-| **Trained LLM (Ours)** | {fair_reward:.3f} | {fair_fairness:.3f} | ✅ **Prioritizes high-vulnerability populations under pressure.** |
-
-> **Key Insight**: While the greedy model seems fast, its "Efficiency" is an illusion built on socioeconomic exclusion. Our **Fair-GRPO-RLVR** agent learns that true recovery must be equitable to be sustainable.
-"""
-
-    # ── Custom Simplified Gradio UI ──────────────────────────────────────────
-    with gr.Blocks(title="FairRecovery++ Simulator", theme=gr.themes.Soft()) as gradio_app:
-        gr.Markdown("# 🏗️ FairRecovery++: Adaptive Multi-Agent Disaster Recovery Environment")
-        
+    with gr.Blocks(title="FairRecovery++", theme=gr.themes.Soft()) as gradio_app:
+        gr.Markdown("# 🏗️ FairRecovery++")
         with gr.Tabs():
-            with gr.Tab("Interactive Simulation"):
-                gr.Markdown("> **An OpenEnv environment that teaches LLMs to make fair resource allocation decisions under adversarial pressure, multi-agent dynamics, and long-horizon planning constraints.**")
-                
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        gr.Markdown("### 🚨 The Fairness Trap (Hard Scenario)")
-                        gr.Markdown("""
-After a disaster, AI systems optimizing for **efficiency alone** consistently neglect vulnerable populations. 
-A greedy planner fixes wealthy **Zone 0** (low damage, easy win) while **Zone 4** (96% vulnerable, 92% damaged) gets nothing. This is the **Fairness Trap** -- and current LLMs fall right into it.
+            with gr.Tab("Simulation"):
+                policy = gr.Dropdown(choices=["Baseline (Greedy)", "Trained LLM"], value="Baseline (Greedy)")
+                btn = gr.Button("Run Simulation")
+                logs = gr.Markdown()
+                results = gr.Markdown()
+                btn.click(run_simulation, inputs=[policy], outputs=[logs, results])
+            with gr.Tab("README"):
+                readme_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "README.md")
+                with open(readme_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+                    if content.startswith("---"):
+                        content = re.sub(r"^---.*?---", "", content, flags=re.DOTALL)
+                    gr.Markdown(content)
 
-**FairRecovery++** teaches the LLM to escape this trap by balancing efficiency vs fairness under tight budgets.
-""")
-                        
-                        policy_dropdown = gr.Dropdown(
-                            choices=["Baseline (Greedy)", "Trained LLM (FairRecovery++)"],
-                            value="Baseline (Greedy)",
-                            label="Choose AI Policy"
-                        )
-                        run_btn = gr.Button("▶ Run Episode", variant="primary")
-                        
-                        gr.Markdown("---")
-                        compare_btn = gr.Button("🔁 Run Both & Compare Impact", variant="secondary")
-                        
-                    with gr.Column(scale=2):
-                        gr.Markdown("### 📜 Mission Log")
-                        log_output = gr.Markdown("*Select a policy and click 'Run Episode' to view the ground-truth simulation.*")
-                        result_output = gr.Markdown("")
-                        
-                run_btn.click(
-                    fn=run_simulation,
-                    inputs=[policy_dropdown],
-                    outputs=[log_output, result_output]
-                )
-                
-                compare_btn.click(
-                    fn=compare_policies,
-                    inputs=[],
-                    outputs=[result_output]
-                )
-
-            with gr.Tab("Project README"):
-                try:
-                    readme_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "README.md")
-                    with open(readme_path, "r", encoding="utf-8") as f:
-                        # Strip the YAML frontmatter for cleaner display in UI
-                        content = f.read()
-                        if content.startswith("---"):
-                            content = re.sub(r"^---.*?---", "", content, flags=re.DOTALL)
-                        gr.Markdown(content)
-                except Exception as e:
-                    gr.Markdown(f"### Error loading README.md\n{e}")
-
-    # Re-add root redirects but ensure trailing slash for /ui/ 
-    # to prevent relative asset 404s (the cause of the blank screen).
     @app.get("/")
-    async def root_redirect():
-        from fastapi.responses import RedirectResponse
-        return RedirectResponse(url="/ui/")
-        
-    @app.get("/web")
-    async def old_web_redirect():
-        from fastapi.responses import RedirectResponse
-        return RedirectResponse(url="/ui/")
+    async def root(): return RedirectResponse(url="/ui/")
 
     return gr.mount_gradio_app(app, gradio_app, path="/ui")
 
-
 app = _build_app()
-
-
-def main(host: str = "0.0.0.0", port: int = 8000) -> None:
-    import uvicorn
-    uvicorn.run(app, host=host, port=port)
-
-
 if __name__ == "__main__":
-    main()
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

server/fairrecovery_environment.py

@@ -1,175 +1,75 @@
 """
-FairRecovery++ — Core Environment Implementation.
+FairRecovery++ — Core Environment.
 
-Adaptive multi-agent OpenEnv environment for post-disaster recovery.
-Implements Environment base class with step(), reset(), state().
-
-Design:
-  * Multi-step protocol: analyze -> allocate -> execute -> adapt -> submit
-  * Multi-agent system: citizens, NGOs, adversaries interact each step
-  * Behavior analysis: patterns extracted from interaction logs
-  * Predictive engine: forecasts events for proactive planning
-  * Dense rewards at every step with 5-component breakdown
-  * Composable rubrics wired via RFC 004 pattern
-  * Safety shield validates before state mutation
-
-Themes Hit:
-  * Theme 3.1 — Real-World Professional Tasks (core)
-  * Theme 2   — Long-Horizon Planning (strong)
-  * Theme 1   — Multi-Agent Interaction (added)
+Perfectly aligned with OpenEnv patterns and reference project structure.
 """
 
 from __future__ import annotations
-
 import uuid
-from typing import Any, Dict, List, Optional
-
+from typing import Any, Optional
 import structlog
 
-try:
-    from openenv.core.env_server.interfaces import Environment
-    from openenv.core.env_server.types import Action, Observation, State
-except ImportError:
-    class Environment:  # type: ignore
-        pass
-    Action = object  # type: ignore
-    Observation = object  # type: ignore
-    State = object  # type: ignore
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import Action, Observation, State
 
 from fairrecovery_env.constants import (
-    ActionType, MAX_DAYS, MAX_STEPS_SAFETY_CAP,
-    MIN_STEPS, CURRICULUM_MAX_STEPS, EARLY_SUBMIT_PENALTY,
-    FINAL_BONUS_WEIGHT_UTILITY, FINAL_BONUS_WEIGHT_FAIRNESS,
-    PENALTY_INVALID_ACTION, PENALTY_WRONG_STAGE,
-    PENALTY_PATTERN_IGNORED, PENALTY_ADVERSARIAL_FAILURE,
+    MAX_STEPS_PER_EPISODE, MAX_DAYS, TaskID, ActionType, ResourceType,
+    COST_MEDICAL, COST_WATER, COST_POWER
 )
 from fairrecovery_env.models import (
-    AgentEvent, FairRecoveryAction, FairRecoveryObservation,
-    FairRecoveryState, ZoneObservation,
+    FairRecoveryAction, FairRecoveryObservation, FairRecoveryState, ZoneState
 )
-from fairrecovery_env.rewards import RewardEngine, compute_fairness_reward
-from fairrecovery_env.rubrics import CompositeRubric
-from fairrecovery_env.shield import validate, check_timeout
-from fairrecovery_env.state import CityState
-from fairrecovery_env.tasks import ScenarioConfig, get_task
-from fairrecovery_env.agents import MultiAgentManager
-from fairrecovery_env.behavior_analyzer import BehaviorAnalyzer
-from fairrecovery_env.predictor import Predictor, Prediction
+from fairrecovery_env.rewards import RewardEngine
+from fairrecovery_env.tasks import get_task, TaskDefinition
 
 logger = structlog.get_logger(__name__)
 
-
-def _clip01(value: float) -> float:
-    return float(max(0.0, min(1.0, value)))
-
-
-def _signed_to_unit(value: float) -> float:
-    # Convert [-1, 1] style component into [0, 1] for transparent dashboards.
-    return _clip01((value + 1.0) / 2.0)
-
-
 class FairRecoveryEnvironment(Environment):
-    """
-    Post-disaster city recovery RL environment with multi-agent dynamics.
-
-    The agent must:
-      1. Analyse zone status (damage, service, vulnerability, agent events)
-      2. Queue resource allocations respecting budget constraints
-      3. Execute allocations and receive dense rewards
-      4. Adapt strategy based on predictions and agent behavior
-      5. Repeat for MAX_DAYS days, then submit for final score
-
-    Rewards balance 5 components:
-      R_exec   — service improvement (utility)
-      R_fair   — service parity between vulnerable/non-vulnerable zones
-      R_adapt  — success against predicted events
-      R_stable — system balance (citizen satisfaction variance)
-      R_safe   — constraint satisfaction
-    """
+    """OpenEnv environment for fair disaster recovery."""
 
     SUPPORTS_CONCURRENT_SESSIONS: bool = True
 
     def __init__(self) -> None:
         super().__init__()
-        self._city: Optional[CityState] = None
-        self._task: Optional[ScenarioConfig] = None
+        self._state: Optional[FairRecoveryState] = None
+        self._task: Optional[TaskDefinition] = None
         self._reward_engine: Optional[RewardEngine] = None
-        self._rubrics: CompositeRubric = CompositeRubric()
-        self._episode_id: Optional[str] = None
-        self._step_count: int = 0
         self._action_history: list[str] = []
-        self._current_difficulty: str = "medium"
-
-        # Multi-agent system
-        self._agent_manager: Optional[MultiAgentManager] = None
-        self._behavior_analyzer: Optional[BehaviorAnalyzer] = None
-        self._predictor: Optional[Predictor] = None
-        self._last_prediction: Optional[Prediction] = None
-        self._last_events: List[AgentEvent] = []
-        self._adaptation_scores: List[float] = []
 
-    # ──────────────────────────────────��───────────────────────────────────────
-    # reset()
-    # ──────────────────────────────────────────────────────────────────────────
-    def reset(self, seed: Optional[int] = None, difficulty: str = "medium",
-              episode_id: Optional[str] = None, **kwargs: Any) -> FairRecoveryObservation:
-        """Reset environment with multi-agent system initialization."""
-        self._task = get_task(difficulty)
-        self._current_difficulty = difficulty
-        self._city = CityState(self._task.__dict__)
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        task_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> FairRecoveryObservation:
+        """Reset the environment."""
+        resolved_task_id = TaskID(task_id) if task_id else TaskID.FLOOD_EASY
+        self._task = get_task(resolved_task_id)
         self._reward_engine = RewardEngine(self._task)
-        self._episode_id = episode_id or str(uuid.uuid4())
-        self._step_count = 0
         self._action_history = []
-        self._city.step_stage = "analyze"
-        self._adaptation_scores = []
 
-        # Initialize multi-agent system
-        vuln_ratios = [z.vulnerable_ratio for z in self._city.zones]
-        actual_seed = seed or 42
-        self._agent_manager = MultiAgentManager(
-            num_zones=len(self._city.zones),
-            zone_vulnerabilities=vuln_ratios,
-            seed=actual_seed,
+        ep_id = episode_id or str(uuid.uuid4())
+        
+        # Deep copy initial zones
+        zones = [ZoneState(**z.model_dump()) for z in self._task.initial_zones]
+        
+        self._state = FairRecoveryState(
+            episode_id=ep_id,
+            step_count=0,
+            day=1,
+            budget_remaining=self._task.budget_limit,
+            zones=zones,
+            task_id=resolved_task_id,
+            difficulty=self._task.difficulty,
         )
-        self._behavior_analyzer = BehaviorAnalyzer(num_zones=len(self._city.zones))
-        self._predictor = Predictor(self._behavior_analyzer)
-        self._last_prediction = None
-        self._last_events = []
 
-        # Rubrics reset
-        self._rubrics.reset()
-        initial_fairness = compute_fairness_reward(self._city.zones)
-        self._rubrics.fairness.set_initial_fairness(initial_fairness)
+        return self._build_observation("Environment reset. Begin disaster recovery.")
 
-        logger.info("environment_reset", episode_id=self._episode_id,
-                     difficulty=difficulty, num_zones=len(self._city.zones),
-                     initial_budget=self._city.budget_left,
-                     initial_fairness=round(initial_fairness, 4))
-
-        obs = self._build_observation(reward=0.0, done=False, r_exec=0.0, r_fair=0.0, r_safe=0.0)
-        obs.step_feedback = (
-            f"Episode started. Difficulty: {difficulty}. "
-            f"Zones: {len(self._city.zones)}. Budget: {self._city.budget_left}. "
-            f"Active agents: {self._agent_manager.get_active_agent_count()}. "
-            f"{self._task.hint}"
-        )
-        return obs
-
-    # ──────────────────────────────────────────────────────────────────────────
-    # step()
-    # ──────────────────────────────────────────────────────────────────────────
-    def step(self, action: Action, timeout_s: Optional[float] = None,
-             **kwargs: Any) -> FairRecoveryObservation:
-        """Execute one agent action with multi-agent dynamics."""
-        if self._city is None or self._reward_engine is None:
-            return self._error_obs("Call reset() before step().")
-
-        if check_timeout(self._step_count):
-            return self._build_observation(
-                reward=PENALTY_INVALID_ACTION, done=True,
-                r_exec=0.0, r_fair=0.0, r_safe=-0.5,
-                feedback="Episode terminated: safety step cap exceeded.")
+    def step(self, action: Action, **kwargs: Any) -> FairRecoveryObservation:
+        """Execute a recovery step."""
+        if self._state is None or self._reward_engine is None:
+            return FairRecoveryObservation(done=True, reward=0.0, step_feedback="Reset first.")
 
         # Parse action
         try:
@@ -179,434 +79,90 @@ class FairRecoveryEnvironment(Environment):
                 typed_action = FairRecoveryAction(**action)
             else:
                 typed_action = FairRecoveryAction(**action.model_dump())
-        except Exception as exc:
-            logger.warning("malformed_action", error=str(exc))
-            typed_action = FairRecoveryAction(action_type="noop",
-                                               reasoning=f"Parse error: {exc}")
-
-        self._step_count += 1
-        city = self._city
-        action_type = str(typed_action.action_type).replace("ActionType.", "").lower()
-
-        # ── HARD GATE: kill early-submit exploitation (structural, not band-aid) ──
-        # If the agent tries to submit before MIN_STEPS, we apply a small penalty,
-        # do NOT end the episode, and return immediately with a clear feedback msg.
-        # This removes the "submit on step 1 to exit fast" exploit entirely.
-        if action_type == "submit" and self._step_count < MIN_STEPS:
-            obs = self._build_observation(
-                reward=-EARLY_SUBMIT_PENALTY, done=False,
-                r_exec=0.0, r_fair=0.0, r_safe=-EARLY_SUBMIT_PENALTY,
-                r_adapt=0.0, r_stable=0.0,
-                feedback=(
-                    f"early_submit_blocked: step {self._step_count} < MIN_STEPS={MIN_STEPS}. "
-                    f"Penalty {EARLY_SUBMIT_PENALTY}. Continue with analyze/allocate/execute."
-                ),
-            )
-            self._action_history.append(
-                f"Step {self._step_count}: submit_blocked (penalty={-EARLY_SUBMIT_PENALTY:+.3f})"
-            )
-            obs.action_history = list(self._action_history[-8:])
-            logger.info("early_submit_blocked", episode_id=self._episode_id,
-                        step=self._step_count, min_steps=MIN_STEPS)
-            return obs
-
-        # Retrieve allocations for shield
-        alloc_dicts = None
-        if typed_action.allocations:
-            alloc_dicts = [a.model_dump() for a in typed_action.allocations]
-
-        # Shield validation
-        is_valid, violations = validate(
-            action_type=action_type, current_stage=city.step_stage,
-            step_count=self._step_count, city=city, allocations=alloc_dicts)
-
-        # Run multi-agent system to generate events
-        agent_events = self._run_agents()
-
-        # Generate prediction for next step
-        prediction = self._generate_prediction()
-
-        # Dispatch by action type
-        reward = r_exec = r_fair = r_safe = r_adapt = r_stable = 0.0
-        done = False
-        feedback = ""
-
-        if action_type == "analyze":
-            components = self._reward_engine.compute_analysis_step(
-                chosen_zones=typed_action.critical_zones or [], city=city)
-            reward = components.R_total
-            r_exec = components.R_analysis
-            feedback = components.feedback
-            city.record(f"analyzed zones={typed_action.critical_zones}")
-            city.step_stage = "allocate"
-
-        elif action_type == "allocate":
-            city.pending_allocations = alloc_dicts or []
-            city.record(f"queued {len(city.pending_allocations)} allocations")
-            city.step_stage = "execute"
-            feedback = (f"Queued {len(city.pending_allocations)} allocation(s). "
-                        f"Budget: {city.budget_left}. Call execute next.")
-
-        elif action_type == "execute":
-            city.snapshot_services()
-            exec_violations = city.apply_allocations()
-            all_violations = violations + exec_violations
-
-            # Evaluate adaptation to predictions
-            adapt_score = self._evaluate_adaptation(typed_action)
-
-            # Apply multi-agent effects
-            self._apply_agent_effects()
-
-            components = self._reward_engine.compute_execute_step(
-                city=city, violations=all_violations)
-            reward = components.R_total
-            r_exec = components.R_exec
-            r_fair = components.R_fair
-            r_safe = components.R_safe
-            feedback = components.feedback
-            city.record(f"executed | {feedback}")
-            city.step_stage = "analyze"
-
-            if self._should_end_episode(action_type=action_type):
-                done = True
-
-        elif action_type == "adapt":
-            # Adaptation step: agent explicitly responds to predictions/events
-            adapt_score = self._evaluate_adaptation(typed_action)
-            self._adaptation_scores.append(adapt_score)
-            reward = 0.05 * adapt_score  # Small reward for adaptation
-            r_adapt = adapt_score
-            feedback = (f"Adaptation score: {adapt_score:.3f}. "
-                        f"Strategy: {typed_action.adaptation_strategy or 'none'}")
-            city.record(f"adapted | {feedback}")
-            city.step_stage = "analyze"
-
-        elif action_type == "submit":
-            components = self._reward_engine.compute_submit_reward(city=city)
-            reward = components.R_total
-            r_fair = components.R_fair
-            feedback = components.feedback
-            done = True
-
-        elif action_type == "noop":
-            reward = -0.02
-            feedback = "No action taken (noop). This wastes a step."
-
-        else:
-            reward = PENALTY_INVALID_ACTION
-            feedback = f"Unknown action_type '{action_type}'."
-            violations.append(f"unknown_action_type:{action_type}")
-
-        # Stage violations
-        if violations and not is_valid:
-            reward += PENALTY_WRONG_STAGE
-            feedback += f" | Violations: {violations}"
-
-        # Global termination check (covers analyze/allocate/adapt/noop paths too).
-        # NOTE: submit pre-MIN_STEPS already short-circuited earlier.
-        if not done and self._should_end_episode(action_type=action_type):
-            done = True
-            feedback += " | episode_end_reached"
-
-        # Build observation with multi-agent data
-        obs = self._build_observation(
-            reward=reward, done=done, r_exec=r_exec, r_fair=r_fair,
-            r_safe=r_safe, feedback=feedback, agent_events=self._last_events,
-            predictions=prediction.to_dict() if prediction else None)
-
-        # Rubric scoring (RFC 004)
-        rubric_score = self._rubrics.forward(typed_action, obs)
-        if rubric_score != 0.0:
-            obs.cumulative_reward += rubric_score
-
-        # Action history
-        action_summary = f"Step {self._step_count}: {action_type} (reward={reward:+.3f})"
-        self._action_history.append(action_summary)
-        obs.action_history = list(self._action_history[-8:])
-
-        # Final grader score
-        if done:
-            obs.grader_score = self._reward_engine.get_final_grader_score(city)
-
-        logger.info("step_executed", episode_id=self._episode_id,
-                     step=self._step_count, action_type=action_type,
-                     day=city.day, reward=round(reward, 4), done=done)
-
-        return obs
-
-    # ──────────────────────────────────────────────────────────────────────────
-    # state property
-    # ──────────────────────────────────────────────────────────────────────────
-    @property
-    def state(self) -> FairRecoveryState:
-        if self._city is None:
-            return FairRecoveryState()
-        city = self._city
-        zones_obs = [
-            ZoneObservation(
-                zone_id=z.zone_id, damage=round(z.damage, 3),
-                service=round(z.service, 3),
-                vulnerable_ratio=round(z.vulnerable_ratio, 3),
-                citizen_satisfaction=round(z.citizen_satisfaction, 3),
-                risk_level=0.0)
-            for z in city.zones
-        ]
-        fairness = compute_fairness_reward(city.zones)
-        adapt_avg = (sum(self._adaptation_scores) / len(self._adaptation_scores)
-                     if self._adaptation_scores else 0.0)
-        return FairRecoveryState(
-            episode_id=self._episode_id, difficulty=self._current_difficulty,
-            day=city.day, budget_left=round(city.budget_left, 2),
-            step_stage=city.step_stage, step_count=self._step_count,
-            cumulative_reward=round(self._reward_engine.cumulative_reward if self._reward_engine else 0.0, 4),
-            fairness_score=round(fairness, 4),
-            is_done=(
-                city.day >= MAX_DAYS
-                or self._average_recovery() >= 0.95
-                or city.budget_left <= 0
-            ),
-            violations_total=city.violations_total, zones=zones_obs,
-            active_agents=self._agent_manager.get_active_agent_count() if self._agent_manager else 0,
-            adversarial_events=self._agent_manager.get_adversarial_event_count() if self._agent_manager else 0,
-            adaptation_score=round(adapt_avg, 4))
-
-    def close(self) -> None:
-        self._city = None
-        self._task = None
-        self._reward_engine = None
-        self._agent_manager = None
-        self._behavior_analyzer = None
-        self._predictor = None
-        self._action_history.clear()
-        logger.info("environment_closed", episode_id=self._episode_id)
-
-    # ──────────────────────────────────────────────────────────────────────────
-    # Multi-Agent Helpers
-    # ──────────────────────────────────────────────────────────────────────────
-    def _run_agents(self) -> List[AgentEvent]:
-        """Run all agents and collect events."""
-        if not self._agent_manager or not self._city:
-            return []
-        city = self._city
-        events = self._agent_manager.step(
-            zone_services=city.current_services,
-            zone_damages=city.current_damages,
-            zone_vulnerabilities=city.current_vulnerabilities,
-            planner_target_zones=city.planner_target_zones,
-            day=city.day)
-
-        # Ingest into behavior analyzer
-        if self._behavior_analyzer:
-            self._behavior_analyzer.ingest(events)
-
-        # Convert to AgentEvent models for observation
-        self._last_events = [
-            AgentEvent(agent_type=e.agent_type, event_type=e.event_type,
-                       zone_id=e.zone_id, intensity=e.intensity,
-                       message=e.message, timestamp=e.timestamp)
-            for e in events
-        ]
-        return self._last_events
-
-    def _generate_prediction(self) -> Optional[Prediction]:
-        """Generate prediction for next events."""
-        if not self._predictor or not self._city:
-            return None
-        city = self._city
-        prediction = self._predictor.predict_next(
-            zone_services=city.current_services,
-            zone_damages=city.current_damages,
-            zone_vulnerabilities=city.current_vulnerabilities,
-            day=city.day)
-        self._last_prediction = prediction
-        return prediction
+        except Exception as e:
+            typed_action = FairRecoveryAction(action_type=ActionType.NOOP, reasoning=str(e))
+
+        self._state.step_count += 1
+        
+        # 1. Update Day Counter
+        # Sequence: Analyze -> Allocate -> Execute -> Day++
+        if typed_action.action_type == ActionType.EXECUTE:
+            self._execute_phase(typed_action)
+            self._state.day += 1
+        elif typed_action.action_type == ActionType.ALLOCATE:
+            self._allocate_phase(typed_action)
+
+        # 2. Compute Reward
+        reward, feedback = self._reward_engine.compute_reward(typed_action, self._state)
+        self._state.cumulative_reward = self._reward_engine.cumulative_reward
+
+        # 3. Check Termination
+        is_done = (
+            typed_action.action_type == ActionType.SUBMIT or
+            self._state.day > MAX_DAYS or
+            self._state.step_count >= MAX_STEPS_PER_EPISODE
+        )
+        self._state.is_done = is_done
 
-    def _evaluate_adaptation(self, action: FairRecoveryAction) -> float:
-        """Evaluate how well the agent adapted to predictions."""
-        if not self._predictor or not self._last_prediction:
-            return 0.5
-        planner_zones = []
-        if action.critical_zones:
-            planner_zones = action.critical_zones
-        elif action.allocations:
-            planner_zones = [a.zone for a in action.allocations]
-        return self._predictor.evaluate_adaptation(
-            prediction=self._last_prediction,
-            actual_events=self._agent_manager.interaction_log[-5:] if self._agent_manager else [],
-            planner_actions=planner_zones)
+        return self._build_observation(feedback)
 
-    def _apply_agent_effects(self) -> None:
-        """Apply multi-agent effects to the world state."""
-        if not self._agent_manager or not self._city:
+    def _allocate_phase(self, action: FairRecoveryAction):
+        """Process resource allocations."""
+        if not action.allocations:
             return
-        city = self._city
-        # Apply disruptions from adversarial agents
-        for event in self._last_events:
-            if event.event_type == "disruption":
-                idx = event.zone_id
-                if 0 <= idx < len(city.zones):
-                    city.zones[idx].apply_disruption(event.intensity)
-            elif event.event_type in ("cooperation", "aid_delivery"):
-                idx = event.zone_id
-                if 0 <= idx < len(city.zones):
-                    city.zones[idx].service = min(1.0, city.zones[idx].service + 0.05)
-        # Update citizen satisfaction
-        for citizen in self._agent_manager.citizens:
-            idx = citizen.zone_id
-            if 0 <= idx < len(city.zones):
-                city.zones[idx].citizen_satisfaction = citizen.satisfaction
 
-    # ──────────────────────────────────────────────────────────────────────────
-    # Observation Builders
-    # ──────────────────────────────────────────────────────────────────────────
-    def _build_observation(self, reward: float, done: bool, r_exec: float,
-                           r_fair: float, r_safe: float, feedback: str = "",
-                           agent_events: Optional[List[AgentEvent]] = None,
-                           predictions: Optional[dict] = None) -> FairRecoveryObservation:
-        city = self._city
-        zones_obs = [
-            ZoneObservation(
-                zone_id=z.zone_id, damage=round(z.damage, 3),
-                service=round(z.service, 3),
-                vulnerable_ratio=round(z.vulnerable_ratio, 3),
-                citizen_satisfaction=round(z.citizen_satisfaction, 3),
-                risk_level=0.0)
-            for z in city.zones
-        ] if city else []
-
-        fairness = compute_fairness_reward(city.zones) if city else 0.0
-        steps_remaining = (MAX_DAYS - city.day) * 3 if city else 0
-        cumulative = self._reward_engine.cumulative_reward if self._reward_engine else 0.0
-
-        # Attach risk levels from behavior analyzer
-        if self._behavior_analyzer and zones_obs:
-            risk_levels = self._behavior_analyzer.get_risk_levels()
-            for i, zo in enumerate(zones_obs):
-                if i < len(risk_levels):
-                    zo.risk_level = round(risk_levels[i], 3)
-
-        info = self._build_reward_info(
-            reward=reward,
-            r_exec=r_exec,
-            r_fair=r_fair,
-            r_safe=r_safe,
-            done=done,
-            fairness_score=fairness,
-        )
-
-        return FairRecoveryObservation(
-            zones=zones_obs, day=city.day if city else 0,
-            budget_left=round(city.budget_left, 2) if city else 0.0,
-            step_stage=city.step_stage if city else "analyze",
-            fairness_score=round(fairness, 4), step_feedback=feedback,
-            steps_remaining=max(0, steps_remaining),
-            cumulative_reward=round(cumulative, 4),
-            r_exec=round(r_exec, 4), r_fair=round(r_fair, 4),
-            r_safe=round(r_safe, 4),
-            done=done, reward=round(reward, 4),
-            info=info,
-            agent_events=agent_events or [], predictions=predictions)
+        for alloc in action.allocations:
+            cost = {
+                ResourceType.MEDICAL: COST_MEDICAL,
+                ResourceType.WATER: COST_WATER,
+                ResourceType.POWER: COST_POWER
+            }.get(alloc.resource, 0.0)
+
+            if self._state.budget_remaining >= cost:
+                self._state.budget_remaining -= cost
+                zone = self._state.zones[alloc.zone]
+                # Resources reduce damage and increase service level
+                zone.damage = max(0.0, zone.damage - 0.05)
+                zone.service_level = min(1.0, zone.service_level + 0.1)
+            else:
+                self._state.violations_total += 1
+
+    def _execute_phase(self, action: FairRecoveryAction):
+        """Natural environment progression (deterioration if no service)."""
+        for zone in self._state.zones:
+            # Deterioration
+            if zone.service_level < 0.2:
+                zone.damage = min(1.0, zone.damage + 0.02)
+            # Service decay
+            zone.service_level = max(0.0, zone.service_level - 0.05)
+
+    def _build_observation(self, feedback: str) -> FairRecoveryObservation:
+        """Construct a FairRecoveryObservation from current state."""
+        
+        # Calculate Equity for observation
+        services = [z.service_level for z in self._state.zones]
+        avg_svc = sum(services) / len(services)
+        mad = sum(abs(s - avg_svc) for s in services) / len(services)
+        equity = 1.0 - mad
+
+        grader_score = self._reward_engine.get_final_score(self._state) if self._state.is_done else None
 
-    def _error_obs(self, msg: str) -> FairRecoveryObservation:
         return FairRecoveryObservation(
-            step_feedback=f"Error: {msg}", done=False,
-            reward=PENALTY_INVALID_ACTION,
-            info=self._build_reward_info(
-                reward=PENALTY_INVALID_ACTION,
-                r_exec=0.0,
-                r_fair=0.0,
-                r_safe=-1.0,
-            ),
-        )
-
-    def _build_reward_info(
-        self,
-        reward: float,
-        r_exec: float,
-        r_fair: float,
-        r_safe: float,
-        done: bool = False,
-        fairness_score: float = 0.0,
-    ) -> Dict[str, Any]:
-        """Curriculum-weighted reward in [0,1] + trajectory-level final bonus.
-
-        Curriculum:
-          progress = min(step_count, CURRICULUM_MAX_STEPS) / CURRICULUM_MAX_STEPS
-          step_r   = (0.6 + 0.4*p)*utility + (0.2 + 0.3*p)*fairness + 0.2*safety
-
-        => Early steps reward utility; late steps reward fairness.
-        => Prevents "fair but useless" policies and "instant-submit" exploits.
-
-        Final bonus (only when done=True) teaches long-horizon planning:
-          bonus = 0.5*final_utility + 0.5*final_fairness
-        """
-        utility = _signed_to_unit(r_exec)
-        fairness = _signed_to_unit(r_fair)
-        safety = _signed_to_unit(r_safe)
-        raw = _signed_to_unit(reward)
-
-        progress = min(self._step_count, CURRICULUM_MAX_STEPS) / float(CURRICULUM_MAX_STEPS)
-        w_u = 0.6 + 0.4 * progress
-        w_f = 0.2 + 0.3 * progress
-        w_s = 0.2
-
-        step_r = w_u * utility + w_f * fairness + w_s * safety
-        # Normalize back to [0,1] (max possible weights ≈ 1.0+0.5+0.2 = 1.7)
-        step_r = _clip01(step_r / (w_u + w_f + w_s))
-
-        # Trajectory-level long-horizon bonus on episode end.
-        bonus = 0.0
-        if done:
-            final_utility = self._average_recovery()
-            final_fairness = _signed_to_unit(fairness_score)
-            bonus = (
-                FINAL_BONUS_WEIGHT_UTILITY * final_utility
-                + FINAL_BONUS_WEIGHT_FAIRNESS * final_fairness
-            )
-            bonus = _clip01(bonus)
-
-        # Blended reward: 0.7 * curriculum step reward + 0.3 * final bonus (if any).
-        # If not done, blend = step_r alone.
-        if done:
-            blended = _clip01(0.6 * step_r + 0.4 * bonus)
-        else:
-            blended = step_r
-
-        return {
-            "reward": round(blended, 4),
-            "reward_step": round(step_r, 4),
-            "reward_raw": round(raw, 4),
-            "final_bonus": round(bonus, 4),
-            "progress": round(progress, 4),
-            "utility": round(utility, 4),
-            "fairness": round(fairness, 4),
-            "safety": round(safety, 4),
-        }
-
-    def _should_end_episode(self, action_type: str) -> bool:
-        """Episode-end logic with structural trajectory shaping.
-
-        done = (step_count >= MIN_STEPS AND action == "submit")
-            OR step_count >= CURRICULUM_MAX_STEPS
-            OR day >= MAX_DAYS
-            OR recovery >= 0.95
-            OR budget_left <= 0
-        """
-        if self._city is None:
-            return False
-        recovery = self._average_recovery()
-        valid_submit = action_type == "submit" and self._step_count >= MIN_STEPS
-        hit_max_steps = self._step_count >= CURRICULUM_MAX_STEPS
-        return (
-            self._city.day >= MAX_DAYS
-            or recovery >= 0.95
-            or valid_submit
-            or hit_max_steps
-            or self._city.budget_left <= 0
+            done=self._state.is_done,
+            reward=0.0, # Per-step reward is handled by cumulative
+            day=self._state.day,
+            budget_left=self._state.budget_remaining,
+            zones=self._state.zones,
+            fairness_score=equity,
+            step_stage="dynamic", # Could be more granular
+            steps_remaining=MAX_STEPS_PER_EPISODE - self._state.step_count,
+            cumulative_reward=self._state.cumulative_reward,
+            action_history=list(self._action_history),
+            grader_score=grader_score,
+            step_feedback=feedback,
+            metadata={"grader_score": grader_score}
         )
 
-    def _average_recovery(self) -> float:
-        if self._city is None or not self._city.zones:
-            return 0.0
-        return float(sum(z.service for z in self._city.zones) / len(self._city.zones))
+    @property
+    def state(self) -> FairRecoveryState:
+        return self._state