caseload 0.1.0__py3-none-any.whl

caseload/__init__.py

@@ -0,0 +1,46 @@
+"""Budgeted fraud investigation under selective labels and regime drift."""
+
+from .mdp import (
+    EXPLORE_GRID,
+    OBS_DIM,
+    OBS_NAMES,
+    PACE_GRID,
+    Episode,
+    InvestigationMDP,
+    RolloutResult,
+    Round,
+    StepRecord,
+)
+from .policies import (
+    EpsilonExplore,
+    RandomGrid,
+    RandomPolicy,
+    TopK,
+    YieldTriggered,
+    oracle_ceiling,
+    rollout,
+)
+from .scorers import GradientBoostScorer, LogisticScorer
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "EXPLORE_GRID",
+    "OBS_DIM",
+    "OBS_NAMES",
+    "PACE_GRID",
+    "Episode",
+    "EpsilonExplore",
+    "GradientBoostScorer",
+    "InvestigationMDP",
+    "LogisticScorer",
+    "RandomGrid",
+    "RandomPolicy",
+    "RolloutResult",
+    "Round",
+    "StepRecord",
+    "TopK",
+    "YieldTriggered",
+    "oracle_ceiling",
+    "rollout",
+]

caseload/agents/__init__.py

@@ -0,0 +1,5 @@
+"""Learned policies. Importing this needs the ``agents`` extra (torch)."""
+
+from .ppo import PPO, ActorCritic, PPOConfig, PPOPolicy
+
+__all__ = ["PPO", "ActorCritic", "PPOConfig", "PPOPolicy"]

caseload/agents/ppo.py

@@ -0,0 +1,197 @@
+"""A compact PPO for the investigation MDP, in one file.
+
+Deliberately small. The action space is two grids of five, the observation is
+fifteen numbers and an episode is a few tens of rounds, so the policy is a
+two-layer MLP and the whole thing trains on a laptop CPU. At these batch sizes
+Apple's MPS backend was slower than CPU when I tried it on an M1 Pro; that
+comparison is not scripted here, so it is a note rather than a result. There is no
+device flag.
+
+The only non-obvious part is that the environment is expensive relative to the
+policy: a round refits a gradient boosting detector, so almost all wall-clock is
+environment, not gradients. Updates are therefore large and infrequent.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.distributions import Categorical
+
+from ..mdp import EXPLORE_GRID, OBS_DIM, PACE_GRID
+
+
+@dataclass
+class PPOConfig:
+    hidden: int = 64
+    lr: float = 1e-3
+    gamma: float = 0.995
+    gae_lambda: float = 0.95
+    clip: float = 0.2
+    epochs: int = 10
+    minibatch: int = 32
+    entropy_coef: float = 0.02
+    value_coef: float = 0.5
+    max_grad_norm: float = 0.5
+    episodes_per_update: int = 8
+    updates: int = 60
+    # the environment refits a detector every round, so it dominates wall-clock and
+    # a batch is only a few hundred transitions. small minibatches and more epochs
+    # are what turn that into enough gradient steps to actually move the policy:
+    # at minibatch 256 a batch is one minibatch, which is ~4 steps per update
+
+    seed: int = 0
+
+
+class ActorCritic(nn.Module):
+    """Two independent categorical heads, one per action grid, and a value head."""
+
+    def __init__(self, obs_dim: int = OBS_DIM, hidden: int = 64) -> None:
+        super().__init__()
+        self.body = nn.Sequential(
+            nn.Linear(obs_dim, hidden),
+            nn.Tanh(),
+            nn.Linear(hidden, hidden),
+            nn.Tanh(),
+        )
+        self.pace_head = nn.Linear(hidden, len(PACE_GRID))
+        self.explore_head = nn.Linear(hidden, len(EXPLORE_GRID))
+        self.value_head = nn.Linear(hidden, 1)
+        # small final-layer init keeps the initial policy close to uniform, which
+        # matters here because a confidently wrong opening policy wastes the budget
+        for head in (self.pace_head, self.explore_head):
+            nn.init.orthogonal_(head.weight, gain=0.01)
+            nn.init.zeros_(head.bias)
+        nn.init.orthogonal_(self.value_head.weight, gain=1.0)
+        nn.init.zeros_(self.value_head.bias)
+
+    def forward(self, obs: torch.Tensor) -> tuple[Categorical, Categorical, torch.Tensor]:
+        h = self.body(obs)
+        return (
+            Categorical(logits=self.pace_head(h)),
+            Categorical(logits=self.explore_head(h)),
+            self.value_head(h).squeeze(-1),
+        )
+
+    @torch.no_grad()
+    def act(self, obs: np.ndarray, greedy: bool = False) -> tuple[int, int, float, float]:
+        o = torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0)
+        dp, de, v = self(o)
+        if greedy:
+            a_p = int(dp.probs.argmax())
+            a_e = int(de.probs.argmax())
+        else:
+            a_p = int(dp.sample())
+            a_e = int(de.sample())
+        lp = float(dp.log_prob(torch.tensor([a_p])) + de.log_prob(torch.tensor([a_e])))
+        return a_p, a_e, lp, float(v)
+
+
+@dataclass
+class Batch:
+    obs: list[np.ndarray] = field(default_factory=list)
+    pace: list[int] = field(default_factory=list)
+    explore: list[int] = field(default_factory=list)
+    logp: list[float] = field(default_factory=list)
+    value: list[float] = field(default_factory=list)
+    reward: list[float] = field(default_factory=list)
+    done: list[bool] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.obs)
+
+    def tensors(self) -> dict[str, torch.Tensor]:
+        return {
+            "obs": torch.as_tensor(np.asarray(self.obs), dtype=torch.float32),
+            "pace": torch.as_tensor(self.pace, dtype=torch.long),
+            "explore": torch.as_tensor(self.explore, dtype=torch.long),
+            "logp": torch.as_tensor(self.logp, dtype=torch.float32),
+            "value": torch.as_tensor(self.value, dtype=torch.float32),
+        }
+
+
+def gae(
+    rewards: list[float],
+    values: list[float],
+    dones: list[bool],
+    gamma: float,
+    lam: float,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generalised advantage estimation over a batch of concatenated episodes."""
+    n = len(rewards)
+    adv = np.zeros(n, dtype=np.float32)
+    last = 0.0
+    for i in reversed(range(n)):
+        nonterminal = 0.0 if dones[i] else 1.0
+        next_value = values[i + 1] if i + 1 < n else 0.0
+        delta = rewards[i] + gamma * next_value * nonterminal - values[i]
+        last = delta + gamma * lam * nonterminal * last
+        adv[i] = last
+    return adv, adv + np.asarray(values, dtype=np.float32)
+
+
+class PPO:
+    def __init__(self, cfg: PPOConfig | None = None) -> None:
+        self.cfg = cfg or PPOConfig()
+        torch.manual_seed(self.cfg.seed)
+        self.net = ActorCritic(hidden=self.cfg.hidden)
+        self.opt = torch.optim.Adam(self.net.parameters(), lr=self.cfg.lr, eps=1e-5)
+        self.history: list[dict[str, float]] = []
+
+    def update(self, batch: Batch) -> dict[str, float]:
+        c = self.cfg
+        adv, ret = gae(batch.reward, batch.value, batch.done, c.gamma, c.gae_lambda)
+        t = batch.tensors()
+        adv_t = torch.as_tensor(adv)
+        adv_t = (adv_t - adv_t.mean()) / (adv_t.std() + 1e-8)
+        ret_t = torch.as_tensor(ret)
+
+        n = len(batch)
+        idx = np.arange(n)
+        stats: dict[str, float] = {}
+        for _ in range(c.epochs):
+            np.random.shuffle(idx)
+            for start in range(0, n, c.minibatch):
+                mb = idx[start : start + c.minibatch]
+                dp, de, v = self.net(t["obs"][mb])
+                logp = dp.log_prob(t["pace"][mb]) + de.log_prob(t["explore"][mb])
+                ratio = torch.exp(logp - t["logp"][mb])
+                a = adv_t[mb]
+                unclipped = ratio * a
+                clipped = torch.clamp(ratio, 1 - c.clip, 1 + c.clip) * a
+                pg_loss = -torch.min(unclipped, clipped).mean()
+                v_loss = ((v - ret_t[mb]) ** 2).mean()
+                ent = (dp.entropy() + de.entropy()).mean()
+                loss = pg_loss + c.value_coef * v_loss - c.entropy_coef * ent
+                self.opt.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm_(self.net.parameters(), c.max_grad_norm)
+                self.opt.step()
+                with torch.no_grad():
+                    stats = {
+                        "pg_loss": pg_loss.item(),
+                        "value_loss": v_loss.item(),
+                        "entropy": ent.item(),
+                        "approx_kl": (t["logp"][mb] - logp).mean().abs().item(),
+                    }
+        self.history.append(stats)
+        return stats
+
+
+class PPOPolicy:
+    """Wraps a trained network so it plugs into ``caseload.policies.rollout``."""
+
+    def __init__(self, net: ActorCritic, greedy: bool = True, name: str = "ppo") -> None:
+        self.net = net
+        self.greedy = greedy
+        self.name = name
+
+    def act(self, obs: np.ndarray) -> tuple[int, int]:
+        p, e, _, _ = self.net.act(obs, greedy=self.greedy)
+        return p, e
+
+    def reset(self) -> None:
+        return None

caseload/envs/__init__.py

@@ -0,0 +1,5 @@
+"""Episode sources: a procedural drift simulator for training, Elliptic for evaluation."""
+
+from .drift import DriftConfig, DriftScenario, make_episode
+
+__all__ = ["DriftConfig", "DriftScenario", "make_episode"]

caseload/envs/drift.py

@@ -0,0 +1,177 @@
+"""A procedural fraud-drift simulator, used for training.
+
+The real dataset in this package, the Elliptic transaction graph, contains exactly
+one regime break. One break is one trajectory, and you cannot fit a policy to one
+trajectory without simply memorising it. So policies are trained here, on randomly
+generated drift scenarios, and evaluated on the real break they have never seen.
+
+Fraud is a mixture of modes in feature space. At a break, some modes stop emitting
+and new ones start. Nothing about the feature vectors announces this: the only
+visible symptom is that the detector's yield falls, and the only way to recover is
+to spend budget on cases the detector does not rank highly, which costs recall now
+in exchange for a detector that still works later.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from ..mdp import Episode, Round
+
+
+@dataclass
+class DriftConfig:
+    n_rounds: int = 20
+    warm_rounds: int = 8
+    pool_lo: int = 400
+    pool_hi: int = 1200
+    dim: int = 24
+    base_rate_lo: float = 0.02
+    base_rate_hi: float = 0.10
+    n_modes: int = 3
+    mode_scale: float = 1.15
+    mode_spread: float = 0.55
+    break_lo_frac: float = 0.45
+    break_hi_frac: float = 0.75
+    n_breaks: int = 1
+    mode_turnover: float = 1.0
+    label_noise: float = 0.0
+    adversarial_break: bool = True
+    adversarial_candidates: int = 40
+
+
+class DriftScenario:
+    """One randomly drawn drift world. Call :meth:`episode` to materialise it."""
+
+    def __init__(self, cfg: DriftConfig, seed: int) -> None:
+        self.cfg = cfg
+        self.rng = np.random.default_rng(seed)
+        self.seed = seed
+        c = cfg
+        self.base_rate = float(self.rng.uniform(c.base_rate_lo, c.base_rate_hi))
+        total = c.warm_rounds + c.n_rounds
+        lo = c.warm_rounds + int(c.break_lo_frac * c.n_rounds)
+        hi = c.warm_rounds + int(c.break_hi_frac * c.n_rounds)
+        self.breaks = sorted(
+            int(self.rng.integers(lo, max(lo + 1, hi + 1))) for _ in range(max(c.n_breaks, 0))
+        )
+        # a bank of fraud modes; regimes pick disjoint-ish subsets of them
+        n_bank = c.n_modes * (len(self.breaks) + 1)
+        self.modes = self.rng.normal(0.0, c.mode_spread, size=(n_bank, c.dim)) * c.mode_scale
+        if c.adversarial_break and len(self.breaks) > 0:
+            self._hide_later_modes(n_bank)
+        self.regimes: list[np.ndarray] = []
+        for r in range(len(self.breaks) + 1):
+            if c.mode_turnover >= 1.0:
+                idx = np.arange(r * c.n_modes, (r + 1) * c.n_modes) % n_bank
+            else:
+                keep = int(round((1 - c.mode_turnover) * c.n_modes))
+                prev = self.regimes[-1] if self.regimes else np.arange(c.n_modes)
+                fresh = (np.arange(c.n_modes - keep) + (r + 1) * c.n_modes) % n_bank
+                idx = np.concatenate([prev[:keep], fresh])
+            self.regimes.append(np.asarray(idx, dtype=int))
+        self.total_rounds = total
+
+    def _hide_later_modes(self, n_bank: int) -> None:
+        """Move post-break modes to where the incumbent detector scores lowest.
+
+        Without this the new fraud mode lands somewhere the old detector still gives
+        a middling score, so a pure top-k policy stumbles into a few cases every
+        round, feeds them to the refit, and quietly heals itself. That recovers a
+        large share of post-break fraud with no exploration at all, which leaves
+        exploration nothing to contribute and makes the environment a poor model of
+        the real thing: on Elliptic the frozen detector finds 2 of the 169 illicit
+        transactions in steps 43-49 (``results/collapse.log``).
+
+        Screening candidate centres against a detector fitted on pre-break traffic
+        reproduces that. It is also the more realistic story, since an adversary who
+        adapts moves to where the current model is not looking.
+        """
+        from sklearn.ensemble import HistGradientBoostingClassifier
+
+        c = self.cfg
+        pre = self.modes[: c.n_modes]
+        # a small pre-break sample, enough to fit a screening detector
+        xs, ys = [], []
+        for _ in range(6):
+            k = int(self.rng.binomial(600, self.base_rate))
+            x = self.rng.normal(0.0, 1.0, size=(600, c.dim))
+            y = np.zeros(600, dtype=int)
+            if k:
+                pick = self.rng.choice(len(pre), size=k, replace=True)
+                x[:k] = pre[pick] + self.rng.normal(0.0, 0.42, size=(k, c.dim))
+                y[:k] = 1
+            xs.append(x)
+            ys.append(y)
+        x_all, y_all = np.concatenate(xs), np.concatenate(ys)
+        if int((y_all == 1).sum()) < 2:
+            return
+        screen = HistGradientBoostingClassifier(
+            max_iter=20, max_leaf_nodes=8, random_state=0, early_stopping=False
+        ).fit(x_all, y_all)
+
+        n_new = n_bank - c.n_modes
+        if n_new <= 0:
+            return
+        cand = (
+            self.rng.normal(0.0, c.mode_spread, size=(c.adversarial_candidates, c.dim))
+            * c.mode_scale
+        )
+        # score each candidate centre by how suspicious the old detector finds it
+        s = screen.predict_proba(cand)[:, 1]
+        # also require distance from the pre-break modes, so "hidden" does not mean
+        # "sitting on top of an old mode the detector happens to score low"
+        far = np.sqrt(((cand[:, None, :] - pre[None]) ** 2).sum(-1)).min(1)
+        rank = np.argsort(s - 0.05 * far)
+        self.modes[c.n_modes :] = cand[rank[:n_new]]
+
+    def _regime_at(self, t: int) -> int:
+        return int(np.searchsorted(self.breaks, t, side="right"))
+
+    def _draw(self, t: int, n: int) -> tuple[np.ndarray, np.ndarray]:
+        c = self.cfg
+        k = int(self.rng.binomial(n, self.base_rate))
+        x = self.rng.normal(0.0, 1.0, size=(n, c.dim))
+        y = np.zeros(n, dtype=int)
+        if k > 0:
+            modes = self.regimes[self._regime_at(t)]
+            pick = self.rng.choice(modes, size=k, replace=True)
+            x[:k] = self.modes[pick] + self.rng.normal(0.0, 0.42, size=(k, c.dim))
+            y[:k] = 1
+        if c.label_noise > 0:
+            flip = self.rng.random(n) < c.label_noise
+            y[flip] = 1 - y[flip]
+        order = self.rng.permutation(n)
+        return x[order], y[order]
+
+    def episode(self) -> Episode:
+        c = self.cfg
+        sizes = self.rng.integers(c.pool_lo, c.pool_hi + 1, size=self.total_rounds)
+        warm_x, warm_y = [], []
+        for t in range(c.warm_rounds):
+            x, y = self._draw(t, int(sizes[t]))
+            warm_x.append(x)
+            warm_y.append(y)
+        rounds = []
+        for t in range(c.warm_rounds, self.total_rounds):
+            x, y = self._draw(t, int(sizes[t]))
+            rounds.append(Round(x=x, y=y))
+        return Episode(
+            rounds=rounds,
+            warm_x=np.concatenate(warm_x),
+            warm_y=np.concatenate(warm_y),
+            name=f"drift-{self.seed}",
+        )
+
+    @property
+    def break_rounds(self) -> list[int]:
+        """Break positions expressed as indices into the episode's rounds."""
+        return [b - self.cfg.warm_rounds for b in self.breaks]
+
+
+def make_episode(seed: int, cfg: DriftConfig | None = None) -> tuple[Episode, list[int]]:
+    """Draw one drift episode and report where its breaks fall."""
+    sc = DriftScenario(cfg or DriftConfig(), seed)
+    return sc.episode(), sc.break_rounds

caseload/envs/elliptic.py

@@ -0,0 +1,131 @@
+"""The Elliptic Bitcoin transaction graph, as a held-out investigation episode.
+
+The data set is from Weber et al., "Anti-Money Laundering in Bitcoin: Experimenting
+with Graph Convolutional Networks for Financial Forensics", KDD '19 Workshop on
+Anomaly Detection in Finance (arXiv:1908.02591): 203,769 transactions over 49 time
+steps, each with 166 features, the first of which is the time step itself. The
+archive keeps the other 165 as ``x``.
+
+Weber et al. also document the regime break this environment is built around. A dark
+market closed at time step 43, and every model they tried, including one retrained
+after each step, performed poorly on the illicit transactions that followed.
+``scripts/measure_collapse.py`` reproduces that failure at an investigation budget
+and writes the per-step counts to ``results/collapse.log``: a detector frozen on
+steps 1 to 34 finds a large share of the illicit cases per step up to step 42 and
+almost none from 43 on, while the same detector refitted on labels up to step 44
+finds many times more of the fraud in steps 45 to 49, so the information needed to
+recover exists. It is only obtainable by spending budget on cases the broken detector
+ranks low.
+
+Most transactions carry no label (the exact share is in ``results/collapse.log``).
+Unlabelled cases are real cases whose status nobody established; investigating one
+costs budget and reveals nothing. That is left in deliberately, because it is what
+the operational problem actually looks like.
+
+The data is CC BY-NC-ND 4.0 and is not shipped here. ``scripts/fetch_elliptic.py``
+says where to get it and converts the Kaggle CSVs into ``elliptic_parsed.npz`` with
+arrays ``x``, ``node_time``, ``labels``, ``src``, ``dst`` and ``txid``. The archive
+is looked for in ``~/.cache/caseload``. Two older locations are still read if that
+is empty: ``~/.cache/auditgym`` (this package's former name) and ``~/.cache/graphspot``
+(the graphspot library, which writes the same arrays in the same row order). The
+graphspot parser rounds some features differently in the last bit, and the post-break
+counts are small enough for that to move them, so convert with
+``scripts/fetch_elliptic.py`` to reproduce the committed results exactly.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import pathlib
+
+import numpy as np
+
+from ..mdp import Episode, Round
+
+CANONICAL_DIR = pathlib.Path.home() / ".cache" / "caseload"
+LEGACY_DIRS = (
+    pathlib.Path.home() / ".cache" / "auditgym",
+    pathlib.Path.home() / ".cache" / "graphspot",
+)
+ARCHIVE = "elliptic_parsed.npz"
+# the dark-market shutdown reported by Weber et al., in the dataset's own numbering
+BREAK_TIME = 43
+DEFAULT_WARM_UNTIL = 34
+
+
+def find_archive(path: str | pathlib.Path | None = None) -> pathlib.Path:
+    if path is not None:
+        p = pathlib.Path(path).expanduser()
+        if p.is_dir():
+            p = p / ARCHIVE
+        if p.exists():
+            return p
+        raise FileNotFoundError(p)
+    for base in (CANONICAL_DIR, *LEGACY_DIRS):
+        p = base / ARCHIVE
+        if p.exists():
+            return p
+    raise FileNotFoundError(
+        f"{ARCHIVE} not found in {CANONICAL_DIR}. "
+        "Run scripts/fetch_elliptic.py, which prints where to get the data."
+    )
+
+
+def load_raw(path: str | pathlib.Path | None = None) -> dict[str, np.ndarray]:
+    d = np.load(find_archive(path), allow_pickle=True)
+    return {k: d[k] for k in d.files}
+
+
+def fingerprint(x: np.ndarray, t: np.ndarray, y: np.ndarray) -> str:
+    """sha256 over the feature, time-step and label arrays, to tell archives apart."""
+    h = hashlib.sha256()
+    for a in (x, t, y):
+        h.update(np.ascontiguousarray(a).tobytes())
+    return h.hexdigest()
+
+
+def archive_fingerprint(path: str | pathlib.Path | None = None) -> str:
+    raw = load_raw(path)
+    return fingerprint(
+        raw["x"], raw["node_time"].astype(np.int64), raw["labels"].astype(np.int64)
+    )
+
+
+def load_episode(
+    path: str | pathlib.Path | None = None,
+    warm_until: int = DEFAULT_WARM_UNTIL,
+    drop_unlabelled: bool = False,
+    labelled_only_rewards: bool = True,
+) -> tuple[Episode, int]:
+    """Build the Elliptic episode.
+
+    Returns the episode and the index of the break within its rounds.
+
+    ``drop_unlabelled`` removes cases nobody adjudicated. That makes the problem
+    easier and less faithful; it exists so the cost of the unlabelled majority can
+    be measured rather than assumed. ``labelled_only_rewards`` maps the unlabelled
+    class to 0, meaning investigating one wastes budget, which is the honest
+    reading: you looked and learned nothing.
+    """
+    raw = load_raw(path)
+    x, t, y = raw["x"], raw["node_time"].astype(int), raw["labels"].astype(int)
+    if drop_unlabelled:
+        keep = y >= 0
+        x, t, y = x[keep], t[keep], y[keep]
+    y_eff = np.where(y == 1, 1, 0)
+    if not labelled_only_rewards:
+        y_eff = np.where(y < 0, 0, y_eff)
+
+    warm = t <= warm_until
+    # the detector warm-starts only from cases that were actually adjudicated
+    warm_labelled = warm & (y >= 0)
+    steps = sorted(int(s) for s in np.unique(t[t > warm_until]))
+    rounds = [Round(x=x[t == s], y=y_eff[t == s]) for s in steps]
+    break_index = steps.index(BREAK_TIME) if BREAK_TIME in steps else -1
+    ep = Episode(
+        rounds=rounds,
+        warm_x=x[warm_labelled],
+        warm_y=y_eff[warm_labelled],
+        name=f"elliptic(t>{warm_until})",
+    )
+    return ep, break_index