1#!/usr/bin/env python2"""Descriptive analysis of counterfactual labels: gain distribution, harm rate, variance decomposition of the3exhaustive 24-action audit into state and action components, per-action mean gain, oracle vs random vs STOP,4and the dependence on the prefix depth. Usage: python -m evaluation.label_analysis <labels_dir> [--out results/analyzed/labels_<name>.json]5"""6from __future__ import annotations78import argparse, json, sys9from pathlib import Path1011import numpy as np12import torch131415def analyse(path: Path) -> dict:16 d = torch.load(path, weights_only=False)17 recs = d["records"]18 g_all, r_pre_all = [], []19 aud = []20 for r in recs:21 v = r["valid"]22 g_all.append(r["gain"][v].numpy().ravel()); r_pre_all.append(np.full(int(v.sum()) * 4, r["r_pre"]))23 if "all_gains" in r:24 aud.append((r["all_gains"][v].numpy(), r["base_loss"][v].numpy(), r["r_pre"]))25 g = np.concatenate(g_all); rp = np.concatenate(r_pre_all)26 out = {"n_labels": int(g.size), "mean": float(g.mean()), "std": float(g.std()), "median": float(np.median(g)),27 "frac_pos": float((g > 0).mean()), "frac_neg": float((g < 0).mean()), "frac_abs_gt_0.01": float((np.abs(g) > 0.01).mean()),28 "frac_abs_gt_0.1": float((np.abs(g) > 0.1).mean()), "q05": float(np.quantile(g, 0.05)), "q95": float(np.quantile(g, 0.95)),29 "by_prefix_depth": {int(k): {"mean": float(g[rp == k].mean()), "frac_pos": float((g[rp == k] > 0).mean()), "n": int((rp == k).sum())} for k in np.unique(rp)}}30 if aud:31 G = np.concatenate([a for a, _, _ in aud]) # [n_states, A]32 base = np.concatenate([b for _, b, _ in aud])33 A = G.shape[1]34 state_mean = G.mean(1, keepdims=True)35 total_var = G.var()36 var_state = state_mean.var() # between-state variance of the mean gain37 var_action_within = (G - state_mean).var() # within-state (action-specific) variance38 best = G.max(1); best_clip = np.clip(best, 0, None)39 rand = G.mean(1)40 out["audit"] = {41 "n_states": int(G.shape[0]), "n_actions": int(A),42 "var_total": float(total_var), "var_between_states": float(var_state), "var_within_state_actions": float(var_action_within),43 "within_share": float(var_action_within / max(total_var, 1e-12)),44 "oracle_one_step_mean": float(best_clip.mean()), "random_action_mean": float(rand.mean()),45 "stop_mean": 0.0, "frac_states_any_positive": float((best > 0).mean()),46 "frac_states_all_positive": float((G > 0).all(1).mean()), "frac_states_all_negative": float((G < 0).all(1).mean()),47 "mean_rank_corr_between_actions": None,48 "per_action_mean": G.mean(0).tolist(), "per_action_frac_best": np.bincount(G.argmax(1), minlength=A).tolist(),49 "base_loss_mean": float(base.mean()),50 "oracle_gain_relative_to_base": float(best_clip.mean() / max(base.mean(), 1e-9)),51 "single_vs_pair_mean": {"single": float(G[:, :16].mean()), "pair": float(G[:, 16:].mean())},52 "winners_curse_check_split_half": None,53 }54 # winner's-curse check: pick the best action on half the... not possible per-state (one label per action);55 # instead measure how often the best action of state i is also the best of a *different* state with the56 # same prefix depth (action consistency) vs chance 1/A.57 am = G.argmax(1)58 consist = (am[:-1] == am[1:]).mean()59 out["audit"]["consecutive_state_best_action_agreement"] = float(consist)60 # decomposition of the one-step oracle: continue/stop on the state mean vs choosing the best action61 state_level = np.clip(G.mean(1), 0, None)62 out["audit"]["oracle_state_level_mean"] = float(state_level.mean()) # STOP unless the average action helps63 out["audit"]["oracle_action_selection_increment"] = float(best_clip.mean() - state_level.mean())64 out["audit"]["random_with_perfect_stop_mean"] = float(state_level.mean())65 out["audit"]["chance_agreement"] = 1.0 / A66 # relation between base loss and oracle gain67 hi = base > np.median(base)68 out["audit"]["oracle_gain_high_vs_low_base_loss"] = {"high": float(best_clip[hi].mean()), "low": float(best_clip[~hi].mean())}69 # where the opportunity is, is a random action still good? (mean over the 24 actions vs their max)70 out["audit"]["random_gain_high_vs_low_base_loss"] = {"high": float(rand[hi].mean()), "low": float(rand[~hi].mean())}71 out["audit"]["frac_actions_positive_high_vs_low"] = {"high": float((G[hi] > 0).mean()), "low": float((G[~hi] > 0).mean())}72 out["audit"]["within_state_std_high_vs_low"] = {"high": float((G[hi] - G[hi].mean(1, keepdims=True)).std()), "low": float((G[~hi] - G[~hi].mean(1, keepdims=True)).std())}73 # top decile of base loss74 q = base >= np.quantile(base, 0.9)75 out["audit"]["top_decile_base_loss"] = {"oracle": float(best_clip[q].mean()), "random": float(rand[q].mean()), "frac_positive": float((G[q] > 0).mean()), "n": int(q.sum())}76 return out777879def main(argv=None):80 ap = argparse.ArgumentParser(description=__doc__)81 ap.add_argument("labels_dir"); ap.add_argument("--out", default=None)82 args = ap.parse_args(argv)83 res = {}84 for part in ["train", "calib"]:85 p = Path(args.labels_dir) / f"labels_{part}.pt"86 if p.exists():87 res[part] = analyse(p)88 print(part, json.dumps({k: v for k, v in res[part].items() if k != "audit"}, indent=None)[:600])89 if "audit" in res[part]:90 a = res[part]["audit"]91 print(" audit:", {k: (round(v, 5) if isinstance(v, float) else v) for k, v in a.items() if k not in ("per_action_mean", "per_action_frac_best")})92 out = Path(args.out or f"results/analyzed/labels_{Path(args.labels_dir).name}.json")93 out.parent.mkdir(parents=True, exist_ok=True); out.write_text(json.dumps(res, indent=1))94 print("wrote", out)959697if __name__ == "__main__":98 main()99