Robot_Arm/ml/eval/check_r2.py

"""方案 R2 数据侧检查：近邻样本的标签是否自洽（时间律不可观测成分会让相同输入对应不同动作）。
  python -m ml.eval.check_r2 --samples ml/data/samples --out ml/data/r2.json
做法：特征逐维标准化；随机取 Q 个查询样本，在 P 个样本池里找最近邻（排除同一任务，避免同轨迹相邻帧冒充"近邻"）；
取距离最近的 5% 近邻对，比较前 3 步（实际提交的动作）标签：std ≈ RMS(ΔY)/√2，单位为每步最大步长 STEP 的比例。
判据：label_std < 0.10（即 < 10% 步长）。起步失败率在闭环仿真结果的 start_fail 字段（acceptance 汇总）。
2026-10-18 第五轮（审查 N-2：判据与搜索池耦合）：最近 5% 的近邻距离随池大小变，label_std 随之系统偏大/偏小。
  改为按"近邻距离 → 标签差"做线性回归，取距离 → 0 的截距 label_std_norm 作判据（与池大小无关：池大只是近邻更近、点更多，
  截距估计的是"同一输入下的固有标签分歧"）；旧口径 label_std_mean 保留作参考。pipeline 默认只警告，R2_STRICT=1 才拦。
  自检：同一份数据 --pool 50000 与 --pool 200000 两跑，label_std_norm 应一致（差 < 0.02）。
2026-10-18 新增（审查：R2 无执行点）。默认 5000 查询 × 50000 池，分块 256：GPU 峰值 < 0.5 GB，CPU 约 1–2 分钟；近邻对为空时输出 pass=null 而不报错。"""
import argparse, json, math
import torch

from ml.data import Samples


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--samples', nargs='+', default=['ml/data/samples']); ap.add_argument('--out', default='ml/data/r2.json')
    ap.add_argument('--queries', type=int, default=5000); ap.add_argument('--pool', type=int, default=50000)
    ap.add_argument('--frac', type=float, default=0.05); ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    a = ap.parse_args()
    d = Samples(a.samples, device=a.device, split='train')
    g = torch.Generator(device=a.device); g.manual_seed(a.seed)
    P = min(a.pool, d.n); Q = min(a.queries, d.n)
    pi = torch.randperm(d.n, device=a.device, generator=g)[:P]; qi = torch.randperm(d.n, device=a.device, generator=g)[:Q]
    X = d.X.float(); mu, sd = X[pi].mean(0), X[pi].std(0).clamp(min=1e-6)
    Xp, Xq = (X[pi] - mu) / sd, (X[qi] - mu) / sd
    tp, tq = d.TID[pi], d.TID[qi]
    best_d = torch.empty(Q, device=a.device); best_j = torch.empty(Q, dtype=torch.long, device=a.device)
    for s in range(0, Q, 256):
        D = torch.cdist(Xq[s:s + 256], Xp)
        D.masked_fill_(tq[s:s + 256, None] == tp[None], float('inf'))
        v, j = D.min(1); best_d[s:s + 256] = v; best_j[s:s + 256] = j
    fin = torch.isfinite(best_d)
    if fin.sum() < 20:
        out = {'pairs': 0, 'label_std_mean': None, 'pass': None, 'pool': P, 'rule': 'label_std_norm < 0.10', 'note': '近邻对不足（样本太少），未判定'}
        print(json.dumps(out)); json.dump(out, open(a.out, 'w'), indent=1); return
    # 与池无关的口径：全部有效近邻对（取距离较近的一半，排除远处非线性段）上回归 std ~ a + b·dist，截距 a
    dd = best_d[fin]; jj = best_j[fin]; qf = torch.nonzero(fin).squeeze(1)
    half = dd <= torch.quantile(dd, 0.5)
    Ya, Yb = d.Y[qi][qf[half], :3].float(), d.Y[pi][jj[half], :3].float()
    sa = (Ya - Yb).pow(2).mean(-1).mean(-1).sqrt() / math.sqrt(2)
    A = torch.stack([torch.ones_like(dd[half]), dd[half]], 1)
    coef = torch.linalg.lstsq(A.cpu(), sa.cpu()[:, None]).solution.squeeze(1)
    std_norm = max(0.0, coef[0].item())
    thr = torch.quantile(best_d[fin], a.frac)
    m = fin & (best_d <= thr)
    Yq, Yp = d.Y[qi][m, :3].float(), d.Y[pi][best_j[m], :3].float()
    rms = (Yq - Yp).pow(2).mean(-1).mean(-1).sqrt()            # 每对、6 轴 × 3 步平均
    std = (rms / math.sqrt(2))
    out = {'pairs': int(m.sum()), 'pool': P, 'queries': Q, 'nn_dist_threshold': thr.item(),
           'label_std_norm': std_norm, 'label_std_slope': coef[1].item(),
           'label_std_mean': std.mean().item(), 'label_std_p90': torch.quantile(std, 0.9).item(),
           'pass': std_norm < 0.10,
           'rule': 'label_std_norm（近邻距离→0 的回归截距）< 0.10（10% 步长）；label_std_mean 为旧口径（随池大小变），仅参考'}
    print(json.dumps(out, indent=1))
    json.dump(out, open(a.out, 'w'), indent=1)


if __name__ == '__main__':
    main()