129 lines
6.0 KiB
Python
129 lines
6.0 KiB
Python
#!/usr/bin/env python
|
||
"""闭环末位误差诊断（只读，手动运行；不在流水线里）。
|
||
|
||
回答一个问题：策略"差一点没进 1 mm 门"还是"根本没朝目标走"。
|
||
输出每题：
|
||
- 最终位置误差 err_mm / 最终姿态误差 oerr_deg
|
||
- 全程最小位置误差 min_err_mm、最小综合误差时是否进过 5 mm & 5° 窗口（DLS 接管门槛）
|
||
- done（是否到达）、碰撞步数、用时
|
||
以及按 cat/mode 分组的汇总与分位数。
|
||
|
||
用法（仓库根目录）：
|
||
PYTHONUTF8=1 python ml/eval/diag_tail.py --ckpt ml/runs/base/last.pt --n 512 --out ml/runs/diag_tail.json
|
||
"""
|
||
import argparse, json, math, os, sys
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||
|
||
import torch
|
||
|
||
|
||
def pct(v, p):
|
||
if not v:
|
||
return None
|
||
v = sorted(v)
|
||
return v[min(len(v) - 1, int(len(v) * p))]
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument('--ckpt', required=True)
|
||
ap.add_argument('--samples', default='ml/data/samples')
|
||
ap.add_argument('--tasks', default=None, help='任务 jsonl（覆盖 --samples/--split）')
|
||
ap.add_argument('--split', default='val')
|
||
ap.add_argument('--n', type=int, default=512)
|
||
ap.add_argument('--bs', type=int, default=64, help='每批题数（显存小就调小）')
|
||
ap.add_argument('--K', type=int, default=8)
|
||
ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
|
||
ap.add_argument('--out', default=None)
|
||
a = ap.parse_args()
|
||
|
||
from ml.common.kin_torch import Kin, D2R
|
||
from ml.data import load_tasks
|
||
from ml.model.policy import build
|
||
from ml.sim.closed_loop import rollout, goal_err
|
||
|
||
from ml.common.ckpt import load_policy # 2026-10-19：EMA 权重 + act_scale
|
||
model, ck = load_policy(a.ckpt, a.device)
|
||
kin = Kin(device=a.device)
|
||
if ck.get('fp') and ck['fp'] != kin.fp:
|
||
sys.exit('checkpoint 的碰撞模型指纹 %s 与当前 kin.json %s 不符' % (ck['fp'], kin.fp))
|
||
|
||
if a.tasks:
|
||
tasks = [json.loads(l) for l in open(a.tasks, encoding='utf-8') if l.strip()]
|
||
else:
|
||
tasks = [t for t in load_tasks([a.samples]) if t['split'] == a.split]
|
||
tasks = tasks[:a.n]
|
||
|
||
per_task = []
|
||
with torch.no_grad():
|
||
for i in range(0, len(tasks), a.bs):
|
||
chunk = tasks[i:i + a.bs]
|
||
res, per = rollout(model, kin, chunk, a.device, K=a.K, safety=False, seed=i, return_traj=True)
|
||
traj = per['traj'] # [B, T, 6]
|
||
B, Tn, _ = traj.shape
|
||
Tt = torch.tensor([t['T'] for t in chunk], dtype=torch.float32, device=a.device)
|
||
Rt = torch.tensor([t['R'] if t.get('R') else [0.] * 9 for t in chunk], dtype=torch.float32, device=a.device)
|
||
min_e = torch.full((B,), float('inf'), device=a.device)
|
||
min_o = torch.full((B,), float('inf'), device=a.device)
|
||
entered = torch.zeros(B, dtype=torch.bool, device=a.device) # 曾同时满足 e<5mm & (无姿态要求或 oe<5°)
|
||
step = 256
|
||
for t0 in range(0, Tn, step):
|
||
q = traj[:, t0:t0 + step].reshape(-1, 6)
|
||
n = q.shape[0] // B
|
||
e, oe = goal_err(kin, q, Tt.repeat_interleave(n, 0), Rt.repeat_interleave(n, 0))
|
||
e = e.view(B, n); oe = oe.view(B, n)
|
||
min_e = torch.minimum(min_e, e.min(1).values)
|
||
min_o = torch.minimum(min_o, oe.min(1).values)
|
||
pose = (Rt.abs().sum(-1) > 0)[:, None]
|
||
window = (e < 5) & (~pose | (oe < 5))
|
||
entered |= window.any(1)
|
||
done = per['done'].to(a.device); col = per['col'].to(a.device); td = per['t_done'].to(a.device)
|
||
end_e = per['err'].to(a.device)
|
||
# 末态姿态误差：traj 最后一帧
|
||
qend = traj[:, -1]
|
||
ee, eo = goal_err(kin, qend, Tt, Rt)
|
||
for j, t in enumerate(chunk):
|
||
per_task.append({
|
||
'tid': t.get('tid'), 'cat': t.get('cat'), 'mode': t.get('mode'),
|
||
'err_mm': float(end_e[j]), 'oerr_deg': float(eo[j]),
|
||
'min_err_mm': float(min_e[j]), 'min_oerr_deg': float(min_o[j]),
|
||
'entered_5mm5deg': bool(entered[j]),
|
||
'done': bool(done[j]), 'collision_steps': float(col[j]),
|
||
't_done_s': None if math.isnan(float(td[j])) else float(td[j]),
|
||
})
|
||
print('[diag] %d/%d done=%d entered=%d min_err<=5mm=%d' % (
|
||
min(i + a.bs, len(tasks)), len(tasks), int(done.sum()), int(entered.sum()),
|
||
sum(1 for p in per_task if p['min_err_mm'] <= 5)), flush=True)
|
||
|
||
def summ(rows):
|
||
if not rows:
|
||
return {'n': 0}
|
||
e = [r['err_mm'] for r in rows]; me = [r['min_err_mm'] for r in rows]
|
||
o = [r['oerr_deg'] for r in rows]
|
||
return {'n': len(rows),
|
||
'done': sum(r['done'] for r in rows),
|
||
'entered_5mm5deg': sum(r['entered_5mm5deg'] for r in rows),
|
||
'min_err_mm<=5': sum(1 for v in me if v <= 5),
|
||
'min_err_mm<=10': sum(1 for v in me if v <= 10),
|
||
'err_mm_p10/p50/p90': [pct(e, .1), pct(e, .5), pct(e, .9)],
|
||
'min_err_mm_p10/p50/p90': [pct(me, .1), pct(me, .5), pct(me, .9)],
|
||
'oerr_deg_p50/p90': [pct(o, .5), pct(o, .9)],
|
||
'collision_steps_mean': sum(r['collision_steps'] for r in rows) / len(rows)}
|
||
|
||
out = {'ckpt': a.ckpt, 'step': ck.get('step'), 'fp': ck.get('fp'), 'n': len(per_task),
|
||
'overall': summ(per_task),
|
||
'by_cat_mode': {k: summ([r for r in per_task if '%s/%s' % (r['cat'], r['mode']) == k])
|
||
for k in sorted({('%s/%s' % (r['cat'], r['mode'])) for r in per_task})},
|
||
'tasks': per_task}
|
||
s = json.dumps(out, indent=1, ensure_ascii=False)
|
||
print(s[:2000])
|
||
if a.out:
|
||
open(a.out + '.tmp', 'w', encoding='utf-8').write(s)
|
||
os.replace(a.out + '.tmp', a.out)
|
||
print('[diag] 已写', a.out)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|