停止策略成本与风险评估

"""Exact offline evaluation of stopping policies on dense-graded trajectories.

A policy sees decision rows in step order and may stop at any row; the result is
that row's `now` grade and output tokens so far. Never stopping gives the
agent's own final grade and total tokens. Scores come from task-grouped
cross-fitting (train/dev) or from a model frozen before test labels are read.
"""
import argparse
import json
import random
from collections import defaultdict
from pathlib import Path
import numpy as np
from .common import write_json

COST = 'out'  # 'out' = output tokens, 'in' = cumulative prompt tokens re-read per call
FEATURES = ['step', 'out_tokens_at', 'ckpt_index', 'steps_since_ckpt', 'is_ckpt_step', 'prod_files',
            'lines_added', 'lines_removed', 'is_test_cmd', 'rc', 'n_tests_run', 'n_tests_pass',
            'last_test_rc', 'tests_since_ckpt', 'passes_since_ckpt']
REPOS = ['astropy', 'django', 'matplotlib', 'pydata', 'pytest-dev', 'scikit-learn', 'sphinx-doc', 'sympy',
         'psf', 'pylint-dev', 'mwaskom']


def load(path):
    rows = [json.loads(l) for l in open(path)]
    trajs = [json.loads(l) for l in open(Path(path).with_suffix('.trajs.jsonl'))]
    return rows, trajs


EXTRA = []  # optional semantic features merged from files (e.g. p_self)


def featurize(r):
    x = [(-99 if r.get(f) is None else float(r[f])) for f in FEATURES + EXTRA]
    return x + [float(r['repo'] == p) for p in REPOS]


def by_traj(rows):
    g = defaultdict(list)
    for r in rows: g[r['traj']].append(r)
    for v in g.values(): v.sort(key=lambda r: r['step'])
    return g


def run_policy(trajs, groups, stop):
    """stop(row) -> bool. Returns per-trajectory (task, resolved, tokens, stopped)."""
    out = []
    for t in trajs:
        res = (t['final'], t[COST + '_tokens_total'], False)
        for r in groups.get(t['traj'], []):
            if stop(r): res = (r['now'], r[COST + '_tokens_at'], True); break
        out.append((t['task_id'],) + res)
    return out


def summarize(results):
    per = defaultdict(list)
    for task, ok, tok, st in results: per[task].append((ok, tok, st))
    tasks = sorted(per)
    res = np.mean([np.mean([a for a, _, _ in per[k]]) for k in tasks])
    tok = np.mean([np.mean([b for _, b, _ in per[k]]) for k in tasks])
    return {'resolve': float(res), 'tokens': float(tok), 'stop_rate': float(np.mean([s for *_, s in results])), 'n': len(results)}


def paired_bootstrap(a, b, n=2000, seed=0):
    """Task-clustered bootstrap of (resolve, tokens) differences a - b."""
    pa, pb = defaultdict(list), defaultdict(list)
    for x in a: pa[x[0]].append(x)
    for x in b: pb[x[0]].append(x)
    tasks = sorted(pa); rng = random.Random(seed); dr, dt = [], []
    for _ in range(n):
        s = [rng.choice(tasks) for _ in tasks]
        dr.append(np.mean([np.mean([x[1] for x in pa[k]]) - np.mean([x[1] for x in pb[k]]) for k in s]))
        dt.append(np.mean([np.mean([x[2] for x in pa[k]]) - np.mean([x[2] for x in pb[k]]) for k in s]))
    q = lambda v: [float(np.percentile(v, 2.5)), float(np.percentile(v, 97.5))]
    return {'resolve_diff_ci': q(dr), 'tokens_diff_ci': q(dt)}


def headroom(rows, trajs):
    g = by_traj(rows); stats = defaultdict(int); wasted = []; tot = []
    for t in trajs:
        rs = g.get(t['traj'], []); nows = [r['now'] for r in rs]
        ever = any(nows); stats['n'] += 1; stats['final'] += t['final']; stats['ever_correct'] += ever
        stats['correct_then_broken'] += ever and not t['final']
        stats['broken_then_fixed'] += (not nows or not nows[0]) and t['final']
        if ever:
            first = next(r for r in rs if r['now'])
            wasted.append(t['out_tokens_total'] - first['out_tokens_at']); tot.append(t['out_tokens_total'])
    stats['tokens_after_first_correct_frac'] = float(sum(wasted) / max(1, sum(tot)))
    return dict(stats)


def oracle(trajs, groups):
    """Earliest stop achieving the best reachable outcome (hopeless runs stop at their first decision row)."""
    out = []
    for t in trajs:
        rs = groups.get(t['traj'], []); best = max([t['final']] + [r['now'] for r in rs])
        hit = next((r for r in rs if r['now'] == best), None)
        out.append((t['task_id'], best, hit[COST + '_tokens_at'], True) if hit else (t['task_id'], t['final'], t[COST + '_tokens_total'], False))
    return out


def fixed_policies():
    return {
        'native': lambda r: False,
        'oracle_first_correct': lambda r: r['now'] == 1,
        'first_ckpt': lambda r: True,
        'first_test_pass_after_edit': lambda r: r['is_test_cmd'] and r['rc'] == 0 and r['steps_since_ckpt'] > 0,
        'two_test_passes_after_edit': lambda r: r['passes_since_ckpt'] >= 2 and r['is_test_cmd'] and r['rc'] == 0,
    }


def gbdt_oof(rows, label, folds=5, seed=0):
    from sklearn.ensemble import HistGradientBoostingClassifier
    tasks = sorted({r['task_id'] for r in rows}); rng = random.Random(seed); rng.shuffle(tasks)
    fold = {t: i % folds for i, t in enumerate(tasks)}
    X = np.array([featurize(r) for r in rows]); y = np.array([r[label] for r in rows]); p = np.zeros(len(rows))
    for k in range(folds):
        tr = np.array([fold[r['task_id']] != k for r in rows]); te = ~tr
        if len(set(y[tr])) < 2: p[te] = y[tr].mean(); continue
        m = HistGradientBoostingClassifier(max_iter=200, learning_rate=0.05, max_leaf_nodes=15, min_samples_leaf=40, random_state=seed)
        m.fit(X[tr], y[tr]); p[te] = m.predict_proba(X[te])[:, 1]
    return p


def frontier(trajs, groups, key, thresholds):
    pts = []
    for th in thresholds:
        pts.append({'threshold': float(th), **summarize(run_policy(trajs, groups, lambda r: r[key] >= th))})
    return pts


def main(a):
    global COST
    COST = a.cost
    rows, trajs = load(a.data)
    if a.probe:
        P = {(q['traj'], q['step']): q['p_self'] for q in json.load(open(a.probe))}
        for r in rows: r['p_self'] = P.get((r['traj'], r['step']))
        EXTRA.append('p_self')
    g = by_traj(rows)
    report = {'headroom': headroom(rows, trajs), 'fixed': {}}
    base = run_policy(trajs, g, lambda r: False)
    o = oracle(trajs, g); report['fixed']['oracle'] = {**summarize(o), **paired_bootstrap(o, base)}
    for name, f in fixed_policies().items():
        res = run_policy(trajs, g, f)
        report['fixed'][name] = {**summarize(res), **(paired_bootstrap(res, base) if name != 'native' else {})}
    if a.probe:
        for r in rows: r['self_score'] = r['p_self'] if r['p_self'] is not None else -9
        report['self_probe'] = frontier(trajs, g, 'self_score', np.linspace(0, 1, 41))
    if a.gbdt:
        p_now = gbdt_oof(rows, 'now'); p_fin = gbdt_oof(rows, 'final')
        for r, x, y in zip(rows, p_now, p_fin):
            dec = r['is_ckpt_step'] or r['is_test_cmd']  # same decision rows as the small model
            r['p_now'] = x if dec else -9; r['p_final'] = y if dec else -9; r['gain'] = x - y if dec else -9
        ths = np.linspace(0, 1, 41)
        report['gbdt_oof'] = {'now_head': frontier(trajs, g, 'p_now', ths), 'final_head': frontier(trajs, g, 'p_final', ths),
                              'gain': frontier(trajs, g, 'gain', np.linspace(-0.5, 0.5, 41))}
    if a.pred:
        for name in a.pred:
            P = {(q['traj'], q['step']): q for q in json.load(open(name))}
            for r in rows:
                q = P.get((r['traj'], r['step']))
                r['m_now'] = q['p_now'] if q else -9; r['m_final'] = q['p_final'] if q else -9
                r['m_gain'] = r['m_now'] - r['m_final'] if q else -9
            report.setdefault('model', {})[name] = {'now_head': frontier(trajs, g, 'm_now', np.linspace(0, 1, 41)),
                'final_head': frontier(trajs, g, 'm_final', np.linspace(0, 1, 41)),
                'gain': frontier(trajs, g, 'm_gain', np.linspace(-0.5, 0.5, 41))}
    nat = report['fixed']['native']
    def best(pts, tol):
        ok = [p for p in pts if p['resolve'] >= nat['resolve'] - tol]
        b = min(ok, key=lambda p: p['tokens']) if ok else None
        return b and {'threshold': b['threshold'], 'resolve': round(b['resolve'], 3), 'saving': round(1 - b['tokens'] / nat['tokens'], 3)}
    fr = {'self_probe': report.get('self_probe')}
    for k in ('gbdt_oof',): fr.update({f'gbdt.{h}': v for h, v in (report.get(k) or {}).items()})
    for m, d in (report.get('model') or {}).items(): fr.update({f'{Path(m).parent.name}.{h}': v for h, v in d.items()})
    report['best_within'] = {k: {f'tol{t}': best(v, t) for t in (0.0, 0.01, 0.02)} for k, v in fr.items() if v}
    write_json(a.out, report)
    for k, v in report['best_within'].items(): print('BEST', k, v)
    print(json.dumps({'headroom': report['headroom'], 'fixed': {k: {kk: round(vv, 3) if isinstance(vv, float) else vv for kk, vv in v.items()} for k, v in report['fixed'].items()}}, indent=1))


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--data', required=True)
    ap.add_argument('--out', required=True)
    ap.add_argument('--gbdt', action='store_true')
    ap.add_argument('--cost', choices=['out', 'in'], default='out')
    ap.add_argument('--probe', help='self_probe output; adds p_self policy and GBDT feature')
    ap.add_argument('--pred', nargs='+', help='stop_train predictions.json files')
    main(ap.parse_args())