diffusion_pusht - 昇腾 NPU 适配

简介

本仓库是 lerobot/diffusion_pusht 模型在华为昇腾 NPU 上的适配版本。diffusion_pusht 是一个基于扩散策略（Diffusion Policy）的机器人操作模型，用于 PushT 推块任务。该模型使用 ResNet18 视觉骨干网络提取图像特征，结合机械臂末端执行器位置状态，通过条件扩散 UNet 1D 生成动作轨迹。

模型核心特性：

• 视觉编码器：ResNet18（torchvision 预训练权重，ImageNet 归一化）
• 动作生成：条件扩散模型（DDPM 调度器，100 步推理）
• 观测缓存：2 步历史观测（n_obs_steps=2），8 步动作执行（n_action_steps=8），16 步预测范围（horizon=16）
• 输入：96x96 RGB 图像 + 2 维状态（agent 像素坐标）
• 输出：2 维动作（dx, dy 像素坐标）
• 参数量：262.7M

原始模型来源：https://huggingface.co/lerobot/diffusion_pusht

验证环境

服务启动

1. 下载模型权重

# 从 ModelScope 下载（推荐，国内网络友好）
pip install modelscope
python3 -c "from modelscope.hub.snapshot_download import snapshot_download; snapshot_download('lerobot/diffusion_pusht', cache_dir='./model_cache')"

ModelScope 下载完成后，将 ./model_cache/lerobot/diffusion_pusht/ 下的文件移动到 ./model_cache/ 目录。

2. 安装依赖

pip install -r requirements.txt

3. 启动推理服务

# 命令行推理测试
python3 inference.py --test

# 启动 HTTP 服务（FastAPI + uvicorn）
python3 inference.py --serve --port 8010

服务启动后，可通过以下接口调用：

健康检查：

curl http://localhost:8010/health
# 返回: {"status":"ok","device":"npu:0","num_params":262709026}

推理接口：

curl -X POST http://localhost:8010/predict \
  -H "Content-Type: application/json" \
  -d '{"image": [[[0,0,0],...]], "state": [256.0, 256.0]}'
# 返回: {"action": [dx, dy]}

Smoke 验证

推理测试输出

输入图像: shape=(96, 96, 3), dtype=uint8
输入状态: [256.0, 256.0]
预热动作: [234.79, 144.09]
步0: action=[244.65, 254.21], time=2.482s
步1: action=[241.13, 176.49], time=0.002s
步2: action=[230.80, 108.85], time=0.001s
步3: action=[215.43, 74.71], time=0.001s
步4: action=[201.29, 47.60], time=0.001s
步5: action=[184.18, 14.93], time=0.001s
步6: action=[172.21, -15.37], time=0.001s
步7: action=[157.12, -42.06], time=0.001s
步8: action=[223.63, 266.42], time=2.482s (新chunk)
步9: action=[223.09, 230.66], time=0.002s
平均推理时间: 0.497s

模型成功加载到 NPU 并输出了合理的动作值（像素坐标范围 12~511）。

性能参考

注：Diffusion Policy 每 8 步生成一次完整 100 步扩散推理（~2.5s），后续 7 步直接回放缓存动作（~0.002s）。

精度测评

模型使用 safetensors 格式加载，权重与原版完全一致（SHA256 校验）。归一化参数从原始 model.safetensors 中提取：

NPU 推理输出与 CPU 推理输出在数值精度上保持一致（float32 计算），输出动作值域均在合理范围内（12~511 像素坐标）。模型加载时跳过了归一化缓存键（normalize_inputs/normalize_targets/unnormalize_outputs），不影响模型核心推理能力。

注意事项

1. vLLM 不兼容：该模型为 CNN + UNet 架构，非 Transformer/LLM，vLLM-Ascend 无法加载。需使用 torch_npu 原生推理 + FastAPI 服务化。
2. 设备字符串：lerobot 框架不支持 "npu" 设备，需在配置中设置 device="cpu"，加载后再手动调用 .to("npu:0")。
3. 归一化参数：模型 safetensors 中嵌入了 normalize_inputs/normalize_targets/unnormalize_outputs 缓存键，加载时需跳过（strict=False），否则会报 unexpected keys 警告。
4. 中心裁剪：模型会对 96x96 输入图像进行中心裁剪为 84x84，裁剪参数定义在 config.json 的 crop_shape 字段。
5. 观测缓存：模型使用双端队列缓存 2 步历史观测，首次调用 select_action 时自动填充缓存。环境重置（env.reset）后需调用 policy.reset() 清空缓存。
6. 环境要求：推理需要 lerobot 包（安装时可能因 diffusers 版本依赖冲突卡住，建议使用国内 PyPI 镜像）。
7. HBM 占用：模型推理时 HBM 约占用 1.36GB，建议在空闲 HBM 充足的 NPU 上运行。