適應多形態多任務，最強開源機器人學習系統「八爪魚」誕生

作者：機器之心 2024-05-28 13:03:50

一位優秀的相聲演員需要吹拉彈唱樣樣在行，類似地，一個優秀的機器人模型也應能適應多樣化的機器人形態和不同的任務，但目前大多數機器人模型都只能控制一種形態的機器人執行一類任務。現在 Octo（八爪魚）來了！

在機器人學習方面，一種常用方法是收集針對特定機器人和任務的數據集，然后用其來訓練策略。但是，如果使用這種方法來從頭開始學習，每一個任務都需要收集足夠數據，并且所得策略的泛化能力通常也不佳。

原理上講，從其它機器人和任務收集的經驗能提供可能的解決方案，能讓模型看到多種多樣的機器人控制問題，而這些問題也許能提升機器人在下游任務上的泛化能力和性能。但是，即便現在已經出現了能處理多種自然語言和計算機視覺任務的通用模型，構建「通用機器人模型」依然困難重重。

要為機器人訓練一個統一的控制策略非常困難，其中涉及諸多難點，包括操作不同的機器人機體、傳感器配置、動作空間、任務規范、環境和計算預算。

為了實現這一目標，已經出現了一些「機器人基礎模型」相關研究成果；它們的做法是直接將機器人觀察映射成動作，然后通過零樣本或少樣本方式泛化至新領域或新機器人。這些模型通常被稱為「通才機器人策略（generalist robot policy）」，簡稱 GRP，這強調了機器人跨多種任務、環境和機器人系統執行低階視覺運動控制的能力。

舉些例子：GNM（General Navigation Model，通用導航模型）適用于多種不同的機器人導航場景，RoboCat 可針對任務目標操作不同的機器人機體，RT-X 能通過語言操控五種不同的機器人機體。盡管這些模型確實是重要進展，但它們也存在多方面的局限：它們的輸入觀察通常是預定義的且通常很有限（比如單相機輸入視頻流）；它們難以有效地微調至新領域；這些模型中最大型的版本都沒有提供人們使用（這一點很重要）。

近日，加州大學伯克利分校、斯坦福大學、卡內基梅隆大學和谷歌 DeepMind 的 18 位研究者組成的 Octo Model Team 發布了他們的開創性研究成果：Octo 模型。該項目有效地克服了上述局限。

論文標題：Octo: An Open-Source Generalist Robot Policy
論文地址：https://arxiv.org/pdf/2405.12213
開源項目：https://octo-models.github.io/

他們設計了一個系統，能讓 GRP 更輕松地應對下游機器人應用的接口多樣化問題。

該模型的核心是 Transformer 架構，其可將任意輸入 token（根據觀察和任務創建）映射成輸出 token（然后編碼成動作），而且該架構可使用多樣化的機器人和任務數據集進行訓練。該策略無需額外訓練就能接受不同的相機配置，也能控制不同的機器人，還能通過語言命令或目標圖像進行引導 —— 所有這些只需通過改變輸入模型的 token 即可實現。

最重要的是，該模型還能適應傳感器輸入、動作空間或機器人形態不同的新機器人配置，所需的只是采用適當的適配器（adapter）并使用一個小的目標領域數據集和少量計算預算進行微調。

不僅如此，Octo 還已經在迄今為止最大的機器人操控數據集上完成了預訓練 —— 該數據集包含來自 Open X-Embodiment 數據集的 80 萬個機器人演示。Octo 不僅是首個可有效微調至新觀察和動作空間的 GRP，也是首個完全開源（訓練工作流程、模型檢查點和數據）的通才機器人操控策略。該團隊也在論文中強調了其組合 Octo 各組件的獨特創新性。

Octo 模型

下面我們來看看 Octo 這個開源的通才機器人策略是如何構建的。總體而言，Octo 的設計目標是讓其成為一個靈活且廣泛適用的通才機器人策略，可被大量不同的下游機器人應用和研究項目使用。

架構

Octo 的核心是基于 Transformer 的策略 π。其包含三個關鍵部分：輸入 token 化器、Transformer 骨干網絡和讀出頭。

如圖 2 所示，其中輸入 token 化器的作用是將語言指令、目標和觀察序列轉換成 token，Transformer 骨干會把這些 token 處理成嵌入，讀出頭則是得出所需的輸出，即動作。

任務和觀察 token 化器

為了將任務定義（比如語言指令和目標圖像）與觀察（比如相機視頻流）轉換成常用的已 token 化的格式，該團隊針對不同模態使用了不同的 token 化器：

對于語言輸入，先 token 化，然后通過一個預訓練的 Transformer 將其處理成一個語言嵌入 token 序列。具體而言，他們使用的模型是 t5-base (111M)。

對于圖像觀察和目標，則是通過一個較淺的卷積堆棧來處理，然后再拆分成平展后圖塊構成的序列。

最后，通過向任務和觀察 token 添加可學習的位置嵌入并按一定順序排列它們來構建 Transformer 的輸入序列。

Transformer 骨干和讀出頭

將輸入處理成一種統一化的 token 序列之后，就能交給 Transformer 處理了。這與之前的研究工作類似：基于觀察和動作系列來訓練基于 Transformer 的策略。

Octo 的注意力模式是逐塊掩碼式：觀察 token 只能按照因果關系關注來自同一或之前時間步驟的 token 以及任務 token。對應于不存在觀察的 token 會被完全掩蔽掉（比如沒有語言指令的數據集）。這種模塊化設計很方便，可在微調階段添加或移除觀察或任務。

除了這些輸入 token 模塊，該團隊還插入了已學習完成的讀出 token。讀出 token 會關注其之前的觀察和任務 token，但不會被任何觀察或任務 token 關注。因此，讀出 token 只能讀取和處理內部嵌入，而無法影響內部嵌入。讀出 token 的作用類似于 BERT 中的 [CLS] token，充當截至目前的觀察序列的緊湊向量嵌入。針對讀出 token 的嵌入，會使用一個輕量的實現擴散過程的「動作頭」。這個動作頭會預測多個連續動作構成的一個「塊（chunk）」。

這樣的設計可讓用戶在下游微調時向模型靈活地添加新的任務和觀察輸入或動作輸出頭。當在下游添加新的任務、觀察或損失函數時，可以在整體上保留 Transformer 的預訓練權重，僅添加新的位置嵌入、一個新的輕量編碼器、或由于規范變化而必需的新頭的參數。這不同于之前的架構 —— 對于之前的架構，如果添加或移除圖像輸入或改變任務規范，就需要重新初始化或重新訓練預訓練模型的大量組件。

要讓 Octo 成為真正的「通才」模型，這種靈活性至關重要：由于我們不可能在預訓練階段覆蓋所有可能的機器人傳感器和動作配置，因此，如果能在微調階段調整 Octo 的輸入和輸出，便能讓其成為機器人社區的一種多功能工具。另外，之前使用標準 Transformer 骨干或融合使用視覺編碼器與 MLP 輸出頭的模型設計固定了模型輸入的類型和順序。相較之下，切換 Octo 的觀察或任務并不需要對大部分模型進行重新初始化。

訓練數據

該團隊從 Open X-Embodiment 中取用了包含 25 個數據集的混合數據集。圖 3 給出了數據集的組成。

有關訓練目標和訓練硬件配置等更多細節請參閱原論文。

模型檢查點和代碼

重點來了！該團隊不僅發了 Octo 的論文，還完全開源了所有資源，其中包括：

預訓練完成的 Octo 檢查點，包括 2700 萬參數的 Octo-Small 和 9300 萬參數的 Octo-Base。
用于 Octo 模型的微調腳本，基于 JAX。
用于在 Open X-Embodiment 數據集上預訓練 Octo 的模型預訓練工作流程，基于 JAX。用于 Open X-Embodiment 數據的數據加載器，兼容 JAX 和 PyTorch。