NVIDIA OmniDreams 完整教學

從概念理解到 post-training 實作：如何用 NVIDIA 的 world model 為自駕模擬生成擬真影像。

1. 專案定位

什麼是 OmniDreams？

OmniDreams 是 NVIDIA Research 開發的 world model（世界模型），專為 autonomous-driving simulation（自動駕駛模擬）設計。它能從一張真實照片出發，搭配道路資訊和行車軌跡，即時生成逼真的多鏡頭駕駛影片。

用一個比喻理解

想像你是電影導演。你有：

• 📸 一張街景照片（第一幀）
• 📝 場景描述（「晴天市區道路，有行人和車輛」）
• 🗺️ 精密地圖（告訴模型路在哪裡、車道標線怎麼畫）
• 🚗 行車路線（車要往哪裡開、怎麼轉彎）

OmniDreams 就像一個 AI 攝影團隊，根據這些資訊自動拍出整段行車影片——而且每一幀都是擬真的（photorealistic），可以從多個攝影機角度看。

核心數據

兩個 Repo 的分工

flowchart LR
    subgraph OmniDreams ["omni-dreams（本 repo）"]
        PT["Post-Training
模型微調
3 種實驗"]
    end
    
    subgraph FlashDreams ["flashdreams（配套 repo）"]
        INF["Inference
離線批次推論"]
        LIVE["Interactive Drive
即時互動駕駛 demo"]
    end
    
    PT -->|"產出微調後的
模型權重"| INF
    PT -->|"產出微調後的
模型權重"| LIVE

2. 核心概念淺解

2.1 什麼是 World Model？

World model（世界模型） 是一種 AI 模型，能夠「想像」未來會發生什麼事。在自駕場景中：

flowchart TD
    A["輸入"] --> B["World Model
OmniDreams"]
    A --> A1["單一 RGB 影像
（第一幀）"]
    A --> A2["Text Prompt
（場景描述）"]
    A --> A3["HD Map
（高精地圖影像）"]
    A --> A4["Trajectory Pose
（行車軌跡姿態）"]
    
    B --> C["輸出"]
    C --> C1["多鏡頭擬真影片
（chunk by chunk）"]
    C --> C2["自回歸生成
（前一段影片
喂入下一段）"]

2.2 Cosmos2 SV-HDMap 是什麼？

Cosmos2 是 NVIDIA 的基礎 video generation model（影片生成模型）。OmniDreams 在 Cosmos2 之上加了兩個關鍵特化：

• SV（Structured Video）— 理解影片的時空結構
• HDMap（High-Definition Map）— 理解高精地圖作為 conditioning（條件輸入）

flowchart TD
    Cosmos["Cosmos2
通用影片生成基礎模型"] --> SV["+ SV 結構化影片理解"]
    SV --> HDMap["+ HDMap 高精地圖條件"]
    HDMap --> OD["= OmniDreams
自駕專用 world model"]

2.3 三種 Post-Training 實驗

OmniDreams 提供三種 fine-tuning 實驗，代表模型訓練的三個階段：

flowchart TD
    E1["E1: Causal Student-Init
（因果學生初始化）"] --> E1D["16N causal，chunk2
先學會「逐步生成」影片
每次生成 2 latent frames"]
    
    E2["E2: Bidirectional Teacher
（雙向教師）"] --> E2D["189-frame teacher
先訓練一個「看全局」的老師模型
能看前後文理解場景"]
    
    E3["E3: Self-Forcing Distillation
（自強化蒸餾）"] --> E3D["學生模型從教師模型蒸餾
學會用更少步驟
生成一樣好的影片"]
    
    E2 -->|"知識蒸餾"| E3
    E1 -->|"學生骨幹"| E3

白話解釋：

1. E1（Student）：訓練一個「一步步往前看」的學生模型
2. E2（Teacher）：訓練一個「能看全局」的教師模型
3. E3（Distillation）：讓學生從教師學習，用更少計算生成一樣好的結果

2.4 FSDP + Context Parallelism

因為模型和影片都很大，訓練需要分散到多個 GPU：

flowchart LR
    subgraph Node ["8x H100 80GB Node"]
        GPU1["GPU 0"]
        GPU2["GPU 1"]
        GPU3["GPU 2"]
        GPU4["GPU 3"]
        GPU5["GPU 4"]
        GPU6["GPU 5"]
        GPU7["GPU 6"]
        GPU8["GPU 7"]
    end
    
    FSDP["FSDP
模型參數分片
每個 GPU 只存
1/8 的參數"] --> Node
    CP["Context Parallelism
影片幀分片
每個 GPU 只處理
部分時間步"] --> Node

• FSDP（Fully Sharded Data Parallel）：把模型參數切成 8 份，每個 GPU 只存 1/8
• Context Parallelism：把影片的時間維度切開，每個 GPU 只處理部分 frame

3. 安裝與環境設定

3.1 硬體需求

3.2 安裝方式 A：Virtual Environment

 1# 安裝系統依賴
 2sudo apt update && sudo apt -y install curl ffmpeg libx11-dev tree wget
 3
 4# 安裝 uv
 5curl -LsSf https://astral.sh/uv/install.sh | sh
 6source $HOME/.local/bin/env
 7
 8# Clone repo
 9git clone https://github.com/nv-tlabs/omni-dreams.git
10cd omni-dreams/post-training
11
12# 安裝 Python 環境
13uv python install
14uv sync --extra=cu128
15source .venv/bin/activate

3.3 安裝方式 B：Docker

 1cd omni-dreams/post-training
 2
 3# Ampere - Hopper GPU
 4image_tag=$(docker build -f Dockerfile -q .)
 5
 6# 執行容器
 7docker run -it --runtime=nvidia --ipc=host --rm \
 8  -v .:/workspace \
 9  -v /workspace/.venv \
10  -v /root/.cache:/root/.cache \
11  -e HF_TOKEN="$HF_TOKEN" \
12  $image_tag

3.4 HuggingFace 設定

模型權重和資料集需要 HuggingFace 授權：

 1# 1. 建立 HF token
 2# https://huggingface.co/settings/tokens/new
 3
 4# 2. 設定 token
 5export HF_TOKEN=<YOUR-HF-TOKEN>
 6
 7# 3. 需要存取的 gated model（需先在 HF 頁面同意 license）
 8# - nvidia/Cosmos-Predict2-2B-Video2World（tokenizer）
 9# - nvidia/Cosmos-Reason1-7B（text encoder，~16GB）
10
11# 4. 資料集
12# - nvidia/PhysicalAI-Autonomous-Vehicles-NuRec（PAI-NuRec 26.01 sample scenes）

3.5 關鍵環境變數

4. 核心功能詳解 — Post-Training 實驗

4.1 三種實驗的執行方式

 1# E1: Causal Student-Init（16N causal chunk2）
 2torchrun --nproc_per_node=8 --master_port=12341 -m scripts.train \
 3  --config=omnidreams/_src/omnidreams/configs/causal_cosmos2/config.py \
 4  -- experiment="causal_cosmos2_2B_single_view_chunk2_t24_hdmap_vae"
 5
 6# E2: Bidirectional Teacher（189-frame teacher）
 7torchrun --nproc_per_node=8 --master_port=12341 -m scripts.train \
 8  --config=omnidreams/_src/omnidreams/configs/causal_cosmos2/config.py \
 9  -- experiment="teacher_cosmos2_2B_single_view_t24_hdmap_vae"
10
11# E3: Self-Forcing Distillation
12torchrun --nproc_per_node=8 --master_port=12341 -m scripts.train \
13  --config=omnidreams/_src/omnidreams/configs/self_forcing/config.py \
14  -- experiment="cosmos_v2_2b_SF_res720p_fps30_i2v_hdmap_chunk2_vae_encode_loc6_release"

4.2 訓練參數解析

4.3 資料流架構

flowchart TD
    subgraph Input ["輸入資料"]
        RGB["RGB Frame
（真實街景照片）"]
        TXT["Text Prompt
（場景描述）"]
        HDMAP["HD Map Image
（車道線 + 路況）"]
        TRAJ["Trajectory Poses
（車輛行進路徑）"]
    end
    
    subgraph Encode ["編碼"]
        VAE["VAE Tokenizer
Cosmos-Predict2-2B
影像 → latent space"]
        TE["Text Encoder
Cosmos-Reason1-7B
文字 → embedding"]
        HE["HDMap VAE Encoding
地圖 → 16-ch latent"]
    end
    
    subgraph Model ["World Model"]
        DIT["DiT Backbone
（Diffusion Transformer）
+ RoPE 3D"]
        RF["Rectified Flow
（去噪擴散生成）"]
    end
    
    subgraph Output ["輸出"]
        VID["Generated Video Chunk
（擬真影片區段）"]
        AR["Autoregressive Loop
（前一 chunk 喂入下一 chunk）"]
    end
    
    RGB --> VAE
    HDMAP --> HE
    TXT --> TE
    VAE --> DIT
    HE --> DIT
    TE --> DIT
    TRAJ --> DIT
    DIT --> RF
    RF --> VID
    VID --> AR
    AR -->|"下一 chunk 輸入"| DIT

4.4 模型 Checkpoint Registry

OmniDreams 使用 checkpoint registry 管理模型權重繼承：

flowchart TD
    L0["L0: Cosmos2 Distilled
（基礎蒸餾模型）"] --> L1b["L1b: Teacher
（189-frame 雙向教師）"]
    L0 --> L2a["L2a: Student-Init
（因果學生初始化）"]
    L1b --> SF["Self-Forcing
（net_real_score_ckpt
+ net_fake_score_ckpt）"]
    L2a --> SF
    SF --> Final["最終推論模型
（低步驟高品質）"]

5. 應用場景

場景 A：自駕模擬測試

用 OmniDreams 生成的擬真影片替代昂貴的實車路測，測試自駕演算法在各種場景（天氣、路況、行人）下的反應。

場景 B：Corner Case 生成

修改 text prompt 和 HD map，生成罕見場景（如行人突然衝出、大雨中的十字路口），用於壓力測試自駕系統。

場景 C：多鏡頭一致性驗證

OmniDreams 支援 multi-camera generation（多鏡頭生成），可驗證自駕系統的多鏡頭融合模組。

場景 D：研究與論文

用 post-training 實驗探索不同的 world model training strategy（如 self-forcing vs teacher-student）。

6. 資安掃描報告

掃描範圍

掃描結論

🟡 中風險 — NVIDIA 官方研究工具。credentials/*.secret 是內部 S3 存取用（公開版本不含實際 key），WANDB_API_KEY 走環境變數。公開使用者只需設定 HF_TOKEN。主要風險來自大型 GPU 訓練工作負載的基礎設施安全（Docker runtime、shared memory、multi-node networking）。

7. FAQ

Q1：我沒有 8x H100，能跑嗎？

A： 目前不行。最低需求是 8x Ampere/Hopper GPU（NPROC=8），更小的 GPU 數量不被支援。可以考慮使用雲端 GPU（AWS p5、Azure NC H100、GCP a3）。

Q2：OmniDreams 跟 FlashDreams 的差別？

A： OmniDreams = 模型訓練（post-training）。FlashDreams = 模型推論（inference + interactive drive）。訓練完的模型權重餵給 FlashDreams 做推論。

Q3：我需要自己準備訓練資料嗎？

A： 不需要。OmniDreams 隨附 PAI-NuRec 26.01 sample scenes（在 HuggingFace 上，需要同意 license 後下載）。如果你有自己的自駕資料集，可以用類似格式準備。

Q4：training 要跑多久？

A： E1/E2 設定 max_iter=150,000，E3 設定 max_iter=10,000。在 8x H100 上，E3 約數小時完成；E1/E2 則需要數天。實際時間取決於資料集大小和 checkpoint save 頻率。

Q5：HDMap VAE encoding 是什麼？

A： 將高精地圖影像（車道線、路況等）通過 VAE（Variational Autoencoder）編碼成 16 channel 的 latent representation，然後 concat 到模型的輸入上，讓模型理解道路結構。

8. 進階技巧

8.1 Slurm 叢集提交

1# 用 QUICKSTART 中的 Slurm wrapper
2sbatch --nodes=1 --ntasks-per-node=8 --gres=gpu:8 \
3  samples/post-training/launch_slurm.sh \
4  --config=... --experiment=...

8.2 離線環境設定

在無網路的 compute node 上預先 stage 模型：

1# Login node（有網路）
2export HF_HOME=/shared/hf_cache
3hf download nvidia/Cosmos-Predict2-2B-Video2World --repo-type model tokenizer/tokenizer.pth
4hf download nvidia/Cosmos-Reason1-7B --repo-type model --revision 3210bec0495fdc7a8d3dbb8d58da5711eab4b423
5
6# Compute node（無網路）
7export HF_HOME=/shared/hf_cache
8export HF_HUB_OFFLINE=1

8.3 Triton Cache 競爭問題

多 rank 同時編譯 Triton kernel 會在 NFS 上產生 race condition：

1# 解法：per-rank 隔離 cache
2export TRITON_CACHE_DIR="/tmp/triton_${SLURM_JOB_ID:-$$}_${SLURM_LOCALID:-0}"

8.4 justfile 快速指令

1just install          # 安裝環境
2just test-cpu         # CPU 測試
3just test-gpu         # GPU 測試
4just lint             # Linting
5just docker-cu128     # Docker (CUDA 12.8)
6just docker-cu130     # Docker (CUDA 13.0, Blackwell)

9. 整合進其他工作流

與 FlashDreams 的接力

flowchart LR
    A["OmniDreams
Post-Training"] -->|"微調後權重"| B["FlashDreams
Inference"]
    B --> C["離線批次推論
mp4 影片"]
    B --> D["即時互動駕駛
Live Demo"]

與 NVIDIA Omniverse 生態系

OmniDreams 是 NVIDIA 自駕模擬生態系的一部分，可搭配 Omniverse、Isaac Sim 等工具。

與自駕開發 Pipeline

flowchart TD
    A["真實路測資料"] --> B["OmniDreams
Post-Training
微調 world model"]
    B --> C["FlashDreams
生成模擬場景"]
    C --> D["自駕演算法
在模擬中測試"]
    D --> E["發現問題
修正演算法"]
    E --> F["再次模擬驗證"]
    F -->|"通過"| G["實車路測"]

10. 重點摘要 Checklist

• 理解定位：OmniDreams = post-training，FlashDreams = inference
• World Model 概念：單幀 + prompt + HD map + trajectory → 擬真影片
• 3 種實驗：E1 Student-Init / E2 Teacher / E3 Self-Forcing Distillation
• 硬體需求：最低 8x H100 80GB，150+ GB 磁碟
• HuggingFace 設定：token + 同意 gated model license
• 環境變數：HF_HOME / HF_HUB_OFFLINE / CUDA_HOME / TRITON_CACHE_DIR
• 安裝方式：uv venv 或 Docker
• FSDP + CP：模型分片 + 時間維度分片 → 多 GPU 訓練
• Cosmos2 SV-HDMap：基礎模型 + 結構化影片 + 高精地圖
• Apache-2.0：開源可商用（但 gated model weight 需另外同意 NVIDIA license）

11. 進一步閱讀

官方資源

• Research Blog
• White Paper (PDF)
• Model Weights (HuggingFace)
• FlashDreams (Inference)
• FlashDreams OmniDreams Docs

相關技術

• Cosmos2 Foundation Model
• FSDP (Fully Sharded Data Parallel)
• Rectified Flow
• Self-Forcing Distillation

社群

• NVIDIA Omniverse Discord
  • #omnidreams / #flashdreams / #world-model-chit-chat