Robot_Arm/ml/mining/gate.py

"""门禁（方案 §7.2 第 6 步）：新模型须 ①门禁集到达率不降 ②A 类降幅 ≤ 0.5 个百分点 ③上一轮修好的错题 95% 仍对
④R7：B/C 类（需绕行）左右比例变化 ≤ 5 个百分点 ⑤R7：两者都到达的题行程中位数比 ≤ 1.05。
门禁集为无老师题，到达按定长 6 s 口径（见 closed_loop.NOREF_SEC）。
  python -m ml.mining.gate --old cur.pt --new cand.pt --tasks ml/runs/gate_tasks.jsonl --out gate.json [--fixed prev_fixed.jsonl]
通过则把 cand 复制为 --current（默认 ml/runs/current.pt）。gate.json 含 accept / gain_pp，loop.sh 据此决定继续或停止。"""
import argparse, json, shutil
import torch

from ml.common.kin_torch import Kin
from ml.model.policy import build
from ml.sim.closed_loop import evaluate


def load(path, dev):
    ck = torch.load(path, map_location=dev)
    m = build(ck['cfg']).to(dev); m.load_state_dict(ck['model']); return m.eval()


def per_task_success(pers):
    return torch.cat([p['done'] & (p['col'] == 0) for _, p in pers]).tolist()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--old', required=True); ap.add_argument('--new', required=True); ap.add_argument('--tasks', required=True)
    ap.add_argument('--fixed', default=None); ap.add_argument('--out', required=True); ap.add_argument('--current', default='ml/runs/current.pt')
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    a = ap.parse_args()
    kin = Kin(device=a.device)
    tasks = [json.loads(l) for l in open(a.tasks) if l.strip()]
    ro, po = evaluate(load(a.old, a.device), kin, tasks, a.device)
    rn, pn = evaluate(load(a.new, a.device), kin, tasks, a.device)
    so, sn = per_task_success(po), per_task_success(pn)
    acc = lambda s, c: (sum(x for x, t in zip(s, tasks) if t['cat'] == c) / max(1, sum(t['cat'] == c for t in tasks)))
    gain = (sum(sn) - sum(so)) / len(tasks) * 100
    a_drop = (acc(so, 'A') - acc(sn, 'A')) * 100
    keep = None
    if a.fixed:
        ft = [json.loads(l) for l in open(a.fixed) if l.strip()]
        if ft:
            r2, p2 = evaluate(load(a.new, a.device), kin, ft, a.device)
            keep = sum(per_task_success(p2)) / len(ft)
    # R7：左右比例与行程
    cat = [t['cat'] for t in tasks]
    so_side = torch.cat([p['side'] for _, p in po]); sn_side = torch.cat([p['side'] for _, p in pn])
    det = torch.tensor([c in ('B', 'C') for c in cat])
    # 2026-10-18：左右比例按全部 B/C 题的首次 J1 方向统计，不以到达为条件（APG 把一侧绕行全做失败时，到达子集比例可纹丝不动）
    left = lambda sd: (((sd > 0) & det).sum() / det.sum().clamp(min=1)).item()   # side 为首次 J1 方向（closed_loop，第五轮改）
    ok_o, ok_n = torch.tensor(so, dtype=torch.bool), torch.tensor(sn, dtype=torch.bool)
    lr_shift = abs(left(sn_side) - left(so_side)) * 100
    both = ok_o & ok_n
    tro = torch.cat([p['travel'] for _, p in po]); trn = torch.cat([p['travel'] for _, p in pn])
    travel_ratio = (trn[both].median() / tro[both].median().clamp(min=1e-6)).item() if both.any() else None
    r7 = lr_shift <= 5.0 and (travel_ratio is None or travel_ratio <= 1.05)
    accept = gain >= 0 and a_drop <= 0.5 and (keep is None or keep >= 0.95) and r7
    out = {'old': ro, 'new': rn, 'gain_pp': gain, 'A_drop_pp': a_drop, 'fixed_kept': keep,
           'R7_left_share_shift_pp': lr_shift, 'R7_travel_median_ratio': travel_ratio, 'R7_ok': r7, 'accept': accept}
    # 本轮新修好的题（旧错新对）留给下一轮做"仍对"检查
    fixed_now = [t for t, x, y in zip(tasks, so, sn) if (not x) and y]
    with open(a.out.replace('.json', '_fixed.jsonl'), 'w') as f:
        for t in fixed_now: f.write(json.dumps(t) + '\n')
    json.dump(out, open(a.out, 'w'), indent=1)
    if accept: shutil.copy(a.new, a.current)
    print(json.dumps({k: out[k] for k in ('gain_pp', 'A_drop_pp', 'fixed_kept', 'R7_left_share_shift_pp', 'R7_travel_median_ratio', 'accept')}))


if __name__ == '__main__':
    main()