Files
Robot_Arm/ml/common/kin_torch.py
T
cbb 6de085fd37 exp+fix: 轮次31——确定界连续认证抓到段内碰撞,但单段P50 22.9ms进不了20ms;自适应cap收益经交错复测撤回
- Kin.cert_lever 用三角不等式给全姿态位移界(替代 JS 随机采样×1.05),界量级 5.93–48.86 mm/整步
- Kin.certify_segments 批量前沿二分+每根查询预算:48 固定段中预算32判安全37/碰撞1/未知拒10,
  第180题1.08–1.10s漏检段以5次查询判碰撞;未知只拒不放
- 修 seg_seg_dist 第二条退化成点不投影(模型内退化胶囊686受影响);但7586个加密点的组最小间隙
  逐位不变、2000样本JS对拍仍 PASS → 不解释轮次30漏检、不改既有成绩
- 负结论:B=1 单段认证 P50 22.93/P90 58.43/P99 101.57 ms,432次中225次>20ms;
  自适应宽相位 cap 交错复测耗时无收益,先前1.6×为GPU顺序漂移,已在日志撤回
- 成本模型 L=ceil(log2(bound/(ga+gb)))、查询=2^L+1 与实测对齐(安全段2–17次,未知段33–65次)
- 新增8项段认证测试+1项回放测试(14→23项通过);变异实测证明张量 max_depth 回绕会误放行
2026-10-20 19:13:40 +08:00

334 lines
18 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, adaptive_cap=True):
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)
# 子段的界是父段一半;宽相位只需覆盖到能被证明的最小间隙,层越深筛得越狠。
qc = min(cap, float(bound.max()) / 2 + eps) if adaptive_cap else cap
gm = self.gaps(mid, qc)['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}
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))