from __future__ import annotations import argparse import json import random import re from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Tuple import matplotlib.pyplot as plt import numpy as np import requests from fairrecovery_env.models import AllocationItem, FairRecoveryAction, FairRecoveryObservation from server.fairrecovery_environment import FairRecoveryEnvironment OLLAMA_URL = "http://localhost:11434/api/generate" @dataclass class RolloutResult: total_reward: float transparent_reward: float final_fairness: float breakdown: List[Dict[str, float]] def _safe_json_extract(text: str) -> Dict: m = re.search(r"\{.*\}", text, flags=re.DOTALL) if not m: return {"action_type": "noop"} try: return json.loads(m.group(0)) except Exception: return {"action_type": "noop"} def _build_prompt(obs: FairRecoveryObservation, strategy_hint: str) -> str: zones = [ { "zone_id": z.zone_id, "damage": round(z.damage, 3), "service": round(z.service, 3), "vulnerable_ratio": round(z.vulnerable_ratio, 3), "risk_level": round(z.risk_level, 3), } for z in obs.zones ] state = { "day": obs.day, "step_stage": obs.step_stage, "budget_left": round(obs.budget_left, 2), "fairness_score": round(obs.fairness_score, 4), "zones": zones, } return ( "You are a disaster recovery planner.\n" "Return only one JSON object. No markdown.\n" "Valid action_type: analyze, allocate, execute, adapt, submit, noop.\n" "Rules:\n" "- Always follow stage protocol analyze->allocate->execute.\n" "- Avoid noop unless absolutely necessary.\n" "- Favor vulnerable high-damage zones fairly.\n" f"- Strategy hint: {strategy_hint}\n\n" f"Observation:\n{json.dumps(state, indent=2)}\n\n" "Output JSON action:" ) def _ask_ollama(model: str, prompt: str, temperature: float) -> Dict: payload = { "model": model, "prompt": prompt, "stream": False, "options": {"temperature": temperature}, } r = requests.post(OLLAMA_URL, json=payload, timeout=120) r.raise_for_status() txt = r.json().get("response", "") return _safe_json_extract(txt) def _to_action(payload: Dict, obs: FairRecoveryObservation) -> FairRecoveryAction: try: if payload.get("action_type") == "allocate": allocs = payload.get("allocations") or [] if not allocs: payload["allocations"] = [{"zone": 0, "resource": "power"}] clean = [] for a in payload["allocations"][:1]: zone = max(0, min(int(a.get("zone", 0)), len(obs.zones) - 1)) resource = a.get("resource", "power") if resource not in {"power", "water", "medical"}: resource = "power" clean.append(AllocationItem(zone=zone, resource=resource)) return FairRecoveryAction(action_type="allocate", allocations=clean) return FairRecoveryAction(**payload) except Exception: if obs.step_stage == "analyze": return FairRecoveryAction(action_type="analyze", critical_zones=[0, 1], reasoning="fallback") if obs.step_stage == "allocate": return FairRecoveryAction(action_type="allocate", allocations=[AllocationItem(zone=0, resource="power")]) if obs.step_stage == "execute": return FairRecoveryAction(action_type="execute") return FairRecoveryAction(action_type="submit") def run_episode(env: FairRecoveryEnvironment, model: str, strategy_hint: str, temperature: float) -> RolloutResult: obs = env.reset(difficulty="hard") total_reward = 0.0 rows: List[Dict[str, float]] = [] noop_streak = 0 step_cap = 24 for _ in range(step_cap): prompt = _build_prompt(obs, strategy_hint) payload = _ask_ollama(model=model, prompt=prompt, temperature=temperature) action = _to_action(payload, obs) if action.action_type == "noop": noop_streak += 1 else: noop_streak = 0 if noop_streak >= 4: action = FairRecoveryAction(action_type="submit") obs = env.step(action) total_reward += obs.reward if obs.info: rows.append( { "reward": float(obs.info.get("reward", 0.0)), "utility": float(obs.info.get("utility", 0.0)), "fairness": float(obs.info.get("fairness", 0.0)), } ) if obs.done: break if not obs.done: obs = env.step(FairRecoveryAction(action_type="submit")) total_reward += obs.reward if obs.info: rows.append( { "reward": float(obs.info.get("reward", 0.0)), "utility": float(obs.info.get("utility", 0.0)), "fairness": float(obs.info.get("fairness", 0.0)), } ) avg_transparent = float(np.mean([r["reward"] for r in rows])) if rows else 0.0 return RolloutResult( total_reward=float(total_reward), transparent_reward=avg_transparent, final_fairness=float(obs.fairness_score), breakdown=rows, ) def _moving_avg(xs: List[float], w: int = 6) -> List[float]: out = [] for i in range(len(xs)): s = max(0, i - w + 1) out.append(float(np.mean(xs[s : i + 1]))) return out def train_local(model: str, episodes: int, seed: int) -> Dict: random.seed(seed) np.random.seed(seed) env = FairRecoveryEnvironment() strategies = [ "strict fairness: prioritize vulnerable zones first", "balanced fairness + utility with low-noop behavior", "utility-first but keep fairness parity each day", ] scores = np.zeros(len(strategies), dtype=float) counts = np.zeros(len(strategies), dtype=float) baseline_hint = "generic planning with no fairness emphasis" baseline_runs = [run_episode(env, model, baseline_hint, 0.1) for _ in range(4)] baseline_avg = float(np.mean([r.transparent_reward for r in baseline_runs])) step_rewards: List[float] = [] step_fairness: List[float] = [] step_utility: List[float] = [] # Lightweight bandit tuning over prompting strategy. for ep in range(episodes): eps = max(0.15, 0.6 - (ep / max(1, episodes))) if random.random() < eps or np.all(counts == 0): idx = random.randint(0, len(strategies) - 1) else: ucb = np.where( counts > 0, scores / np.maximum(counts, 1.0) + np.sqrt(2 * np.log(ep + 2) / np.maximum(counts, 1.0)), 1e9, ) idx = int(np.argmax(ucb)) rr = run_episode(env, model, strategies[idx], 0.15) reward = rr.transparent_reward scores[idx] += reward counts[idx] += 1 for row in rr.breakdown: step_rewards.append(row["reward"]) step_fairness.append(row["fairness"]) step_utility.append(row["utility"]) best_idx = int(np.argmax(np.where(counts > 0, scores / np.maximum(counts, 1.0), -1e9))) best_hint = strategies[best_idx] trained_runs = [run_episode(env, model, best_hint, 0.1) for _ in range(4)] trained_avg = float(np.mean([r.transparent_reward for r in trained_runs])) x = np.arange(1, len(step_rewards) + 1) plt.figure(figsize=(9, 4.5)) plt.plot(x, _moving_avg(step_rewards), linewidth=2) plt.title("Reward vs Steps (Ollama qwen2.5)") plt.xlabel("Step") plt.ylabel("Reward (0-1)") plt.ylim(0, 1) plt.grid(alpha=0.25) plt.tight_layout() plt.savefig("reward_vs_steps.png", dpi=140) plt.close() plt.figure(figsize=(9, 4.5)) plt.plot(x, _moving_avg(step_fairness), linewidth=2, color="#1f77b4") plt.title("Fairness vs Steps (Ollama qwen2.5)") plt.xlabel("Step") plt.ylabel("Fairness (0-1)") plt.ylim(0, 1) plt.grid(alpha=0.25) plt.tight_layout() plt.savefig("fairness_vs_steps.png", dpi=140) plt.close() plt.figure(figsize=(6, 6)) plt.scatter(step_utility, step_fairness, alpha=0.5, s=16) plt.title("Utility vs Fairness") plt.xlabel("Utility (0-1)") plt.ylabel("Fairness (0-1)") plt.xlim(0, 1) plt.ylim(0, 1) plt.grid(alpha=0.25) plt.tight_layout() plt.savefig("utility_vs_fairness.png", dpi=140) plt.close() summary = { "engine": "ollama", "model": model, "episodes": episodes, "baseline_avg_reward": round(baseline_avg, 3), "trained_avg_reward": round(trained_avg, 3), "best_strategy": best_hint, "strategy_counts": counts.tolist(), } Path("ollama_training_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") return summary def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--model", default="qwen2.5:7b") parser.add_argument("--episodes", type=int, default=10) parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() summary = train_local(model=args.model, episodes=args.episodes, seed=args.seed) print(json.dumps(summary, indent=2)) if __name__ == "__main__": main()