超
gcw_g59KKkes/edge-inference-smolvla-so101-pick-orange
模型介绍
文件和版本
Pull Requests
讨论
分析

edge-inference/smolvla-so101-pick-orange · 昇腾 NPU 视觉-语言-动作策略推理

本仓库将 edge-inference/smolvla-so101-pick-orange(SmolVLA 视觉-语言-动作策略,微调自 lerobot/smolvla_base,任务为在 LeIsaac/Isaac-Sim 的 SO101 仿真环境中"拾取橙子并放到盘子上")适配到昇腾 (Ascend) NPU(torch-npu,设备 npu:0)。权重在首次使用时自动从镜像下载,无需手动干预。

  • HuggingFace 原始仓库:https://huggingface.co/edge-inference/smolvla-so101-pick-orange
  • GitCode 权重镜像(声明主源,实际未镜像,自动回退 hf-mirror):https://ai.gitcode.com/hf_mirrors/edge-inference/smolvla-so101-pick-orange
  • 模型架构:SmolVLA —— SmolVLM2-500M-Video-Instruct 视觉-语言主干(截断为 16 层)+ 16 层动作专家头(cross-attention),flow matching 采样(10 步去噪),输出 50 步 × 6D 关节角动作块。

模型信息

字段值
模型edge-inference/smolvla-so101-pick-orange
策略类型lerobot.SmolVLAPolicy("type": "smolvla")
任务SO101 仿真(LeIsaac)"拾取橙子并放到盘子上"(pick-orange)
视觉主干SmolVLM2-500M-Video-Instruct,16 层(从 32 层截断),SigLIP 视觉编码器 12 层,bf16
动作专家16 层 × hidden 720(0.75×),cross-attention(self_attn_every_n_layers=2)
输入2 路相机(front/wrist,各 3×480×640,策略内部 resize 到 512×512 并做 padding)+ 6D 关节状态(MEAN_STD 归一化)+ 语言指令(tokenizer max_length=48)
输出50 步动作块(chunk_size=50)× 6D 关节角;返回**真实(去归一化)**动作值
权重文件config.json、model.safetensors(约 1.20 GB)、policy_preprocessor*.json/safetensors、policy_postprocessor*.json/safetensors、train_config.json

图像预处理(来自模型卡):2 张 480×640 相机图 → resize 至 512×512 并 padding → 16×16 patch 化(每图 1024 tokens)→ SigLIP 视觉编码器。像素从 [0,1] 映射到 [-1,1]。VISUAL 特征归一化模式为 IDENTITY(不做 ImageNet 均值方差归一化,policy_preprocessor.json 中 norm_map.VISUAL=IDENTITY)。

昇腾 NPU 适配说明

  1. 无任何 CUDA 设备调用 — 模型与所有输入 Tensor 全部在 npu:0 上运行。项目中的 model_code/(vendored lerobot 0.4.4 smolvla 策略源码)已逐文件清除 CUDA 引用并做 NPU 适配。
  2. NPU 初始化 包裹在 init_npu() 中,失败时输出 JSON 错误并退出。
  3. 权重自动下载 — 策略权重(7 个 lerobot 格式文件)与 SmolVLM2 tokenizer/处理器文件自动下载到 model_weights/,本地缓存优先;GitCode 镜像 429/未镜像时自动回退到 hf-mirror.com。下载后对 model.safetensors 做 8 字节 header + JSON + 数据偏移 + NaN/Inf 校验。
  4. model.safetensors 自包含完整 VLM — checkpoint 内含 model.vlm_with_expert.vlm.*(345 个张量)+ lm_expert.*(145 个张量)+ 5 个投影模块,共 500 个张量 / 450M 参数,无需单独下载 SmolVLM2 权重;仅需其 tokenizer/config 文件用于指令 token 化。
  5. 设备回退 — checkpoint config.json 里 "device": "cuda"(训练记录),在 NPU 环境自动识别为不可用并切换到 npu:0(见 model_code/lerobot/utils/utils.py 中 is_torch_device_available / auto_select_torch_device 的 NPU 分支)。
  6. torch.compile 跳过 — config.compile_model=false,不触发编译。

Vendored 模型源码(model_code/)

model_code/lerobot/ 是从 lerobot==0.4.4(PyPI wheel)逐文件取得的 smolvla 策略源码(316 个 .py,约 4.4 MB),并做了以下 NPU 适配:

  • policies/smolvla/smolvlm_with_expert.py — 在 eager_attention_forward 中把 attention mask 强制转为 bool(torch.where 在 torch 2.9 上要求 bool 条件),并清除 CUDA 注释。
  • utils/utils.py — auto_select_torch_device/get_safe_torch_device/is_torch_device_available/is_amp_available 增加 Ascend NPU(torch.npu)分支;device="cuda" 字符串在检测时视为不可用并回退。
  • 顶层 policies/__init__.py、processor/__init__.py、envs/__init__.py、teleoperators/__init__.py、motors/__init__.py、robots/__init__.py、cameras/__init__.py 被裁剪为纯推理所需的最小导出(避免引入 diffusers/gym/串口/电机驱动等无关重依赖)。
  • 其余 CUDA 引用(torch.cuda.*、"cuda:0" 默认值、文档串等)在 vendored 源码内全部改写为 NPU 等价形式。

说明:model_code 仅保留策略前向推理所需路径,非本模型相关的策略(act/pi0/groot/xvla 等)与训练/仿真脚本不在运行依赖中,但仍随源码保留以供参考。

权重

权重在首次使用时自动从镜像下载,缓存于 model_weights/edge-inference/smolvla-so101-pick-orange/(策略权重)与 model_weights/HuggingFaceTB/SmolVLM2-500M-Video-Instruct/(tokenizer/config)。后续运行优先加载本地缓存。

  • GitCode 镜像(主,实际未镜像):https://ai.gitcode.com/hf_mirrors/edge-inference/smolvla-so101-pick-orange
  • 备用镜像:https://hf-mirror.com/edge-inference/smolvla-so101-pick-orange
  • VLM 处理器镜像:https://hf-mirror.com/HuggingFaceTB/SmolVLM2-500M-Video-Instruct

使用方法

命令行(单图推理)

python3 inference.py --img examples/demo_orange.png

或指定指令文本:

python3 inference.py --img examples/demo_orange.png --instr "Pick up the orange and place it on the plate"

输出 JSON 到 stdout(设备 npu:0),action 为**真实(去归一化)**的 6D 关节角动作(action_chunk 为 50×6 动作块)。

FastAPI 服务

python3 inference.py --mode server --port 8784
  • GET / — 最小化 HTML 上传页面
  • POST /predict — multipart 上传一张相机图(可选表单字段 instr 覆盖指令)→ 与 CLI 相同的 JSON

示例:

curl -s -F file=@examples/demo_orange.png http://127.0.0.1:8784/predict

输出 JSON

{
  "model": "edge-inference/smolvla-so101-pick-orange",
  "source": "https://ai.gitcode.com/hf_mirrors/edge-inference/smolvla-so101-pick-orange",
  "device": "npu:0",
  "status": "success",
  "elapsed_ms": 608.87,
  "instruction": "Pick up the orange and place it on the plate",
  "action_dim": 6,
  "action_horizon": 50,
  "action": [8.87, 3.82, 1.65, 76.76, 17.48, 33.10],
  "action_chunk": [[...50 x 6 真实动作...]],
  "raw_action": [0.0004, 0.0215, -0.0142, 0.0132, -0.0836, 0.0429],
  "default_state": [8.84, 4.43, 3.70, 75.60, 18.46, 33.34],
  "state_mean_std": true,
  "camera_views": ["observation.images.front", "observation.images.wrist"],
  "input_image": "demo_orange.png"
}

字段说明:

  • action — 真实(去归一化) 的第一帧动作(6D 关节角,由 policy_postprocessor 的 MEAN_STD unnormalizer 还原:x_real = x_norm * std + mean)。
  • action_chunk — 真实去归一化的 50×6 动作块(chunk_size=50)。
  • raw_action — 模型原始(归一化)的第一帧动作。
  • default_state — 无真实状态输入时使用的默认 6D 关节状态(取数据集状态均值作为中性位姿)。

示例输入

examples/demo_orange.png(640×480 RGB)是合成的演示帧:PIL 绘制的桌面上有一个橙色的球体与白色盘子(so101/pick-orange 场景的示意)。因单图输入契约,该帧同时作为 front 与 wrist 两路相机特征输入。未能从网络获取真实 so101/pick-orange 相机帧(LeIsaac 数据集为 parquet+视频,下载成本高且带宽受限),故按模板回退为合成帧,特此说明。

环境

  • Python 3.11.14,PyTorch(昇腾 NPU 版 2.9.0),torch_npu 2.9.0
  • 昇腾 Ascend 910(64 GB HBM)
  • 依赖:lerobot==0.4.4(vendored 于 model_code/)、transformers>=4.57、torch、torch-npu、safetensors、fastapi、uvicorn、python-multipart、Pillow、draccus、num2words 等(见 requirements.txt)

文件说明

  • inference.py — 双模式推理入口(CLI + FastAPI,NPU 适配)
  • model_code/ — vendored lerobot 0.4.4 smolvla 策略源码(CUDA 已清除)
  • make_screenshots.py — 生成验证截图
  • examples/demo_orange.png — 合成演示帧
  • assets/ — 验证截图(agent_workflow / npu_device_call / model_result)

引用

@article{shukor2025smolvla,
  title={SmolVLA: A Vision-Language-Action Model for Affordable and Efficient Robotics},
  author={Shukor, Mustafa and Aubakirova, Dana and Capuano, Francesco and Kooijmans, Pepijn and Palma, Steven and Zouitine, Adil and Aractingi, Michel and Pascal, Caroline and Russi, Martino and Marafioti, Andres and Alibert, Simon and Cord, Matthieu and Wolf, Thomas and Cadene, Remi},
  journal={arXiv preprint arXiv:2506.01844},
  year={2025}
}