z
zhangyuge147/robotics-diffusion-transformer-rdt-1b-NPU
模型介绍
文件和版本
Pull Requests
讨论
分析

robotics-diffusion-transformer/rdt-1b on Ascend NPU

agent workflow npu device call model result

1. 模型简介

  • 模型名称: robotics-diffusion-transformer/rdt-1b (RDT-1B)
  • 权重URL: https://huggingface.co/robotics-diffusion-transformer/rdt-1b (revision eb09036cc64ca4945051acbd1bd581d30a1d7711)
  • 任务类型: 强化学习 / 具身智能 — Vision-Language-Action (VLA) 双臂操作
  • 架构: Diffusion Transformer (RDT) — 1B 参数, hidden 2048, depth 28, 32 heads, 支持多模态条件 (语言 1024×4096, 图像 4374×1152, 状态 128)
  • 输入契约: 语言指令 embeddings (T5-XXL 4096-dim, 1024 tokens), 多视角 RGB 图像 tokens (SigLIP-so400m 1152-dim, 4374 tokens = 2 历史 ×3 相机 ×729 patch), proprioception 统一动作向量 128-dim (包含左右臂 6-DoF + gripper + 底盘), 控制频率 25Hz, 动作掩码 128-dim
  • 输出契约: 连续动作块 action_chunk [1, 64, 128] (64 步 ×128 维统一空间, 覆盖关节位置/速度、EEF 位姿/速度、轮式底盘速度), 经 diffusion 去噪 5 步 (DPMSolver), 值为有限浮点, 幅值约 [-1.2,1.2]
  • 预训练数据: 46 数据集, 1M+ 多机器人轨迹 (RT-1, RH20T, DROID, BridgeData V2, RoboSet, Open X-Embodiment)
  • 权重: pytorch_model.bin 2.46GB, OrderedDict 618 keys, FP32 (可转 BF16)
  • 许可证: mit
  • 论文: https://arxiv.org/pdf/2410.07864 (RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation)
  • 代码: https://github.com/thu-ml/RoboticsDiffusionTransformer

2. 数据契约

  • 验证数据: 合成可复现观测 (与官方一致的随机种子):
    • lang_tokens [1,1024,4096] N(0,1) seed 0 (模拟 T5-XXL embeddings)
    • img_tokens [1,4374,1152] N(0,1) seed 0 (模拟 SigLIP tokens = 2×3×729)
    • state_tokens [1,1,128] N(0,0.1) seed 0 (proprio 14 维映射到 128 统一空间)
    • action_mask [1,1,128] 全 1 (所有 128 维有效)
    • ctrl_freqs [1] 25.0
    • lang_attn_mask [1,1024] 全 True
  • 预处理: 无额外归一化, 直接使用适配器 mlp2x_gelu (语言/图像) 与 mlp3x_gelu (状态) 映射到 hidden 2048
  • 后处理: Diffusion 采样后乘 action_mask (掩无效维度), 再按 STATE_VEC_IDX_MAPPING 解包到 14 维关节动作 (本验证保留 128 维原始输出)
  • 形状: 输入同上, 输出 [1,64,128] (horizon 64, action_dim 128)
  • 采样: DPMSolverMultistepScheduler, num_inference_timesteps=5, beta_schedule=squaredcos_cap_v2, prediction_type=sample, 初始噪声 seed 42 (CPU 与 NPU 共享噪声时严格一致)

3. 验证环境

  • NPU: Ascend910_9362 (Ascend910B, 2 卡, HBM 64GB/卡), CANN 8.5.1, driver 25.5.5
  • PyTorch: 2.9.0+cpu, torch_npu: 2.9.0.post1+gitee7ba04
  • diffusers: 0.27.2, transformers: 4.41.0, timm: 1.0.28
  • Python: 3.11.14
  • 权重本地: /opt/atomgit/.cache/huggingface/hub/models--robotics-diffusion-transformer--rdt-1b/snapshots/eb09036cc64ca4945051acbd1bd581d30a1d7711 (pytorch_model.bin 2.456G + config.json + head.mp4)
  • 输入: 合成可复现, 非真实机器人图像, 但契约与官方 scripts/agilex_model.py 完全一致

4. 安装依赖

pip install -r requirements.txt
# torch==2.9.0
# torch-npu==2.9.0.post1
# diffusers==0.27.2
# transformers==4.41.0
# timm==1.0.28
# Pillow==12.2.0
# huggingface_hub==0.36.2
# numpy==1.26.4

5. NPU 推理

python inference.py                 # 自动 NPU (若可用)
python inference.py --device npu    # 强制 npu:0
python inference.py --device cpu    # 强制 cpu
python inference.py --mode validate   # CPU vs NPU 一致性 (共享噪声, 严格 FP32)
python inference.py --mode benchmark  # 预热 3 + 10 次计时 (synchronize)

核心逻辑 (inference.py:1):

  • 本地加载 RDTRunner (1.228B 参数) + pytorch_model.bin → model.to("npu:0"), 验证 next(model.parameters()).device == npu:0
  • 构造相同观测 (固定种子 0 的 lang/img/state), 在 CPU 与 NPU 均使用固定种子 42 的共享初始噪声 (通过 patch torch.randn) 保证对比公平
  • 计时 torch.npu.synchronize() 前后, 预热 3 次, 正式 10 次

6. 真实推理结果

执行命令: python inference.py (在 npu:0 真实执行, 权重来自本地快照, torch.float32)

Model: robotics-diffusion-transformer/rdt-1b
Revision: eb09036cc64ca4945051acbd1bd581d30a1d7711
Route: manipulation-vla (Diffusion Transformer, horizon=64, action_dim=128, hidden 2048 depth 28)
Weight: https://huggingface.co/robotics-diffusion-transformer/rdt-1b
Backend: torch + torch_npu + diffusers DDPM/DPMSolver (pure PyTorch), dtype torch.float32
[npu-smi]
  Ascend910_9362  OK  169.7W  45C  3135/65536 MB
torch_npu_available: True
device_count: 2, device_name: Ascend910_9362
Requested device: auto -> actual: npu:0
Snapshot: .../snapshots/eb09036cc64ca4945051acbd1bd581d30a1d7711
Config: pred_horizon=64 action_dim=128 hidden=2048 depth=28 lang_dim=4096 img_dim=1152
Weight file: .../pytorch_model.bin (2.46 GB)
Diffusion params: 1.228320e+09
Model first param device (cpu): cpu -> npu:0
Param count: 1.228B
Input: lang (1, 1024, 4096) img (1, 4374, 1152) state (1, 1, 128) mask (1, 1, 128) ctrl (1,)
Input sample: lang mean -0.0002 img mean 0.0002 state sample [0.00096, 0.0435, -0.0104]
Input device NPU: lang npu:0 img npu:0 state npu:0 mask npu:0
[RESULT] Model: robotics-diffusion-transformer/rdt-1b rev eb09036... route manipulation-vla horizon 64 dim 128
[RESULT] Backend: npu:0, dtype torch.float32, device_name Ascend910_9362
[RESULT] Action shape: (1, 64, 128), device npu:0, dtype torch.float32
[RESULT] First action (5 dims): [-0.01267491, 0.21815849, 0.25815707, -1.12188291, -0.08900752]
[RESULT] Last action (5 dims): [0.02416026, 0.22093831, 0.20951719, -1.11252224, -0.09642471]
[RESULT] Finite: True, max 1.1219, mean 0.0157
[PERF] latency 462.12 ms (synchronized with torch.npu.synchronize)  # 首轮含编译
[PERF] actions_per_s 138.49 horizon 64
[RESULT] PASS
  • NPU 首轮含编译 ~462 ms, 稳态 257 ms/块 (见性能节), CPU 同任务需 133.6 秒, 加速比 ~519 倍
  • 输出动作块已验证有限性、形状 [1,64,128] 与维度 128, 可直接映射到 AgileX 双臂 14 维控制

7. CPU-NPU 一致性验证

  • 方法: 同一权重、同观测、同适配器、同种子 0 (lang/img/state)、同 action_mask 与 ctrl_freqs、同 dtype float32、同 5 步 DPMSolver、共享初始噪声 (seed 42 的 torch.randn([1,64,128]) 在 CPU 生成后搬运至 NPU, 通过 mock torch.randn 确保两端同噪声), 分别在 cpu 与 npu:0 执行 RDTRunner.predict_action, 保存原始动作 .json 后比较
  • 命令: 
    python inference.py --mode validate
# 内部等价于:
python .opencode/skills/npu_adapt_robotics/scripts/validate_policy_outputs.py --cpu /tmp/cpu_shared.json --npu /tmp/npu_shared.json --atol 1e-4 --rtol 1e-3
  • 指标 (共享噪声):
    • count: 8192 (64*128)
    • finite: true
    • max_abs: 2.56e-06
    • max_ref: 1.126
    • tolerance (strict FP32 atol 1e-4 rtol 1e-3): 0.001226
    • passed: true (strict)
    • relaxed (atol 5e-3 rtol 5e-2): tolerance 0.0613 passed true
  • 非共享噪声对照 (原始 torch.randn 各自生成, CPU 与 NPU RNG 实现差异):
    • max_abs: 0.04428, 需 relaxed 阈值才 PASS (说明差异来自 RNG 而非数值误差), 共享噪声后严格 PASS 证明数值一致
  • 说明: 为保证同噪声对比, 已将 conditional_sample 的 torch.randn(..., device=device) 改为 torch.randn(..., device="cpu").to(device) 思路, 本仓库通过 unittest.mock.patch("torch.randn") 实现; 剩余差异 2e-06 来自 Ascend NPU 与 CPU 在 Linear/GELU/Attention 的微小累积误差, 远低于阈值, 不影响策略语义

8. 性能测试

条件: npu:0, float32, batch=1, lang 1024×4096 + img 4374×1152 + state 1×128, horizon 64, steps 5, torch.npu.synchronize() 计时, 预热 3 次, 正式 10 次

warmup done (3 runs)
run 0: 257.23 ms
run 1: 257.19 ms
run 2: 257.19 ms
run 3: 257.09 ms
run 4: 257.19 ms
run 5: 257.19 ms
run 6: 257.25 ms
run 7: 257.49 ms
run 8: 258.31 ms
run 9: 258.66 ms
avg 257.48 ms min 257.09 max 258.66 p50 257.23 p90 258.31 p95 258.66
actions_per_s 248.56  chunks_per_s 3.88
  • 首轮编译: 462 ms (含图编译), 稳定后 257 ms/块
  • 吞吐: 248.6 actions/s, 3.88 chunks/s (chunks/s = 1000/avg_ms, actions/s = horizon*1000/avg_ms)
  • 对比 CPU: 133610 ms/块 → 加速比 519×
  • 资源: 单卡 Ascend910B 64GB 显存占用 ~6GB (权重 2.46GB FP32 + 激活 ~3GB), 符合单卡门禁 (无须删减视觉/语言塔)

9. 限制与说明

  • 仅适配 manipulation-vla 路由, 不含 T5-XXL 与 SigLIP 的端到端编码, 语言/图像 tokens 为预计算 embeddings (与官方 scripts/agilex_model.py 中加载预存 text_embeds 方式一致)
  • 真实机器人闭环 rollout 未执行, 仅单步动作块验证 (符合技能要求: 环境 reward 不替代数值一致性)
  • 权重未包含在仓库, 通过本地快照加载; 若重新下载需保证 revision eb09036cc64ca4945051acbd1bd581d30a1d7711
  • trust_remote_code: 不需要, 全部本地可审计 inference.py 与 models/rdt_runner.py

10. 许可证

mit (同上游 robotics-diffusion-transformer/rdt-1b)