318 lines
14 KiB
Markdown
318 lines
14 KiB
Markdown
# 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），还是你的机器跑更快。
|