File size: 4,792 Bytes
f27fd6a 1c03487 0092607 7537292 1c03487 0092607 1c03487 f27fd6a 7537292 1c03487 7537292 1c03487 f27fd6a 1c03487 f27fd6a 1c03487 f27fd6a 1c03487 0092607 f27fd6a 1c03487 f27fd6a 1c03487 7537292 f27fd6a 1c03487 7537292 1c03487 7537292 1c03487 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 | import os, sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from typing import List
from models import ContainmentAction, CityObservation
class EpisodicMemory:
"""
Stores high-reward steps from past rollouts and retrieves similar
past decisions to guide the next rollout via prompt injection.
Key design choice: actions are stored as RELATIVE ranks rather than
absolute district IDs. Rank 0 = highest-infected district at the time
of the action, rank 1 = second-highest, and so on.
This matters because each episode randomises spread rates and densities,
so "allocate D1" from rollout 1 may refer to a completely different
epidemiological situation in rollout 2. Storing rank instead means
memory encodes the strategy ("target the worst district") rather than
an accident of episode initialisation ("target district 1").
On retrieval, ranks are resolved back to actual current district IDs
so the injected prompt text is immediately actionable.
"""
def __init__(self, max_size: int = 20):
self.memories: List[dict] = []
self.max_size = max_size
def store(self, obs: CityObservation, action: ContainmentAction, reward: float):
if reward < -0.3:
return
phase = "early" if obs.current_step <= obs.max_steps // 3 else \
"mid" if obs.current_step <= 2 * obs.max_steps // 3 else "late"
# Sort districts by infection rate descending to get current rankings
sorted_by_infection = sorted(
obs.districts,
key=lambda d: d.reported_infection_rate,
reverse=True
)
id_to_rank = {d.district_id: rank for rank, d in enumerate(sorted_by_infection)}
# Store rank rather than absolute ID
district_rank = id_to_rank.get(action.district_id, 0)
target_infection = next(
(d.reported_infection_rate for d in obs.districts if d.district_id == action.district_id),
0.0
)
highest_infection = sorted_by_infection[0].reported_infection_rate if sorted_by_infection else 0.0
self.memories.append({
"infection_profile": [round(d.reported_infection_rate, 2) for d in obs.districts],
"resources": obs.available_resources,
"phase": phase,
"action_type": action.action_type,
"district_rank": district_rank, # 0 = highest infected
"target_infection": round(target_infection, 2),
"highest_infection": round(highest_infection, 2),
"reward": round(reward, 4),
})
self.memories.sort(key=lambda m: m["reward"], reverse=True)
self.memories = self.memories[:self.max_size]
def retrieve(self, obs: CityObservation, top_k: int = 5) -> str:
if not self.memories:
return ""
current = [round(d.reported_infection_rate, 2) for d in obs.districts]
phase = "early" if obs.current_step <= obs.max_steps // 3 else \
"mid" if obs.current_step <= 2 * obs.max_steps // 3 else "late"
def score(memory: dict) -> float:
profile = memory["infection_profile"]
if len(profile) != len(current):
return float("inf")
l1 = sum(abs(a - b) for a, b in zip(profile, current))
phase_bonus = 0.0 if memory.get("phase") == phase else 0.1
return l1 + phase_bonus
ranked = sorted(self.memories, key=score)
top = ranked[:top_k]
# Resolve stored ranks back to current district IDs for this episode
current_sorted = sorted(
obs.districts,
key=lambda d: d.reported_infection_rate,
reverse=True
)
lines = ["Past decisions that earned positive reward (use as guidance):"]
for m in top:
rank = m["district_rank"]
if rank < len(current_sorted):
resolved_id = current_sorted[rank].district_id
rank_label = f"rank-{rank} district (currently D{resolved_id})"
else:
resolved_id = current_sorted[0].district_id if current_sorted else 0
rank_label = f"rank-0 district (currently D{resolved_id})"
lines.append(
f" Phase={m.get('phase','?')} resources={m['resources']} "
f"highest={m['highest_infection']:.2f} target={m['target_infection']:.2f}: "
f"'{m['action_type']}' {rank_label} → reward {m['reward']:+.2f}"
)
return "\n".join(lines)
def clear(self):
self.memories = []
def __len__(self) -> int:
return len(self.memories)
|