137 lines
6.7 KiB
Python
137 lines
6.7 KiB
Python
#!/usr/bin/env python
|
||
"""闭环碰撞定位：第一次碰撞发生在第几秒、碰撞时最近的是哪两条连杆（J几-J几）。
|
||
|
||
用法（仓库根目录）：
|
||
PYTHONUTF8=1 python -m ml.eval.diag_collision --ckpt ml/runs/base/last.pt --n 512 --out ml/runs/diag_collision.json
|
||
|
||
口径与 closed_loop.rollout 一致：每步同时检查当前状态 q 与相邻状态中点 mid 的间隙（cap=0.5 mm）；
|
||
间隙组对索引 = gA*7+gB（gA<gB，g 为连杆组 0..5 → J1..J6）。
|
||
"""
|
||
import argparse, json, math, os, sys
|
||
from collections import Counter
|
||
|
||
import torch
|
||
|
||
from ml.common.kin_torch import Kin
|
||
from ml.data import load_tasks
|
||
from ml.model.policy import build
|
||
from ml.sim.closed_loop import DT, rollout
|
||
|
||
|
||
def pair_names():
|
||
# kin.json 有 7 个连杆组：g=0 是固定基座/立柱，g=1..6 是 J1..J6
|
||
def nm(g):
|
||
return '基座' if g == 0 else f'J{g}'
|
||
return {a * 7 + b: f'{nm(a)}-{nm(b)}' for a in range(7) for b in range(a + 1, 7)}
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument('--ckpt', required=True)
|
||
ap.add_argument('--samples', default='ml/data/samples')
|
||
ap.add_argument('--split', default='val')
|
||
ap.add_argument('--n', type=int, default=512)
|
||
ap.add_argument('--bs', type=int, default=64)
|
||
ap.add_argument('--chunk', type=int, default=64, help='回放轨迹分块步数（显存小就调小）')
|
||
ap.add_argument('--K', type=int, default=8)
|
||
ap.add_argument('--cap', type=float, default=0.5)
|
||
ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
|
||
ap.add_argument('--out', default=None)
|
||
a = ap.parse_args()
|
||
|
||
ck = torch.load(a.ckpt, map_location=a.device)
|
||
model = build(ck['cfg']).to(a.device)
|
||
model.load_state_dict(ck['model'])
|
||
model.eval()
|
||
kin = Kin(device=a.device)
|
||
if ck.get('fp') and ck['fp'] != kin.fp:
|
||
sys.exit(f'checkpoint 指纹 {ck["fp"]} 与 kin.json {kin.fp} 不符')
|
||
tasks = [t for t in load_tasks([a.samples]) if t['split'] == a.split][:a.n]
|
||
PN = pair_names()
|
||
rows = []
|
||
tau = [0, 0.5, 1.0, 2.0, 3.0, 5.0, 1e9]
|
||
hist = Counter()
|
||
pair_first = Counter(); pair_steps = Counter()
|
||
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'].to(a.device) # [B,T,6]
|
||
B, Tn, _ = traj.shape
|
||
first_t = torch.full((B,), float('nan'), device=a.device)
|
||
first_pair = torch.full((B,), -1, dtype=torch.long, device=a.device)
|
||
first_gap = torch.full((B,), float('nan'), device=a.device)
|
||
min_gap = torch.full((B,), float('inf'), device=a.device)
|
||
col_steps = torch.zeros(B, device=a.device)
|
||
for k in range(Tn):
|
||
qk = traj[:, k]
|
||
mid = 0.5 * (traj[:, max(0, k - 1)] + qk)
|
||
pts = torch.stack([qk, mid], 1).reshape(2 * B, 6) # 每帧只喂 2B 个点（与 rollout 同量级；一次 2048 点会爆显存）
|
||
gm = kin.gaps(pts, a.cap)['gm'].view(B, 2, 49)
|
||
gmin, gpair = gm.min(-1) # [B,2]
|
||
g1 = gmin.min(1).values
|
||
min_gap = torch.minimum(min_gap, g1)
|
||
newly = (g1 <= 0) & torch.isnan(first_t)
|
||
if bool(newly.any()):
|
||
for b in newly.nonzero().flatten().tolist():
|
||
pos = 0 if float(gmin[b, 0]) <= 0 else 1
|
||
first_t[b] = k * DT
|
||
first_pair[b] = gpair[b, pos]
|
||
first_gap[b] = gmin[b, pos]
|
||
if not bool(torch.isnan(first_t).any()):
|
||
break
|
||
for b, t in enumerate(chunk):
|
||
ft = float(first_t[b])
|
||
fp = int(first_pair[b])
|
||
name = PN.get(fp, '—')
|
||
if not math.isnan(ft):
|
||
pair_first[name] += 1
|
||
for j in range(len(tau) - 1):
|
||
if tau[j] <= ft < tau[j + 1]:
|
||
hist[f'{tau[j]}–{tau[j + 1]} s' if tau[j + 1] < 1e8 else '> 5 s'] += 1
|
||
break
|
||
rows.append({
|
||
'tid': t.get('tid'), 'cat': t.get('cat'), 'mode': t.get('mode'),
|
||
'collided': bool(per['col'][b] > 0),
|
||
'first_collision_s': None if math.isnan(ft) else round(ft, 2),
|
||
'first_pair': None if name == '—' else name, 'first_pair_idx': None if fp < 0 else fp,
|
||
'first_gap_mm': None if math.isnan(float(first_gap[b])) else round(float(first_gap[b]), 3),
|
||
'min_gap_mm': round(float(min_gap[b]), 3),
|
||
'collision_steps': int(per['col'][b]),
|
||
})
|
||
print(f'[diag] {min(i + a.bs, len(tasks))}/{len(tasks)} 已回放', flush=True)
|
||
|
||
collided = [r for r in rows if r['collided']]
|
||
def pct(v, p):
|
||
v = sorted(v); return v[min(len(v) - 1, int(len(v) * p))] if v else None
|
||
times = [r['first_collision_s'] for r in collided]
|
||
out = {
|
||
'ckpt': a.ckpt, 'step': ck.get('step'), 'n': len(rows), 'n_collided': len(collided),
|
||
'collision_rate': len(collided) / max(1, len(rows)),
|
||
'first_time_s': {'p10': pct(times, .1), 'p50': pct(times, .5), 'p90': pct(times, .9)},
|
||
'first_time_hist': dict(hist),
|
||
'first_pair_top': pair_first.most_common(),
|
||
'min_gap_mm_p10/p50': [pct([r['min_gap_mm'] for r in rows], .1), pct([r['min_gap_mm'] for r in rows], .5)],
|
||
'by_cat': {},
|
||
'tasks': rows,
|
||
}
|
||
for k in sorted({f"{r['cat']}/{r['mode']}" for r in rows}):
|
||
g = [r for r in rows if f"{r['cat']}/{r['mode']}" == k]
|
||
out['by_cat'][k] = {'n': len(g), 'first_t_p50': pct([r['first_collision_s'] for r in g if r['collided']], .5),
|
||
'first_pairs': Counter([r['first_pair'] for r in g if r['first_pair']]).most_common(3)}
|
||
s = json.dumps(out, indent=1, ensure_ascii=False)
|
||
if a.out:
|
||
open(a.out + '.tmp', 'w', encoding='utf-8').write(s); os.replace(a.out + '.tmp', a.out)
|
||
print('== 汇总 ==')
|
||
print('题数 %d，撞过 %d（%.1f%%）' % (len(rows), len(collided), 100 * out['collision_rate']))
|
||
print('首次碰撞时间 p10/p50/p90 = %s / %s / %s s' % (
|
||
out['first_time_s']['p10'], out['first_time_s']['p50'], out['first_time_s']['p90']))
|
||
print('首次碰撞时间分布：', out['first_time_hist'])
|
||
print('首次碰撞的连杆对 Top：', pair_first.most_common(8))
|
||
print('全程最小间隙 p10/p50 = %s / %s mm' % tuple(out['min_gap_mm_p10/p50']))
|
||
if a.out: print('已写', a.out)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|