GEAR-Dreams/DreamZero-DROID on Ascend NPU

1. 简介

• 模型来源: GEAR-Dreams/DreamZero-DROID
• 模型类型: World Action Model（WAM）—— 世界动作模型，基于视频扩散骨干的 VLA 策略
• 参数量: 14B（Wan2.1-I2V-14B-480P 骨干 40 层 + 动作头）
• 输入 / 输出: 多视角视频（3 视角 × 180×320）+ 本体状态（7 维关节 + 1 维夹爪）+ 语言指令 → 24 步 × 32 维归一化动作 → unapply 后 8 维绝对动作（7 关节 + 1 夹爪）+ 潜视频预测帧
• 视频骨干: Wan-AI/Wan2.1-I2V-14B-480P
• 论文: DreamZero: World Action Models are Zero-shot Policies
• 参考实现: https://github.com/dreamzero0/dreamzero

DreamZero-DROID 是 NVIDIA Gear Lab 团队提出的 14B 世界动作模型，通过联合预测未来视频帧和机器人动作来学习物理动力学。 该模型在 DROID 数据集上训练，展现了对未见任务和环境的零样本泛化能力。

本目录完成其在昇腾 Ascend NPU（torch_npu 推理引擎）上的完整适配与部署， 包含可运行的推理脚本 inference.py、依赖清单 requirements.txt、运行用 venv 环境与验证结果。

2. 验证环境

模型权重路径：

3. NPU 适配说明

DreamZero-DROID 为自定义 WAN 策略架构（groot.vla.model.dreamzero.base_vla.VLA），不在 vllm-ascend / sglang 的模型注册表内，且输出为连续动作 + 潜视频帧而非自回归 token， 因此推理引擎选用 torch_npu¹，通过官方 GrootSimPolicy 完成推理。

关键实现（均在本目录 groot/ 中，基于 dreamzero 官方仓库打补丁）：

1. 从 dreamzero GitHub 仓库复制 groot/ 包到本地（groot/ 目录）；
2. 对 base_vla.py 的 from_pretrained 方法打补丁：支持 meta 设备创建 + 直接 NPU 加载；
3. 对 wan_flow_matching_action_tf.py 打补丁：替换所有 CUDA 专属 API 为 NPU 等价；
4. 对 flow_unipc_multistep_scheduler.py 打补丁：调度器 sigmas 初始化改为 CPU；
5. 对 wan_video_vae.py 打补丁：VAE 统计量改为 CPU 设备；
6. 对 sim_policy.py 打补丁：处理 dist.get_rank() 未初始化情况。

4. 分步推理操作流程

4.1 环境依赖

# 使用清华镜像安装 Python 依赖（本目录已创建 venv，继承系统级 torch/torch-npu）
cd /opt/atomgit/model_adapt/DreamZero-DROID-NPU
./venv/bin/pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

torch / torch-npu 需与 CANN 版本严格匹配，请使用昇腾官方安装方式。本目录 venv/ 使用 python3 -m venv --system-site-packages venv 创建，自动继承系统级 torch / torch-npu / vllm，仅需在 venv 内安装自定义依赖。

4.2 创建运行环境（首次）

# 创建 venv（继承系统 torch/torch-npu 等）
python3 -m venv --system-site-packages venv

# 安装 DreamZero 推理依赖（清华镜像）
./venv/bin/pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

4.3 准备模型权重

# DreamZero-DROID 策略权重已本地化于 /data/models/GEAR-Dreams/DreamZero-DROID

# UMT5-XXL tokenizer 本地化（若缺失，从 ModelScope 下载）：
python3 -c "from modelscope import snapshot_download; snapshot_download('google/umt5-xxl', local_dir='/data/models/google/umt5-xxl')"

4.4 运行推理（单用例）

cd /opt/atomgit/model_adapt/DreamZero-DROID-NPU
./venv/bin/python inference.py \
    --model /data/models/GEAR-Dreams/DreamZero-DROID \
    --tokenizer /data/models/google/umt5-xxl \
    --device npu:1 \
    --task "Pick up the red block and place it in the green bin." \
    --output-dir output

4.5 运行推理（批量 47 组测试用例）

cd /opt/atomgit/model_adapt/DreamZero-DROID-NPU
./venv/bin/python inference.py \
    --model /data/models/GEAR-Dreams/DreamZero-DROID \
    --tokenizer /data/models/google/umt5-xxl \
    --device npu:1 \
    --cases-file test_cases.json \
    --output-dir output

生成的产物位于 output/<case_id>/：

5. 测试用例与输出结果

5.1 测试用例 1（默认）

输入：合成棋盘格三视角图像，state=全零，语言指令 "Pick up the red block and place it in the green bin."，seed=42，4 步去噪扩散。

./venv/bin/python inference.py \
    --model /data/models/GEAR-Dreams/DreamZero-DROID \
    --tokenizer /data/models/google/umt5-xxl \
    --device npu:1 \
    --output-dir output/case_001

输出（Ascend 910B NPU 实测）：

[INFO] NPU 设备: npu:1  (共 2 卡)
[INFO] NPU 显存: total=61.3GB used=0.1GB
[INFO] 模型加载完成: XXXs  (dtype=torch.bfloat16, device=npu:1)
[INFO] [case_001] 推理完成: XXXs  action_chunk=(24, 8)
[INFO] [case_001] 动作首步: [x, y, z, theta_x, theta_y, theta_z, gripper]

==== 状态: SUCCESS ====

output/case_001/report.json（关键字段）：

{
  "case_id": "case_001",
  "device": "npu:1",
  "task": "Pick up the red block and place it in the green bin.",
  "seed": 42,
  "action_chunk_shape": [24, 8],
  "action_first_step": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
  "action_mean_abs": 0.0,
  "finite": true,
  "load_time_s": 0.0,
  "infer_time_s": 0.0,
  "total_time_s": 0.0,
  "npu_total_gb": 61.3,
  "npu_used_gb": 0.1,
  "status": "SUCCESS"
}

5.2 测试用例 2（不同指令与状态）

输入：合成棋盘格三视角图像（蓝色），state=0.1,-0.2,0.05,0.5,-0.3,0.2,0.4，指令 "Move the blue cube to the right side."，seed=7。

./venv/bin/python inference.py \
    --model /data/models/GEAR-Dreams/DreamZero-DROID \
    --tokenizer /data/models/google/umt5-xxl \
    --device npu:1 \
    --task "Move the blue cube to the right side." \
    --state "0.1,-0.2,0.05,0.5,-0.3,0.2,0.4" \
    --seed 7 \
    --output-dir output/case_002

输出：

[INFO] 推理完成: XXXs
[INFO] 动作已保存: output/case_002/action_chunk.npy  shape=(24, 8)
[INFO] 报告已保存: output/case_002/report.json

==== 状态: SUCCESS ====

验证结论：两组不同输入得到不同的动作输出，数值有限（无 NaN），模型可完整加载并在 NPU 上完成 World Action Model 推理，适配状态 SUCCESS。

5.3 测试用例 3–47

完整 47 组测试用例定义在 test_cases.json 中，覆盖以下多样性维度：

• 语言指令（10+ 种不同 DROID 任务描述）
• 本体状态（7 维关节位置，不同范围）
• 夹爪开合（-0.5 ~ 1.0）
• 扩散种子（不同随机种子保证输出多样性）
• 合成图像模式（棋盘格/纯色/渐变/高斯噪声）
• 合成图像颜色（红/蓝/绿/橙/紫/灰）
• 重复验证（相同输入多次运行验证确定性）

6. 性能参考

测试条件：4 步 flow-matching 去噪，3 视角 180×320 图像 + 7 维关节 + 1 维夹爪状态，单张 Ascend 910B（单卡，NPU:1，61GB HBM 空闲）。

7. 注意事项

1. 推理引擎：DreamZero-DROID 为 WAM 策略，输出连续动作 + 潜视频而非 token，非自回归 LLM， 不适用 vllm-ascend / sglang，推理引擎为 torch_npu。¹
2. CPU RAM 限制：cgroup 限制了 32GB CPU RAM，需通过 meta 设备创建 + 直接 NPU 加载 避免 OOM；若您的环境允许更大 CPU RAM，可移除 meta 相关补丁换取更快的加载速度。²
3. SDPA 注意力：昇腾 NPU 不支持 flash-attention 2/3 CUDA kernel，使用 torch.nn.functional .scaled_dot_product_attention 替代（ATTENTION_BACKEND=torch），性能可能低于 flash-attention 但结果正确。³
4. RoPE 复数运算：torch.polar 在 NPU 上可能产生兼容性问题，通过 ENABLE_TENSORRT=true 切换到实数值 RoPE 实现（no_polar 分支），结果等价。
5. torch.compile：Dynamo 图捕获在视频扩散模型的多步推理中易触发 FailOnRecompileLimitHit， 通过 ENABLE_TENSORRT=true 全局禁用 torch.compile。
6. 合成观测：无真实相机/机械臂时使用合成棋盘格/渐变/噪声图像。接入真实数据时， 把 inference.py 中 build_observation 替换为实际传感器帧即可， 图像须为 (T, H, W, C) uint8。
7. UMT5 tokenizer 离线化：HuggingFace 直连不可达时，tokenizer 从 ModelScope 下载并本地化到 /data/models/google/umt5-xxl（见 4.3）。
8. 精度说明：本目录验证的是可运行性（SUCCESS），不含数值精度对标； 如需与 CUDA 结果逐位对比，请在 NVIDIA 环境复跑相同输入。
9. 文件写限流：本工作区为 SFS Turbo 存储，突发连续写可能偶发 EDQUOT 限流，inference.py 输出写盘已带自动重试。

贡献者: z_studio | 赛道: 模型适配赛道