87 lines
3.7 KiB
Python
87 lines
3.7 KiB
Python
"""无老师出题（PyTorch）：A 均匀 / B 直线被挡 / C 贴立柱 / D 工作空间外；pos/pose 各半。
|
|
2026-10-18：无老师题没有参考时长与行程，exec_s/travel 置 None（此前写死 3.0/0.0 造成假时长预算）；
|
|
闭环仿真对这类题按定长 6 s 口径判到达，time_ok_1p3 / travel_ratio 只在有老师参考的题上统计。"""
|
|
import math
|
|
import torch
|
|
|
|
from ml.common.features import wrap180
|
|
|
|
|
|
def min_gap_chunked(kin, q, cap=0.5, chunk=512):
|
|
"""分块算最小间隙：随机位形每个约 2800 个幸存胶囊对、峰值约 1.5 MB/位形，一次 4 万位形会爆内存/显存（2026-10-18 实测出题 OOM），每块 512"""
|
|
return torch.cat([kin.gaps(q[i:i + chunk], cap)['gm'].min(-1).values for i in range(0, q.shape[0], chunk)])
|
|
|
|
|
|
def rand_q(kin, n, g):
|
|
u = torch.rand(n, 6, device=kin.lo.device, generator=g)
|
|
return kin.lo + u * (kin.hi - kin.lo)
|
|
|
|
|
|
def free_q(kin, n, g, chunk=1024):
|
|
out = []
|
|
while sum(x.shape[0] for x in out) < n:
|
|
q = rand_q(kin, chunk, g)
|
|
ok = min_gap_chunked(kin, q) > 0
|
|
out.append(q[ok])
|
|
return torch.cat(out)[:n]
|
|
|
|
|
|
def straight_blocked(kin, a, b, n=24):
|
|
d = torch.where(kin.cont, wrap180(b - a), b - a)
|
|
s = torch.linspace(0, 1, n, device=a.device)
|
|
pts = (a[:, None] + s[None, :, None] * d[:, None]).reshape(-1, 6)
|
|
return (min_gap_chunked(kin, pts) <= 0).view(a.shape[0], n).any(-1)
|
|
|
|
|
|
def column_gap(kin, q):
|
|
gm = kin.gaps(q, 40)['gm']
|
|
return torch.stack([gm[:, g] for g in (2, 3, 4, 5, 6)], -1).min(-1).values # 组 0（立柱/底座）× 组 2..6
|
|
|
|
|
|
@torch.no_grad()
|
|
def sample_tasks(kin, n, seed, mix=(0.3, 0.3, 0.3, 0.1)):
|
|
g = torch.Generator(device=kin.lo.device); g.manual_seed(seed)
|
|
want = [int(round(n * m)) for m in mix]; want[0] += n - sum(want)
|
|
tasks = []
|
|
def emit(cat, q0, b, T=None):
|
|
pose = torch.rand(q0.shape[0], device=q0.device, generator=g) < 0.5
|
|
if cat == 'D': pose[:] = False
|
|
xf = kin.cum_xf(b)
|
|
Tb = kin.fk(b, xf) if T is None else T
|
|
Rb = kin.orient(b, xf).reshape(-1, 9)
|
|
for i in range(q0.shape[0]):
|
|
tasks.append({'cat': cat, 'mode': 'pose' if pose[i] else 'pos', 'q0': q0[i].tolist(), 'T': Tb[i].tolist(),
|
|
'R': Rb[i].tolist() if pose[i] else None, 'exec_s': None, 'travel': None})
|
|
if want[0]:
|
|
emit('A', free_q(kin, want[0], g), free_q(kin, want[0], g))
|
|
for cat, k in (('B', want[1]), ('C', want[2])):
|
|
got_a, got_b, have = [], [], 0
|
|
while have < k:
|
|
a, b = free_q(kin, 2048, g), free_q(kin, 2048, g)
|
|
m = straight_blocked(kin, a, b)
|
|
if cat == 'C': m &= column_gap(kin, b) < 8
|
|
got_a.append(a[m]); got_b.append(b[m]); have += int(m.sum())
|
|
if k: emit(cat, torch.cat(got_a)[:k], torch.cat(got_b)[:k])
|
|
if want[3]:
|
|
a, b = free_q(kin, want[3], g), rand_q(kin, want[3], g)
|
|
o = kin.fk(torch.zeros(1, 6, device=a.device))
|
|
T = o + (kin.fk(b) - o) * 1.6
|
|
emit('D', a, b, T)
|
|
return tasks
|
|
|
|
|
|
def neighbours(t, k, g):
|
|
"""±5 mm / ±3° 邻居"""
|
|
out = []
|
|
for _ in range(k):
|
|
u = torch.randn(3, generator=g); u = u / u.norm() * 5.0
|
|
nt = dict(t); nt['T'] = [x + float(d) for x, d in zip(t['T'], u)]
|
|
if t.get('R'):
|
|
ax = torch.randn(3, generator=g); ax = ax / ax.norm(); ang = math.radians(3.0) * float(torch.rand(1, generator=g))
|
|
K = torch.tensor([[0, -ax[2], ax[1]], [ax[2], 0, -ax[0]], [-ax[1], ax[0], 0]])
|
|
Rr = torch.eye(3) + math.sin(ang) * K + (1 - math.cos(ang)) * K @ K
|
|
nt['R'] = (Rr @ torch.tensor(t['R']).view(3, 3)).flatten().tolist()
|
|
nt['hard'] = True
|
|
out.append(nt)
|
|
return out
|