ml: 全流程代码与运行手册（数据打包/可微运动学与对拍/MeanFlow+WTA 模型/训练/闭环仿真/APG/挖错循环/ONNX 导出）；webapp 策略运行时 ik-policy.js + ?policy=1 接入；tests/policy_bench.cjs；ik-solver 增 pairList/clusterExport/criticalPairs/pairGap/leverTable
This commit is contained in:
Hark
@@ -0,0 +1,59 @@
|
||||
"""策略输入特征（83 维）—— 与 webapp/js/ik-policy.js features() 逐项一致。改动须两边同步并升 FEAT_VER。"""
|
||||
import math
|
||||
import torch
|
||||
|
||||
FEAT_VER = 1
|
||||
VMAX = [60., 60., 90., 90., 120., 180.]
|
||||
DT = 0.02
|
||||
STEP = [v * DT for v in VMAX]
|
||||
ACC = 5.0
|
||||
H, NJ, PREFIX = 16, 6, 3
|
||||
GP = [(0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (2, 4), (2, 5), (3, 5)]
|
||||
GAP_CAP = 40.0
|
||||
F_DIM = 83
|
||||
D2R = math.pi / 180.0
|
||||
ROT6 = [0, 3, 6, 1, 4, 7] # 行主序 3×3 的前两列
|
||||
|
||||
|
||||
def wrap180(x):
|
||||
return torch.remainder(x + 180.0, 360.0) - 180.0
|
||||
|
||||
|
||||
def features(kin, q, q_prev, T, R, prefix, prefix_len, sub, sub_flag, gm=None, xf=None):
|
||||
"""全部为 batch 张量：
|
||||
q, q_prev [B,6] 度；T [B,3] mm；R [B,9]（pos 模式全 0）；pose 标志由 R 是否全 0 判定
|
||||
prefix [B,3,6] 增量（度），prefix_len [B]；sub [B,6]，sub_flag [B]（0/1）
|
||||
gm：kin.gaps(q, 40)['gm'] 可传入复用。"""
|
||||
B = q.shape[0]
|
||||
dt = q.dtype
|
||||
step = torch.tensor(STEP, dtype=dt, device=q.device)
|
||||
xf = xf if xf is not None else kin.cum_xf(q)
|
||||
out = []
|
||||
rad = q * D2R
|
||||
out.append(torch.stack([torch.sin(rad), torch.cos(rad)], -1).reshape(B, 12))
|
||||
lo, hi = kin.lo.to(dt), kin.hi.to(dt)
|
||||
idx = [1, 3, 4]
|
||||
out.append(2 * (q[:, idx] - lo[idx]) / (hi[idx] - lo[idx]) - 1)
|
||||
d = torch.where(kin.cont, wrap180(q - q_prev), q - q_prev)
|
||||
out.append(d / step)
|
||||
p = kin.fk(q, xf); Rq = kin.orient(q, xf).reshape(B, 9)
|
||||
out.append(p / 600)
|
||||
out.append(Rq[:, ROT6])
|
||||
out.append(T / 600)
|
||||
pose = (R.abs().sum(-1) > 0).to(dt)
|
||||
out.append(R[:, ROT6] * pose[:, None])
|
||||
out.append((T - p) / 600)
|
||||
out.append(((T - p) / 50).clamp(-1, 1))
|
||||
out.append(torch.stack([1 - pose, pose], -1))
|
||||
m = (torch.arange(PREFIX, device=q.device)[None, :] < prefix_len[:, None]).to(dt) # [B,3]
|
||||
out.append((prefix / step * m[:, :, None]).reshape(B, 18))
|
||||
out.append(m)
|
||||
sd = torch.where(kin.cont, wrap180(sub - q), sub - q) / 180 * sub_flag[:, None].to(dt)
|
||||
out.append(sd)
|
||||
out.append(sub_flag[:, None].to(dt))
|
||||
if gm is None:
|
||||
gm = kin.gaps(q, GAP_CAP, xf)['gm']
|
||||
out.append(torch.stack([gm[:, a * 7 + b].clamp(max=GAP_CAP) / GAP_CAP for a, b in GP], -1))
|
||||
f = torch.cat(out, -1)
|
||||
assert f.shape[1] == F_DIM, f.shape
|
||||
return f
|
||||
@@ -0,0 +1,160 @@
|
||||
"""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)
|
||||
# 退化（线段长度≈0）
|
||||
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'])
|
||||
|
||||
# ---------------- 运动学 ----------------
|
||||
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 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
|
||||
|
||||
# ---------------- 间隙 ----------------
|
||||
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）
|
||||
b, i, j, gap：幸存胶囊对（gap 可微）
|
||||
"""
|
||||
xf = xf if xf is not None else 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] - cc[:, self.cp_b]).pow(2).sum(-1) # [B,M]
|
||||
keep = d2 <= (self.cp_lim + cap).pow(2)
|
||||
bb, mm = keep.nonzero(as_tuple=True)
|
||||
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)
|
||||
return {'gm': gm.view(B, 49), 'b': rep_b, 'i': pi, 'j': pj, 'gap': gap}
|
||||
|
||||
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))
|
||||
Reference in New Issue
Block a user