Initial commit
This commit is contained in:
277
app/services/ml/strategy_bandit.py
Normal file
277
app/services/ml/strategy_bandit.py
Normal file
@@ -0,0 +1,277 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user