Files
Robot_Arm/ml/prep/build_dataset.cjs
T

76 lines
5.0 KiB
JavaScript
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
/* 原始老师轨迹 → 闭环训练样本(方案 §4.5 / §5)
* node ml/prep/build_dataset.cjs ml/data/shard_0001.jsonl [outDir=ml/data/samples]
* 每条 keep=true 的轨迹按 50 Hz 时间律重放,每 3 步取一个样本:
* X [N,83] 特征(webapp/js/ik-policy.js features(),与部署同一份代码)
* Y [N,16,6] 未来 16 步每步增量 / STEP(到终点后补 0)
* V [N] 剩余时间(s)/ 5
* SG [N,13] 子目标标签:剩余两个途经点相对当前 q 的差 / 180(不足用终点补)+ 是否需要
* TID[N] 任务编号(tasks.jsonl 行号所在的全局 id,用于划分与闭环验证)
* 增广:30% 样本对 q_t 加 σ=1° 噪声,标签改为 8 步内线性收回老师轨迹(噪声态须无干涉且间隙 > 1 mm);
* 已提交前缀长度 r ∈ {0..3} 均匀随机,前缀取老师接下来 r 步的增量。 */
'use strict';
const fs = require('fs'), path = require('path');
const IN = process.argv[2], OUTDIR = process.argv[3] || path.join(__dirname, '../data/samples');
if (!IN) { console.error('usage: build_dataset.cjs <shard.jsonl> [outDir]'); process.exit(1); }
const { S } = require('../../tests/route_bench_setup.cjs')();
const POL = require('../../webapp/js/ik-policy.js');
const { timeParam } = require('../gen/timelaw.cjs');
const { writeNpy } = require('./npy.cjs');
const { F_DIM, STEP, H } = POL;
const L = S.LIM, O = S.ORDER;
const base = path.basename(IN, '.jsonl');
let st = 0; for (const ch of base) st = (st * 31 + ch.charCodeAt(0)) >>> 0;
const rnd = () => { st = (st * 1664525 + 1013904223) >>> 0; return st / 4294967296; };
const gauss = () => Math.sqrt(-2 * Math.log(rnd() + 1e-12)) * Math.cos(2 * Math.PI * rnd());
const hash = s => { let h = 2166136261; for (const c of s) h = Math.imul(h ^ c.charCodeAt(0), 16777619); return h >>> 0; };
const X = [], Y = [], V = [], SG = [], TID = [], tasks = [];
const lines = fs.readFileSync(IN, 'utf8').split('\n').filter(Boolean);
let nAug = 0, nAugRej = 0;
for (const line of lines) {
const r = JSON.parse(line);
if (!r.keep) continue;
const tid = hash(r.id) & 0x7fffffff;
const split = (hash(r.id + '#split') % 100) < 5 ? 'val' : 'train';
const P = r.path;
const tp = timeParam(P);
const traj = tp.traj, taus = tp.taus, N = traj.length - 1;
tasks.push({ tid, id: r.id, cat: r.cat, mode: r.mode, q0: r.q0, T: r.T, R: r.R, exec_s: N * POL.DT, travel: r.travel_deg, reached: r.reached, err: r.err, split });
const vt = tp.cum; // 顶点 τ
const inter = P.slice(1, -1).map((p, i) => ({ p, tau: vt[i + 1] }));
for (let t = 0; t <= N; t += 3) {
let q = traj[t].slice(), qPrev = (t > 0 ? traj[t - 1] : traj[0]).slice();
let off = [0, 0, 0, 0, 0, 0], aug = false;
if (rnd() < 0.3) {
const qn = q.map((v, j) => { let x = v + gauss(); const t2 = L[O[j]]; if (!t2.cont) x = Math.max(t2.lo, Math.min(t2.hi, x)); return x; });
nAug++;
if (S.collideFree(qn) && S.clearance(qn, 5) > 1) { off = qn.map((v, j) => v - q[j]); q = qn; qPrev = qPrev.map((v, j) => v + off[j]); aug = true; }
else nAugRej++;
}
const pos = k => { const i = Math.min(N, t + k); const w = aug ? Math.max(0, 1 - k / 8) : 0; return traj[i].map((v, j) => v + off[j] * w); };
const y = new Float32Array(H * 6);
let prevP = q;
for (let k = 1; k <= H; k++) { const pk = pos(k); for (let j = 0; j < 6; j++) y[(k - 1) * 6 + j] = Math.max(-1, Math.min(1, (pk[j] - prevP[j]) / STEP[j])); prevP = pk; }
const rr = Math.floor(rnd() * 4), prefix = [];
for (let k = 0; k < rr; k++) { const row = []; for (let j = 0; j < 6; j++) row.push(y[k * 6 + j] * STEP[j]); prefix.push(row); }
const remain = inter.filter(v => v.tau > taus[t] + 0.02);
const sub = remain.length && rnd() < 0.8 ? remain[0].p : null;
const gm = S.gapMatrix(q, POL.GAP_CAP);
const f = POL.features(S, { q, qPrev, T: r.T, R: r.R, prefix, sub }, gm);
const v1 = remain[0] ? remain[0].p : P[P.length - 1], v2 = remain[1] ? remain[1].p : P[P.length - 1];
const sg = new Float32Array(13);
for (let j = 0; j < 6; j++) { sg[j] = (v1[j] - q[j]) / 180; sg[6 + j] = (v2[j] - q[j]) / 180; }
sg[12] = remain.length ? 1 : 0;
X.push(f); Y.push(y); V.push((N - t) * POL.DT / 5); SG.push(sg); TID.push(tid);
}
}
fs.mkdirSync(OUTDIR, { recursive: true });
const n = X.length;
const cat = (arrs, w) => { const o = new Float32Array(n * w); arrs.forEach((a, i) => o.set(a, i * w)); return o; };
writeNpy(path.join(OUTDIR, base + '.X.npy'), cat(X, F_DIM), [n, F_DIM]);
writeNpy(path.join(OUTDIR, base + '.Y.npy'), cat(Y, H * 6), [n, H, 6]);
writeNpy(path.join(OUTDIR, base + '.V.npy'), Float32Array.from(V), [n]);
writeNpy(path.join(OUTDIR, base + '.SG.npy'), cat(SG, 13), [n, 13]);
writeNpy(path.join(OUTDIR, base + '.TID.npy'), TID, [n], 'int32');
fs.writeFileSync(path.join(OUTDIR, base + '.tasks.jsonl'), tasks.map(t => JSON.stringify(t)).join('\n') + '\n');
console.log(JSON.stringify({ shard: base, tasks: tasks.length, samples: n, aug: nAug, augRejected: nAugRej, featVer: POL.FEAT_VER }));