""" Thompson Sampling Contextual Bandit — Strategy Selector ========================================================= Replaces the greedy SQL auto-tuner with proper online RL. Problem ------- The old greedy tuner picks the strategy with the highest *average* quality across all history. It never explores alternatives once one strategy leads, and it ignores context (the same strategy isn't best at all times of day and all load levels). Solution -------- A contextual multi-armed bandit with Thompson Sampling: Context = (time_band, load_band) — up to 7 × 4 = 28 states Arms = 4 strategies — balanced / fuel_saver / aggressive_speed / zone_strict Reward = quality_score / 100 — 0..1 float, already logged by MLDataCollector How Thompson Sampling works --------------------------- Each (context, arm) pair has a Beta(α, β) posterior where: α = sum of rewards seen so far (high quality calls push α up) β = sum of "anti-rewards" (low quality calls push β up) To SELECT a strategy: 1. For each arm, sample θ ~ Beta(α, β) 2. Pick arm with highest θ → Naturally balances exploration (uncertain arms get sampled often) with exploitation (well-known good arms dominate when confident) To UPDATE after an assignment: reward = quality_score / 100 α += reward β += (1 - reward) Bootstrap --------- On first startup the entire historical SQLite DB is replayed to warm up posteriors so the bandit starts informed, not blank. If saved state already exists it is loaded from disk instead. Persistence ----------- ml_data/strategy_bandit.json — saved every 10 updates. """ import json import logging import os import threading from typing import Dict, List, Optional, Tuple import numpy as np logger = logging.getLogger(__name__) _SAVE_PATH = os.getenv("BANDIT_PATH", "ml_data/strategy_bandit.json") _DB_PATH = os.getenv("ML_DB_PATH", "ml_data/ml_store.db") ARMS = ["balanced", "fuel_saver", "aggressive_speed", "zone_strict"] class ContextualBandit: """ Thompson Sampling bandit for ml_strategy selection. Context key : "{time_band}|{load_band}" Arms : ARMS list (4 strategies) Prior : Beta(1, 1) — uniform, no initial preference """ def __init__(self): self._lock = threading.Lock() # _posteriors[ctx_key][arm] = [alpha, beta] self._posteriors: Dict[str, Dict[str, List[float]]] = {} self._update_count = 0 self._total_pulls = 0 self._load() self._bootstrap_from_db() # ── Public API ─────────────────────────────────────────────────────────── def select(self, time_band: str, load_band: str) -> str: """ Sample from Beta posteriors and return the strategy with the highest sample. Explores uncertain arms naturally; exploits known-good arms as confidence grows. """ ctx = self._ctx(time_band, load_band) with self._lock: samples = { arm: float(np.random.beta(*self._ab(ctx, arm))) for arm in ARMS } self._total_pulls += 1 chosen = max(samples, key=samples.__getitem__) logger.debug( f"[Bandit] ctx={ctx} " f"samples={{{', '.join(f'{k}:{v:.3f}' for k, v in samples.items())}}} " f"→ {chosen}" ) return chosen def update(self, time_band: str, load_band: str, strategy: str, quality_score: float) -> None: """Update the Beta posterior for (context, arm) after observing quality.""" if strategy not in ARMS: return ctx = self._ctx(time_band, load_band) reward = min(1.0, max(0.0, float(quality_score) / 100.0)) with self._lock: ab = self._ab(ctx, strategy) ab[0] += reward # alpha ← quality adds to success mass ab[1] += (1.0 - reward) # beta ← (1-quality) adds to failure mass self._update_count += 1 should_save = self._update_count % 10 == 0 if should_save: self._save() def best_arm(self, time_band: str, load_band: str) -> Tuple[str, float]: """ Return the arm with the highest posterior mean (pure exploitation, no sampling noise). Used for logging and dashboard, not for live selection. """ ctx = self._ctx(time_band, load_band) with self._lock: means = { arm: self._ab(ctx, arm)[0] / sum(self._ab(ctx, arm)) for arm in ARMS } best = max(means, key=means.__getitem__) return best, round(means[best], 4) def get_stats(self) -> Dict: """Return per-context arm statistics for the ML admin dashboard.""" with self._lock: stats: Dict[str, Dict] = {} for ctx, arms in sorted(self._posteriors.items()): stats[ctx] = {} for arm in ARMS: if arm not in arms: stats[ctx][arm] = {"mean_reward": 0.5, "observations": 0, "alpha": 1.0, "beta": 1.0} continue alpha, beta = arms[arm] n = alpha + beta - 2.0 # subtract the Beta(1,1) prior mass mean = alpha / (alpha + beta) stats[ctx][arm] = { "mean_reward": round(mean, 4), "observations": round(max(0, n), 1), "alpha": round(alpha, 2), "beta": round(beta, 2), } return { "total_pulls": self._total_pulls, "total_updates": self._update_count, "context_count": len(self._posteriors), "arms": ARMS, "contexts": stats, } # ── Persistence ────────────────────────────────────────────────────────── def _save(self) -> None: try: with self._lock: snapshot = { "posteriors": { ctx: {arm: list(ab) for arm, ab in arms.items()} for ctx, arms in self._posteriors.items() }, "update_count": self._update_count, "total_pulls": self._total_pulls, } os.makedirs(os.path.dirname(_SAVE_PATH) or ".", exist_ok=True) with open(_SAVE_PATH, "w", encoding="utf-8") as f: json.dump(snapshot, f, indent=2) except Exception as e: logger.warning(f"[Bandit] Save failed: {e}") def _load(self) -> None: try: if not os.path.exists(_SAVE_PATH): logger.info("[Bandit] No saved state — will bootstrap from DB.") return with open(_SAVE_PATH, "r", encoding="utf-8") as f: data = json.load(f) self._posteriors = data.get("posteriors", {}) self._update_count = data.get("update_count", 0) self._total_pulls = data.get("total_pulls", 0) logger.info( f"[Bandit] Loaded from disk — " f"{len(self._posteriors)} contexts, {self._update_count} updates" ) except Exception as e: logger.warning(f"[Bandit] Load failed (starting fresh): {e}") # ── Bootstrap ──────────────────────────────────────────────────────────── def _bootstrap_from_db(self) -> None: """ Warm-up posteriors by replaying every historical assignment event. Skipped if the saved JSON already reflects current DB data. """ if self._update_count > 0: logger.info( f"[Bandit] Already warmed ({self._update_count} updates) — " "skipping DB bootstrap." ) return try: import sqlite3 from app.services.ml.behavior_analyzer import ( time_band as _tb, load_band as _lb, ) conn = sqlite3.connect(_DB_PATH) rows = conn.execute( "SELECT timestamp, avg_load, ml_strategy, quality_score " "FROM assignment_ml_log " "WHERE ml_strategy IS NOT NULL AND quality_score IS NOT NULL " "ORDER BY id ASC" ).fetchall() conn.close() if not rows: logger.info("[Bandit] No historical data — starting with uniform priors.") return for ts, avg_load, strategy, quality in rows: t_band = _tb(str(ts)) if ts else "UNKNOWN" l_band = _lb(float(avg_load or 0)) self.update(t_band, l_band, strategy, float(quality or 50)) # Reset counter so we don't confuse bootstrap updates with live ones with self._lock: self._update_count = len(rows) logger.info( f"[Bandit] Bootstrapped from {len(rows)} historical events — " f"{len(self._posteriors)} contexts, {len(ARMS)} arms." ) self._save() except Exception as e: logger.warning(f"[Bandit] Bootstrap failed (non-fatal): {e}") # ── Helpers ────────────────────────────────────────────────────────────── @staticmethod def _ctx(time_band: str, load_band: str) -> str: return f"{time_band}|{load_band}" def _ab(self, ctx: str, arm: str) -> List[float]: """Get or create Beta(1, 1) uniform prior for (ctx, arm). NOT thread-safe alone.""" if ctx not in self._posteriors: self._posteriors[ctx] = {} if arm not in self._posteriors[ctx]: self._posteriors[ctx][arm] = [1.0, 1.0] return self._posteriors[ctx][arm] # ── Singleton ──────────────────────────────────────────────────────────────── _bandit_instance: Optional[ContextualBandit] = None _bandit_lock = threading.Lock() def get_bandit() -> ContextualBandit: """Return the process-level singleton ContextualBandit.""" global _bandit_instance if _bandit_instance is None: with _bandit_lock: if _bandit_instance is None: _bandit_instance = ContextualBandit() return _bandit_instance