RFMEP02：Octo

在Robot Foundation Model（RFM）系列第一期RFM EP01：Pi和π0具身基礎模型，我們透過Sergey在CoRL 2024的兩個Talk瞭解了美國具身智慧獨角獸PI推出的機器人大模型π0。作為RFM的第二期，我們會關注更早一個版本的機器人大模型Octo，以及從Octo到π0的最佳化過程。

這篇文章是Octo和π0的一作Kevin Black在CoRL 2024，the 1st Workshop on X-Embodiment Robot Learning workshop上題為“From Octo to π0：How to Train Your Generalist Robot Policy”的報告。

這個報告深入淺出的介紹了Octo的成功與失敗，以及π0如何在Octo的經驗基礎上進一步提高。非常難得的一點是，Kevin不吝將他在這兩個機器人模型工作的失敗經驗分享出來。除了詳細的介紹了Octo，Kevin也提到了π0訓練的一些獨到經驗。

值得注意的是，這篇文章主要來自於Kevin Black報告的翻譯總結，關於他表述的一些細節，我也有很多不理解的地方，所以如果你在閱讀文章的過程中，看到一些事實錯誤，歡迎在文末留言糾正。

以下為本文目錄

👇

1. Octo的提出

2. Octo模型設計

3. Octo模型評估

4. Octo的效果

5. Octo值得借鑑的經驗

6. Octo的一些問題

資料

模型設計

7. 在Octo基礎上π0如何改進

擴充套件模型的規模+改進預訓練

改進Decoder

解決分佈漂移問題：Action Experts

訓練與結果

π0有效嗎？

1. Octo的提出

Open X Embodiment 資料集，很多領域的人可能都已經聽說過這個資料集。它在一年前釋出，對於機器人學者來說，這是一項巨大的里程碑。這個資料集的規模遠超我們之前使用過的任何資料集，涵蓋了來自 21 個機構的貢獻、22 臺機器人產生的資料，總計約 4000 小時的記錄活動。在資料驅動的機器人學中，最關鍵的因素就是資料，而這個資料集提供了我們一直在等待的數量和多樣性。

谷歌並沒有止步於此——他們還發布了使用這個資料集的令人印象深刻的成果，他們的模型 RT-2，使用 Open X Embodiment 資料集進行訓練，在不同實驗室和任務之間表現出了令人矚目的泛化能力。這是一項不可思議的成就。當我第一次評估這個模型時，我曾經說過“它不可阻擋”，這是 Sergey 喜歡在所有幻燈片中引用的句子，雖然這讓我一直很尷尬。

然而，如果我對RT-2提出批評，唯一的問題就是那個“550億”——指的是它使用的語言模型的規模。這個數字帶來了兩個重大問題。首先，它使得這個模型本質上成為私有的，這意味著它不太可能被開源。其次，即使它開源了，大多數機器人學者也沒有資源來部署，更不用說訓練如此龐大的模型了。我們大多數人不習慣啟動足夠的GPU來處理如此龐大的計算量。

僅僅為了在一個 550 億引數的模型上進行推理，就需要足夠的GPU，更不用說訓練了，這可不是一件小事。谷歌確實釋出了一個名為RT-1的較小模型，但它也有自己的侷限性。例如，訓練基礎設施並沒有開源，並且它並不特別適合微調。

就在這個時候，我和一些在 Berkeley 的朋友看到了一個機會。我們想：“為什麼不自己用這些驚人的資料訓練一個模型呢？”這個想法並不新穎——其實也沒什麼深奧的，但值得一試。於是，我們開始了一個叫做 Octo 的專案。

Octo 的主要目標是靈活性。正如我們所知，Cross-Embodiment Learning需要在多種系統和任務中具備適應性，我們希望確保我們的模型能夠支援這一點。專案的另一個基石是將整個基礎設施開源。這是我們的一大創意：創造出真正能被人們使用的東西。對於很多人來說，這將是第一次直接使用一個在如此大規模、多樣化資料集上訓練的模型。

我們希望人們能夠真正使用這個模型——這就是我們的核心想法。有一個類比提到，Octo 就像機器人領域的 GPT-2。如果你不熟悉，GPT-2 是一個相對較小的模型，雖然它對於zero-shot 任務並不特別有用——它能生成連貫的內容，但僅此而已。然而，它對於微調非常有用。GPT-2 有大約 1 億個引數，你可以輕鬆下載並微調它，應用到生成詩歌等任務上。

Octo 的理念大致也是如此。它可能有些雄心勃勃，但我們的希望是，如果人們能夠使用它並取得良好的表現，它將像 GPT-2 對機器人領域的貢獻一樣，具有類似的實用性。

2. Octo模型設計

Octo 的設計並不複雜。它基本上是一個大 Transformer。包含一個language encoder處理語言指令，observation encoders處理感測器輸入和其他觀察資料，然後一切都輸入到 Transformer 中。主要創新點是action heads——這些小模組，如 MLPs，用來輸出動作。這個設計旨在增加模型的靈活性。如果機器人有一個不在預訓練資料中出現的動作空間，你可以簡單地替換掉action heads。我們的想法是，儘管action heads會發生變化，龐大的共享encoder仍然能夠提供知識轉移，而action heads則能輕鬆適應新的動作表示。

3. Octo模型評估

我們進行了很多評估。對於微調，我們在具有新動作空間的任務上進行了實驗，例如 Berkeley Pickup 和 Berkeley Bimanual。這些動作空間在預訓練資料中沒有出現。預訓練使用的是經典的末端執行器位置控制，而這些任務涉及的是關節位置和關節扭矩。

我們還進行了zero-shot 評估，意味著沒有微調，而是直接在訓練資料中出現的場景上測試。這並不是為了展示泛化能力，而是為了確認模型能夠適應它在訓練時見過的資料。

我們付出了很多努力讓 Octo 更易獲取，我們釋出了開原始碼庫；提供了完整的基礎設施，包括資料載入器。

4. Octo的效果

那麼，它有效嗎？人們真的在使用 Octo 嗎？嗯……算是吧。但並不是為了微調。大多數人將它作為基準進行對比。坦白說，這也沒關係。這仍然是一種影響力，表明我們讓系統足夠簡單，便於人們進行推理。

但如果說實話，人們並沒有真正使用 Octo 進行微調，因為它的表現並沒有那麼好。微調帶來的效能提升不足以彌補投入的努力。格式化資料、設定訓練基礎設施和學習如何使用——這些都需要付出努力。而微調 Octo 所帶來的收益並不足以讓這些努力變得值得。

在很多方面，Octo 是失敗的。我們沒有取得我們所期待的突破性成果。但這並不意味著它在各個方面都是失敗的。在談論哪裡出了問題之前，我想先花點時間反思一下，我們做對了什麼，和在這個過程中我們學到了什麼。

5. Octo值得借鑑的經驗

我們學到的一件事是，擴散模型（diffusion）非常有效。在我們開始做這個專案時，擴散策略（diffusion policies）已經相對為人所知，但對我個人來說，這個專案真正讓我深刻理解了這一點。一開始什麼都不行，但在截止日期前大約兩週，我們添加了一個diffusion head——突然間，一切開始順利。這次經歷加強了擴散（diffusion）在資料擴充套件（scale up data）過程中尤其有用的觀點，這可能是因為它能很好地處理多模態資料。

另一個重要的啟示是action chunking（動作分塊）。這並不是一個新想法——像diffusion policy和 ACT 這樣的方式，已經使用了某種形式的action chunking。然而，這個專案讓我深刻意識到它的重要性。在Berkeley，我們傳統上沒有在某些任務中使用action chunking，但我們發現，當action chunking的大小在半秒到一秒之間時，效能顯著提高。

還有一些看似平凡但同樣重要的經驗教訓——這些東西可能不夠光鮮亮麗，不值得寫進論文，但在訓練這些模型時卻至關重要。例如，資料洗牌（shuffling data）證明是極其重要的。機器人學資料有一個不幸的特性——它們高度相關。每條軌跡與自身顯然高度相關，甚至同一場景的軌跡也常常非常相似。如果沒有適當的洗牌，損失曲線的改善將不如預期。我們是在試錯中學到這一點的。一旦我們更有效地進行了資料洗牌，效能得到了顯著提升——有時甚至翻倍。

另一個看似微不足道但影響深遠的細節是動作表示（action representation）。如何表示gripper的狀態？是絕對位置還是相對位置？對於關節，使用位置、增量位置還是扭矩？對於末端執行器控制，使用什麼參考框架？這些問題可能看起來微不足道，但它們實際非常重要，尤其是在cross-embodiment訓練中。在 PI，我們花了無數小時討論這些問題。看起來這些問題隨著資料規模的擴大可能會變得不那麼重要，也許將來是這樣，但目前這些細節差異很大。

6. Octo的一些問題

現在，談談為什麼 Octo 沒有成為機器人領域的 GPT-2。我要把問題分為兩大類：資料和設計。

資料

在資料驅動的機器人學中，資料始終是最重要的因素。這一點在更廣泛的機器學習領域也是如此。關於資料已經有很多討論了，所以我決定採取一種獨特的方式來闡明這一點。我做了一件有些傻的事情：我嘗試將文字資料集轉換成“小時”，以便與機器人資料集進行比較。

如果你將平均人類閱讀速度與文字資料集中的標記數相乘，你可以估算一個普通人閱讀完這個資料集所需的時間。例如，GPT-2 的訓練資料集，在當時被認為是一個較小的資料集，但它的規模遠遠超過了任何機器人資料集。

需要明確的是，這種比較有很多侷限性。資料集的多樣性和質量遠比資料的原始大小更為重要，但我認為這個練習仍然突出了一個關鍵點：我們在機器人學中擁有的資料遠遠不夠。即使是我們現在認為相對較弱的 GPT-2，它在訓練時也接觸到了比任何機器人模型現在所能獲得的更多的資料。

在 Octo 開始時，我們高估了 Open X Embodiment 資料集的規模，認為它非常龐大。但實際上，它與語言資料集相比仍然很小。即使在 PI，我們全天候收集資料，至今也只有大約 10,000 小時的資料——這幾乎只是語言模型資料的一小部分。做個對比，現代語言訓練資料集包含大約 7.9 億小時的資料。為了說明這一點，我做了一個圖表，展示了它們之間的規模差異，真是令人震驚。

這個比較並不是說我們永遠無法在機器人領域達到這樣的資料量，但它確實意味著我們需要調整預期。我們不會看到像語言模型中那樣的湧現性質，因為我們根本沒有足夠的資料量、資料多樣性或狀態空間的覆蓋。

設計

雖然資料是最重要的因素，但設計同樣非常關鍵——尤其是在我們擴充套件規模時（scale up）。如果我們假設機器人資料集會隨著時間增長，那麼探索合適的設計以及它們如何隨著更多資料的引入而擴充套件，仍然是非常有價值的。

我們在 Octo 中犯了一個大錯誤，就是沒有使用視覺預訓練。這與資料密切相關，因為利用大規模的網際網路資料集是克服機器人資料集侷限性的一個有效途徑。在 Octo 中，我們自認為 Open X Embodiment 資料集已經足夠大，所以沒有對observation encoder進行預訓練，只對語言編碼器進行了預訓練。結果證明這是一個錯誤。即使是在相對較大的資料集上，視覺預訓練也能帶來巨大的改進。

在 PI，我們早期進行了一些內部實驗，甚至在 ImageNet 或類似資料集上進行簡單的預訓練（例如，使用 ResNet 或 ViT），也顯示出顯著的效能提升。這並不是關於泛化能力或抗干擾能力，而僅僅是提高了訓練資料中任務的表現。預訓練帶來的幫助遠遠超出了我們的預期。

另一個問題是與action head（動作頭）相關的。最初的設計是將action head做得很小——佔總引數的不到 5%，這樣更換新的動作空間時不會丟失太多資訊。然而，這證明是一個問題，尤其是在使用基於擴散的目標（diffusion-based objectives）時。隨著資料變得更加多模態和複雜，需要一個更大的解碼器來有效捕捉資料分佈。在 Octo 中，龐大的 Transformer 在單次前向傳播中處理輸入，而action head (或decoder)則在去噪過程中進行了多次前向傳播。更強大的解碼器會更適合處理複雜的分佈。

我們遇到的另一個問題是高頻控制設定，例如當時越來越流行的 ALOHA 型別配置。這些設定建議將動作塊大小設為 50 或更多，但 Octo 無法處理這些設定。損失值沒有下降，甚至無法擬合數據。在 Octo 釋出後，我加入了 PI。在我的第一週，我提出了一種新的架構。這個架構與 Octo 相似，但具有一個顯著更大的解碼器——大約佔總引數的 30%，並且可能在計算成本上佔據更大的比例（雖然我沒有計算具體數字）。

這個架構仍然基於 Transformer，因為我們知道 Transformer 非常有效，所以沒有理由不使用它。該設計受到了“Diffusion Transformer”論文的啟發，這篇論文在計算機視覺領域取得了很大的成功。此外，我們透過使用 ImageNet 對視覺骨幹進行預訓練，解決了預訓練的問題。

這個架構表現得非常好。在π0論文中，這個版本被稱為“π0 Small”，但從時間順序來看，這個模型是在 Octo 之前的，我們在擴充套件之前使用了它很長時間。它的表現異常出色。有趣的是，幾乎完全相同的架構也被其他幾個專案獨立開發出來，例如 ALOHA Unleashed和 DIT Block Policy。多個團隊的嘗試說明這是一個非常強大的設計。

7. 在PI Scale Up—π0如何改進

擴充套件模型的規模+改進預訓練

在此之後，我們在 PI 的下一個目標是透過兩個方面來推動進展：(1) 擴充套件模型的規模和 (2) 改進預訓練。雖然 ImageNet 預訓練確實有幫助，但我們很快意識到，使用VLM模型作為骨幹網路——即我們現在所稱的VLA——會更加有效。其基本思想是，利用一個VLM，並對其進行微調，以輸出動作。這種方法已顯示出巨大的潛力。

為了在保持擴散目標和整體設計的同時，將現有架構轉變為 VLA，我開始進行一些簡單的修改。首先，我用一個語言模型替換了原來的大 Transformer 編碼器，如GLaM。其次，我將視覺骨幹網路升級為更強大的模型，如 SimCLIP。

通常，VLM使用的是decoder-only Transformers，因此語言輸入透過cross-attention 和 self-attention mechanisms進行新增——這是VLM研究中的常見做法。此時，模型看起來類似於 PaLM-E，它將 SimCLIP 視覺編碼器與GLaM結合起來。PaLM-E 是由 Google 開發的一個優秀模型，專門設計用於在新任務上進行微調。

我們將這種方法整合到現有架構中的初步嘗試涉及使用相同的解碼器，但將編碼器替換為 PaLM-E。它的表現還可以——能夠擬合數據，但並沒有比我們之前的小模型好多少，而且存在一些問題，比如偶爾的損失不穩定。我嘗試了許多修復方法，例如凍結編碼器或對編碼器使用較慢的預熱，但都沒有取得很好的效果。

那時，我意識到問題所在：GLaM 是一個decoder-only Transformer，但我們將它當作encoder使用。問題出在其架構上。decoder-only 模型的最終層經過自迴歸目標訓練，並不設計用來生成豐富的表示——它最佳化的目的是預測下一個token。從理論上講，如果模型經過完美訓練，最終層只會告訴你下一個詞，並不會攜帶有意義的上下文表示。

下一步：改進Decoder

為了解決這個問題，下一步的邏輯是確保Decoder能夠從語言模型的每一層獲取表示（Representation），而不僅僅是從最後一層。這一改變使得整個系統看起來更像是一個大Decoder。觀察輸入（如影像）被送入模型，同時動作輸入也被加入。

新架構的一個顯著特點是將原本的encoder和decoder之間進行了分離。在推理過程中，我們透過encoder部分進行一次前向傳播，然後在去噪過程中透過decoder部分進行多次前向傳播。為了提高效率，我們使用了 KV 快取，這使我們能夠儲存來自encoder的鍵和值。因此，解碼器的計算速度大大提高，因為它不需要重新計算整個字首的注意力層。

這個設計與“Transfusion”論文中提出的某些概念非常相似，這篇論文大約在同一時間釋出。當我看到那篇論文時，我想：“哇，他們有同樣的想法！”這通常是判斷一個概念好壞的標誌——它在不同團隊之間產生共鳴。

新的架構比前一個版本表現得要好得多。它在泛化能力上有所提升，特別是在利用VLM預訓練時。例如，它能夠識別一些之前沒有見過的新物體，並展現出較強的魯棒性。得益於擴散目標，它還能夠處理高頻控制任務。

然而，仍然存在一個主要問題：收斂非常緩慢。訓練時間遠比預期要長。我的假設是，這與分佈漂移（distribution shift）有關。VLM 的預訓練只見過某些型別的輸入，因此在微調過程中引入未初始化的線性投影層可能會使模型產生困惑。

解決分佈漂移問題：動作專家（Action Experts）

為了解決分佈漂移問題，我提出了一個靈感來自專家混合架構的解決方案。核心思路是為新的輸入使用不同的權重。概念上，你可以將其視為一個 Transformer 序列，其中不同的 token 使用不同的權重集合。例如：原始輸入（例如來自 PaLM-E 的視覺和語言）將使用預訓練的權重，保持與模型在預訓練期間所見的一致性。新輸入（例如機器人狀態和噪聲動作）將使用一組新的權重，從頭開始初始化，以避免分佈漂移。

這種分離也使我們能夠最佳化效率。由於自注意力層是共享的，我們可以大大縮小新部分（即“動作專家”部分）的 MLP 塊。這節省了引數，減少了訓練時間，並大大加快了推理速度，因為動作專家部分在去噪過程中會多次執行。

結果非常有希望。不僅訓練變得更快，模型的收斂性也得到了顯著改善，導致了更高的效能。這種架構最終成為了π0模型，如我們在π0論文中所描述的那樣。唯一的不同是，釋出版本中的插圖更精美。