Files
Robot_Arm/ml/common/kin_torch.py
T
cbb c6b46b8bb1 exp: 轮次32——均匀细分认证一次调用可证段内碰撞,界经720项检验未证伪,但P50≤20ms的档只能证出54~63%
- certify_uniform:固定 n 个子段一次 gaps 调用,判据与二分同源(子段界=整段界/n),证不出即拒
- 确定界反证:48段×121点×15组对=720项检验 0 违例;有信息量仅 61 项,最紧仍高出 0.035 mm
  (是"未能在真实运动上证伪",不是证明;浮点误差仍未形式化)
- 第180题 1.08-1.10s 漏检段在每个 n>=8 均判碰撞;二次加密(2n+1)反证 0 误放行
- 更正轮次31的"自适应cap无收益":旧公式 bound/2+eps 会把本可通过的子段钳成不可证;
  改为 bound+2eps 后交错实测决策不变、快 1.36~1.65x(n=8 13.13→9.68,n=128 98.46→59.76 ms)
- 负结论:只有 n<=16 能把认证压进 20 ms,但该档仅证出 26-30/48 根,其余必须刹停;
  成本=log(bound/间隙) 决定,故下一步要在反馈后限幅并让安全层要求正余量(现 cap=0.5 会截掉余量)
- 测试 23→26 项:n=5拒/n=6受的解析对照、cap_required 公式、n 张量类型门
2026-10-20 19:30:02 +08:00

365 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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 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,gA<gB;其余为 cap)
pairs=True 时另有 b, i, j, gap:幸存胶囊对(gap 可微;只有 APG 的 hinge 与 check_parity 需要)
2026-10-19(用户实测 4096 点 cap 40 峰值 4.46 GiB 全在拼接 8.6e7 对的 b/i/j/gap):
默认只返回 gm,每块算完就丢掉胶囊对,峰值只剩单块。
"""
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, pairs)
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]), pairs) 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, 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):
"""只算指定组对 gp=[(a,b),...] 的最小间隙 [B,len(gp)](特征用,避免 49 组全量)。
数值与 gaps(q,cap)['gm'][:, a*7+b] 一致;组对在碰撞粗对里不存在时返回 cap。"""
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.cum_xf(q); 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')
# 预算给成张量时 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')
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))