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

机器之心發表於2024-05-28

一位優秀的相聲演員需要吹拉彈唱樣樣在行,類似地,一個優秀的機器人模型也應能適應多樣化的機器人形態和不同的任務,但目前大多數機器人模型都只能控制一種形態的機器人執行一類任務。現在 Octo(八爪魚)來了!這個基於 Transformer 的模型堪稱當前最強大的開源機器人學習系統,無需額外訓練就能完成多樣化的機器人操控任務並能在一定程度適應新機器人形態和新任務,就像肢體靈活的八爪魚

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

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

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

為了實現這一目標,已經出現了一些「機器人基礎模型」相關研究成果;它們的做法是直接將機器人觀察對映成動作,然後透過零樣本或少樣本方式泛化至新領域或新機器人。這些模型通常被稱為「通才機器人策略(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。

實驗

該團隊也透過實驗對 Octo 進行了實證分析,在多個維度上評估了其作為機器人基礎模型的效能:

  1. 能否直接使用 Octo 控制多臺機器人機體並解決語言和目標任務?

  2. Octo 權重可否作為優質的初始化基礎,支援面向新任務和機器人的資料高效型微調,以及它是否優於從頭開始訓練的方法和常用的預訓練表徵?

  3. 在構建通才機器人策略方面,Octo 中的哪種設計決策最重要?

圖 4 展示了評估 Octo 的 9 種任務。

圖片

直接使用 Octo 控制多臺機器人

該團隊比較了 Octo、RT-1-X、RT-2-X 的零樣本操控能力,結果見圖 5。

圖片

可以看到,Octo 的成功率比 RT-1-X(3500 萬引數)高 29%。而在 WidowX 和 RT-1 Robot 評估上,Octo 與 550 億引數的 RT-2-X 效能相當。

此外,RT-1-X 和 RT-2-X 僅支援語言指令,而 Octo 還支援以目標影像為條件。該團隊還發現,在 WidowX 任務上,如果使用目標影像為條件,成功率比使用語言為條件高 25%。這可能是因為目標影像能提供更多有關任務完成的資訊。

Octo 能高效地使用資料來適應新領域

表 1 給出了資料高效型微調的實驗結果。

圖片

可以看到,相比於從頭開始訓練或使用預訓練的 VC-1 權重進行預訓練,微調 Octo 得到的結果更好。在 6 種評估設定上,Octo 相較於第二名基準的平均優勢為 52%!

並且不得不提的是:針對所有這些評估任務,微調 Octo 時使用的配方和超引數全都一樣,由此可見該團隊找到了一個非常好的預設配置。

通才機器人策略訓練的設計決策

上面的結果表明 Octo 確實能作為零樣本多機器人控制器,也能作為策略微調的初始化基礎。接下來,該團隊分析了不同設計決策對 Octo 策略效能的影響。具體而言,他們關注的重點是以下方面:模型架構、訓練資料、訓練目標、模型規模。為此,他們進行了消融研究。

表 2 給出了在模型架構、訓練資料和訓練目標的消融研究結果。

圖片

圖 6 則展現了模型規模對零樣本成功率的影響,可以看出來更大的模型有更好的視覺場景感知能力。

圖片

整體而言,Octo 各元件的有效性得到了證明。

相關文章