Robot_Arm/ml/export_onnx.py

"""导出部署包（方案 §8.4）：
  python -m ml.export_onnx --ckpt ml/runs/current.pt --out webapp/models/traj_policy.onnx [--fp16]
输出 traj_policy.onnx（输入 feat[1 或 K,83]、z[K,16,6]；输出 act[K,16,6] 归一化增量、value[K]、sub[K,13]）
与 traj_policy.meta.json（碰撞模型指纹、特征版本、维度、enc_once）。导出后用 onnxruntime 与 PyTorch 对拍。
2026-10-18 第五轮：改为"编码一次、解码 K 次"（Policy.sample_shared），meta.enc_once=true；运行时据此只送 1 行特征。
2026-10-18 第六轮：feat 轴改为动态（1 行 = 主推理编码一次；K 行 = 前瞻逐行编码），meta.feat_dyn=true。
  第五轮的 enc_once 模型 feat 固定 1 行，运行时前瞻只能逐行推（ik-policy.js ortInfer 已兼容），重新导出即可。
  --random 导出的是延迟探针（meta.random_init=true）：查看器 / bench / 验收都会拒绝把它当策略用。"""
import argparse, json, os
import numpy as np
import torch

from ml.common.kin_torch import Kin
from ml.common.features import F_DIM, H, FEAT_VER, STEP
from ml.model.policy import build


class Wrap(torch.nn.Module):
    def __init__(self, m):
        super().__init__(); self.m = m

    def forward(self, feat, z):
        return self.m.sample_shared(feat, z)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt'); ap.add_argument('--random', metavar='CFG', help='R4 测延迟：按配置随机初始化导出，不需要训练'); ap.add_argument('--out', default='webapp/models/traj_policy.onnx')
    ap.add_argument('--fp16', action='store_true'); ap.add_argument('--opset', type=int, default=17)
    a = ap.parse_args()
    if a.random:   # 2026-10-18：R4（浏览器 batch=8 P99）提前到训练前测，延迟只取决于结构不取决于权重
        torch.manual_seed(0); ck = {'cfg': json.load(open(a.random)), 'random_init': True}
        m = build(ck['cfg']); m.eval()
    elif a.ckpt:
        ck = torch.load(a.ckpt, map_location='cpu')
        m = build(ck['cfg']); m.load_state_dict(ck['model']); m.eval()
    else: raise SystemExit('需要 --ckpt 或 --random <配置>')
    kin = Kin()
    if ck.get('fp') and ck['fp'] != kin.fp: raise SystemExit(f"checkpoint 指纹 {ck['fp']} ≠ 当前碰撞模型 {kin.fp}：须重训")
    os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True)
    feat, z = torch.randn(1, F_DIM), torch.randn(8, H, 6)
    torch.onnx.export(Wrap(m), (feat, z), a.out, input_names=['feat', 'z'], output_names=['act', 'value', 'sub'],
                      dynamic_axes={'feat': {0: 'B'}, 'z': {0: 'K'}, 'act': {0: 'K'}, 'value': {0: 'K'}, 'sub': {0: 'K'}},
                      opset_version=a.opset, dynamo=False)
    if a.fp16:
        from onnxconverter_common import float16
        import onnx
        mo = float16.convert_float_to_float16(onnx.load(a.out), keep_io_types=True)
        onnx.save(mo, a.out)
    import onnxruntime as ort
    s = ort.InferenceSession(a.out, providers=['CPUExecutionProvider'])
    o = s.run(None, {'feat': feat.numpy(), 'z': z.numpy()})
    with torch.no_grad():
        t = m.sample_shared(feat, z)
        t8 = m.sample(feat.expand(8, -1), z)               # 与旧的"复制 8 行"口径逐元素一致
    err8 = max(float((t[i] - t8[i]).abs().max()) for i in range(3))
    assert err8 < 1e-4, f'sample_shared 与 sample 不一致：{err8}'
    err = max(float(np.abs(o[i] - t[i].numpy()).max()) for i in range(3))
    f8 = torch.randn(8, F_DIM)                              # 前瞻口径：8 行不同特征，须与 sample() 逐行一致
    o8 = s.run(None, {'feat': f8.numpy(), 'z': z.numpy()})
    with torch.no_grad(): r8 = m.sample(f8, z)
    err_rows = max(float(np.abs(o8[i] - r8[i].numpy()).max()) for i in range(3))
    assert err_rows < 1e-3, f'feat 为 K 行时与 sample 不一致：{err_rows}'
    err = max(err, err_rows)
    meta = {'fp': kin.fp, 'tcp0': kin.tcp0.tolist(), 'head': m.head, 'feat_ver': FEAT_VER, 'f_dim': F_DIM, 'h': H, 'step': STEP, 'fp16': a.fp16,
            'ckpt': os.path.abspath(a.ckpt) if a.ckpt else None, 'random_init': bool(a.random), 'enc_once': True, 'feat_dyn': True, 'eval': ck.get('eval'), 'max_abs_err_vs_torch': err}
    json.dump(meta, open(a.out.replace('.onnx', '.meta.json'), 'w'), indent=1)
    print(json.dumps({'onnx': a.out, 'MB': os.path.getsize(a.out) / 1e6, 'max_abs_err_vs_torch': err}))


if __name__ == '__main__':
    main()