Diffusion Policy for PushT — 扩散策略模型 (昇腾 NPU 版)

Diffusion Policy 是基于扩散模型的机器人操控策略, 在 PushT 任务上预训练。本仓库为 昇腾 Ascend NPU (torch_npu) 适配版, 全流程无 CUDA。

模型架构

• 视觉编码器: ResNet18 backbone + SpatialSoftmax (32 关键点)
• 扩散模型: ConditionalUnet1D (ConditionalResnetBlock1D), 3 层下采样 + 2 层中间 + 2 层上采样
• 噪声调度: DDPM, 100 训练步, squaredcos_cap_v2 beta 调度, epsilon 预测
• 输入: 图像 (3x96x96) + 机器人状态 (2 维)
• 输出: 动作 chunk (2 维, horizon=16, n_action_steps=8)

快速开始

# 命令行单图推理 (需要样本输入图像)
python3 inference.py --img sample_input.png

# FastAPI 服务 (POST /predict 上传图片返回动作 chunk JSON)
python3 inference.py --server

输出格式

{
  "model": "lerobot/diffusion_pusht",
  "device": "npu:0",
  "action_dim": 2,
  "horizon": 16,
  "n_action_steps": 8,
  "denoise_steps": 100,
  "inference_time_s": 0.45,
  "action_chunk": [[0.01, -0.02, ...], ...],
  "action_mean": 0.003,
  "action_std": 0.015
}

权重

Diffusion Policy 权重自动从 GitCode 镜像下载, 缓存 ./model_weights/。