60 lines
2.2 KiB
Python
60 lines
2.2 KiB
Python
"""策略输入特征（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
|