"""Blame attribution for one test case: Shapley values and simpler baselines. The cooperative game -------------------- Players are pipeline steps that have a "fix" (an alternative implementation). For a coalition S, ``v(S)`` is the success metric when the steps in S run their fix and all others run as-run, *under the same exogenous noise* as the factual run (or averaged over a set of noise worlds shared by all coalitions). ``v(empty)`` is the factual outcome; ``v(all)`` is the outcome with every fix applied. Shapley values split ``v(all) - v(empty)`` among the steps. """ from __future__ import annotations import hashlib import math from dataclasses import dataclass, field from statistics import NormalDist from typing import Any, Callable, Mapping, Sequence import numpy as np from .noise import NoiseSource, abduct from .pipeline import Case, Pipeline, Trace, _fingerprint __all__ = [ "BlameGame", "CaseBlame", "ShapleyResult", "analyze_case", "bernstein_radius", "exact_shapley", "hoeffding_radius", "last_touch", "leave_one_out", "mc_shapley", "single_swap", ] # --------------------------------------------------------------------------- # The game # --------------------------------------------------------------------------- class BlameGame: """Coalition value function ``v(S)`` for one case, with memoization. Parameters ---------- pipeline, inputs: What to run. fixes: ``{step: variant}`` -- the alternative each player switches to. Defaults to each step's ``default_fix``. players: Steps to attribute over (default: all steps that have a fix). world: A single :class:`NoiseSource` shared by all coalitions (deterministic counterfactual replay). world_factory, n_worlds, crn: Alternatively, ``v(S)`` is the mean over ``n_worlds`` worlds ``world_factory(j)``. With ``crn=True`` every coalition uses the same worlds ``j = 0..n_worlds-1`` (common random numbers). With ``crn=False`` each coalition gets its own independent worlds -- the naive "just re-run it" estimator, kept for comparison. """ def __init__(self, pipeline: Pipeline, inputs: Mapping[str, Any], fixes: Mapping[str, str] | None = None, players: Sequence[str] | None = None, world: NoiseSource | None = None, world_factory: Callable[[int], NoiseSource] | None = None, n_worlds: int = 1, crn: bool = True): self.pipeline = pipeline self.inputs = dict(inputs) self.fixes = dict(fixes) if fixes is not None else pipeline.default_fixes() self.players = list(players) if players is not None else [ s for s in pipeline.step_names if s in self.fixes] missing = [p for p in self.players if p not in self.fixes] if missing: raise ValueError(f"players without a fix variant: {missing}") if world_factory is None: w = world if world is not None else NoiseSource(0) world_factory, n_worlds = (lambda j: w), 1 self.world_factory = world_factory self.n_worlds = n_worlds self.crn = crn self._values: dict[int, float] = {} self._crn_worlds: list[NoiseSource] | None = None @property def n(self) -> int: return len(self.players) @property def full_mask(self) -> int: return (1 << self.n) - 1 def mask_of(self, names) -> int: m = 0 for nm in names: m |= 1 << self.players.index(nm) return m def names_of(self, mask: int) -> list[str]: return [p for i, p in enumerate(self.players) if mask >> i & 1] def intervention(self, mask: int) -> dict[str, str]: return {p: self.fixes[p] for p in self.names_of(mask)} def worlds_for(self, mask: int) -> list[NoiseSource]: if self.crn: if self._crn_worlds is None: self._crn_worlds = [self.world_factory(j) for j in range(self.n_worlds)] return self._crn_worlds base = self.n_worlds * (mask + 1) return [self.world_factory(base + j) for j in range(self.n_worlds)] def value(self, mask: int) -> float: """``v(S)`` for the coalition encoded by bitmask ``mask`` (memoized).""" v = self._values.get(mask) if v is None: iv = self.intervention(mask) v = float(np.mean([self.pipeline.run(self.inputs, iv, w).score for w in self.worlds_for(mask)])) self._values[mask] = v return v def value_of(self, names) -> float: return self.value(self.mask_of(names)) def trace(self, mask: int, world_index: int = 0) -> Trace: """One concrete trace for a coalition, in a reference world shared by all masks.""" w = self.worlds_for(0)[world_index] # same reference world for every mask return self.pipeline.run(self.inputs, self.intervention(mask), w) @property def coalitions_evaluated(self) -> int: return len(self._values) # --------------------------------------------------------------------------- # Concentration bounds # --------------------------------------------------------------------------- def hoeffding_radius(n: int, value_range: float, delta: float) -> float: """Two-sided Hoeffding radius for the mean of n i.i.d. variables with given range.""" return value_range * math.sqrt(math.log(2.0 / delta) / (2.0 * n)) def bernstein_radius(sample_var: float, n: int, value_range: float, delta: float) -> float: """Two-sided empirical-Bernstein radius (Maurer & Pontil 2009, Thm 4). With probability >= 1 - delta, ``|mean - E X| <= sqrt(2 V ln(4/delta) / n) + 7 R ln(4/delta) / (3 (n - 1))`` where V is the unbiased sample variance and R the range. Much tighter than Hoeffding when the variance is small (e.g. most marginal contributions are exactly 0). """ if n < 2: return float("inf") L = math.log(4.0 / delta) return math.sqrt(2.0 * sample_var * L / n) + 7.0 * value_range * L / (3.0 * (n - 1)) def _radius(bound: str, xs: np.ndarray, value_range: float, delta: float) -> float: n = len(xs) if bound == "hoeffding": return hoeffding_radius(n, value_range, delta) if bound == "bernstein": return bernstein_radius(float(np.var(xs, ddof=1)) if n > 1 else 0.0, n, value_range, delta) if bound == "normal": # asymptotic 95%-style CI, *not* a finite-sample guarantee if n < 2: return float("inf") z = NormalDist().inv_cdf(1.0 - delta / 2.0) return z * float(np.std(xs, ddof=1)) / math.sqrt(n) raise ValueError(f"unknown bound {bound!r}") def _separated(means: np.ndarray, rad: np.ndarray, rule: str) -> bool: order = np.argsort(-means) lo, hi = means - rad, means + rad if rule == "top1": return bool(lo[order[0]] > np.max(hi[order[1:]])) if len(means) > 1 else True if rule == "ranking": return all(lo[a] > hi[b] for a, b in zip(order[:-1], order[1:])) raise ValueError(f"unknown stopping rule {rule!r}") # --------------------------------------------------------------------------- # Shapley values # --------------------------------------------------------------------------- @dataclass class ShapleyResult: values: dict[str, float] method: str ci: dict[str, tuple[float, float]] | None = None n_permutations: int = 0 coalitions_evaluated: int = 0 cost: float = 0.0 stop_reason: str = "" v_empty: float = 0.0 v_full: float = 0.0 history: list[dict] = field(default_factory=list) # checkpoints for MC def exact_shapley(game: BlameGame) -> ShapleyResult: """Exact Shapley values by enumerating all 2^n coalitions (memoized). ``phi_i = sum_{S not containing i} |S|! (n-|S|-1)! / n! * (v(S+i) - v(S))``. """ n = game.n c0 = game.pipeline.meter.total w = [math.factorial(s) * math.factorial(n - s - 1) / math.factorial(n) for s in range(n)] vals = [game.value(m) for m in range(1 << n)] phi = {} for i, p in enumerate(game.players): bit = 1 << i tot = 0.0 for m in range(1 << n): if not m & bit: tot += w[bin(m).count("1")] * (vals[m | bit] - vals[m]) phi[p] = tot return ShapleyResult(phi, "exact", coalitions_evaluated=game.coalitions_evaluated, cost=game.pipeline.meter.total - c0, stop_reason="exhaustive", v_empty=vals[0], v_full=vals[-1]) def mc_shapley(game: BlameGame, max_permutations: int = 2000, *, delta: float = 0.05, bound: str = "bernstein", stop: str | None = None, width: float | None = None, budget: float | None = None, min_permutations: int = 8, growth: float = 1.5, metric_range: tuple[float, float] = (0.0, 1.0), sample_worlds: bool = False, seed: int = 0) -> ShapleyResult: """Monte Carlo permutation Shapley with simultaneous confidence intervals. Each random permutation yields one marginal contribution per player; their average is unbiased for the Shapley value, and because each permutation's marginals telescope to ``v(all) - v(empty)``, the estimates satisfy efficiency exactly. Coalition values are memoized, so late permutations are often free. Confidence intervals: marginals lie in ``[-(hi-lo), hi-lo]``. At checkpoint k (geometric schedule ``min_permutations * growth^k``) we use per-player radius with ``delta_k = delta / (n * k * (k+1))``; summing over players and checkpoints gives a union bound <= delta, so the intervals hold *simultaneously* and remain valid under adaptive stopping (for the finite-sample bounds "bernstein"/"hoeffding"; "normal" is heuristic). With ``stop=None`` there is one checkpoint at the end (``delta / n``). Stopping (first one met wins): * ``stop="top1"``: top player's interval is disjoint from all others (best-arm identification); ``stop="ranking"``: all adjacent pairs disjoint. * ``width``: every radius <= width. * ``budget``: cost spent in this call reaches ``budget``. * ``max_permutations`` reached. ``sample_worlds=True`` (for games built with a ``world_factory``) draws a fresh noise world per permutation instead of using the averaged, memoized ``v(S)``: the estimate is then unbiased for the Shapley value of the world-averaged game while each permutation costs only n+1 runs. With ``game.crn`` all n+1 evaluations of permutation t share world t (common random numbers); otherwise every evaluation gets its own world, which turns each marginal into a difference of two independent outcomes. """ rng = np.random.default_rng(seed) n = game.n R = 2.0 * (metric_range[1] - metric_range[0]) c0 = game.pipeline.meter.total adaptive = stop is not None or width is not None marg: list[np.ndarray] = [] history: list[dict] = [] k = 0 next_check = min_permutations if adaptive else max_permutations reason = "max_permutations" rad = np.full(n, np.inf) means = np.zeros(n) t = 0 def evaluate(mask: int, t: int, pos: int) -> float: if not sample_worlds: return game.value(mask) j = t if game.crn else 10**7 + t * (n + 1) + pos if j not in world_cache: world_cache.clear() world_cache[j] = game.world_factory(j) return game.pipeline.run(game.inputs, game.intervention(mask), world_cache[j]).score world_cache: dict[int, NoiseSource] = {} while t < max_permutations: perm = rng.permutation(n) row = np.empty(n) mask, prev = 0, evaluate(0, t, 0) for pos, i in enumerate(perm, start=1): mask |= 1 << int(i) cur = evaluate(mask, t, pos) row[i] = cur - prev prev = cur marg.append(row) t += 1 spent = game.pipeline.meter.total - c0 over_budget = budget is not None and spent >= budget if t >= next_check or over_budget or t == max_permutations: k += 1 dk = delta / (n * k * (k + 1)) if adaptive else delta / n X = np.array(marg) means = X.mean(axis=0) rad = np.array([_radius(bound, X[:, j], R, dk) for j in range(n)]) history.append({"permutations": t, "cost": spent, "means": means.tolist(), "radius": rad.tolist(), "coalitions": game.coalitions_evaluated}) if stop is not None and _separated(means, rad, stop): reason = f"separated ({stop})" break if width is not None and np.all(rad <= width): reason = f"width <= {width}" break if over_budget: reason = "budget" break next_check = max(t + 1, int(math.ceil(min_permutations * growth ** k))) X = np.array(marg) means = X.mean(axis=0) phi = {p: float(means[i]) for i, p in enumerate(game.players)} ci = {p: (float(means[i] - rad[i]), float(means[i] + rad[i])) for i, p in enumerate(game.players)} if sample_worlds: # v(empty)/v(all) were never computed exactly; don't spend runs on them v0 = vf = float("nan") else: v0, vf = game.value(0), game.value(game.full_mask) return ShapleyResult(phi, f"mc-{bound}", ci, n_permutations=t, coalitions_evaluated=game.coalitions_evaluated, cost=game.pipeline.meter.total - c0, stop_reason=reason, v_empty=v0, v_full=vf, history=history) # --------------------------------------------------------------------------- # Baselines # --------------------------------------------------------------------------- def single_swap(game: BlameGame) -> dict[str, float]: """"Fix-alone" necessity: ``v({i}) - v(empty)``. Does fixing only this step help?""" v0 = game.value(0) return {p: game.value(1 << i) - v0 for i, p in enumerate(game.players)} def leave_one_out(game: BlameGame) -> dict[str, float]: """``v(all) - v(all \\ {i})``: does the full fix still need this step? Zero for every member of an over-determined set of causes (any one of which suffices), which is why it under-blames redundant failures. """ vf = game.value(game.full_mask) return {p: vf - game.value(game.full_mask & ~(1 << i)) for i, p in enumerate(game.players)} def last_touch(game: BlameGame) -> str | None: """Naive heuristic: the *last* step (in execution order) whose output differs between the factual run and the all-fixed run -- "who last touched the wrong value". Uses the first noise world. Returns None if nothing differs. """ f = game.trace(0) g = game.trace(game.full_mask) last = None for s in game.pipeline.steps: if _fingerprint(f.outputs[s.output]) != _fingerprint(g.outputs[s.output]): last = s.name return last # --------------------------------------------------------------------------- # One-case analysis bundle # --------------------------------------------------------------------------- @dataclass class CaseBlame: case_id: str players: list[str] shapley: ShapleyResult fix_alone: dict[str, float] leave_one_out: dict[str, float] last_touch: str | None factual_score: float v_empty: float v_full: float coalition_values: dict[str, float] | None factual_outputs: dict fixed_outputs: dict def make_game(pipeline: Pipeline, case: Case, *, fixes=None, players=None, mode: str = "recorded", n_worlds: int = 1, crn: bool = True, seed: int = 0, factual: Trace | None = None) -> BlameGame: """Build the game for a case under one of three noise modes. * ``"recorded"``: replay the factual world exactly (needs the noise, e.g. because it was logged/cached or the run is simulated). Deterministic. * ``"posterior"``: only the factual *choices* are known; average over ``n_worlds`` worlds abducted with :func:`~agentblame.noise.abduct`. * ``"prior"``: ignore what happened; average over fresh worlds. This is the interventional question "how often would the fix help on a new sample of this input", not a counterfactual about this failure. """ if mode == "recorded": return BlameGame(pipeline, case.inputs, fixes, players, world=NoiseSource(case.seed)) if mode == "posterior": if factual is None: factual = pipeline.run(case.inputs, noise=NoiseSource(case.seed)) # The abduction RNG must depend on the case too; otherwise every case would # reuse the same posterior draws and batch averages would be correlated. cid = int.from_bytes(hashlib.blake2b(repr(case.seed).encode(), digest_size=4).digest(), "little") fac = lambda j: abduct(factual, seed=[seed, j, cid], name=f"post[{case.seed!r},{seed},{j}]") elif mode == "prior": fac = lambda j: NoiseSource(("prior", case.seed, seed, j)) else: raise ValueError(f"unknown mode {mode!r}") return BlameGame(pipeline, case.inputs, fixes, players, world_factory=fac, n_worlds=n_worlds, crn=crn) def analyze_case(pipeline: Pipeline, case: Case, *, fixes=None, players=None, mode: str = "recorded", n_worlds: int = 1, crn: bool = True, method: str = "auto", mc: Mapping | None = None, seed: int = 0, factual: Trace | None = None) -> CaseBlame: """Shapley + baselines for one case. ``method``: "exact", "mc", or "auto" (exact when there are at most 10 players).""" if factual is None: factual = pipeline.run(case.inputs, noise=NoiseSource(case.seed)) game = make_game(pipeline, case, fixes=fixes, players=players, mode=mode, n_worlds=n_worlds, crn=crn, seed=seed, factual=factual) use_exact = method == "exact" or (method == "auto" and game.n <= 10) sh = exact_shapley(game) if use_exact else mc_shapley(game, seed=seed, **dict(mc or {})) coal = None if game.n <= 6 and use_exact: coal = {"{" + ", ".join(game.names_of(m)) + "}": game.value(m) for m in range(1 << game.n)} full = game.trace(game.full_mask) return CaseBlame(case.id, game.players, sh, single_swap(game), leave_one_out(game), last_touch(game), factual.score, game.value(0), game.value(game.full_mask), coal, factual.outputs, full.outputs)