π0 和 π0-FAST：用於通用機器人控制的視覺-語言-動作模型

釋出於 2025 年 2 月 4 日

在 GitHub 上更新

贊

168

我們已經將首批 機器人基礎模型 移植到了 Hugging Face LeRobot！由 Physical Intelligence 開發的 π0 和 π0-FAST 現已在 LeRobot 程式碼庫 中提供，為 Hugging Face 生態系統帶來了通用機器人智慧。如果你對視覺-語言-動作 (VLA) 模型與視覺-語言模型 (VLM) 的區別以及動作是如何表示的感到好奇，那就深入閱讀這篇博文一探究竟吧！

在我們的程式碼庫中探索模型集和模型的 PyTorch 版本：Huggingface Pi0 模型合集 | Huggingface Pi0+FAST 模型合集 | LeRobot 程式碼庫

引言

羅伯特·海因萊因（Robert Heinlein）提出，一個全面發展的人應當能夠處理廣泛的任務——無論是智力上的還是體力上的——而不是侷限於某一特定領域。將全面發展的人與機器智慧進行類比：人工智慧系統千差萬別，但人類智慧的卓越之處在於其多功能性——能夠適應各種任務、環境和意外情況。儘管大型語言模型和視覺-語言模型 (LLM, VLM) 顯示出巨大潛力，但它們缺乏與物理世界的互動。為了彌合這一差距，我們需要基於機器人資料訓練的模型。通用機器人模型可以增強適應性，利用多樣化的資料來提高泛化能力和魯棒性。與在孤立任務上進行訓練不同，在多樣化的機器人資料上進行預訓練——類似於 LLM 的做法——能夠提升效率和效能。

開發通用機器人策略，或稱機器人基礎模型，面臨三大關鍵挑戰：

需要大規模研究 才能充分利用預訓練的優勢。
設計能夠整合不同資料來源同時捕捉複雜物理互動的模型架構。這方面的一個關鍵挑戰是 跨形態訓練 (cross-embodiment training)，即模型必須從具有不同配置、控制空間和動作表示的多種機器人型別中學習。現有方法透過以下方式解決這一問題：
- 結合來自不同機器人平臺的多模態資料集 以增強泛化能力。
- 使用共享表示 來彌合不同機器人形態（如單臂、雙臂和移動機械臂）之間的差距。
制定有效的訓練方案，因為自然語言處理和視覺領域的最新進展在很大程度上依賴於精心的預訓練和後訓練策略。

在本文中，我們介紹由 Physical Intelligence 開發的原型模型和學習框架 π0 與 π0-FAST，旨在克服這些挑戰。

🔍 π0 是什麼？

論文 | Jax 程式碼

π0 (Pi-Zero) 是一個由 Physical Intelligence 團隊開發的 視覺-語言-動作 (VLA) 模型，專為 通用機器人控制 設計。它建立在 大規模預訓練 和 基於流匹配的動作生成 的基礎上，使機器人能夠在不同形態下執行 靈巧的操作任務。

π0 在來自 7 個機器人平臺 和 68 個獨特任務 的資料上進行訓練，在複雜的真實世界任務（如 疊衣服、收拾餐桌、打包雜貨、組裝盒子和取物）中表現出強大的 零樣本 和 微調效能。

與標準的機器人策略不同，π0 採用流匹配 來生成 流暢、即時的 50Hz 動作軌跡，使其在實際部署中 高效、精確且適應性強。流匹配曾用於連續歸一化流，並提高了擴散模型中的生成質量。π0 使用的去噪過程以同樣的方式工作，從一個隨機噪聲開始，逐步收斂到一個有意義的運動動作序列。

如何在 LeRobot 中使用 π0？

首先，你需要升級你的 lerobot 安裝，它現在依賴於 transformers！只需在克隆 git 倉庫後執行：

pip install -e ".[pi0]"

π0 模型是基礎模型，就像 PaliGemma 一樣，旨在適應各種框架、環境和場景輸入。這裡的基座模型可以直接使用，尤其是 π0。

在 π0 預訓練模型上進行推理

python lerobot/scripts/eval.py \
--pretrained_policy.path=/path/to/pretrained/pi0

然而，由於這是從 jax 到 torch 的轉換以及從特定環境的轉換，效能有所下降。我們建議您像下面這樣，將您自己的 π0 微調到您自己的環境中。

微調 π0 預訓練模型

要使用來自 openpi 的 pi0_base 檢查點微調 π0 模型，請執行以下命令：

python lerobot/scripts/train.py \
--policy.path=lerobot/pi0 \
--dataset.repo_id=danaaubakirova/koch_test

要在 PaliGemma 和 Expert Gemma 上微調 π0 神經網路（這些網路在 π0 微調之前使用了 VLM 預設引數進行預訓練），請執行：

python lerobot/scripts/train.py \
--policy.type=pi0 \
--dataset.repo_id=danaaubakirova/koch_test

您也可以獨立於 LeRobot 訓練框架，使用以下程式碼來使用預訓練的 π0 模型：

policy = Pi0Policy.from_pretrained("lerobot/pi0")

VLM 和 VLA 之間有什麼區別？

視覺-語言模型 (VLM) 和視覺-語言-動作模型 (VLA) 共享一個共同的基礎：transformer。然而，關鍵區別在於動作表示。VLM 處理和生成多模態表示（影像和文字），而 VLA 則透過加入動作和觀察狀態詞元來擴充套件這一點。有了這些額外的詞元，下一個挑戰就是理解注意力是如何計算的。

機器人策略中的注意力機制

讓我們擴充套件一下詞彙，介紹一些關鍵術語：

狀態詞元 (State Token)

它是一個單一的詞元，代表機器人的 當前環境狀態（例如，關節角度、感測器數值或其他相關觀測值）。
掩碼規則允許該詞元 關注字首中的影像和文字，這意味著狀態詞元可以“看到”決策所需的所有視覺或文字線索。
它還以 三角方式關注之前的狀態。如果使用多個狀態詞元，每個新的狀態詞元可以看到舊的詞元，但反之不行。

動作詞元 (Action Tokens)

代表 電機命令序列。
對除填充區域外的所有內容都具有 完全可見性。這意味著每個動作詞元可以關注：
- 所有非填充的影像詞元（整個場景），
- 所有非填充的文字詞元（指令或描述），
- 狀態詞元（當前和之前的），
- 其他動作詞元.

字首詞元 (Prefix Tokens)

代表 整個場景，並完全相互關注，類似於 PaliGemma。

核心思想

這些詞元封裝了：

機器人對環境的 內部表示 (狀態)，
機器人發出的 命令或控制 (動作)，
時間或步驟索引的編碼 (時間嵌入)。

它們被附加在字首部分（影像+文字）之後，因此字首作為上下文（例如，場景影像、語言指令如 “做個好機器人” 或 “移動方塊”），而後綴則捕捉特定於策略的特徵。

⚡ 邁向 π0 中更快的注意力機制

然而，在 π0 中高效地處理注意力機制也帶來了一系列挑戰。其注意力掩碼的獨特形狀影響了注意力的計算方式——讓我們深入瞭解細節！

處理二維注意力掩碼

由此產生的 二維因果掩碼 表現出強烈的 塊稀疏性，但定義每個塊的邊界——尤其是在一批樣本中——有點棘手。我們習慣於自迴歸建模中的三角結構因果掩碼，但這並非如此。

正如您在下面的示例中所見：影像（第一個元素）有一些填充詞元，代表空置的攝像頭。然後，新增文字詞元和狀態詞元。這個“字首”部分形成了一個完全非因果的注意力，就像在 PaliGemma 中一樣。然後，“字尾”（狀態 + 動作/時間詞元）具有一個因果塊結構。簡單的 eager 實現方式執行矩陣乘法並在整個輸入上應用 softmax，這使其效率極低。

圖 1：VLA 注意力掩碼的視覺化

我們能使用 FlashAttention2 嗎？

FlashAttention2 提供了一個 varlen 介面，但 cu_seqlens（累積字首長度）必須手動計算。它專為具有統一查詢和鍵長度的 連續（或嚴格因果）注意力模式 而設計。
它 無法自然處理不規則的塊掩碼 或每個詞元任意的“允許”位置，而這正是我們所需要的。
所以，雖然以一定的實現成本使用它是可能的，但我們決定轉向……

在 PyTorch 中使用 FlexAttention

這看起來像是 FlexAttention 的活兒！它有一個純 PyTorch 介面，我們探索了兩種方案：

在我們因果掩碼中注意力被關閉的位置 新增一個 score_mod。然而，即使是一個標量加法也會 顯著降低 FlexAttention 的效能。這是因為在我們的案例中，score_mod 是在最佳化的 CUDA 核心之外新增的。
正確的選項是 對我們的因果掩碼進行索引，並將得到的簽名傳遞給 block mask 的建立。 這個 block mask 能有效地指示哪些地方需要計算注意力，哪些地方可以完全跳過。

# Example of indexing the causal mask and using mask_mod
causal_mask = generate_causal_mask(your_samples) # should be[batch, head, q_len, kv_len]
# Now we need to wrap this mask 
def precomputed_mask_factory(precomputed_mask: torch.Tensor) -> _mask_mod_signature:
    def mask_mod(b, h, q_idx, kv_idx):
        return precomputed_mask[b][h][q_idx][kv_idx]
    return mask_mod
flex_attention_output = flex_attention(query, key, value, mask_mod=mask_mod)

mask_mod = precomputed_mask_factory(causal_mask)
# create a block mask with that signature
block_mask = create_block_mask(
    mask_mod=mask_mod,
    # ...
)

# Call flex attention now!
attn_output, attention_weights = flex_attention(
    query,
    key,
    value,
    block_mask=block_mask,
)

目前的實現可以執行，但仍在進行中，目標是讓它支援 torch.compile 並充分利用其優勢！

如何有效表示“動作”？

既然我們知道動作只是可以被分詞化的 n 維 向量，我們就可以探討在視覺-語言-動作 (VLA) 模型中動作表示的挑戰。動作的表示方式直接影響效率、泛化能力和執行保真度。

一種方法是 語義化動作表示，其中動作被描述為高階概念，如子任務或關鍵點。雖然這允許小樣本和零樣本學習，但它通常依賴於手動設計的底層控制器，限制了在不同機器人間的靈活性。相比之下，底層控制表示直接將動作對映到電機指令，能夠實現精確運動，但訓練過程不太穩定且難以擴充套件。

大多數現有的 VLA 使用 離散動作分詞，將連續動作轉換為自迴歸生成的離散詞元。最常用的方法——逐維度、逐時間步分箱——在處理高頻控制任務時表現不佳，導致表示有損且訓練效率低下。向量量化 (VQ) 和時間序列壓縮等替代方案有所幫助，但 VQ 對超引數敏感，使其在多樣化的機器人設計中可靠性較低。

為了解決這些問題，頻域動作序列分詞 (FAST) 引入了一種新穎的時間序列壓縮方法，使用離散餘弦變換 (DCT)。FAST 減少了冗餘，提高了效率，並增強了動作的保真度。

基於此，我們推出了 π0-FAST，這是 π0 的一個更快、自迴歸的版本，也已在 Lerobot 程式碼庫中提供。它是 π0 的擴充套件，利用了這種新的分詞器來更好地表示動作。

🚀 π0-FAST 是什麼？

論文 | Jax 程式碼 | 我們在 Lerobot 中的實現

π0-FAST 是 π0 的一個 自迴歸版本，引入了 **FAST (頻域動作序列分詞)**——一種新的分詞方案，可提高效率和效能。

π0-FAST 的主要優勢：

與基於擴散的 VLA 相比，訓練速度快 5 倍。
改進的動作表示，減少了動作序列中的冗餘。
在未見過的環境和機器人形態上具有 更強的泛化能力。

🔗 π0-FAST 分詞器 可在此處獲取：FAST Tokenizer

🔗 預訓練權重可在此處獲取：Pi0+FAST

FAST 是如何工作的？

FAST 使用離散餘弦變換 (DCT) 將連續的動作序列壓縮為離散的詞元。如圖 2 所示，該過程首先對原始機器人動作進行歸一化，將每個動作維度的第 1 和第 99 百分位數對映到 [-1, 1] 範圍內。這種歸一化用於確保不同機器人系統間的一致性，並提高對異常值的魯棒性。

然後，每個動作維度使用 DCT 獨立進行變換，將時域訊號轉換為頻域。為了減少冗餘，透過縮放和四捨五入操作去除不重要的係數，其中一個超引數平衡了壓縮率和重建精度。得到的 DCT 係數矩陣通常是稀疏的，被展平成一個一維整數序列，首先跨維度交錯低頻分量以保留關鍵資訊。

為了進一步壓縮序列，應用了位元組對編碼 (BPE)。像往常一樣，BPE 合併跨維度頻繁出現的模式，同時保持固定大小的詞彙表。

圖 2：FAST 動作分詞流程

由於所有操作都是可逆的，因此可以從詞元中高效且無損地重建動作。該分詞流程只有兩個超引數：四捨五入前應用的縮放係數和 BPE 詞彙表大小。這兩個引數在不同資料集上都保持魯棒性。

此外，一個名為 FAST+ 的通用版本 FAST 已經在一個包含單臂、雙臂和移動操作機器人的一百萬個動作序列上進行了訓練，使其適用於各種機器人設定。FAST+ 作為 Hugging Face AutoProcessor 提供，允許使用者僅用幾行程式碼就能對動作序列進行分詞。

為實現最佳壓縮效果，輸入動作在分詞前應進行分位數歸一化至 [-1, 1]。藉助 AutoProcessor 模組，使用者可以在自己的資料集上訓練自定義的 FAST 分詞器。

如何使用 FAST 分詞器？

🔗 有關自定義動作分詞器的使用和訓練程式碼，請參見官方 FAST 程式碼庫

FAST 已整合到 Hugging Face Transformers 中，可以輕鬆用於編碼和解碼機器人動作序列。

通用機器人智慧的下一步是什麼？

藉助 π0 和 π0-FAST，我們向通用機器人智慧邁出了重要一步，將可擴充套件、高效且多功能的視覺-語言-動作 (VLA) 模型引入 LeRobot。透過利用 FAST 分詞，我們增強了動作表示，使機器人能夠以更高的效率和適應性執行各種任務。這些進展為未來的多形態、即時機器人策略打開了大門，推動了機器人在現實世界中所能達到的極限。🚀

其他資源

參考文獻

@book{heinlein2021time,
  title={Time enough for love},
  author={Heinlein, Robert A},
  year={2021},
  publisher={Penguin}
}

@article{black2024pi_0,
  title={$$\backslash$pi\_0 $: A Vision-Language-Action Flow Model for General Robot Control},
  author={Black, Kevin and Brown, Noah and Driess, Danny and Esmail, Adnan and Equi, Michael and Finn, Chelsea and Fusai, Niccolo and Groom, Lachy and Hausman, Karol and Ichter, Brian and others},
  journal={arXiv preprint arXiv:2410.24164},
  year={2024}
}

@article{pertsch2025fast,
  title={FAST: Efficient Action Tokenization for Vision-Language-Action Models},
  author={Pertsch, Karl and Stachowicz, Kyle and Ichter, Brian and Driess, Danny and Nair, Suraj and Vuong, Quan and Mees, Oier and Finn, Chelse|a and Levine, Sergey},
  journal={arXiv preprint arXiv:2501.09747},
  year={2025}
}

更多部落格文章