cds-interpret 0.1.0__py3-none-any.whl

cds_interpret/__init__.py

@@ -0,0 +1,106 @@
+"""cds-interpret: look inside any PyTorch Rocket League bot.
+
+The interpretability toolkit from the Cat Distribution System, on its own. It answers three questions about a bot:
+what does it *know* (can "the ball is about to go in my goal" be read from its network?), what *features* has its
+network learned, and what *happens* when you change one of them - and it can show the answers live, in the real
+game, while the bot plays.
+
+    import cds_interpret as ci
+
+    policy = ci.Policy(net, obs=my_obs, actions=my_table)     # bring your own network (or ci.demo.policy())
+    moments = ci.record.rocketsim(policy, games=20)           # or record the real game, or your RLGym env
+    study = ci.Study(policy, moments)
+    print(study.probes().summary())                           # what each layer knows, against baselines
+    study.features("mlp.3")                                   # a sparse autoencoder's features of a layer
+    study.report("report.html")
+
+Three tiers, by what you have of the bot - every tier works for any bot:
+
+1. **Watch** (game state only): concept labels, habits, drawings - for any bot you can watch.
+2. **Open** (its PyTorch network): probes, sparse-autoencoder features, edits, reports, the live overlay.
+3. **Run** (it plays in RocketSim or your RLGym env): record its own play at will, causal tests in whole games,
+   concepts followed across training snapshots.
+
+Start with docs/quickstart.md.
+"""
+
+from . import (
+    analysis,
+    behavior,
+    curriculum,
+    demo,
+    emergence,
+    frames,
+    hooks,
+    intervene,
+    labels,
+    names,
+    overlay,
+    probes,
+    record,
+    report,
+    rlgym,
+    sae,
+    store,
+    tracking,
+    viz,
+)
+from ._version import __version__
+from .errors import CDSInterpretError, LayerNotFound, MeshesNotFound, MissingDependency, ShapeError
+from .frames import Frames, from_arrays, from_rlbot_packet, from_rlgym
+from .hooks import LayerTap, list_layers, run_layers, sequential_layers
+from .labels import LABEL_KIND, concept_labels, describe_label, future_labels, register_label, short_label
+from .moments import Moments, Tracks
+from .policy import Policy, describe_action, eager_copy, lookup_table
+from .study import FeatureSet, ProbeResults, Study
+
+__all__ = [
+    "LABEL_KIND",
+    "CDSInterpretError",
+    "FeatureSet",
+    "Frames",
+    "LayerNotFound",
+    "LayerTap",
+    "MeshesNotFound",
+    "MissingDependency",
+    "Moments",
+    "Policy",
+    "ProbeResults",
+    "ShapeError",
+    "Study",
+    "Tracks",
+    "__version__",
+    "analysis",
+    "behavior",
+    "concept_labels",
+    "curriculum",
+    "demo",
+    "describe_action",
+    "describe_label",
+    "eager_copy",
+    "emergence",
+    "frames",
+    "from_arrays",
+    "from_rlbot_packet",
+    "from_rlgym",
+    "future_labels",
+    "hooks",
+    "intervene",
+    "labels",
+    "list_layers",
+    "lookup_table",
+    "names",
+    "overlay",
+    "probes",
+    "record",
+    "register_label",
+    "report",
+    "rlgym",
+    "run_layers",
+    "sae",
+    "sequential_layers",
+    "short_label",
+    "store",
+    "tracking",
+    "viz",
+]

cds_interpret/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from .cli import main
+
+sys.exit(main())

cds_interpret/_rlbot.py

@@ -0,0 +1,186 @@
+"""What the RLBot recorder and the thoughts overlay share: finding your bot's network, read-only taps on it, the
+events between two packets, the ball's path from RLBot's prediction, and running code after the controls are sent.
+
+Nothing here imports rlbot at import time.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import torch
+
+from .errors import CDSInterpretError
+from .frames import from_rlbot_packet
+from .labels import PATH_T
+
+
+def find_network(bot: Any, network: str | Callable[[Any], torch.nn.Module] | torch.nn.Module | None) -> torch.nn.Module:
+    """Your bot's torch network: ``network`` as a dotted attribute path of the bot (``"agent.policy"``), a function
+    of the bot, a module, or None to search the bot's attributes (two levels deep) for the largest nn.Module."""
+    if isinstance(network, torch.nn.Module):
+        return network
+    if callable(network) and not isinstance(network, str):
+        net = network(bot)
+    elif isinstance(network, str):
+        net = bot
+        for part in network.split("."):
+            if not hasattr(net, part):
+                raise CDSInterpretError(f"the bot has no attribute {network!r} (stopped at {part!r}): pass network= the "
+                                        "dotted name of the attribute that holds your torch network, as set in "
+                                        "initialize()")
+            net = getattr(net, part)
+    else:
+        found: list[tuple[int, str, torch.nn.Module]] = []
+
+        def look(obj: Any, prefix: str, depth: int) -> None:
+            for k, v in list(getattr(obj, "__dict__", {}).items()):
+                if isinstance(v, torch.nn.Module):
+                    found.append((sum(p.numel() for p in v.parameters()), prefix + k, v))
+                elif depth > 0 and not k.startswith("_") and hasattr(v, "__dict__") and not isinstance(v, type):
+                    look(v, f"{prefix}{k}.", depth - 1)
+        look(bot, "", 1)
+        if not found:
+            raise CDSInterpretError("no torch network found on the bot: pass network= (the attribute that holds it, "
+                                    "e.g. network=\"policy\", or a function bot -> module)")
+        found.sort(key=lambda t: -t[0])
+        net = found[0][2]
+    if isinstance(net, torch.jit.ScriptModule):
+        raise CDSInterpretError("the bot's network is a TorchScript module, whose layers PyTorch cannot hook. Load it "
+                                "as an eager module in your bot (cds_interpret.eager_copy(net) for MLP blocks).")
+    if not isinstance(net, torch.nn.Module):
+        raise CDSInterpretError(f"network {network!r} is a {type(net).__name__}, not a torch.nn.Module")
+    return net
+
+
+class PacketEvents:
+    """The events of the step between two packets: touches (a car's latest touch changed), demolitions (a car's
+    demolition count went up; a car became demolished) and goals (a team's score went up)."""
+
+    def __init__(self) -> None:
+        self.prev: Any = None
+
+    def update(self, packet: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
+        c = max(1, len(packet.players))
+        touch, made, demoed = np.zeros(c, np.float32), np.zeros(c, np.float32), np.zeros(c, np.float32)
+        goal = -1
+        prev = self.prev
+        if prev is not None and len(prev.players) == len(packet.players):
+            for j, (a, b) in enumerate(zip(prev.players, packet.players)):
+                ta, tb = getattr(a, "latest_touch", None), getattr(b, "latest_touch", None)
+                if tb is not None and (ta is None or tb.game_seconds != ta.game_seconds):
+                    touch[j] = 1.0
+                da = getattr(getattr(a, "score_info", None), "demolitions", 0)
+                db = getattr(getattr(b, "score_info", None), "demolitions", 0)
+                made[j] = float(db > da)
+                demoed[j] = float(b.demolished_timeout > 0 and not a.demolished_timeout > 0)
+            for t, (ta2, tb2) in enumerate(zip(prev.teams, packet.teams)):
+                if tb2.score > ta2.score:
+                    goal = t
+        self.prev = packet
+        return touch, made, demoed, goal
+
+
+def packet_row(packet: Any) -> dict[str, np.ndarray]:
+    """One packet as a row of Frames fields (with prev_controls from each car's last input)."""
+    fr = from_rlbot_packet(packet)
+    row = {k: v[0] for k, v in fr.data.items()}
+    ctl = np.zeros((fr.n_cars, 8), np.float32)
+    for j, p in enumerate(packet.players):
+        li = getattr(p, "last_input", None)
+        if li is not None:
+            ctl[j] = [li.throttle, li.steer, li.pitch, li.yaw, li.roll, float(li.jump), float(li.boost),
+                      float(li.handbrake)]
+    row["prev_controls"] = ctl
+    return row
+
+
+def prediction_path(prediction: Any, now: float) -> np.ndarray | None:
+    """RLBot's ball prediction at labels.PATH_T seconds after ``now`` (nearest slices), [K, 3]; None without one."""
+    slices = getattr(prediction, "slices", None)
+    if not slices:
+        return None
+    ts = np.array([s.game_seconds for s in slices], np.float64) - float(now)
+    pos = np.array([[s.physics.location.x, s.physics.location.y, s.physics.location.z] for s in slices], np.float32)
+    if ts[-1] < PATH_T[-1] - 0.05:
+        return None
+    idx = np.clip(np.searchsorted(ts, PATH_T), 0, len(ts) - 1)
+    return pos[idx]
+
+
+def landing_point(prediction: Any, step: int = 4, below: float = 120.0) -> tuple[float, float, float] | None:
+    """Where RLBot's prediction says the ball next comes down (first slice under ``below`` uu, every ``step``-th)."""
+    slices = getattr(prediction, "slices", None)
+    if not slices:
+        return None
+    for s in slices[::step]:
+        loc = s.physics.location
+        if loc.z < below:
+            return (loc.x, loc.y, loc.z)
+    return None
+
+
+def wrap_after_decision(base: Any, name: str, after: Callable[[Any, Any], None],
+                        before: Callable[[Any, Any], None] | None = None,
+                        init: Callable[[Any], None] | None = None, retire: Callable[[Any], None] | None = None) -> type:
+    """A subclass of an RLBot v5 Bot / Hivemind class whose decisions are the parent's own, with ``before(bot,
+    packet)`` run just before the decision and ``after(bot, packet)`` run once its controls have been sent. Errors in
+    the added code are logged and switch it off; they never reach the bot."""
+    if not hasattr(base, "_packet_processor"):
+        raise CDSInterpretError(f"{base.__name__} does not look like an RLBot v5 Bot or Hivemind class (no "
+                                "_packet_processor). cds-interpret works with the rlbot package 2.x (RLBot v5).")
+
+    def initialize(self: Any) -> None:
+        base.initialize(self)
+        self._cds_on = True
+        if init is not None:
+            try:
+                init(self)
+            except Exception as e:                    # the added code must never stop the bot
+                _log(self, f"cds-interpret: {name} off ({e})")
+                self._cds_on = False
+
+    def _packet_processor(self: Any, packet: Any) -> None:
+        on = getattr(self, "_cds_on", False)
+        if on and before is not None:
+            try:
+                before(self, packet)
+            except Exception as e:                    # (as above)
+                _log(self, f"cds-interpret: {name} stopped ({e})")
+                self._cds_on = on = False
+        base._packet_processor(self, packet)          # the decision and its controls go out first
+        if on:
+            try:
+                after(self, packet)
+            except Exception as e:                    # (as above)
+                _log(self, f"cds-interpret: {name} stopped ({e})")
+                self._cds_on = False
+
+    def retire_(self: Any) -> None:
+        if retire is not None and getattr(self, "_cds_on", False):
+            try:
+                retire(self)
+            except Exception as e:
+                _log(self, f"cds-interpret: {name} could not finish ({e})")
+        base.retire(self)
+
+    return type(f"{base.__name__}With{name.title().replace(' ', '')}", (base,),
+                {"initialize": initialize, "_packet_processor": _packet_processor, "retire": retire_,
+                 "__doc__": f"{base.__name__}, playing exactly as it does, with cds-interpret's {name}."})
+
+
+def _log(bot: Any, msg: str) -> None:
+    logger = getattr(bot, "logger", None) or getattr(bot, "_logger", None)
+    if logger is not None:
+        logger.warning(msg)
+    else:
+        print(msg)
+
+
+def our_indices(bot: Any) -> list[int]:
+    """The car indices the bot controls (a Bot: [index]; a Hivemind: its indices)."""
+    if hasattr(bot, "indices") and getattr(bot, "indices", None):
+        return [int(i) for i in bot.indices]
+    return [int(bot.index)]

cds_interpret/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

cds_interpret/analysis.py

@@ -0,0 +1,244 @@
+"""Reading a feature: how it relates to the labels, its top moments, and a plain description of a moment.
+
+    rel = label_relations(act, labels, kinds)          # per label: AUROC (binary) or Pearson r (regression)
+    top = top_moments(act, k=12, groups=game_ids, min_gap=8, steps=step_ids)   # the strongest, at most one a stretch
+    describe(frame, me)                                # "ours at (x, y) 1840 uu/s on the ground, 34 boost; ball ..."
+    guess_name(rel)                                    # the best-matching label, as a GUESS with its score
+    causal_targets(features)                           # a causal test's target, group and controls, by a fixed rule
+
+All game-independent except describe(), which reads a Frames moment (any bot's).
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .frames import Frames, canon, team_sign
+from .probes import auroc
+
+
+def label_relations(act: np.ndarray, labels: dict[str, np.ndarray], kinds: dict[str, str],
+                    min_pos: int = 30) -> dict[str, dict]:
+    """For a feature's activation act [N] (0 where it is off): per label, AUROC of act as a detector (binary labels)
+    or Pearson r (regression labels), on the rows where the label is defined; also, for binary labels, P(label |
+    feature on) against the base rate (the lift)."""
+    act = np.asarray(act, np.float64)
+    on = act > 0
+    out = {}
+    for k, y in labels.items():
+        y = np.asarray(y, np.float64)
+        v = ~np.isnan(y)
+        if v.sum() < 50:
+            continue
+        a, yy = act[v], y[v]
+        if kinds.get(k) == "bin":
+            pos = int((yy > 0.5).sum())
+            if pos < min_pos or pos == len(yy):
+                continue
+            base = pos / len(yy)
+            p_on = float((yy[on[v]] > 0.5).mean()) if on[v].any() else float("nan")
+            out[k] = {"auroc": auroc(a, yy), "base": base, "p_given_on": p_on,
+                      "lift": p_on / base if base > 0 and np.isfinite(p_on) else float("nan"),
+                      "recall": float((on[v] & (yy > 0.5)).sum() / pos), "coverage": float(v.mean())}
+        else:
+            if np.std(a) == 0 or np.std(yy) == 0:
+                continue
+            out[k] = {"r": float(np.corrcoef(a, yy)[0, 1]), "coverage": float(v.mean())}
+    return out
+
+
+def relation_strength(rel: dict) -> float:
+    """One number per label relation, 0 = none, 1 = perfect: |2 AUROC - 1| or |r|."""
+    return abs(2.0 * rel["auroc"] - 1.0) if "auroc" in rel else abs(rel["r"])
+
+
+def guess_name(rel: dict[str, dict], top: int = 3, min_coverage: float = 0.3) -> list[tuple[str, float, str]]:
+    """The labels a feature relates to most: [(label, strength, sign "+" / "-")] - a starting point for a name, not
+    a finding: check the top moments. Labels defined on fewer than ``min_coverage`` of the moments (``threat_t`` has
+    a value only when there is a threat) are left out: a relation measured on a small slice of moments makes a
+    misleading name."""
+    scored = []
+    for k, r in rel.items():
+        if r.get("coverage", 1.0) < min_coverage:
+            continue
+        s = relation_strength(r)
+        sign = "+" if (r.get("auroc", 0.5) >= 0.5 if "auroc" in r else r["r"] >= 0) else "-"
+        scored.append((k, s, sign))
+    return sorted(scored, key=lambda t: -t[1])[:top]
+
+
+def top_moments(act: np.ndarray, k: int = 12, groups: np.ndarray | None = None, steps: np.ndarray | None = None,
+                min_gap: int = 8) -> np.ndarray:
+    """Indices of the k strongest activations, at most one per stretch of min_gap steps of the same game (so the top
+    list is not one moment repeated)."""
+    order = np.argsort(-np.asarray(act), kind="mergesort")
+    picked: list[int] = []
+    taken: dict = {}
+    for i in order:
+        if act[i] <= 0 or len(picked) >= k:
+            break
+        if groups is not None and steps is not None:
+            g, s = groups[i], steps[i]
+            if any(abs(s - t) < min_gap for t in taken.get(g, [])):
+                continue
+            taken.setdefault(g, []).append(s)
+        picked.append(int(i))
+    return np.asarray(picked, np.int64)
+
+
+def _where(xy) -> str:
+    x, y = xy
+    side = "left" if x < -1200 else "right" if x > 1200 else "middle"
+    third = "own third" if y < -1700 else "attacking third" if y > 1700 else "midfield"
+    return f"{third}, {side}"
+
+
+def describe(fr: Frames, me: int, t: int = 0) -> str:
+    """One moment (row t of fr) in plain words, from car slot me's side (its team attacks +y)."""
+    team = float(fr.car_team[t, me])
+    s = float(team_sign(np.array([team]))[0])
+    pos = canon(fr.car_pos[t], s)
+    vel = canon(fr.car_vel[t], s)
+    b = canon(fr.ball_pos[t], s) + 0.0                  # (+ 0.0: no "-0" in the text)
+    bv = canon(fr.ball_vel[t], s) + 0.0
+    pos = pos + 0.0
+    mine = pos[me]
+
+    def car_txt(j):
+        sp = float(np.linalg.norm(vel[j]))
+        where = "demolished" if fr.car_demoed[t, j] > 0.5 else (
+            "in the air" if fr.car_on_ground[t, j] < 0.5 else ("on a wall" if pos[j, 2] > 150 else "on the ground"))
+        return f"({pos[j, 0]:.0f}, {pos[j, 1]:.0f}, {pos[j, 2]:.0f}) {sp:.0f} uu/s {where}, {fr.car_boost[t, j]:.0f} boost"
+    parts = [f"us {car_txt(me)} [{_where(mine[:2])}]"]
+    for j in range(fr.n_cars):
+        if j == me or fr.car_present[t, j] < 0.5:
+            continue
+        who = "mate" if fr.car_team[t, j] == team else "opp"
+        d = float(np.linalg.norm(pos[j] - mine))
+        parts.append(f"{who} {car_txt(j)}, {d:.0f} uu from us")
+    speed = float(np.linalg.norm(bv))
+    heading = ("toward their goal" if bv[1] > 300 else "toward our goal" if bv[1] < -300 else "across the field")
+    moving = "still" if speed < 50 else f"moving {speed:.0f} uu/s {heading}"
+    parts.append(f"ball ({b[0]:.0f}, {b[1]:.0f}, {b[2]:.0f}) {moving}, {np.linalg.norm(b - mine):.0f} uu from us")
+    return "; ".join(parts)
+
+
+# ------------------------------------------------------------------------- many sparse features at once
+def sparse_columns(idx: np.ndarray, val: np.ndarray, m: int):
+    """Top-k codes (idx, val [n, k]) as columns: (rows, values, starts) - feature j fires on rows[starts[j]:starts[j+1]]
+    with values[...] (positive activations only)."""
+    flat_f = idx.reshape(-1).astype(np.int64)
+    flat_v = val.reshape(-1).astype(np.float32)
+    flat_r = np.repeat(np.arange(len(idx), dtype=np.int64), idx.shape[1])
+    keep = flat_v > 0
+    f, r, v = flat_f[keep], flat_r[keep], flat_v[keep]
+    order = np.argsort(f, kind="stable")
+    f, r, v = f[order], r[order], v[order]
+    return r, v, np.searchsorted(f, np.arange(m + 1))
+
+
+def _avg_ranks(x: np.ndarray) -> np.ndarray:
+    order = np.argsort(x, kind="mergesort")
+    _, start, counts = np.unique(x[order], return_index=True, return_counts=True)
+    r = np.empty(len(x))
+    r[order] = np.repeat(start + (counts - 1) / 2.0 + 1.0, counts)
+    return r
+
+
+def sparse_relations(rows_nz: np.ndarray, vals_nz: np.ndarray, labels: dict[str, np.ndarray], kinds: dict[str, str],
+                     n: int, min_pos: int = 30) -> dict[str, dict]:
+    """label_relations for a feature that is 0 except on rows_nz (values vals_nz) of n rows - in time proportional to
+    how often it fires. Same outputs: AUROC / base / P(label | on) / lift / recall, or Pearson r."""
+    out = {}
+    for k, y in labels.items():
+        kind = kinds.get(k)
+        valid = ~np.isnan(y)
+        nv = int(valid.sum())
+        if nv < 50:
+            continue
+        yn = y[rows_nz]
+        okn = ~np.isnan(yn)
+        vv, yy = vals_nz[okn], yn[okn]
+        if kind == "bin":
+            n1 = int((y[valid] > 0.5).sum())
+            n0 = nv - n1
+            if n1 < min_pos or n0 == 0:
+                continue
+            pos = yy > 0.5
+            nz1, nz0 = int(pos.sum()), int((~pos).sum())
+            z1, z0 = n1 - nz1, n0 - nz0
+            u_nz = (_avg_ranks(vv)[pos].sum() - nz1 * (nz1 + 1) / 2.0) if len(vv) else 0.0
+            u = u_nz + nz1 * z0 + 0.5 * z1 * z0
+            base = n1 / nv
+            p_on = nz1 / len(yy) if len(yy) else float("nan")
+            out[k] = {"auroc": float(u / (n1 * n0)), "base": base, "p_given_on": float(p_on),
+                      "lift": float(p_on / base) if len(yy) else float("nan"), "recall": nz1 / n1,
+                      "coverage": nv / max(n, 1)}
+        else:
+            yv = y[valid].astype(np.float64)
+            my, sy = yv.mean(), yv.std()
+            v64 = vv.astype(np.float64)
+            ma = v64.sum() / nv
+            sa = np.sqrt(max((v64 ** 2).sum() / nv - ma * ma, 0.0))
+            if sa == 0 or sy == 0:
+                continue
+            cov = (v64 * yy).sum() / nv - ma * my
+            out[k] = {"r": float(cov / (sa * sy)), "coverage": nv / max(n, 1)}
+    return out
+
+
+# ------------------------------------------------------------------------- choosing a causal test's targets
+def causal_targets(features: dict, label: str = "demo_chance", also: str = "demo_soon", min_freq: float = 5e-4,
+                   min_lift: float = 5.0, min_auroc: float = 0.75, n_group: int = 3, n_controls: int = 3,
+                   seed: int = 0) -> dict:
+    """A causal test's targets, chosen by a rule fixed BEFORE any game is played with an edit (so the result cannot
+    steer the choice).
+
+    features: a dictionary report's {id: {"freq", "relations": {label: {"auroc", "lift", ...}}}} (label_relations per
+    feature, plus how often it fires). Target = the feature that best detects `label` (held-out AUROC) among those
+    firing on >= min_freq of moments with lift >= min_lift and AUROC >= min_auroc; group = the best n_group such.
+    Without any: the cleanest features (best |2 AUROC - 1| on any yes / no label, lift >= min_lift, firing on >= 0.1%
+    of moments), each feature counted once, at its best label. Controls = n_controls features firing about as often
+    as the target (within 30%), outside the group, that detect neither `label` nor `also` (AUROC < 0.6), drawn with
+    `seed`. Returns {"rule", "target", "group", "controls", "candidates" (the top 10), "arms"}: the arms to play -
+    base, the target ablated, the group ablated, the target scaled 3x, the controls ablated."""
+    feats = {int(j): f for j, f in features.items()}
+
+    def score(f, lab):
+        r = f["relations"].get(lab, {})
+        return r.get("auroc", 0.5), r.get("lift", 0.0)
+
+    cands = sorted(((score(f, label)[0], j, label) for j, f in feats.items()
+                    if f["freq"] >= min_freq and score(f, label)[1] >= min_lift and score(f, label)[0] >= min_auroc),
+                   reverse=True)
+    kind = label
+    if not cands:
+        kind = "cleanest"
+        best: dict[int, tuple[float, str]] = {}
+        for j, f in feats.items():
+            if f["freq"] < 1e-3:
+                continue
+            for lab, r in f["relations"].items():
+                if "auroc" in r and r.get("lift", 0) >= min_lift:
+                    s = abs(2 * r["auroc"] - 1)
+                    if s > best.get(j, (-1.0, ""))[0]:
+                        best[j] = (s, lab)
+        cands = sorted(((s, j, lab) for j, (s, lab) in best.items()), reverse=True)
+    target = cands[0][1] if cands else None
+    group = [c[1] for c in cands[:n_group]]
+    controls: list[int] = []
+    if target is not None:
+        ft = feats[target]["freq"]
+        pool = sorted(j for j, f in feats.items() if 0.7 * ft <= f["freq"] <= 1.3 * ft and j not in group
+                      and score(f, label)[0] < 0.6 and score(f, also)[0] < 0.6)
+        if pool:
+            controls = sorted(int(c) for c in np.random.default_rng(seed).choice(
+                pool, size=min(n_controls, len(pool)), replace=False))
+    arms = ["base"]
+    if target is not None:
+        arms += [f"ablate:{target}", "ablate:" + ",".join(map(str, group)), f"scale3:{target}"]
+        if controls:
+            arms.append("ablate:" + ",".join(map(str, controls)))
+    return {"rule": kind, "target": target, "group": group, "controls": controls,
+            "candidates": [{"feature": c[1], "label": c[2], "score": round(float(c[0]), 4),
+                            "score_is": "|2 AUROC - 1|" if kind == "cleanest" else "AUROC"} for c in cands[:10]],
+            "arms": arms}