File size: 4,719 Bytes
4d57df8 1c03487 4d57df8 1c03487 4d57df8 1c03487 0092607 1c03487 0092607 4d57df8 1c03487 4d57df8 1c03487 0092607 4d57df8 1c03487 c5694dd 4d57df8 1c03487 0092607 1c03487 0092607 4d57df8 1c03487 0092607 4d57df8 1c03487 0092607 4d57df8 1c03487 0092607 4d57df8 9eafd01 4d57df8 1c03487 4d57df8 1c03487 4d57df8 0092607 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from typing import List
from dataclasses import dataclass
from models import CityState, ContainmentAction
from server.constants import (
INFECTION_THRESHOLD,
SAFE_THRESHOLD,
HOSPITAL_BREACH_POINT,
TREATMENT_REDUCTION,
TASK_CONFIG,
)
@dataclass
class TrajectoryStep:
step: int
city_state: CityState
action: ContainmentAction
reward: float
done: bool
@dataclass
class GradeResult:
final_score: float
containment_score: float
hospital_score: float
efficiency_score: float
speed_score: float
hospital_breached: bool
districts_contained: int
total_steps: int
def grade_trajectory(trajectory: List[TrajectoryStep], task_name: str) -> GradeResult:
if not trajectory:
return GradeResult(0.0, 0.0, 0.0, 0.0, 0.0, False, 0, 0)
config = TASK_CONFIG[task_name]
num_districts = config["num_districts"]
max_steps = config["max_steps"]
total_steps = len(trajectory)
# Containment: fraction of district-days below the infection threshold.
# First 2 steps are excluded — initial conditions are outside the agent's control.
safe_district_days = 0
for step in trajectory[2:]:
for district in step.city_state.districts:
if district.true_infection_rate <= INFECTION_THRESHOLD:
safe_district_days += 1
total_district_days = max(len(trajectory) - 2, 1) * num_districts
containment_score = safe_district_days / total_district_days
# Hospital: average capacity preserved across all district-days.
# Any breach applies a 0.6 multiplier to the final hospital component.
hospital_breached = False
total_capacity_preserved = 0.0
for step in trajectory:
for district in step.city_state.districts:
if district.hospital_capacity_remaining <= HOSPITAL_BREACH_POINT:
hospital_breached = True
total_capacity_preserved += district.hospital_capacity_remaining
avg_capacity = total_capacity_preserved / (total_steps * num_districts)
hospital_score = round(min(1.0, max(0.0, avg_capacity * (0.6 if hospital_breached else 1.0))), 4)
# Efficiency: fraction of resource actions that targeted the right district.
# Uses pre-action infection state so successful treatments aren't penalised retroactively.
correct_actions = 0
total_resource = 0
for idx, step in enumerate(trajectory):
if step.action.action_type not in {"allocate", "test"}:
continue
total_resource += 1
if idx > 0:
prev_districts = trajectory[idx - 1].city_state.districts
pre_action_rate = prev_districts[step.action.district_id].true_infection_rate
highest_before = max(prev_districts, key=lambda d: d.true_infection_rate).district_id
else:
curr_d = step.city_state.districts[step.action.district_id]
pre_action_rate = curr_d.true_infection_rate + TREATMENT_REDUCTION
highest_before = max(step.city_state.districts, key=lambda d: d.true_infection_rate).district_id
if pre_action_rate > INFECTION_THRESHOLD or step.action.district_id == highest_before:
correct_actions += 1
efficiency_score = correct_actions / max(total_resource, 1)
# Speed: reward finishing before max_steps. Zero if episode ran to the limit.
last_step = trajectory[-1]
speed_score = round(max(0.0, 1.0 - total_steps / max_steps), 4) \
if last_step.done and total_steps < max_steps else 0.0
# Weighted final score.
# Hospital is highest-weighted because system collapse is catastrophic and irreversible.
final_score = round(min(1.0, max(0.0,
containment_score * 0.30 +
hospital_score * 0.45 +
efficiency_score * 0.15 +
speed_score * 0.10
)), 4)
districts_contained = sum(
1 for d in trajectory[-1].city_state.districts
if d.true_infection_rate < SAFE_THRESHOLD
)
return GradeResult(
final_score = final_score,
containment_score = round(containment_score, 4),
hospital_score = hospital_score,
efficiency_score = round(efficiency_score, 4),
speed_score = speed_score,
hospital_breached = hospital_breached,
districts_contained = districts_contained,
total_steps = total_steps,
)
def grade_task(trajectory: List[TrajectoryStep], task_name: str) -> float:
return grade_trajectory(trajectory, task_name).final_score
|