| """
|
| Scene-based disaster response coordination environment.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| from copy import deepcopy
|
| from typing import Any, Dict, List, Optional
|
| from uuid import uuid4
|
|
|
| from openenv.core.env_server.interfaces import Environment
|
| from openenv.core.env_server.types import State
|
|
|
| try:
|
| from ..models import (
|
| DisasterAction,
|
| DisasterObservation,
|
| ResourceStatus,
|
| TargetStatus,
|
| )
|
| from .scene_catalog import DEFAULT_SCENE_ID, SCENE_CATALOG, SceneConfig, ordered_scene_ids
|
| except ImportError:
|
| from models import DisasterAction, DisasterObservation, ResourceStatus, TargetStatus
|
| from server.scene_catalog import DEFAULT_SCENE_ID, SCENE_CATALOG, SceneConfig, ordered_scene_ids
|
|
|
|
|
| class DisasterResponseEnvironment(Environment):
|
| """
|
| Multi-scene disaster response environment with hidden-state reward shaping.
|
|
|
| The agent sees targets, resources, and timing cues, but rewards come from a
|
| latent harm model so the policy cannot self-certify mediocre behavior.
|
| """
|
|
|
| SUPPORTS_CONCURRENT_SESSIONS: bool = True
|
|
|
| def __init__(self) -> None:
|
| self._state = State(episode_id=str(uuid4()), step_count=0)
|
| self._scene: SceneConfig = SCENE_CATALOG[DEFAULT_SCENE_ID]
|
| self._targets: Dict[str, Dict[str, Any]] = {}
|
| self._resources: Dict[str, Dict[str, Any]] = {}
|
| self._metrics: Dict[str, float] = {}
|
| self._turn: int = 0
|
| self._baseline_harm: float = 0.0
|
| self._final_score: Optional[float] = None
|
|
|
| def reset(
|
| self,
|
| seed: Optional[int] = None,
|
| episode_id: Optional[str] = None,
|
| scene_id: Optional[str] = None,
|
| level: Optional[int] = None,
|
| **kwargs: Any,
|
| ) -> DisasterObservation:
|
| self._state = State(
|
| episode_id=episode_id or str(uuid4()),
|
| step_count=0,
|
| )
|
| self._scene = self._select_scene(scene_id=scene_id, level=level)
|
| self._targets = self._init_targets(self._scene)
|
| self._resources = self._init_resources(self._scene)
|
| self._metrics = {
|
| "fatalities": 0.0,
|
| "critical_injuries": 0.0,
|
| "exposure_harm": 0.0,
|
| "service_loss": 0.0,
|
| "invalid_actions": 0.0,
|
| "ineffective_assignments": 0.0,
|
| "deadline_misses": 0.0,
|
| "reassignment_churn": 0.0,
|
| "resolved_targets": 0.0,
|
| "failed_targets": 0.0,
|
| }
|
| self._turn = 0
|
| self._final_score = None
|
| self._baseline_harm = self._simulate_noop_baseline()
|
|
|
| feedback = (
|
| f"Level {self._scene.level}: {self._scene.name}\n"
|
| f"{self._scene.briefing}\n"
|
| f"Why this is hard: {self._scene.why_harder}\n"
|
| "Objective: minimize preventable deaths, critical injuries, exposure, and service collapse.\n"
|
| "Submit assignments as a JSON list of {resource_id, target_id} objects."
|
| )
|
| return self._build_observation(feedback=feedback, reward=0.0, done=False)
|
|
|
| def step(self, action: DisasterAction, **kwargs: Any) -> DisasterObservation:
|
| self._turn += 1
|
| self._state.step_count += 1
|
|
|
| feedback_parts: List[str] = []
|
| prev_potential = self._potential(self._targets)
|
|
|
| assignments_by_target: Dict[str, List[str]] = {tid: [] for tid in self._targets}
|
| used_resources: set[str] = set()
|
| penalty = 0.0
|
|
|
| for assignment in action.assignments:
|
| resource_id = assignment.resource_id
|
| target_id = assignment.target_id
|
|
|
| if resource_id not in self._resources:
|
| penalty += 6.0
|
| self._metrics["invalid_actions"] += 1
|
| feedback_parts.append(f"[ERR] Unknown resource '{resource_id}'")
|
| continue
|
| if target_id not in self._targets:
|
| penalty += 6.0
|
| self._metrics["invalid_actions"] += 1
|
| feedback_parts.append(f"[ERR] Unknown target '{target_id}'")
|
| continue
|
| if resource_id in used_resources:
|
| penalty += 5.0
|
| self._metrics["invalid_actions"] += 1
|
| feedback_parts.append(f"[ERR] Resource '{resource_id}' assigned more than once")
|
| continue
|
| if not self._resource_available(self._resources[resource_id], self._turn):
|
| penalty += 5.0
|
| self._metrics["invalid_actions"] += 1
|
| feedback_parts.append(f"[ERR] Resource '{resource_id}' is unavailable")
|
| continue
|
| if self._targets[target_id]["status"] == "resolved":
|
| penalty += 3.0
|
| self._metrics["ineffective_assignments"] += 1
|
| feedback_parts.append(f"[WARN] Target '{target_id}' already resolved")
|
| continue
|
|
|
| used_resources.add(resource_id)
|
| assignments_by_target[target_id].append(resource_id)
|
|
|
| penalty += self._apply_idle_penalty(used_resources)
|
| penalty += self._advance_system(assignments_by_target, feedback_parts)
|
|
|
| next_potential = self._potential(self._targets)
|
| reward = round((next_potential - prev_potential) / 10.0 - penalty, 3)
|
|
|
| done = self._all_targets_resolved() or self._turn >= self._scene.max_turns |
| if done: |
| self._final_score = self._compute_final_score() |
| feedback_parts.append( |
| f"Episode complete. Final score={self._final_score:.1f}/100." |
| ) |
|
|
| feedback = " | ".join(feedback_parts) if feedback_parts else "Assignments executed."
|
| return self._build_observation(feedback=feedback, reward=reward, done=done)
|
|
|
| @property
|
| def state(self) -> State:
|
| return State(
|
| episode_id=self._state.episode_id,
|
| step_count=self._state.step_count,
|
| scene_id=self._scene.scene_id,
|
| scene_name=self._scene.name,
|
| level=self._scene.level,
|
| )
|
|
|
| def _select_scene(
|
| self,
|
| scene_id: Optional[str],
|
| level: Optional[int],
|
| ) -> SceneConfig:
|
| if scene_id:
|
| if scene_id not in SCENE_CATALOG:
|
| raise ValueError(f"Unknown scene_id '{scene_id}'")
|
| return SCENE_CATALOG[scene_id]
|
| if level is not None:
|
| for candidate in SCENE_CATALOG.values():
|
| if candidate.level == level:
|
| return candidate
|
| raise ValueError(f"Unknown level '{level}'")
|
| return SCENE_CATALOG[DEFAULT_SCENE_ID]
|
|
|
| def _init_targets(self, scene: SceneConfig) -> Dict[str, Dict[str, Any]]:
|
| targets: Dict[str, Dict[str, Any]] = {}
|
| for cfg in scene.targets:
|
| targets[cfg.target_id] = {
|
| "config": cfg,
|
| "status": "active",
|
| "progress": 0.0,
|
| "risk": cfg.initial_risk,
|
| "people_remaining": cfg.people_true,
|
| "time_remaining": cfg.deadline_turns,
|
| "last_assigned_resources": [],
|
| "deadline_missed": False,
|
| "failed": False,
|
| }
|
| return targets
|
|
|
| def _init_resources(self, scene: SceneConfig) -> Dict[str, Dict[str, Any]]:
|
| resources: Dict[str, Dict[str, Any]] = {}
|
| for cfg in scene.resources:
|
| resources[cfg.resource_id] = {
|
| "config": cfg,
|
| "remaining_uses": cfg.max_uses,
|
| "last_target_id": None,
|
| }
|
| return resources
|
|
|
| def _resource_available(self, resource: Dict[str, Any], turn: int) -> bool:
|
| cfg = resource["config"]
|
| if cfg.available_until_turn is not None and turn > cfg.available_until_turn:
|
| return False
|
| if resource["remaining_uses"] is not None and resource["remaining_uses"] <= 0:
|
| return False
|
| return True
|
|
|
| def _apply_idle_penalty(self, used_resources: set[str]) -> float:
|
| penalty = 0.0
|
| critical_targets = [
|
| target
|
| for target in self._targets.values()
|
| if target["status"] != "resolved" and target["time_remaining"] <= 2
|
| ]
|
| if not critical_targets:
|
| return penalty
|
|
|
| for resource_id, resource in self._resources.items(): |
| if resource_id in used_resources or not self._resource_available(resource, self._turn): |
| continue |
| if self._resource_can_help_any_target(resource["config"].capabilities, critical_targets): |
| penalty += 3.0 |
| return penalty |
|
|
| def _resource_can_help_any_target(
|
| self,
|
| capabilities: Dict[str, float],
|
| targets: List[Dict[str, Any]],
|
| ) -> bool:
|
| for target in targets:
|
| weights = target["config"].capability_weights
|
| if any(capability in weights for capability in capabilities):
|
| return True
|
| return False
|
|
|
| def _advance_system(
|
| self,
|
| assignments_by_target: Dict[str, List[str]],
|
| feedback_parts: List[str],
|
| ) -> float:
|
| penalty = 0.0
|
| newly_resolved: List[str] = []
|
| deadline_hits: List[str] = []
|
|
|
| for target_id, target in self._targets.items():
|
| cfg = target["config"]
|
| resource_ids = assignments_by_target.get(target_id, [])
|
| response_power = 0.0
|
| assigned_names: List[str] = []
|
|
|
| for resource_id in resource_ids:
|
| resource = self._resources[resource_id]
|
| resource_cfg = resource["config"]
|
| match = max(
|
| (
|
| resource_cfg.capabilities[capability] * weight
|
| for capability, weight in cfg.capability_weights.items()
|
| if capability in resource_cfg.capabilities
|
| ),
|
| default=0.0,
|
| )
|
| if match <= 0.0:
|
| penalty += 3.0
|
| self._metrics["ineffective_assignments"] += 1
|
| feedback_parts.append(
|
| f"[WARN] {resource_id} does not materially help {target_id}"
|
| )
|
| continue
|
|
|
| if resource["last_target_id"] not in (None, target_id):
|
| penalty += 1.0
|
| self._metrics["reassignment_churn"] += 1
|
| response_power += match
|
| assigned_names.append(resource_id)
|
| resource["last_target_id"] = target_id
|
| if resource["remaining_uses"] is not None:
|
| resource["remaining_uses"] -= 1
|
|
|
| target["last_assigned_resources"] = assigned_names
|
| if target["status"] == "resolved" or target["failed"]:
|
| continue
|
|
|
| progress_gain = cfg.progress_per_power * response_power
|
| protection = min(0.92, target["progress"] * 0.55 + response_power * cfg.protection_per_power)
|
| target["progress"] = min(1.0, target["progress"] + progress_gain)
|
|
|
| target["risk"] = max(
|
| 0.15,
|
| min(
|
| 2.5,
|
| target["risk"] + cfg.escalation_rate - response_power * cfg.risk_reduction_per_power,
|
| ),
|
| )
|
|
|
| time_pressure = 1.0 + max(0, 1 - max(target["time_remaining"], 0) / max(1, cfg.deadline_turns)) * 0.6
|
| if target["time_remaining"] <= 0:
|
| time_pressure += 0.4
|
|
|
| protective_gap = max(0.05, 1.0 - protection)
|
|
|
| deaths_now = target["people_remaining"] * cfg.death_rate * target["risk"] * time_pressure * protective_gap
|
| critical_now = target["people_remaining"] * cfg.critical_rate * target["risk"] * time_pressure * protective_gap
|
| exposure_now = cfg.exposed_population * cfg.exposure_rate * target["risk"] * time_pressure * protective_gap
|
| service_now = cfg.service_scale * cfg.service_rate * target["risk"] * time_pressure * protective_gap
|
|
|
| self._metrics["fatalities"] += deaths_now
|
| self._metrics["critical_injuries"] += critical_now
|
| self._metrics["exposure_harm"] += exposure_now
|
| self._metrics["service_loss"] += service_now
|
|
|
| if target["people_remaining"] > 0.0:
|
| target["people_remaining"] = max(0.0, target["people_remaining"] - deaths_now)
|
|
|
| if target["progress"] >= 1.0 or (target["progress"] >= 0.86 and target["risk"] <= 0.25):
|
| if target["status"] != "resolved":
|
| target["status"] = "resolved"
|
| self._metrics["resolved_targets"] += 1
|
| newly_resolved.append(cfg.name)
|
| continue
|
|
|
| if not target["deadline_missed"] and target["time_remaining"] <= 0 and target["progress"] < 0.60: |
| target["deadline_missed"] = True |
| weighted_miss = cfg.deadline_weight * cfg.vulnerability |
| self._metrics["deadline_misses"] += weighted_miss |
| penalty += 4.0 * weighted_miss |
| deadline_hits.append(cfg.name) |
|
|
| if target["time_remaining"] < -2 and target["progress"] < 0.35 and not target["failed"]: |
| target["failed"] = True |
| target["status"] = "failed" |
| weighted_fail = cfg.deadline_weight * cfg.vulnerability |
| self._metrics["failed_targets"] += weighted_fail |
| penalty += 6.0 * weighted_fail |
| elif target["progress"] >= 0.55: |
| target["status"] = "contained" |
| else: |
| target["status"] = "active" |
|
|
| target["time_remaining"] -= 1
|
|
|
| if newly_resolved:
|
| feedback_parts.append("Resolved: " + ", ".join(newly_resolved))
|
| if deadline_hits:
|
| feedback_parts.append("Critical window missed: " + ", ".join(deadline_hits))
|
|
|
| hot_targets = self._hot_target_summaries(limit=3)
|
| if hot_targets:
|
| feedback_parts.append("Hot targets: " + ", ".join(hot_targets))
|
|
|
| return penalty
|
|
|
| def _hot_target_summaries(self, limit: int) -> List[str]:
|
| active_targets = [
|
| target
|
| for target in self._targets.values()
|
| if target["status"] not in {"resolved", "failed"}
|
| ]
|
| active_targets.sort(
|
| key=lambda target: (
|
| -target["risk"],
|
| target["time_remaining"],
|
| -target["config"].vulnerability,
|
| )
|
| )
|
| summaries: List[str] = []
|
| for target in active_targets[:limit]:
|
| summaries.append(
|
| f"{target['config'].target_id}(risk={target['risk']:.2f}, t={target['time_remaining']})"
|
| )
|
| return summaries
|
|
|
| def _potential(self, targets: Dict[str, Dict[str, Any]]) -> float: |
| total = 0.0 |
| for target in targets.values(): |
| if target["status"] == "resolved": |
| continue |
| cfg = target["config"] |
| if target["failed"]: |
| total += ( |
| 140.0 * max(0.0, target["people_remaining"]) |
| + 24.0 * cfg.exposed_population |
| + 28.0 * cfg.service_scale |
| + 40.0 * cfg.deadline_weight * cfg.vulnerability |
| ) |
| continue |
| urgency = target["risk"] * (1.0 + max(0, 2 - target["time_remaining"]) * 0.35) |
| protective_gap = max(0.05, 1.0 - target["progress"] * 0.75) |
| expected_deaths = target["people_remaining"] * cfg.death_rate * urgency * protective_gap * cfg.vulnerability |
| expected_critical = target["people_remaining"] * cfg.critical_rate * urgency * protective_gap * cfg.vulnerability |
| expected_exposure = cfg.exposed_population * cfg.exposure_rate * urgency * protective_gap |
| expected_service = cfg.service_scale * cfg.service_rate * urgency * protective_gap
|
| equity_gap = cfg.equity_weight * cfg.vulnerability * urgency * protective_gap * (1.0 - cfg.visibility)
|
| deadline_gap = max(0.0, 1.0 - max(target["time_remaining"], 0) / max(1, cfg.deadline_turns))
|
| total += (
|
| 100.0 * expected_deaths
|
| + 35.0 * expected_critical
|
| + 12.0 * expected_exposure
|
| + 18.0 * expected_service
|
| + 10.0 * equity_gap
|
| + 8.0 * deadline_gap * cfg.deadline_weight
|
| )
|
| return -total
|
|
|
| def _simulate_noop_baseline(self) -> float:
|
| targets = deepcopy(self._targets)
|
| resources = deepcopy(self._resources)
|
| metrics = deepcopy(self._metrics)
|
| for turn in range(1, self._scene.max_turns + 1):
|
| empty_assignments = {target_id: [] for target_id in targets}
|
| self._advance_copy(targets, resources, metrics, empty_assignments, turn)
|
| return max(1.0, self._compute_total_harm(metrics))
|
|
|
| def _advance_copy(
|
| self,
|
| targets: Dict[str, Dict[str, Any]],
|
| resources: Dict[str, Dict[str, Any]],
|
| metrics: Dict[str, float],
|
| assignments_by_target: Dict[str, List[str]],
|
| turn: int,
|
| ) -> None:
|
| for target_id, target in targets.items():
|
| cfg = target["config"]
|
| response_power = 0.0
|
| for resource_id in assignments_by_target.get(target_id, []):
|
| resource = resources[resource_id]
|
| resource_cfg = resource["config"]
|
| match = max(
|
| (
|
| resource_cfg.capabilities[capability] * weight
|
| for capability, weight in cfg.capability_weights.items()
|
| if capability in resource_cfg.capabilities
|
| ),
|
| default=0.0,
|
| )
|
| if match <= 0.0:
|
| metrics["ineffective_assignments"] += 1
|
| continue
|
| response_power += match
|
| if resource["remaining_uses"] is not None:
|
| resource["remaining_uses"] -= 1
|
|
|
| if target["status"] in {"resolved", "failed"}:
|
| continue
|
|
|
| progress_gain = cfg.progress_per_power * response_power
|
| protection = min(0.92, target["progress"] * 0.55 + response_power * cfg.protection_per_power)
|
| target["progress"] = min(1.0, target["progress"] + progress_gain)
|
| target["risk"] = max(
|
| 0.15,
|
| min(
|
| 2.5,
|
| target["risk"] + cfg.escalation_rate - response_power * cfg.risk_reduction_per_power,
|
| ),
|
| )
|
|
|
| time_pressure = 1.0 + max(0, 1 - max(target["time_remaining"], 0) / max(1, cfg.deadline_turns)) * 0.6
|
| if target["time_remaining"] <= 0:
|
| time_pressure += 0.4
|
| protective_gap = max(0.05, 1.0 - protection)
|
|
|
| deaths_now = target["people_remaining"] * cfg.death_rate * target["risk"] * time_pressure * protective_gap
|
| critical_now = target["people_remaining"] * cfg.critical_rate * target["risk"] * time_pressure * protective_gap
|
| exposure_now = cfg.exposed_population * cfg.exposure_rate * target["risk"] * time_pressure * protective_gap
|
| service_now = cfg.service_scale * cfg.service_rate * target["risk"] * time_pressure * protective_gap
|
|
|
| metrics["fatalities"] += deaths_now
|
| metrics["critical_injuries"] += critical_now
|
| metrics["exposure_harm"] += exposure_now
|
| metrics["service_loss"] += service_now
|
|
|
| if target["people_remaining"] > 0.0:
|
| target["people_remaining"] = max(0.0, target["people_remaining"] - deaths_now)
|
|
|
| if target["progress"] >= 1.0 or (target["progress"] >= 0.86 and target["risk"] <= 0.25):
|
| target["status"] = "resolved"
|
| metrics["resolved_targets"] += 1
|
| continue
|
|
|
| if not target["deadline_missed"] and target["time_remaining"] <= 0 and target["progress"] < 0.60: |
| target["deadline_missed"] = True |
| metrics["deadline_misses"] += cfg.deadline_weight * cfg.vulnerability |
|
|
| if target["time_remaining"] < -2 and target["progress"] < 0.35 and not target["failed"]: |
| target["failed"] = True |
| target["status"] = "failed" |
| metrics["failed_targets"] += cfg.deadline_weight * cfg.vulnerability |
| elif target["progress"] >= 0.55: |
| target["status"] = "contained" |
| else: |
| target["status"] = "active" |
|
|
| target["time_remaining"] -= 1
|
|
|
| def _compute_total_harm(self, metrics: Dict[str, float]) -> float:
|
| return (
|
| 100.0 * metrics["fatalities"]
|
| + 35.0 * metrics["critical_injuries"]
|
| + 12.0 * metrics["exposure_harm"]
|
| + 18.0 * metrics["service_loss"]
|
| + 18.0 * metrics["deadline_misses"]
|
| + 24.0 * metrics["failed_targets"]
|
| + 4.0 * metrics["invalid_actions"]
|
| + 2.0 * metrics["ineffective_assignments"]
|
| + 1.0 * metrics["reassignment_churn"]
|
| )
|
|
|
| def _compute_final_score(self) -> float:
|
| realized_harm = self._compute_total_harm(self._metrics)
|
| raw = 100.0 * (self._baseline_harm - realized_harm) / self._baseline_harm
|
| return max(0.0, min(100.0, round(raw, 2)))
|
|
|
| def _all_targets_resolved(self) -> bool: |
| return all(target["status"] == "resolved" for target in self._targets.values()) |
|
|
| def _priority_band(self, target: Dict[str, Any]) -> str: |
| cfg = target["config"] |
| if target["failed"]: |
| return "failed" |
| urgency = target["risk"] * cfg.vulnerability |
| if target["time_remaining"] <= 1 or urgency >= 1.6: |
| return "immediate" |
| if target["time_remaining"] <= 2 or urgency >= 1.15: |
| return "high" |
| if target["time_remaining"] <= 3 or urgency >= 0.8: |
| return "medium" |
| return "monitor" |
|
|
| def _build_observation(
|
| self,
|
| feedback: str,
|
| reward: float,
|
| done: bool,
|
| ) -> DisasterObservation:
|
| targets = {
|
| target_id: TargetStatus(
|
| name=target["config"].name,
|
| category=target["config"].category,
|
| status=target["status"],
|
| estimated_people=target["config"].estimated_people,
|
| observed_risk=round( |
| max( |
| 0.05, |
| min( |
| 1.0, |
| target["config"].observed_risk
|
| + (target["risk"] - target["config"].initial_risk) * 0.35,
|
| ),
|
| ), |
| 3, |
| ), |
| critical_now=(target["time_remaining"] <= 1 and target["status"] not in {"resolved", "failed"}), |
| priority_band=self._priority_band(target), |
| vulnerability=target["config"].vulnerability_label, |
| visibility=target["config"].visibility, |
| progress=round(target["progress"], 3), |
| time_remaining=target["time_remaining"],
|
| recommended_capabilities=list(target["config"].recommended_capabilities),
|
| last_assigned_resources=list(target["last_assigned_resources"]),
|
| description=(
|
| f"{target['config'].description} Critical window: {target['config'].deadline_note}"
|
| ),
|
| )
|
| for target_id, target in self._targets.items()
|
| }
|
| resources = {
|
| resource_id: ResourceStatus(
|
| name=resource["config"].name,
|
| capabilities=sorted(resource["config"].capabilities.keys()),
|
| available=self._resource_available(resource, self._turn + 1 if not done else self._turn),
|
| remaining_uses=resource["remaining_uses"],
|
| available_until_turn=resource["config"].available_until_turn,
|
| description=resource["config"].description,
|
| )
|
| for resource_id, resource in self._resources.items()
|
| }
|
| resolved_count = sum(1 for target in self._targets.values() if target["status"] == "resolved")
|
| metadata: Dict[str, Any] = {
|
| "scene_ids": ordered_scene_ids(),
|
| "score_method": "normalized_against_noop_baseline",
|
| }
|
| if done and self._final_score is not None:
|
| metadata["audit_metrics"] = {
|
| key: round(value, 2) for key, value in self._metrics.items()
|
| }
|
| metadata["baseline_harm"] = round(self._baseline_harm, 2)
|
|
|
| return DisasterObservation(
|
| scene_id=self._scene.scene_id,
|
| scene_name=self._scene.name,
|
| level=self._scene.level,
|
| narrative=self._scene.briefing,
|
| targets=targets,
|
| resources=resources,
|
| resolved_count=resolved_count,
|
| turn=self._turn,
|
| max_turns=self._scene.max_turns,
|
| feedback=feedback,
|
| final_score=self._final_score if done else None,
|
| done=done,
|
| reward=reward,
|
| metadata=metadata,
|
| )
|
|
|
