"""PyTorch 批量运动学 + 可微胶囊间隙(与 webapp/js/ik-solver.js 同口径)。 数据来自 ml/assets/kin.json(node ml/prep/export_kin.cjs 生成): - cum_xf / fk / orient:与 JS cumXf / fk(q,0) / orient(q) 一致(角度制输入) - gaps(q, cap):所有 |Δg|≥2 胶囊对的 间隙 = 线段距离 − (r_i + r_j + margin)。 宽相位复刻 JS 的簇球筛选(被筛掉的对间隙必 > cap),故 min(gap, cap) 与 JS gapMatrix 逐位同口径; 筛选在 no_grad 下做,幸存对的精确距离保留梯度(APG / FK 损失用)。 """ import json, math, os import torch HERE = os.path.dirname(os.path.abspath(__file__)) DEFAULT_KIN = os.path.join(HERE, '..', 'assets', 'kin.json') D2R = math.pi / 180.0 def rot_axis(d, a): """d: [3] 单位轴;a: [...] 弧度 → [..., 3, 3](Rodrigues,右手)""" c, s = torch.cos(a), torch.sin(a) x, y, z = float(d[0]), float(d[1]), float(d[2]) C = 1 - c row0 = torch.stack([c + x * x * C, x * y * C - z * s, x * z * C + y * s], -1) row1 = torch.stack([y * x * C + z * s, c + y * y * C, y * z * C - x * s], -1) row2 = torch.stack([z * x * C - y * s, z * y * C + x * s, c + z * z * C], -1) return torch.stack([row0, row1, row2], -2) def seg_seg_dist(p1, q1, p2, q2, eps=1e-12): """两线段最近距离(Ericson, Real-Time Collision Detection §5.1.9),[..., 3] → [...]""" d1, d2, r = q1 - p1, q2 - p2, p1 - p2 a = (d1 * d1).sum(-1); e = (d2 * d2).sum(-1); f = (d2 * r).sum(-1) c = (d1 * r).sum(-1); b = (d1 * d2).sum(-1) denom = a * e - b * b s = torch.where(denom > eps, ((b * f - c * e) / denom.clamp_min(eps)).clamp(0, 1), torch.zeros_like(a)) t = (b * s + f) / e.clamp_min(eps) # t 越界:夹回 [0,1] 并重算 s t_lo, t_hi = t < 0, t > 1 s = torch.where(t_lo, (-c / a.clamp_min(eps)).clamp(0, 1), s) s = torch.where(t_hi, ((b - c) / a.clamp_min(eps)).clamp(0, 1), s) t = t.clamp(0, 1) # 第二条退化成点时,仍须在第一条上投影,而不是固定取其起点。 s = torch.where(e <= eps, (-c / a.clamp_min(eps)).clamp(0, 1), s) s = torch.where(a <= eps, torch.zeros_like(s), s) t = torch.where(e <= eps, torch.zeros_like(t), t) c1 = p1 + d1 * s.unsqueeze(-1) c2 = p2 + d2 * t.unsqueeze(-1) return ((c1 - c2).pow(2).sum(-1) + 1e-18).sqrt() class Kin: def __init__(self, path=DEFAULT_KIN, device='cpu', dtype=torch.float32): with open(path, 'r', encoding='utf-8') as f: K = json.load(f) self.fp = K['fp'] self.device, self.dtype = device, dtype T = lambda x, dt=dtype: torch.tensor(x, dtype=dt, device=device) self.P = T(K['P']); self.D = [list(map(float, d)) for d in K['D']] lim = K['lim'] self.lo = T([l[0] for l in lim]); self.hi = T([l[1] for l in lim]); self.cont = T([l[2] for l in lim]).bool() self.tcp0 = T(K['tcp0']); self.capm = float(K['capm']) caps = K['caps'] self.cap_g = torch.tensor([c['g'] for c in caps], device=device) self.cap_a = T([c['a'] for c in caps]); self.cap_b = T([c['b'] for c in caps]); self.cap_r = T([c['r'] for c in caps]) cl = K['clusters'] self.cl_g = torch.tensor([c['g'] for c in cl], device=device) self.cl_c = T([c['c'] for c in cl]) cp = K['cpairs'] self.cp_a = torch.tensor([p[0] for p in cp], device=device) self.cp_b = torch.tensor([p[1] for p in cp], device=device) self.cp_lim = T([p[2] + p[3] + self.capm for p in cp]) # 每个簇对展开成胶囊对(全组合),按簇对连续存放 ci, cj, start, cnt = [], [], [], [] for p in cp: A, B = cl[p[0]]['caps'], cl[p[1]]['caps'] start.append(len(ci)); cnt.append(len(A) * len(B)) for i in A: for j in B: ci.append(i); cj.append(j) self.pi = torch.tensor(ci, device=device, dtype=torch.long) self.pj = torch.tensor(cj, device=device, dtype=torch.long) self.cp_start = torch.tensor(start, device=device, dtype=torch.long) self.cp_cnt = torch.tensor(cnt, device=device, dtype=torch.long) hl = 0.5 * (self.cap_b - self.cap_a).norm(dim=-1) self.pair_need = self.cap_r[self.pi] + self.cap_r[self.pj] + self.capm self.pair_bl = hl[self.pi] + hl[self.pj] + self.cap_r[self.pi] + self.cap_r[self.pj] + self.capm gi, gj = self.cap_g[self.pi], self.cap_g[self.pj] self.pair_gp = torch.minimum(gi, gj) * 7 + torch.maximum(gi, gj) self.lever = T(K['lever']) table = [[0.0] * 7 for _ in range(6)] for group in range(1, 7): endpoints = [c[key] for c in caps if c['g'] == group for key in ('a', 'b')] if not endpoints: continue pivot, axis = K['P'][group - 1], K['D'][group - 1] unit = [v / math.hypot(*axis) for v in axis] relative = [[e[j] - pivot[j] for j in range(3)] for e in endpoints] radius = max(math.hypot(*v) for v in relative) radial = max(math.hypot(v[1] * unit[2] - v[2] * unit[1], v[2] * unit[0] - v[0] * unit[2], v[0] * unit[1] - v[1] * unit[0]) for v in relative) for joint in range(group): chain = sum(math.dist(K['P'][j], K['P'][j + 1]) for j in range(joint, group - 1)) table[joint][group] = (chain + (radial if joint == group - 1 else radius)) * D2R self.cert_groups = torch.tensor([a * 7 + b for a in range(7) for b in range(a + 2, 7)], device=device) # 三角不等式给出全姿态界,不用随机姿态最大值冒充上界。 self.cert_lever = T([[table[k][b] if a <= k < b else 0.0 for k in range(6)] for a in range(7) for b in range(a + 2, 7)]) # ---------------- 运动学 ---------------- def _checked_xf(self, q, xf): """xf 是按行对应 q 的位姿数组,行数不匹配会静默广播出错(宽相位筛到别的行 ⇒ device assert 或错间隙)。""" if xf is None: return self.cum_xf(q) Rs, Ts = xf if Rs.shape[0] != q.shape[0] or Ts.shape[0] != q.shape[0]: raise ValueError('xf 的行数 %d 与 q 的行数 %d 不符' % (Rs.shape[0], q.shape[0])) return xf def cum_xf(self, q): """q: [B,6] 度 → Rs [B,7,3,3], Ts [B,7,3]""" B = q.shape[0] R = torch.eye(3, dtype=q.dtype, device=q.device).expand(B, 3, 3) Tt = torch.zeros(B, 3, dtype=q.dtype, device=q.device) Rs, Ts = [R], [Tt] for k in range(6): Rk = rot_axis(self.D[k], q[:, k] * D2R) Pk = self.P[k].to(q.dtype) tk = Pk - torch.einsum('bij,j->bi', Rk, Pk) Tt = torch.einsum('bij,bj->bi', R, tk) + Tt R = R @ Rk Rs.append(R); Ts.append(Tt) return torch.stack(Rs, 1), torch.stack(Ts, 1) def fk(self, q, xf=None): Rs, Ts = xf if xf is not None else self.cum_xf(q) return torch.einsum('bij,j->bi', Rs[:, 6], self.tcp0.to(q.dtype)) + Ts[:, 6] def orient(self, q, xf=None): Rs, _ = xf if xf is not None else self.cum_xf(q) return Rs[:, 6] def jac(self, q, xf=None): """几何雅可比(对度):返回 (Jv [B,3,6] mm/deg, Jw [B,3,6] 单位轴/deg)""" Rs, Ts = xf if xf is not None else self.cum_xf(q) p = self.fk(q, (Rs, Ts)) Jv, Jw = [], [] for k in range(6): d = torch.tensor(self.D[k], dtype=q.dtype, device=q.device) a = Rs[:, k] @ d o = torch.einsum('bij,j->bi', Rs[:, k], self.P[k].to(q.dtype)) + Ts[:, k] Jv.append(torch.linalg.cross(a, p - o, dim=-1) * D2R) Jw.append(a * D2R) return torch.stack(Jv, -1), torch.stack(Jw, -1) def caps_world(self, q, xf=None): Rs, Ts = xf if xf is not None else self.cum_xf(q) Rg, Tg = Rs[:, self.cap_g], Ts[:, self.cap_g] # [B,C,3,3], [B,C,3] a = torch.einsum('bcij,cj->bci', Rg, self.cap_a.to(q.dtype)) + Tg b = torch.einsum('bcij,cj->bci', Rg, self.cap_b.to(q.dtype)) + Tg return a, b # ---------------- 间隙 ---------------- # 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): out = {'gm': torch.cat([p.pop('gm') for p in parts], 0)} if 'b' in parts[0]: # 非 in-place:b 会被 autograd 的索引反向图引用(APG 的碰撞铰链梯度路径),+= 会让 backward 报 version 错 for p, s in zip(parts, starts): p['b'] = p['b'] + s for k in ('b', 'i', 'j', 'gap'): out[k] = torch.cat([p.pop(k) for p in parts]) # 逐键拼、拼完即释放 return out def gaps(self, q, cap=40.0, xf=None, max_pairs=None, pairs=False): """返回 dict: gm [B,49]:组对最小间隙 min(gap, cap)(gA*7+gB,gAbki', 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, pairs): """粗筛已在显存预算内:按幸存细对数再切块做精算""" Rs, Ts = xf B = q.shape[0] with torch.no_grad(): 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: cs = torch.cumsum(per, 0).tolist(); bounds = [0]; base = 0 for k, c in enumerate(cs): 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, pairs) return self._cat_gaps([self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e]), pairs) for s, e in zip(bounds[:-1], bounds[1:])], bounds[:-1]) def _gaps_one(self, q, cap, xf, pairs): Rs, Ts = xf B = q.shape[0] with torch.no_grad(): 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) pidx = off a, b = self.caps_world(q, xf) with torch.no_grad(): pi, pj = self.pi[pidx], self.pj[pidx] mid = 0.5 * (a[rep_b, pi] + b[rep_b, pi] - a[rep_b, pj] - b[rep_b, pj]) k2 = mid.pow(2).sum(-1) <= (self.pair_bl[pidx] + cap).pow(2) rep_b, pidx, pi, pj = rep_b[k2], pidx[k2], pi[k2], pj[k2] d = seg_seg_dist(a[rep_b, pi], b[rep_b, pi], a[rep_b, pj], b[rep_b, pj]) gap = d - self.pair_need[pidx].to(q.dtype) gm = torch.full((B * 49,), float(cap), dtype=q.dtype, device=q.device) gm = gm.scatter_reduce(0, rep_b * 49 + self.pair_gp[pidx], gap.clamp(max=cap), reduce='amin', include_self=True) if not pairs: return {'gm': gm.view(B, 49)} return {'gm': gm.view(B, 49), 'b': rep_b, 'i': pi, 'j': pj, 'gap': gap} def gm_groups(self, q, gp, cap=40.0, xf=None): """只算指定组对 gp=[(a,b),...] 的最小间隙 [B,len(gp)](特征用,避免 49 组全量)。 数值与 gaps(q,cap)['gm'][:, a*7+b] 一致;组对在碰撞粗对里不存在时返回 cap。 xf 可传入调用者已算好的 cum_xf(q),避免同一姿态重复算 FK(特征路径每周期省一次)。""" key = tuple(gp) if getattr(self, '_gp_key', None) != key: m = torch.full((49,), -1, dtype=torch.long, device=q.device) for i, (a_, b_) in enumerate(gp): m[min(a_, b_) * 7 + max(a_, b_)] = i ga, gb = self.cl_g[self.cp_a], self.cl_g[self.cp_b] cg = torch.minimum(ga, gb) * 7 + torch.maximum(ga, gb) # 簇对 → 7 连杆组对 sel = torch.isin(cg, torch.tensor([min(a, b) * 7 + max(a, b) for a, b in gp], device=q.device)) self._gp = (m, sel, sel.nonzero(as_tuple=True)[0]); self._gp_key = key self._gp49, sel, idx = self._gp xf = self._checked_xf(q, 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[idx]] - cc[:, self.cp_b[idx]]).pow(2).sum(-1) keep = d2 <= (self.cp_lim[idx] + cap).pow(2) bb, mm = keep.nonzero(as_tuple=True) del keep, d2, cc cnt = self.cp_cnt[idx][mm] rep_b = torch.repeat_interleave(bb, cnt) off = torch.repeat_interleave(self.cp_start[idx][mm] - torch.cumsum(cnt, 0) + cnt, cnt) + torch.arange(int(cnt.sum()), device=q.device) pidx = off a, b = self.caps_world(q, xf) with torch.no_grad(): pi, pj = self.pi[pidx], self.pj[pidx] mid = 0.5 * (a[rep_b, pi] + b[rep_b, pi] - a[rep_b, pj] - b[rep_b, pj]) k2 = mid.pow(2).sum(-1) <= (self.pair_bl[pidx] + cap).pow(2) rep_b, pidx, pi, pj = rep_b[k2], pidx[k2], pi[k2], pj[k2] d = seg_seg_dist(a[rep_b, pi], b[rep_b, pi], a[rep_b, pj], b[rep_b, pj]) gap = d - self.pair_need[pidx].to(q.dtype) gm = torch.full((B * len(gp),), float(cap), dtype=q.dtype, device=q.device) gm = gm.scatter_reduce(0, rep_b * len(gp) + self._gp49[self.pair_gp[pidx]], gap.clamp(max=cap), reduce='amin', include_self=True) return gm.view(B, len(gp)) @torch.no_grad() def certify_segments(self, a, b, max_queries=32, max_depth=12, eps=1e-4, cap=20.0): """逐层二分认证:每根段独立查询预算,未知(超预算/超深度/非有限/越界)一律拒绝。 判据对每条组对:子段两端间隙之和 > 该子段位移界 + 2eps(间隙取 min(gap,cap),钳制只会更保守)。 cap 必须 ≥ 该层子段界 + 2eps 才不会挡住本可通过的判定,故不做自适应收紧——同批其他段会改变 钳制程度、让结果随批次而变(实测也无耗时收益,见探索日志轮次 31)。 """ if self.dtype != torch.float64 or a.dtype != torch.float64 or b.dtype != torch.float64: raise ValueError('segment certification requires float64') if a.shape != b.shape or a.ndim != 2 or a.shape[1] != 6: raise ValueError('a and b must be [count,6] with identical shape') # 预算给成张量时 max_depth + 1 会在 uint8 这类窄类型上回绕成 0,循环整个跳过,把碰撞段当安全返回。 if type(max_queries) is not int or type(max_depth) is not int: raise ValueError('certification budgets must be plain ints') if max_queries < 2 or max_depth < 0 or eps <= 0 or cap <= eps: raise ValueError('invalid certification budget or tolerance') count = a.shape[0] status = torch.zeros(count, dtype=torch.long, device=a.device) inside = (((a >= self.lo) & (a <= self.hi) & (b >= self.lo) & (b <= self.hi)) | self.cont).all(-1) status[~inside] = 3 root = torch.arange(count, device=a.device) gaps = self.gaps(torch.cat([a, b]), cap)['gm'][:, self.cert_groups] ga, gb = gaps[:count], gaps[count:] queries = torch.full((count,), 2, dtype=torch.long, device=a.device) bound = (b - a).abs() @ self.cert_lever.T for depth in range(max_depth + 1): finite = torch.isfinite(ga).all(-1) & torch.isfinite(gb).all(-1) collision = (ga <= 0).any(-1) | (gb <= 0).any(-1) status[root[~finite]] = 2 status[root[collision]] = 1 positive = (ga > eps).all(-1) & (gb > eps).all(-1) proven = positive & (ga + gb > bound + 2 * eps).all(-1) pending = (status[root] == 0) & ~proven if not bool(pending.any()): break a, b, ga, gb, root, bound = a[pending], b[pending], ga[pending], gb[pending], root[pending], bound[pending] if depth == max_depth: status[root] = 2 break requested = torch.bincount(root, minlength=count) status[(status == 0) & (queries + requested > max_queries)] = 2 keep = status[root] == 0 a, b, ga, gb, root, bound = (t[keep] for t in (a, b, ga, gb, root, bound)) if not root.numel(): break mid = a + 0.5 * (b - a) gm = self.gaps(mid, cap)['gm'][:, self.cert_groups] queries += torch.bincount(root, minlength=count) a, b = torch.cat([a, mid]), torch.cat([mid, b]) ga, gb = torch.cat([ga, gm]), torch.cat([gm, gb]) root = torch.cat([root, root]) bound = torch.cat([bound, bound]) / 2 return {'safe': status == 0, 'collision': status == 1, 'unresolved': status == 2, 'outside_box': status == 3, 'queries': queries} @torch.no_grad() def certify_uniform(self, a, b, n=64, eps=1e-4, cap=None): """固定层数的均匀细分认证:n 个子段只需一次 gaps 调用(n+1 个点),没有逐层依赖的往返。 判据与 certify_segments 完全相同;界对 Δq 线性,所以子段界 = 整段界 / n。未通过即 unproven (拒绝),不存在“未知当安全”。cap=None 时取“最小不挡住可通过判定的宽相位”= 最大子段界 + 2eps: 小于它钳制会把本可证明的子段压成不可证(偏保守),大于它只是多算对(偏慢),不改正确性。 """ if self.dtype != torch.float64 or a.dtype != torch.float64 or b.dtype != torch.float64: raise ValueError('segment certification requires float64') if type(n) is not int or n < 1 or eps <= 0 or (cap is not None and cap <= eps): raise ValueError('invalid subdivision count or tolerance') if a.shape != b.shape or a.ndim != 2 or a.shape[1] != 6: raise ValueError('a and b must be [count,6] with identical shape') count = a.shape[0] inside = (((a >= self.lo) & (a <= self.hi) & (b >= self.lo) & (b <= self.hi)) | self.cont).all(-1) bound = (b - a).abs() @ self.cert_lever.T / n # [count,15] 子段界 required = float(bound.max()) + 2 * eps if count else 2 * eps qc = required if cap is None else cap alpha = torch.arange(n + 1, dtype=self.dtype, device=a.device) / n points = (a[:, None] + (b - a)[:, None] * alpha[None, :, None]).reshape(-1, 6) gm = self.gaps(points, qc)['gm'][:, self.cert_groups].view(count, n + 1, -1) adjacent = gm[:, :-1] + gm[:, 1:] > bound[:, None, :] + 2 * eps collision = (gm <= 0).any((1, 2)) safe = inside & (gm > eps).all((1, 2)) & adjacent.all((1, 2)) & ~collision return {'safe': safe, 'collision': collision, 'unproven': (~safe & ~collision & inside), 'outside_box': ~inside, 'points': n + 1, 'queries': n + 1, 'cap_required_mm': required, 'cap_used_mm': qc} def min_gap(self, q, cap=40.0): return self.gaps(q, cap)['gm'].min(-1).values def unwrap_to(self, a, b): """连续轴取离 a 最近的等价角(与 JS unwrapTo 同口径)""" d = b - a dw = torch.remainder(d + 180.0, 360.0) - 180.0 return torch.where(self.cont, a + dw, b) def clamp_box(self, q): return torch.where(self.cont, q, torch.maximum(torch.minimum(q, self.hi), self.lo))