170 lines
11 KiB
Python
170 lines
11 KiB
Python
#!/usr/bin/env python
|
||
"""闭环失败定位（只读，手动运行，约 5–8 分钟；训练进行中也可以另开终端跑，显存约 2–3 GB）。
|
||
|
||
2026-10-19 新增：第二次预训练（979a581）开环已经正常（候选0 末端误差 ~6 mm、K 中最好 ~2.9 mm、多样性 0.65），
|
||
闭环却仍 100% 碰撞、0% 到达。开环用的是 JS 生成的特征 X，闭环用的是 Python features() 现算的特征，
|
||
两者之间只要有一处口径不同，模型在闭环里看到的就是没见过的输入。本脚本回答三件事：
|
||
A. 特征口径：val 任务起点（非增广、老师状态）上，Python features() 与数据集 X 的逐组最大差；
|
||
B. 同一噪声下，用两份特征得到的动作差（16 步末端关节差，度）——A 有差但 B 很小说明差异无害；
|
||
C. 闭环失败形态：首次碰撞时刻分布、碰撞步数、首次碰撞的组对、全程最小目标误差、目标误差进度比；
|
||
以及同一批题开安全层（safety=True，近似部署）的到达率 / 刹停次数。
|
||
F. 一步采样残留噪声：开环比较 1/2/4/8 段采样与零噪声的幅度、抖动、残差与噪声相关；闭环用 4 段采样、零噪声各跑一遍。
|
||
用法（仓库根目录）：
|
||
PYTHONUTF8=1 python -m ml.eval.diag_cl --ckpt ml/runs/base/last.pt --n 256 --out ml/runs/diag_cl.json
|
||
"""
|
||
import argparse, json, os
|
||
|
||
import torch
|
||
|
||
from ml.common.kin_torch import Kin
|
||
from ml.common.features import features, STEP, H
|
||
from ml.data import Samples, load_tasks
|
||
from ml.model.policy import build
|
||
from ml.sim.closed_loop import rollout, goal_err, DT
|
||
|
||
# 83 维特征分组（与 features.py 拼接顺序一致）
|
||
GROUPS = [('sincos', 0, 12), ('joint_norm', 12, 15), ('dq', 15, 21), ('p', 21, 24), ('Rq', 24, 30), ('T', 30, 33),
|
||
('R', 33, 39), ('T-p', 39, 42), ('T-p_clamp', 42, 45), ('mode', 45, 47), ('prefix', 47, 65), ('mask', 65, 68),
|
||
('sub', 68, 74), ('sub_flag', 74, 75), ('gaps', 75, 83)]
|
||
|
||
|
||
PROG_T = (0.5, 1.0, 2.0) # 目标误差 / 起点误差 的取样时刻（s）
|
||
|
||
|
||
def pct(v, p):
|
||
v = sorted(v)
|
||
return v[min(len(v) - 1, int(len(v) * p))] if v else None
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument('--ckpt', required=True); ap.add_argument('--samples', default='ml/data/samples')
|
||
ap.add_argument('--n', type=int, default=256); ap.add_argument('--bs', type=int, default=64)
|
||
ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu'); ap.add_argument('--out', default='ml/runs/diag_cl.json')
|
||
a = ap.parse_args()
|
||
dev = a.device
|
||
ck = torch.load(a.ckpt, map_location=dev)
|
||
model = build(ck['cfg']).to(dev); model.load_state_dict(ck['model']); model.eval()
|
||
kin = Kin(device=dev)
|
||
out = {'ckpt': a.ckpt, 'step': ck.get('step')}
|
||
tasks = [t for t in load_tasks([a.samples]) if t['split'] == 'val'][:a.n]
|
||
by_tid = {t['tid']: t for t in tasks}
|
||
|
||
# ---------- A / B：起点特征口径 ----------
|
||
vd = Samples(a.samples, device=dev, split='val')
|
||
tid = vd.TID.long()
|
||
first = torch.ones_like(tid, dtype=torch.bool); first[1:] = tid[1:] != tid[:-1] # 每题第一行 = t=0 样本
|
||
idx = [i for i in first.nonzero().flatten().tolist() if int(tid[i]) in by_tid]
|
||
X0 = vd.X[idx].float()
|
||
tk = [by_tid[int(tid[i])] for i in idx]
|
||
q0 = torch.tensor([t['q0'] for t in tk], dtype=torch.float32, device=dev)
|
||
qx = torch.atan2(X0[:, 0:12:2], X0[:, 1:12:2]) * 180 / torch.pi
|
||
dqx = torch.remainder(qx - q0 + 180, 360) - 180
|
||
clean = dqx.abs().max(-1).values < 1e-2 # 未被噪声增广的起点
|
||
out['A_rows'] = int(clean.sum()); out['A_rows_aug_skipped'] = int((~clean).sum())
|
||
if clean.any():
|
||
Xc = X0[clean]; qc = q0[clean]; tc = [t for t, c in zip(tk, clean.tolist()) if c]
|
||
T = torch.tensor([t['T'] for t in tc], dtype=torch.float32, device=dev)
|
||
R = torch.tensor([t['R'] if t.get('R') else [0.] * 9 for t in tc], dtype=torch.float32, device=dev)
|
||
B = qc.shape[0]
|
||
pre = Xc[:, 47:65].view(B, 3, 6) * torch.tensor(STEP, device=dev) # 前缀 / 子目标直接取数据集里的，只比状态相关部分
|
||
pl = Xc[:, 65:68].sum(-1).round()
|
||
sf = Xc[:, 74]
|
||
sub = qc + Xc[:, 68:74] * 180
|
||
Fp = features(kin, qc, qc, T, R, pre, pl, sub, sf)
|
||
out['A_maxdiff'] = {g: (Fp[:, s:e] - Xc[:, s:e]).abs().max().item() for g, s, e in GROUPS}
|
||
out['A_meandiff'] = {g: (Fp[:, s:e] - Xc[:, s:e]).abs().mean().item() for g, s, e in GROUPS}
|
||
with torch.no_grad():
|
||
z = torch.randn(B * 8, H, 6, device=dev, generator=torch.Generator(device=dev).manual_seed(7))
|
||
a1, _, _ = model.sample_k(Xc, 8, z); a2, _, _ = model.sample_k(Fp, 8, z)
|
||
stepv = torch.tensor(STEP, device=dev)
|
||
d = ((a1.clamp(-1, 1) - a2.clamp(-1, 1)) * stepv).sum(2).abs().max(-1).values # [B,K] 16 步末端关节差（度）
|
||
mag = (a1.clamp(-1, 1) * stepv).sum(2).abs().max(-1).values
|
||
out['B_act_diff_deg'] = {'mean': d.mean().item(), 'p90': d.flatten().kthvalue(max(1, int(0.9 * d.numel()))).values.item(),
|
||
'ref_motion_deg_mean': mag.mean().item()}
|
||
|
||
# ---------- C：闭环失败形态 ----------
|
||
first_t, col_steps, pairs, min_err, reach = [], [], {}, [], 0
|
||
prog = {ts: [] for ts in PROG_T}
|
||
reach_s, brakes_s, col_s = 0, [], 0
|
||
for i in range(0, len(tasks), a.bs):
|
||
ch = tasks[i:i + a.bs]
|
||
with torch.no_grad():
|
||
_, per = rollout(model, kin, ch, dev, K=8, safety=False, seed=i, return_traj=True)
|
||
tr = per['traj'].to(dev) # [B,Tn,6]
|
||
B, Tn, _ = tr.shape
|
||
flat = tr.reshape(-1, 6) # 分块 ≤1024 点（同 dagger/collect.py 的显存修复）
|
||
gm = torch.cat([kin.gaps(flat[s:s + 1024], 0.5)['gm'] for s in range(0, flat.shape[0], 1024)]).view(B, Tn, 49)
|
||
hit = gm.min(-1).values <= 0 # [B,Tn]
|
||
Tt = torch.tensor([t['T'] for t in ch], dtype=torch.float32, device=dev)
|
||
Rt = torch.tensor([t['R'] if t.get('R') else [0.] * 9 for t in ch], dtype=torch.float32, device=dev)
|
||
e, _ = goal_err(kin, tr.reshape(-1, 6), Tt.repeat_interleave(Tn, 0), Rt.repeat_interleave(Tn, 0))
|
||
e = e.view(B, Tn)
|
||
for b in range(B):
|
||
for ts in PROG_T:
|
||
k = min(Tn - 1, int(round(ts / DT))); prog[ts].append((e[b, k] / e[b, 0].clamp(min=1e-3)).item())
|
||
# 到达后轨迹被冻结（d=0），之后不再计
|
||
for b in range(B):
|
||
if per['done'][b]: reach += 1
|
||
min_err.append(e[b].min().item()); col_steps.append(int(per['col'][b]))
|
||
h = hit[b].nonzero().flatten()
|
||
if len(h):
|
||
k = int(h[0]); first_t.append(k * DT)
|
||
pr = int(gm[b, k].argmin()); key = '%d-%d' % (pr // 7, pr % 7); pairs[key] = pairs.get(key, 0) + 1
|
||
rs, ps = rollout(model, kin, ch, dev, K=8, safety=True, seed=i)
|
||
reach_s += int(ps['done'].sum()); col_s += int((ps['col'] > 0).sum()); brakes_s += ps['brakes'].tolist()
|
||
n = len(tasks)
|
||
out['C'] = {'n': n, 'reach': reach / n, 'collided': len(first_t) / n,
|
||
'first_hit_s': {'p10': pct(first_t, .1), 'p50': pct(first_t, .5), 'p90': pct(first_t, .9),
|
||
'within_0.1s': sum(t <= 0.1 for t in first_t) / max(1, len(first_t))},
|
||
'col_steps': {'p50': pct(col_steps, .5), 'p90': pct(col_steps, .9)},
|
||
'first_hit_pair(group a-b)': dict(sorted(pairs.items(), key=lambda kv: -kv[1])[:8]),
|
||
'min_goal_err_mm': {'p10': pct(min_err, .1), 'p50': pct(min_err, .5), 'p90': pct(min_err, .9)},
|
||
'safety_on': {'reach': reach_s / n, 'collided': col_s / n, 'brakes_p50': pct(brakes_s, .5)}}
|
||
out['C']['goal_err_ratio_p50'] = {'%.1fs' % ts: pct(v, .5) for ts, v in prog.items()} # <1 = 在靠近目标
|
||
|
||
# ---------- F：一步采样残留噪声（2026-10-19 加） ----------
|
||
# 前几轮：子目标、抄速度、起步都已排除；但动作幅度在开环、起点、闭环三处一致是老师的 1.4–1.6 倍，
|
||
# 而开环 16 步末端 FK 误差只有 ~6 mm。幅度大、终点却准 = 每步动作里叠着来回抵消的抖动。
|
||
# 嫌疑：MeanFlow 一步生成 x = z − u(z,0,1) 没把噪声 z 去干净（只有 25% 样本学区间平均速度、14k 步），
|
||
# 闭环每次只执行前 3 步，抖动不抵消，表现为「动得多、走得少」，再挤进碰撞。
|
||
# F0：开环 val 随机 4096 行，比较 1/2/4/8 段采样和零噪声（z=0）：幅度比、抖动比、残差与 z 的相关、前 3 步与 16 步末关节误差
|
||
# F1：同一批闭环题，用 4 段采样、零噪声各跑一遍，对比 C 的基线
|
||
stepv = torch.tensor(STEP, device=dev)
|
||
ri = torch.randperm(vd.X.shape[0], generator=torch.Generator().manual_seed(3))[:4096].to(dev)
|
||
Xr, Yr = vd.X[ri].float(), vd.Y[ri].float()
|
||
zr = torch.randn(len(ri), 1, H, 6, device=dev, generator=torch.Generator(device=dev).manual_seed(11))
|
||
ymag = Yr.abs().mean().item(); yjit = (Yr[:, 1:] - Yr[:, :-1]).abs().mean().item()
|
||
f0 = {'teacher_amp': ymag, 'teacher_jitter': yjit}
|
||
for tag, steps, zs in (('1step', 1, 1.0), ('2step', 2, 1.0), ('4step', 4, 1.0), ('8step', 8, 1.0), ('1step_z0', 1, 0.0)):
|
||
with torch.no_grad(): ac, _, _ = model.sample_k(Xr, 1, zr * zs, steps=steps)
|
||
ac = ac[:, 0]; acc = ac.clamp(-1, 1); res = ac - Yr
|
||
zc = zr[:, 0] - zr[:, 0].mean(); rc_ = res - res.mean()
|
||
corr = ((zc * rc_).sum() / (zc.pow(2).sum().sqrt() * rc_.pow(2).sum().sqrt()).clamp(min=1e-9)).item() if zs > 0 else None
|
||
e3 = ((acc[:, :3] - Yr[:, :3]) * stepv).sum(1).abs().max(-1).values
|
||
e16 = ((acc - Yr) * stepv).sum(1).abs().max(-1).values
|
||
f0[tag] = {'amp_ratio': acc.abs().mean().item() / ymag, 'jitter_ratio': (acc[:, 1:] - acc[:, :-1]).abs().mean().item() / yjit,
|
||
'resid_corr_z': corr, 'first3_err_deg_p50': pct(e3.tolist(), .5), 'end16_err_deg_p50': pct(e16.tolist(), .5)}
|
||
out['F0_openloop_sampler'] = f0
|
||
for tag, kw in (('F1_4step', {'steps': 4}), ('F1_z0', {'z_scale': 0.0})):
|
||
rc = cc_ = 0; me = []; r2 = []
|
||
for i in range(0, len(tasks), a.bs):
|
||
ch = tasks[i:i + a.bs]
|
||
with torch.no_grad():
|
||
_, pp = rollout(model, kin, ch, dev, K=8, safety=False, seed=i, return_traj=True, **kw)
|
||
tr = pp['traj'].to(dev); B, Tn, _ = tr.shape
|
||
k2 = min(Tn - 1, int(round(2.0 / DT)))
|
||
Tt = torch.tensor([t['T'] for t in ch], dtype=torch.float32, device=dev)
|
||
Rt = torch.tensor([t['R'] if t.get('R') else [0.] * 9 for t in ch], dtype=torch.float32, device=dev)
|
||
e0, _ = goal_err(kin, tr[:, 0], Tt, Rt); e2, _ = goal_err(kin, tr[:, k2], Tt, Rt)
|
||
r2 += (e2 / e0.clamp(min=1e-3)).tolist()
|
||
rc += int(pp['done'].sum()); cc_ += int((pp['col'] > 0).sum()); me += pp['err'].tolist()
|
||
out[tag] = {'reach': rc / n, 'collided': cc_ / n, 'goal_err_ratio_2s_p50': pct(r2, .5), 'end_goal_err_mm_p50': pct(me, .5)}
|
||
s = json.dumps(out, indent=1, ensure_ascii=False); print(s)
|
||
os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True)
|
||
open(a.out, 'w', encoding='utf-8').write(s)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|