diff --git a/docs/Transformer轨迹拟合训练方案.md b/docs/Transformer轨迹拟合训练方案.md new file mode 100644 index 0000000..0ce3ab7 --- /dev/null +++ b/docs/Transformer轨迹拟合训练方案.md @@ -0,0 +1,317 @@ +# Transformer 拟合避障轨迹规划:训练方案(审查稿) + +> 2026-10-17 起草,基于 `main@3c7c814`(贴立柱修复后的规划器)。本文只是方案,代码尚未编写。 +> 审查时重点看:§0 定位、§2 数据分布、§3 表示方式、§5 训练配置、§7 验收标准。 + +--- + +## 0. 定位与边界 + +| 项目 | 决定 | +|---|---| +| 模型角色 | **提议器**:给出若干条候选关节轨迹,由现有 `collideFree` + 半格复验把关,全部不合法时回落 `routeRT` / `solveDyn` | +| 不做的事 | 不让模型直接驱动电机;不替代碰撞检测;不学习动力学(速度/加速度仍由 `animTo` 时间参数化) | +| 老师 | 当前 JS 规划器(`webapp/js/ik-solver.js`):松手口径用 `S.solveDyn(T, opt)`,拖拽口径用 `S.solveRoute(T, {quick:true,...})` | +| 收益目标 | 实时拖拽重规划 P99 ≈ 321 ms → **< 50 ms**;首次全量 IK 尖峰(最差约 2 s)→ < 100 ms | +| 底线 | 「模型+验证+回落」整体的到达率、路径合法率**不得低于**现有规划器(任何基准上都不允许退化) | + +为什么必须保留验证:碰撞是"硬约束",神经网络只能做到高概率满足。验证器把"偶尔撞一下"变成"偶尔慢一下"。 + +--- + +## 1. 训练环境 + +| 阶段 | 在哪里跑 | 说明 | +|---|---|---| +| 数据生成 | Hark 沙箱(2 vCPU / 7 GB,Node) 或 用户机器 | 纯 JS,无需 GPU;可多进程并行 | +| 训练 | **用户的 GPU 机器** | 沙箱无 GPU、无 PyTorch。最低建议 8 GB 显存(RTX 3060 级别) | +| 部署推理 | 浏览器(onnxruntime-web,WebGPU 优先,wasm 兜底) | 与现有查看器同页运行 | + +Python 环境(训练机): + +```bash +conda create -n armtf python=3.11 -y && conda activate armtf +pip install torch --index-url https://download.pytorch.org/whl/cu121 # 按本机 CUDA 版本选 +pip install numpy pyarrow tqdm tensorboard onnx onnxruntime +``` + +--- + +## 2. 第一步:生成数据(规划器当老师) + +### 2.1 新增文件 + +| 文件 | 作用 | +|---|---| +| `ml/gen/gen_dataset.cjs` | 采样任务 → 调用规划器 → 校验 → 写分片 | +| `ml/gen/sampler.cjs` | 起点/目标采样器(三类分布,见 2.3) | +| `ml/gen/run_parallel.sh` | 按 CPU 核数起 N 个进程,每个进程不同随机种子 | +| `ml/data/shard_XXXX.jsonl` | 原始样本分片(每片 5000 条) | + +复用 `tests/route_bench_setup.cjs` 装配求解器(关节表、`docs/travel.json`、`docs/collide.json`,无需网格)。 + +### 2.2 单条样本的字段 + +```json +{ + "id": "s0003_000127", + "mode": "pos" | "pose", + "teacher": "dyn" | "rt", + "q0": [6 个关节角, 度], + "T": [x, y, z], // 目标 TCP 位置, mm + "R": [[3x3]] | null, // 目标姿态(pose 模式) + "tool": 0, + "path": [[6], [6], ...], // 老师给出的折线: q0 → via... → q + "q_goal": [6], + "ok": true, // 到达 (误差<1mm, pose 另需<0.5°) 且 1° 细验无干涉 + "poseRelaxed": false, + "travel_deg": 123.4, + "plan_ms": 812.0, + "fp": "collide.json 指纹" // modelFingerprint(),碰撞模型变了旧数据作废 +} +``` + +### 2.3 采样分布(目标 20 万条原始任务) + +关节限位(`docs/travel.json`):J1/J3/J6 连续 ±180°;J2 [-92, 9],J4 [-111, 16],J5 [-93, 34]。 + +| 类别 | 占比 | 采样方式 | +|---|---:|---| +| A 均匀 | 50% | 起点 q0、目标位形 b 均在限位内均匀采样且本身无干涉;目标取 `fk(b)`(pose 模式加 `orient(b)`) | +| B 直线被挡 | 30% | 同 A,但要求 q0→b 关节直线 1° 细验有干涉(route_bench 的生成口径) | +| C 贴立柱 | 15% | b 满足 `nearColumn`(J1/J2/J3 任一 ±4° 即撞立柱),q0 在立柱另一侧(near_column_bench 口径) | +| D 不可达/无安全解 | 5% | 目标在工作空间外或所有 IK 分支都撞;标签为"拒绝",教模型说"到不了" | + +- `mode`:pos 与 pose 各 50%。 +- 老师:每条任务用 `solveDyn` 求一次(松手口径,质量更高);其中 30% 额外用 `routeRT` 拖拽口径求一次,作为"快速路径"风格的辅助样本(`teacher:"rt"`),训练时可按需过滤。 +- **测试集隔离**:用固定种子另生成 2000 条留作测试集,且排除与训练目标距离 < 5 mm 的样本;`route_bench` / `near_column_bench` / `rt_drag_bench` 的案例**不进入训练集**。 + +### 2.4 质量过滤 + +1. `ok == false` 的样本:A/B/C 类丢弃;D 类保留为"拒绝"标签。 +2. 整条路径 1° 细验(`segOK`)+ 半格复验,任一段干涉即丢弃。 +3. `poseRelaxed == true` 的样本单独打标签,训练时只当作 pos 模式样本用(避免模型学到"姿态随便放")。 +4. 关节行程超过同任务 `routeRT` 结果 1.5 倍的视为老师绕远,丢弃(可选,先统计比例再决定)。 + +### 2.5 数据增广(不需要重新规划) + +- **后缀切片**:合法路径的任意后半段也是合法路径。从每条路径中按弧长取 2–4 个中间点作为新 q0,得到新样本(目标不变)。约 ×3。 +- **连续轴平移**:J1 绕立柱对称性**不成立**(立柱在固定位置),不做 J1 旋转增广;J6 若与碰撞无关可加 ±360° 等价表示,归一化时统一处理。 +- 增广后预计约 60 万条训练样本。 + +### 2.6 耗时与存储估算 + +- `solveDyn` 约 0.7 s/条,沙箱 2 进程约 1 万条/小时 → 20 万条约 20 小时;用户机器 8 核约 5 小时。 +- 每条约 1–2 KB,全部原始数据约 300 MB;转 Parquet 后约 80 MB。 + +### 2.7 本步检查点(审查时可看) + +- [ ] `ml/data/stats.json`:各类别数量、ok 率、平均行程、路径点数分布、规划耗时分布 +- [ ] 随机抽 20 条在查看器里回放(新增 `?replay=shard_0000.jsonl#127` URL 参数) +- [ ] 测试集与 3 个基准案例和训练集无重叠 + +--- + +## 3. 第二步:预处理(Python) + +文件:`ml/prep/build_dataset.py` → 输出 `ml/data/train.parquet`、`val.parquet`(95/5 划分)、`test.parquet`、`norm.json`。 + +### 3.1 路径重采样 + +1. 连续轴(J1/J3/J6)先做展开(`unwrapTo` 同口径),保证相邻点差 < 180°。 +2. 按**加权关节空间弧长**重采样为固定 **K = 32** 个点,第 0 点 = q0,第 31 点 = q_goal。权重默认各轴 1(可改成按连杆长度加权)。 +3. 记录原始弧长 `L`(模型额外回归它,用于早停与排序)。 + +### 3.2 条件输入(编码器一侧) + +| 量 | 维度 | 编码 | +|---|---:|---| +| q0 | 6 | 每轴 (sin, cos) → 12;限位轴再加归一化值 (q−lo)/(hi−lo) → +3 | +| 目标位置 T | 3 | 除以 600 mm 归一化 | +| 目标姿态 R | 6 | 旋转矩阵前两列(6D 表示);pos 模式填 0 | +| 模式 | 2 | one-hot [pos, pose] | +| 工具号 | 1 | 目前恒 0,留位 | +| 起点 TCP | 3 | `fk(q0)` 归一化(给模型一个笛卡尔线索) | + +### 3.3 输出表示(解码器一侧)——离散化 + +- 每个关节值相对 q0 的**增量** Δq = q_k − q0 离散化:连续轴范围 [−360, 360],限位轴范围 [lo−hi, hi−lo],均分 **512 个桶**(约 1.4°/桶)。 +- 每个桶再回归一个桶内偏移 ∈ [−0.5, 0.5](细化到亚度级)。 +- 序列:K−1 = 31 个路径点 × 6 轴 = **186 个 token**(第 0 点即 q0,不需要预测)。 +- 额外 1 个 `[REJECT]/[OK]` 分类 token 放在序列最前:D 类样本只训练这个 token。 + +**为什么不用直接回归(MSE)**:立柱两侧绕行(左绕/右绕、J3 翻转/不翻转)是多峰分布,MSE 会输出两条路径的平均——恰好穿过立柱。离散化 + 自回归采样能保留多峰。 + +### 3.4 归一化文件 `norm.json` + +记录桶边界、K、权重、碰撞模型指纹 `fp`;部署端必须校验 `fp` 一致,否则拒绝加载模型。 + +--- + +## 4. 第三步:模型结构 + +文件:`ml/model/traj_transformer.py` + +``` +条件向量 (≈37 维) + └─ MLP → 8 个条件 token (d=256) ← 编码器:4 层 Transformer encoder +目标序列: [BOS][OK] (j1,j2,...,j6)×31 + └─ 嵌入 = 桶嵌入 + 轴嵌入(6) + 路径点位置嵌入(32) + └─ 6 层 Transformer decoder(因果自注意力 + 对条件 token 的交叉注意力) +输出头: + ├─ 桶分类头 (512 类 / token) + ├─ 桶内偏移回归头 (1 / token) + ├─ OK/REJECT 头 (2 类) + └─ 路径弧长回归头 (1) +``` + +| 超参 | 值 | +|---|---| +| d_model | 256 | +| 头数 | 8 | +| FFN | 1024,GELU | +| 编码器/解码器层数 | 4 / 6 | +| Dropout | 0.1 | +| 参数量 | 约 7 M(fp32 约 28 MB,fp16 约 14 MB) | + +备选(第一版不做):同一骨干改为 flow matching,直接输出 31×6 连续值,推理 8–16 步;在自回归版精度不够时再试。 + +--- + +## 5. 第四步:训练 + +文件:`ml/train.py`,配置 `ml/configs/base.yaml`。 + +### 5.1 损失 + +``` +L = CE(桶) 权重 1.0(label smoothing 0.05) + + L1(桶内偏移) 权重 0.5 + + CE(OK/REJECT) 权重 0.5 + + L1(弧长 L) 权重 0.1 + + 平滑项 ‖q_{k+1} − 2q_k + q_{k−1}‖²(用期望桶值计算) 权重 0.05 + + 终点项 ‖q_31 − q_goal‖(期望值) 权重 0.5 +``` + +### 5.2 训练配置 + +| 项 | 值 | +|---|---| +| 优化器 | AdamW,β=(0.9, 0.95),weight decay 0.05 | +| 学习率 | 3e-4,前 2000 步线性预热,余弦退火到 3e-6 | +| batch | 256(显存不够用梯度累积) | +| 精度 | bf16/fp16 混合精度 | +| 轮数 | 最多 40 epoch,验证集"合法率"连续 5 次不升即早停 | +| 教师强制 | 全程 teacher forcing;第 20 epoch 后 10% 概率用模型自己的预测(scheduled sampling),减轻曝光偏差 | +| 预计耗时 | RTX 3060:约 1.5–3 小时(估算,跑完第一轮后校正) | + +### 5.3 训练中验证(每个 epoch) + +在 `val` 抽 1000 条,**自回归采样**生成路径后用 Python 版碰撞检查器评估: +- 用 `ml/eval/collide_py.py`(从 `collide.json` 读胶囊,与 JS 同口径;先用 JS 结果对拍 1 万个位形,误差必须为 0)。 +- 指标:合法率(1° 细验无干涉)、终点误差、行程比(模型/老师)、REJECT 准确率。 +- 保存最优 checkpoint 依据 = **合法率**,其次终点误差。 + +### 5.4 一条命令运行 + +```bash +python ml/prep/build_dataset.py --in ml/data --out ml/data +python ml/train.py --config ml/configs/base.yaml --run runs/v1 +tensorboard --logdir runs +``` + +--- + +## 6. 第五步:推理与集成 + +### 6.1 导出 + +`ml/export_onnx.py`:导出带 KV cache 的解码器(opset 17),fp16 量化;输出 `webapp/models/traj_v1.onnx` + `traj_v1.norm.json`。 + +### 6.2 推理流程(新增 `webapp/js/ik-learned.js`,在 `routeRT` 之前调用) + +1. 编码条件 → 一次 batch 采样 **B = 8** 条(温度 {0, 0.5, 0.8, 1.0} 各 2 条,0 为贪心)。 +2. 若 OK 头判 REJECT 概率 > 0.9 → 直接走原规划器(不信任模型的"到不了",只用来省时间排序)。 +3. 每条候选:终点用 `solveTo` 做 1–3 次牛顿精修 → 整条 `collideFree` 1° 细验 + 半格复验。 +4. 选合法且行程最短的一条;**全部不合法 → 原 `routeRT` / `solveDyn` 回落**。 +5. 拖拽中热启动:把上一帧路径剩余部分作为解码前缀,只重新生成后半段。 + +### 6.3 延迟预算(目标) + +| 环节 | 目标 | +|---|---:| +| 模型推理(8 条 × 186 token,WebGPU) | < 15 ms | +| 碰撞验证(8 条) | < 15 ms(已有包围球预筛) | +| 合计(命中时) | < 30 ms | + +--- + +## 7. 第六步:评估与验收 + +### 7.1 测试集合 + +| 集合 | 来源 | +|---|---| +| route_bench | `tests/route_bench.cjs`,60 例 | +| near_column_bench | `tests/near_column_bench.cjs`,30 例(pos / pose / rt) | +| rt_drag_bench | `tests/rt_drag_bench.cjs`,12 例 | +| held-out | §2.3 的 2000 条留出测试集 | + +新增 `tests/learned_bench.cjs`:对同一批案例分别跑 ①原规划器 ②仅模型 ③模型+验证+回落,输出对比表。 + +### 7.2 验收标准 + +| 指标 | 仅模型(参考) | 模型+验证+回落(必须) | +|---|---|---| +| 路径合法率 | ≥ 90% | = 100%(与现状一致) | +| 到达率 | ≥ 95% | ≥ 现状(route 60/60,near pos 30/30,rt 30/30) | +| 关节行程 | ≤ 老师 × 1.1 | ≤ 现状 × 1.05 | +| 拖拽重规划 P99 | — | < 50 ms(现状约 321 ms) | +| 拖拽碰撞帧 | — | 0 | +| 模型命中率(不需回落) | — | ≥ 85% | + +不达标的处理:命中率低 → 加数据(重点补失败类别)/ 加大模型;行程差 → 加大行程 loss 或过滤老师绕远样本;延迟不降 → 减少 B 或改 flow matching。 + +--- + +## 8. 风险与维护 + +1. **碰撞模型/立柱变动 = 重训**:`collide.json` 指纹写进数据和模型,不一致时前端自动禁用模型,回落原规划器,不会出错只会变慢。重训流程 = §2 → §5,约半天。 +2. **老师的缺陷会被学走**:例如 `ci27` 姿态根枚举漏分支(姿态偏差约 97.5°)。数据里 `poseRelaxed` 样本已隔离;老师修好后再重造数据。 +3. **分布外目标**:模型在工作空间边缘可能乱给;有验证+回落兜底,只影响速度。 +4. **J3/J6 连续轴展开**:预处理与部署必须同一套 `unwrapTo` 口径,§3 对拍时一并验证。 + +--- + +## 9. 目录规划 + +``` +ml/ + gen/ gen_dataset.cjs sampler.cjs run_parallel.sh + data/ shard_*.jsonl stats.json train/val/test.parquet norm.json (不入库,.gitignore) + prep/ build_dataset.py + model/ traj_transformer.py + eval/ collide_py.py eval_paths.py + configs/ base.yaml + train.py export_onnx.py README.md +webapp/js/ik-learned.js +webapp/models/traj_v1.onnx traj_v1.norm.json +tests/learned_bench.cjs +``` + +## 10. 里程碑 + +| # | 内容 | 产出 | 在哪跑 | +|---|---|---|---| +| M1 | 数据生成器 + 1000 条试跑 + stats | `ml/gen/*`、stats.json | 沙箱 | +| M2 | Python 碰撞检查器对拍 | `collide_py.py`,1 万位形 0 误差 | 沙箱 | +| M3 | 全量 20 万条数据 | Parquet | 沙箱(约 20 h)或用户机器 | +| M4 | 训练工程 + 小数据过拟合测试(1000 条能背下来) | `train.py` 等 | 沙箱 CPU 验证代码可跑 | +| M5 | 正式训练 | checkpoint、TensorBoard | **用户 GPU** | +| M6 | 导出 ONNX + 前端集成 + learned_bench | 对比表 | 沙箱 + 浏览器 | + +## 11. 需要用户确认 + +1. 训练机显卡型号与显存。 +2. 定位:只做加速器(本文默认),还是希望最终能脱离规划器。 +3. 数据生成在沙箱跑(约 20 h),还是你的机器跑更快。