fix(kin/diag): kin.gaps 粗筛段也按点分块(COARSE_ENTRIES=3e7,≤446 点/块;实测上一版粗筛 4096 点仍 9.3 GiB);diag_cl 加 D 段(子目标头精度 / 不用子目标 / 老师真实途经点)与目标误差进度比;rollout 加 use_sub、sub_override;本机执行清单更新

This commit is contained in:
Hark
2026-10-18 18:32:45 +00:00
parent 95ba43a1c7
commit 0653cf0d8d
4 changed files with 102 additions and 75 deletions
+36 -15
View File
@@ -118,11 +118,22 @@ class Kin:
return a, b
# ---------------- 间隙 ----------------
# 2026-10-19(用户实测 dagger/诊断吃满 24 GB 显存):粗筛后幸存胶囊对数随 点数 × cap 暴涨,
# 一次几千个点、cap=40 mm 时中间张量可达十几 GB(训练里 eval 峰值 15.9 GiB 也是它)。
# 现在先按簇级粗筛估出每个点的幸存胶囊对数,按 PAIR_BUDGET 把 batch 切块逐块算,峰值与调用方的 batch 大小无关。
# 2026-10-19(用户实测 dagger/诊断吃满 24 GB 显存):本模型 1936 条胶囊 / 491 组 / 67256 粗对 / 104.7 万细对。
# 显存有两段峰值,都随点数线性涨:
# 粗筛段 cc[:,cp_a]-cc[:,cp_b] 的广播中间量 ≈ 36 B × 点数 × 粗对数(4096 点实测 9.3 GiB,与 cap 无关);
# 精算段 ≈ 每对几百字节 × 幸存细对数(4096 点 cap 40 幸存 2.4 亿对;旧版不分块实测 36.9 GiB)。
# 所以先按 COARSE_ENTRIES 把点切块(粗筛每块 ≤ ~1 GiB),块内再按幸存细对数 PAIR_BUDGET 切(精算每块 ≤ ~1 GiB)。
# 两段不同时存在,峰值与调用方 batch 大小无关。结果与不分块逐位一致(梯度经 cat 保留)。
COARSE_ENTRIES = 30_000_000
PAIR_BUDGET = 2_000_000
@staticmethod
def _cat_gaps(parts, starts):
return {'gm': torch.cat([p['gm'] for p in parts], 0),
'b': torch.cat([p['b'] + s for p, s in zip(parts, starts)]),
'i': torch.cat([p['i'] for p in parts]), 'j': torch.cat([p['j'] for p in parts]),
'gap': torch.cat([p['gap'] for p in parts])}
def gaps(self, q, cap=40.0, xf=None, max_pairs=None):
"""返回 dict:
gm [B,49]:组对最小间隙 min(gap, cap)(gA*7+gB,gA<gB;其余为 cap)
@@ -131,11 +142,25 @@ class Kin:
xf = xf if xf is not None else self.cum_xf(q)
Rs, Ts = xf
B = q.shape[0]
step = max(1, self.COARSE_ENTRIES // int(self.cp_a.shape[0]))
if B <= step:
return self._gaps_split(q, cap, xf)
st = list(range(0, B, step))
return self._cat_gaps([self._gaps_split(q[s:s + step], cap, (Rs[s:s + step], Ts[s:s + step])) for s in st], st)
def _coarse(self, q, cap, xf):
Rs, Ts = xf
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g]
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1) # [B,M]
del cc
return d2 <= (self.cp_lim + cap).pow(2)
def _gaps_split(self, q, cap, xf):
"""粗筛已在显存预算内:按幸存细对数再切块做精算"""
Rs, Ts = xf
B = q.shape[0]
with torch.no_grad():
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g]
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1)
per = ((d2 <= (self.cp_lim + cap).pow(2)).long() * self.cp_cnt).sum(1) # [B] 每点幸存胶囊对数
del cc, d2
per = (self._coarse(q, cap, xf).long() * self.cp_cnt).sum(1) # [B] 每点幸存细对数
if B <= 1 or int(per.sum()) <= self.PAIR_BUDGET:
bounds = [0, B]
else:
@@ -144,22 +169,18 @@ class Kin:
if c - base > self.PAIR_BUDGET and k > bounds[-1]:
bounds.append(k); base = cs[k - 1]
bounds.append(B)
del per
if len(bounds) == 2:
return self._gaps_one(q, cap, xf)
parts = [self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e])) for s, e in zip(bounds[:-1], bounds[1:])]
return {'gm': torch.cat([p['gm'] for p in parts], 0),
'b': torch.cat([p['b'] + s for p, s in zip(parts, bounds[:-1])]),
'i': torch.cat([p['i'] for p in parts]), 'j': torch.cat([p['j'] for p in parts]),
'gap': torch.cat([p['gap'] for p in parts])}
return self._cat_gaps([self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e])) for s, e in zip(bounds[:-1], bounds[1:])], bounds[:-1])
def _gaps_one(self, q, cap, xf):
Rs, Ts = xf
B = q.shape[0]
with torch.no_grad():
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g]
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1) # [B,M]
keep = d2 <= (self.cp_lim + cap).pow(2)
keep = self._coarse(q, cap, xf)
bb, mm = keep.nonzero(as_tuple=True)
del keep
cnt = self.cp_cnt[mm]
rep_b = torch.repeat_interleave(bb, cnt)
off = torch.repeat_interleave(self.cp_start[mm] - torch.cumsum(cnt, 0) + cnt, cnt) + torch.arange(int(cnt.sum()), device=q.device)
+39
View File
@@ -27,6 +27,9 @@ GROUPS = [('sincos', 0, 12), ('joint_norm', 12, 15), ('dq', 15, 21), ('p', 21, 2
('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
@@ -81,6 +84,7 @@ def main():
# ---------- 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]
@@ -95,6 +99,9 @@ def main():
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
@@ -113,6 +120,38 @@ def main():
'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 = 在靠近目标
# ---------- D:子目标是不是元凶(2026-10-19 加:C 显示全程最小目标误差中位 254 mm,策略根本没朝目标走) ----------
# 训练时特征里的子目标是老师真实途经点(80% 概率给),闭环里却换成子目标头自己的预测;预测偏了,策略就会忠实地朝错的方向走。
sg_by_tid = {}
for i, t in zip(idx, tk):
qrow = torch.atan2(vd.X[i, 0:12:2].float(), vd.X[i, 1:12:2].float()) * 180 / torch.pi
sgr = vd.SG[i].float()
sg_by_tid[t['tid']] = (qrow[None] + sgr[:12].view(2, 6) * 180, bool(sgr[12] > 0.5))
# D0:冷启动子目标头精度(起点特征里抹掉数据集的真子目标,与闭环冷启动同口径)
Xs = X0.clone(); Xs[:, 68:75] = 0
with torch.no_grad(): _, sp = model.heads(model.encode(Xs))
lab = torch.stack([vd.SG[i].float() for i in idx])
need_l = lab[:, 12] > 0.5; need_p = sp[:, 12] > 0.5
via_err = ((sp[:, :6] - lab[:, :6]) * 180).abs().max(-1).values
out['D0_subhead'] = {'n': len(idx), 'need_label_rate': need_l.float().mean().item(), 'need_acc': (need_l == need_p).float().mean().item(),
'via1_maxjoint_err_deg_p50': pct(via_err[need_l].tolist(), .5), 'via1_maxjoint_err_deg_p90': pct(via_err[need_l].tolist(), .9)}
# D1 不用子目标、D2 用老师真实途经点(oracle),同一批题、同一随机种子,对比 C 的基线
def run(tag, **kw):
rc = cc_ = 0; me = []
for i in range(0, len(tasks), a.bs):
ch = tasks[i:i + a.bs]; extra = dict(kw)
if kw.get('sub_override') == 'oracle':
subs = torch.stack([sg_by_tid[t['tid']][0] if t['tid'] in sg_by_tid else torch.zeros(2, 6, device=dev) for t in ch])
need = torch.tensor([sg_by_tid[t['tid']][1] if t['tid'] in sg_by_tid else False for t in ch], device=dev)
extra['sub_override'] = (subs, need)
with torch.no_grad():
_, pp = rollout(model, kin, ch, dev, K=8, safety=False, seed=i, **extra)
rc += int(pp['done'].sum()); cc_ += int((pp['col'] > 0).sum()); me += pp['err'].tolist()
out[tag] = {'reach': rc / n, 'collided': cc_ / n, 'end_goal_err_mm_p50': pct(me, .5)}
run('D1_no_sub', use_sub=False)
run('D2_oracle_sub', sub_override='oracle')
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)
+4 -1
View File
@@ -65,7 +65,8 @@ def dls_steps(kin, q, T, R, n=3, scale=0.3, lam=0.5):
@torch.no_grad()
def rollout(model, kin, tasks, device, K=8, safety=False, seed=0, max_sec=None, return_traj=False):
def rollout(model, kin, tasks, device, K=8, safety=False, seed=0, max_sec=None, return_traj=False, use_sub=True, sub_override=None):
"""use_sub=False:不用子目标(sf=0);sub_override=(subs [B,2,6] 绝对关节角, need [B] bool):冷启动时用给定途经点代替子目标头(诊断用)"""
g = torch.Generator(device=device); g.manual_seed(seed)
q0, T, R, ex, tr = tasks_to_tensors(tasks, device)
B = q0.shape[0]
@@ -108,6 +109,8 @@ def rollout(model, kin, tasks, device, K=8, safety=False, seed=0, max_sec=None,
sub0 = sub
need = sub0[:, 12] > 0.5
subs = q[:, None] + sub0[:, :12].view(B, 2, 6) * 180
if sub_override is not None: subs, need = sub_override[0].to(device), sub_override[1].to(device)
if not use_sub: need = torch.zeros_like(need)
nsub = torch.where(need, torch.full_like(nsub, 2), torch.zeros_like(nsub)); sub_t.fill_(tnow)
fut = qe[:, None, None] + torch.cumsum(inc[:, :, r:], 2) # [B,K,H-r,6]
if last is None: