102 lines
4.6 KiB
Python
102 lines
4.6 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 auto_chunk(kin, mb_per_q, lo=256, hi=8192):
|
||
"""2026-10-18 第五轮（审查 N-6）：分块大小按设备空闲显存定，用掉其 1/4。
|
||
两处每位形峰值不同不是矛盾：kin.gaps 按 (胶囊对下限 + cap)² 筛候选对，cap=0.5 mm 时幸存约 2800 对（≈1.5 MB/位形），
|
||
column_gap 用 cap=40 mm，40 mm 内幸存的对多一个量级（≈20 MB/位形）。CPU 上取下限 256。"""
|
||
if kin.lo.device.type == 'cuda':
|
||
free = torch.cuda.mem_get_info(kin.lo.device)[0]
|
||
return int(max(lo, min(hi, free * 0.25 / (mb_per_q * 1e6))))
|
||
return lo
|
||
|
||
|
||
def min_gap_chunked(kin, q, cap=0.5, chunk=None):
|
||
"""分块算最小间隙：一次 4 万位形会爆内存/显存（2026-10-18 实测出题 OOM）；块大小见 auto_chunk（≈1.5 MB/位形）"""
|
||
chunk = chunk or auto_chunk(kin, 1.5)
|
||
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, chunk=None):
|
||
chunk = chunk or auto_chunk(kin, 20.0) # cap=40 → ≈20 MB/位形
|
||
out = []
|
||
for s in range(0, q.shape[0], chunk): # 与 min_gap_chunked 一样分块防 OOM
|
||
gm = kin.gaps(q[s:s + chunk], 40)['gm']
|
||
out.append(torch.stack([gm[:, g] for g in (2, 3, 4, 5, 6)], -1).min(-1).values) # 组 0（立柱/底座）× 组 2..6
|
||
return torch.cat(out)
|
||
|
||
|
||
@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
|