Robot_Arm/webapp/js/ik-policy.js

/* ============================================================================
 *  ik-policy.js —— Transformer 避障策略的闭环运行时（浏览器 / Node 通用，UMD）
 *  方案文档 §5 / §8：50 Hz 执行，每 3 步（60 ms）推理一次，8 个噪声 → 8 条 16 步动作块，
 *  1 步胶囊梯度投影 → 排序 → 距离下界认证（前 3 步 + 刹停段）→ 提交；全不过则沿已认证刹停段停。
 *  不依赖任何规划器（solveDyn / routeRT / 路网）：只用求解器的 fk / orient / gapMatrix / 胶囊对间隙。
 *  特征定义与 ml/common/features.py 逐项一致（ml/eval 对拍）。改特征必须同时改两边并升 FEAT_VER。
 * ==========================================================================*/
(function (root, factory) {
  if (typeof module === 'object' && module.exports) module.exports = factory();
  else root.ArmPolicy = factory();
}(typeof self !== 'undefined' ? self : this, function () {
'use strict';

const FEAT_VER = 1;
const VMAX = [60, 60, 90, 90, 120, 180];        // °/s（与 ml/gen/timelaw.cjs 同）
const DT = 0.02;
const STEP = VMAX.map(v => v * DT);             // 每步最大增量（°）
const ACC = 5.0;                                 // 归一化加速度 1/s²（与时间律同）
const H = 16, NJ = 6, PREFIX = 3;
const GP = [[0, 2], [0, 3], [0, 4], [0, 5], [0, 6], [2, 4], [2, 5], [3, 5]];   // 特征里的 8 个组对
const GAP_CAP = 40;
const F_DIM = 83;
const D2R = Math.PI / 180;
const wrap180 = a => { a = ((a + 180) % 360 + 360) % 360 - 180; return a; };

/* ---------------- 特征（83 维） ---------------- */
function features(S, st, gm) {
  const q = st.q, L = S.LIM, O = S.ORDER, f = new Float32Array(F_DIM);
  let o = 0;
  for (let j = 0; j < 6; j++) { f[o++] = Math.sin(q[j] * D2R); f[o++] = Math.cos(q[j] * D2R); }
  for (const j of [1, 3, 4]) { const t = L[O[j]]; f[o++] = 2 * (q[j] - t.lo) / (t.hi - t.lo) - 1; }
  const qp = st.qPrev || q;
  for (let j = 0; j < 6; j++) { const d = L[O[j]].cont ? wrap180(q[j] - qp[j]) : q[j] - qp[j]; f[o++] = d / STEP[j]; }
  const p = S.fk(q, 0), R = S.orient(q);
  for (let k = 0; k < 3; k++) f[o++] = p[k] / 600;
  for (const k of [0, 3, 6, 1, 4, 7]) f[o++] = R[k];
  for (let k = 0; k < 3; k++) f[o++] = st.T[k] / 600;
  if (st.R) for (const k of [0, 3, 6, 1, 4, 7]) f[o++] = st.R[k]; else o += 6;
  for (let k = 0; k < 3; k++) f[o++] = (st.T[k] - p[k]) / 600;
  for (let k = 0; k < 3; k++) f[o++] = Math.max(-1, Math.min(1, (st.T[k] - p[k]) / 50));
  f[o++] = st.R ? 0 : 1; f[o++] = st.R ? 1 : 0;
  const pre = st.prefix || [];
  for (let k = 0; k < PREFIX; k++) for (let j = 0; j < 6; j++) f[o++] = k < pre.length ? pre[k][j] / STEP[j] : 0;
  for (let k = 0; k < PREFIX; k++) f[o++] = k < pre.length ? 1 : 0;
  if (st.sub) { for (let j = 0; j < 6; j++) { const d = L[O[j]].cont ? wrap180(st.sub[j] - q[j]) : st.sub[j] - q[j]; f[o++] = d / 180; } f[o++] = 1; }
  else o += 7;
  const G = gm || S.gapMatrix(q, GAP_CAP);
  for (const [a, b] of GP) f[o++] = Math.min(G[a * 7 + b], GAP_CAP) / GAP_CAP;
  if (o !== F_DIM) throw new Error('feature dim ' + o);
  return f;
}

/* ---------------- 距离下界认证（ml/gen/cert.cjs 复用本实现） ---------------- */
function makeCert(S, opt) {
  opt = opt || {};
  const CAP = opt.cap || 40, MIN_DEG = opt.minDeg || 0.002;
  /* 分组杠杆表 TL[k][g]：组对 (gA,gB) 只用关节 gA..gB−1 带动组 gB 的杠杆（比全局上界紧得多） */
  const TL = opt.leverTable || S.leverTable(4000);
  const L = S.LIM, O = S.ORDER;
  const memo = new Map();
  const key = q => q.map(v => v.toFixed(5)).join(',');
  let queries = 0, budgetLeft = Infinity;
  const gm = q => { const k = key(q); let M = memo.get(k); if (!M) { M = S.gapMatrix(q, CAP); queries++; budgetLeft--; if (memo.size > 50000) memo.clear(); memo.set(k, M); } return M; };
  const minGap = M => { let m = Infinity; for (let a = 0; a < 7; a++) for (let b = a + 2; b < 7; b++) m = Math.min(m, M[a * 7 + b]); return m; };
  const inBox = q => O.every((id, i) => L[id].cont || (q[i] >= L[id].lo - 1e-6 && q[i] <= L[id].hi + 1e-6));
  function seg(a, b, depth) {
    if (budgetLeft <= 0) return false;
    const Ma = gm(a), Mb = gm(b);
    const dq = a.map((v, i) => Math.abs(b[i] - v));
    let ok = true;
    for (let ga = 0; ga < 7 && ok; ga++) for (let gb = ga + 2; gb < 7; gb++) {
      const lhs = Ma[ga * 7 + gb] + Mb[ga * 7 + gb];
      let rhs = 0; for (let k = ga; k < gb; k++) rhs += TL[k][gb] * dq[k];
      if (!(lhs > rhs)) { ok = false; break; }
    }
    if (ok) return true;
    if (minGap(Ma) <= 0 || minGap(Mb) <= 0) return false;
    if (Math.max.apply(null, dq) < MIN_DEG || depth > 40) return false;
    const m = a.map((v, i) => 0.5 * (v + b[i]));
    return seg(a, m, depth + 1) && seg(m, b, depth + 1);
  }
  /* 折线（相邻点按短弧展开）→ {ok, queries, minGap, failAt, why}；budget = 本次最多新增距离查询数（超出判不过，why='budget'） */
  function path(pts, budget) {
    const q0 = queries; let mg = Infinity;
    budgetLeft = budget || Infinity;
    const P = [pts[0].slice()];
    for (let i = 1; i < pts.length; i++) P.push(P[i - 1].map((v, k) => L[O[k]].cont ? v + wrap180(pts[i][k] - v) : pts[i][k]));
    for (let i = 0; i < P.length; i++) { if (!inBox(P[i])) return { ok: false, queries: queries - q0, minGap: mg, failAt: i, why: 'box' }; }
    for (let i = 0; i < P.length; i++) mg = Math.min(mg, minGap(gm(P[i])));
    for (let i = 1; i < P.length; i++) if (!seg(P[i - 1], P[i], 0)) return { ok: false, queries: queries - q0, minGap: mg, failAt: i, why: budgetLeft <= 0 ? 'budget' : 'cert' };
    budgetLeft = Infinity;
    return { ok: true, queries: queries - q0, minGap: mg, unwrapped: P };
  }
  return { seg: (a, b) => { budgetLeft = Infinity; return seg(a, b, 0); }, path, leverTable: TL, gapMin: q => minGap(gm(q)), gm, clear: () => memo.clear() };
}

/* ---------------- 刹停段：从 q 以每步增量 v 起，按 ACC 均匀减速到 0 ---------------- */
function brakeSteps(v) {
  let n = 0; for (let j = 0; j < 6; j++) n = Math.max(n, Math.ceil(Math.abs(v[j]) / (STEP[j] * ACC * DT) - 1e-9));
  const out = [];
  for (let k = 1; k <= n; k++) out.push(v.map(x => x * (1 - k / (n + 1))));
  return out;
}

/* ---------------- 1 步胶囊梯度投影（每轴 ≤ maxDeg，仅 gap < margin） ---------------- */
/* pairs：本次推理开始时在 qe 处取的近距胶囊对（S.criticalPairs(qe, margin+10)），候选各步只对这些对求间隙与有限差分（微秒级） */
function project(S, q, margin, maxDeg, pairs) {
  if (!pairs || !pairs.length) return q;
  const gs = pairs.map(p => S.pairGap(q, p[0], p[1]));
  const g0 = Math.min.apply(null, gs);
  if (g0 >= margin) return q;
  pairs = pairs.filter((p, i) => gs[i] < margin + 2).slice(0, 6);
  const h = 0.01, grad = [0, 0, 0, 0, 0, 0];
  for (let j = 0; j < 6; j++) {
    const qa = q.slice(); qa[j] += h;
    const g1 = Math.min.apply(null, pairs.map(p => S.pairGap(qa, p[0], p[1])));
    grad[j] = (g1 - g0) / h;
  }
  const n2 = grad.reduce((a, x) => a + x * x, 0);
  if (n2 < 1e-12) return q;
  const need = margin - g0;
  return q.map((v, j) => v + Math.max(-maxDeg, Math.min(maxDeg, need * grad[j] / n2)));
}

/* ---------------- 局部 DLS 精修（接近目标时） ---------------- */
function refineStep(S, q, T, Rt, maxScale) {
  const e = x => {
    const p = S.fk(x, 0), r = [T[0] - p[0], T[1] - p[1], T[2] - p[2]];
    if (!Rt) return r;
    const R = S.orient(x);   // 旋转误差向量 ≈ 0.5·vee(R_t Rᵀ − R Rtᵀ)，按 100 mm/rad 折算到 mm 量级
    const M = [0, 0, 0, 0, 0, 0, 0, 0, 0];
    for (let i = 0; i < 3; i++) for (let k = 0; k < 3; k++) { let s = 0; for (let m = 0; m < 3; m++) s += Rt[i * 3 + m] * R[k * 3 + m]; M[i * 3 + k] = s; }
    return r.concat([0.5 * (M[7] - M[5]) * 100, 0.5 * (M[2] - M[6]) * 100, 0.5 * (M[3] - M[1]) * 100]);
  };
  const e0 = e(q), m = e0.length, h = 1e-3, J = [];
  for (let j = 0; j < 6; j++) { const qa = q.slice(); qa[j] += h; const e1 = e(qa); J.push(e1.map((v, i) => -(v - e0[i]) / h)); }
  // Δq = Jᵀ (J Jᵀ + λ²I)⁻¹ e，m ≤ 6
  const lam = 0.5, A = [];
  for (let i = 0; i < m; i++) { A.push([]); for (let k = 0; k < m; k++) { let s = i === k ? lam * lam : 0; for (let j = 0; j < 6; j++) s += J[j][i] * J[j][k]; A[i].push(s); } }
  const y = solveLin(A, e0.slice());
  const dq = [0, 0, 0, 0, 0, 0];
  for (let j = 0; j < 6; j++) { let s = 0; for (let i = 0; i < m; i++) s += J[j][i] * y[i]; dq[j] = s; }
  let sc = 1; for (let j = 0; j < 6; j++) sc = Math.max(sc, Math.abs(dq[j]) / (STEP[j] * (maxScale || 0.5)));
  return { dq: dq.map(v => v / sc), err: Math.hypot(e0[0], e0[1], e0[2]) };
}
function solveLin(A, b) {
  const n = b.length;
  for (let c = 0; c < n; c++) {
    let p = c; for (let r = c + 1; r < n; r++) if (Math.abs(A[r][c]) > Math.abs(A[p][c])) p = r;
    [A[c], A[p]] = [A[p], A[c]]; [b[c], b[p]] = [b[p], b[c]];
    for (let r = c + 1; r < n; r++) { const f = A[r][c] / A[c][c]; for (let k = c; k < n; k++) A[r][k] -= f * A[c][k]; b[r] -= f * b[c]; }
  }
  const x = new Array(n).fill(0);
  for (let r = n - 1; r >= 0; r--) { let s = b[r]; for (let k = r + 1; k < n; k++) s -= A[r][k] * x[k]; x[r] = s / A[r][r]; }
  return x;
}

/* ---------------- 正态噪声 ---------------- */
function makeRng(seed) {
  let s = (seed >>> 0) || 1;
  const u = () => { s = (s * 1664525 + 1013904223) >>> 0; return (s + 0.5) / 4294967296; };
  return { u, n: () => Math.sqrt(-2 * Math.log(u())) * Math.cos(2 * Math.PI * u()) };
}

/* ============================================================================
 *  闭环运行器
 *  infer(feat Float32Array[B*F], z Float32Array[B*H*6], B) → Promise<{act[B*H*6], value[B], sub[B*13]}>
 *  用法：r = new Runner(S, infer, opts); r.reset(q0); r.setGoal(T, R);
 *        每 20 ms 调 r.tick()；r.needPlan() 为真时 await r.plan()（tick 与 plan 可并行：plan 只在队列 ≥3 步时发起）
 * ==========================================================================*/
class Runner {
  constructor(S, infer, opts) {
    this.S = S; this.infer = infer;
    this.o = Object.assign({ K: 8, commit: 3, margin: 3,   /* 2026-10-18：投影余量 5→3 mm，与 APG margin_train 一致（老师最小间隙 p50 4.4 mm） */
       projMax: 0.5, refineMm: 5, arriveMm: 1, arriveDeg: 0.5,
      stuckPlans: 5, lockSec: 1.0, subReachDeg: 5, subTimeout: 1.5, seed: 12345, lookahead: true,
      certBudget: 40, safetyMs: 30 }, opts || {});   // 每条候选最多 40 次距离查询；每次推理安全层总预算 30 ms（超出按刹停处理）
    this.cert = makeCert(S, { leverTable: this.o.leverTable, cap: 20 });
    this.rng = makeRng(this.o.seed);
    this.stats = { plans: 0, inferMs: [], safetyMs: [], brakes: 0, refine: 0, lookahead: 0, locks: 0, projected: 0 };
    this.reset(null);
  }
  reset(q0) {
    this.q = q0 ? q0.slice() : null; this.qPrev = this.q ? this.q.slice() : null;
    this.queue = []; this.brakeTail = []; this.lastPlan = null; this.goal = null;
    this.subs = []; this.subT = 0; this.vals = []; this.brakeRun = 0; this.lock = null; this.t = 0; this.cold = true; this.done = false;
  }
  setGoal(T, R) {
    const jump = !this.goal || Math.hypot(T[0] - this.goal.T[0], T[1] - this.goal.T[1], T[2] - this.goal.T[2]) > 50 || (!!R !== !!this.goal.R);
    this.goal = { T: T.slice(), R: R ? R.slice() : null };
    if (jump) { this.needSub = true; this.subs = []; this.cold = true; this.vals = []; this.lock = null; }
    this.done = false;
  }
  qEnd() { let q = this.q.slice(); for (const d of this.queue) q = q.map((v, j) => v + d[j]); return q; }
  err(q) {
    const p = this.S.fk(q, 0), T = this.goal.T;
    const e = Math.hypot(T[0] - p[0], T[1] - p[1], T[2] - p[2]);
    let oe = 0;
    if (this.goal.R) { const R = this.S.orient(q); let tr = 0; for (let i = 0; i < 9; i++) tr += R[i] * this.goal.R[i]; oe = Math.acos(Math.max(-1, Math.min(1, (tr - 1) / 2))) / D2R; }
    return { e, oe };
  }
  needPlan() { return !!this.goal && !this.done && !this.planning && this.queue.length <= this.o.commit; }
  tick() {
    if (!this.q) return null;
    const d = this.queue.shift();
    this.qPrev = this.q.slice();
    if (d) this.q = this.q.map((v, j) => v + d[j]);
    this.t += DT;
    if (this.goal && !this.queue.length) {
      const { e, oe } = this.err(this.q);
      if (e < this.o.arriveMm && oe < this.o.arriveDeg) this.done = true;
    }
    return this.q;
  }
  _cand(act, k, r) {   // 第 k 条候选：增量（°），丢掉模型的前 r 步（已提交前缀），取接下来的步
    const inc = [];
    for (let s = 0; s < H; s++) { const row = []; for (let j = 0; j < 6; j++) row.push(Math.max(-1, Math.min(1, act[(k * H + s) * 6 + j])) * STEP[j]); inc.push(row); }
    return inc;
  }
  _certCommit(qe, steps) {   // 认证：提交步 + 刹停段（先用离散判定快速剔除，再做距离下界认证）
    const pts = [qe.slice()]; let q = qe.slice();
    for (const d of steps) { q = q.map((v, j) => v + d[j]); pts.push(q.slice()); }
    for (let i = 1; i < pts.length; i++) if (!this.S.collideFree(pts[i])) return { ok: false, tail: [] };
    const tail = brakeSteps(steps.length ? steps[steps.length - 1] : [0, 0, 0, 0, 0, 0]);
    /* 刹停段各步同向等比减速 → 关节空间里是一条直线，认证终点一段即可 */
    if (tail.length) { for (const d of tail) q = q.map((v, j) => v + d[j]); pts.push(q.slice()); }
    const c = this.cert.path(pts, this.o.certBudget);
    return { ok: c.ok, tail };
  }
  async plan() {
    if (!this.needPlan()) return;
    this.planning = true;
    try { await this._plan(); } finally { this.planning = false; }
  }
  async _plan() {
    const S = this.S, o = this.o, K = o.K;
    const r = Math.min(this.queue.length, PREFIX), prefix = this.queue.slice(0, r);
    const qe = this.qEnd();
    this.stats.plans++;
    /* 接近目标：局部 DLS 精修代替模型 */
    const ee = this.err(qe);
    if (ee.e < o.refineMm && (!this.goal.R || ee.oe < 5)) {
      if (ee.e < o.arriveMm && ee.oe < o.arriveDeg) return;
      const steps = []; let q = qe.slice();
      for (let k = 0; k < o.commit; k++) { const rs = refineStep(S, q, this.goal.T, this.goal.R, 0.3); steps.push(rs.dq); q = q.map((v, j) => v + rs.dq[j]); }
      const c = this._certCommit(qe, steps);
      if (c.ok) { this.queue.push(...steps); this.brakeTail = c.tail; this.stats.refine++; return; }
    }
    /* 子目标推进 */
    if (this.subs.length) {
      const sg = this.subs[0]; let m = 0; for (let j = 0; j < 6; j++) m = Math.max(m, Math.abs((S.LIM[S.ORDER[j]].cont ? wrap180(sg[j] - qe[j]) : sg[j] - qe[j])));
      if (m < o.subReachDeg || this.t - this.subT > o.subTimeout) { this.subs.shift(); this.subT = this.t; }
    }
    const gm = S.gapMatrix(this.q, GAP_CAP);
    const st = { q: this.q, qPrev: this.qPrev, T: this.goal.T, R: this.goal.R, prefix, sub: this.subs[0] || null };
    const f = features(S, st, gm);
    const feat = new Float32Array(K * F_DIM), z = new Float32Array(K * H * 6);
    for (let k = 0; k < K; k++) feat.set(f, k * F_DIM);
    for (let i = 0; i < z.length; i++) z[i] = this.rng.n();
    const t0 = now();
    const out = await this.infer(feat, z, K);
    this.stats.inferMs.push(now() - t0);
    const t1 = now();
    /* 子目标：新目标/跳变后取价值最好的那条输出 */
    let best = 0; for (let k = 1; k < K; k++) if (out.value[k] < out.value[best]) best = k;
    if (this.needSub) {
      this.needSub = false; this.subT = this.t;
      const sb = out.sub.subarray(best * 13, best * 13 + 13);
      if (sb[12] > 0.5) this.subs = [0, 1].map(v => this.q.map((x, j) => x + sb[v * 6 + j] * 180));
    }
    /* 候选：去掉前缀后的提交步，投影，打分 */
    const cands = [], near = S.criticalPairs(qe, o.margin + 10).slice(0, 64);
    for (let k = 0; k < K; k++) {
      const inc = this._cand(out.act, k, r);
      let steps = inc.slice(r, r + o.commit), q = qe.slice();
      const adj = [];
      for (const d of steps) {
        let qn = q.map((v, j) => v + d[j]);
        const qp = project(S, qn, o.margin, o.projMax, near);
        if (qp !== qn) this.stats.projected++;
        adj.push(qp.map((v, j) => v - q[j])); q = qp;
      }
      steps = adj;
      const fut = []; let qq = qe.slice(); for (let s = r; s < H; s++) { qq = qq.map((v, j) => v + inc[s][j]); fut.push(qq); }
      let j1 = 0; for (let s = 0; s < H; s++) j1 += inc[s][0];
      let score = -2.0 * out.value[k];
      if (this.lastPlan) {   // 与上一条已选计划的未来重叠步对齐（时间平移 3 步）
        let c = 0, w = 1;
        for (let s = 0; s < fut.length && s + o.commit < this.lastPlan.length; s++, w *= 0.85) {
          let d2 = 0; for (let j = 0; j < 6; j++) d2 += Math.pow((fut[s][j] - this.lastPlan[s + o.commit][j]) / STEP[j], 2);
          c += w * d2;
        }
        score -= 0.02 * c;
      }
      cands.push({ k, steps, fut, j1, score, value: out.value[k] });
    }
    /* 卡死检测：价值 0.3 s 内不降 或 连续 2 次刹停 → 锁定另一绕行方向 1 s */
    const vbest = out.value[best]; this.vals.push(vbest); if (this.vals.length > o.stuckPlans) this.vals.shift();
    const stuck = this.brakeRun >= 2 || (this.vals.length >= o.stuckPlans && this.vals[this.vals.length - 1] > this.vals[0] * 0.99 && ee.e > 20);
    if (stuck && !(this.lock && this.t < this.lock.until)) {
      const lastDir = this.lastPlan ? Math.sign(this.lastPlan[this.lastPlan.length - 1][0] - this.lastPlan[0][0]) || 1 : 1;
      this.lock = { side: -lastDir, until: this.t + o.lockSec }; this.stats.locks++; this.vals = [];
    }
    let pool = cands;
    if (this.lock && this.t < this.lock.until) { const f2 = cands.filter(c => Math.abs(c.j1) < 2 || Math.sign(c.j1) === this.lock.side); if (f2.length) pool = f2; }
    /* 短前瞻：冷启动 / 卡死时，各候选块末态再推一块，用其价值打分 */
    if (o.lookahead && (this.cold || stuck)) {
      this.stats.lookahead++;
      const feat2 = new Float32Array(pool.length * F_DIM), z2 = new Float32Array(pool.length * H * 6);
      pool.forEach((c, i) => {
        const qa = c.fut[c.fut.length - 1], qb = c.fut[Math.max(0, c.fut.length - 2)];
        feat2.set(features(S, { q: qa, qPrev: qb, T: this.goal.T, R: this.goal.R, prefix: [], sub: null }), i * F_DIM);
      });
      for (let i = 0; i < z2.length; i++) z2[i] = this.rng.n();
      const o2 = await this.infer(feat2, z2, pool.length);
      pool.forEach((c, i) => { c.score = -2.0 * o2.value[i] - 0.5 * c.value; });
    }
    this.cold = false;
    pool.sort((a, b) => b.score - a.score);
    /* 安全层：依次认证，第一个通过即提交 */
    let picked = null;
    for (const c of pool) {
      if (now() - t1 > o.safetyMs && picked === null && c !== pool[0]) break;
      const v = this._certCommit(qe, c.steps); if (v.ok) { picked = c; this.brakeTail = v.tail; break; }
    }
    this.stats.safetyMs.push(now() - t1);
    if (picked) {
      this.queue.push(...picked.steps); this.lastPlan = [qe.slice()].concat(picked.fut); this.brakeRun = 0;
    } else {
      this.stats.brakes++; this.brakeRun++;
      if (this.brakeTail.length) { this.queue.push(...this.brakeTail); this.brakeTail = []; }
      this.lastPlan = null;
    }
  }
}
const now = () => (typeof performance !== 'undefined' ? performance.now() : Date.now());

/* ---------------- onnxruntime 适配（web: ort = window.ort；node: require('onnxruntime-node')） ---------------- */
function ortInfer(ort, session) {
  return async (feat, z, B) => {
    const r = await session.run({ feat: new ort.Tensor('float32', feat, [B, F_DIM]), z: new ort.Tensor('float32', z, [B, H, 6]) });
    return { act: r.act.data, value: r.value.data, sub: r.sub.data };
  };
}

return { FEAT_VER, VMAX, DT, STEP, ACC, H, F_DIM, GP, GAP_CAP, features, makeCert, brakeSteps, project, refineStep, Runner, ortInfer, wrap180 };
}));