適應多形態(tài)多任務,最強開源機器人學習系統(tǒng)「八爪魚」誕生
在機器人學習方面,一種常用方法是收集針對特定機器人和任務的數(shù)據(jù)集,然后用其來訓練策略。但是,如果使用這種方法來從頭開始學習,每一個任務都需要收集足夠數(shù)據(jù),并且所得策略的泛化能力通常也不佳。
原理上講,從其它機器人和任務收集的經(jīng)驗能提供可能的解決方案,能讓模型看到多種多樣的機器人控制問題,而這些問題也許能提升機器人在下游任務上的泛化能力和性能。但是,即便現(xiàn)在已經(jīng)出現(xiàn)了能處理多種自然語言和計算機視覺任務的通用模型,構(gòu)建「通用機器人模型」依然困難重重。
要為機器人訓練一個統(tǒng)一的控制策略非常困難,其中涉及諸多難點,包括操作不同的機器人機體、傳感器配置、動作空間、任務規(guī)范、環(huán)境和計算預算。
為了實現(xiàn)這一目標,已經(jīng)出現(xiàn)了一些「機器人基礎模型」相關(guān)研究成果;它們的做法是直接將機器人觀察映射成動作,然后通過零樣本或少樣本方式泛化至新領(lǐng)域或新機器人。這些模型通常被稱為「通才機器人策略(generalist robot policy)」,簡稱 GRP,這強調(diào)了機器人跨多種任務、環(huán)境和機器人系統(tǒng)執(zhí)行低階視覺運動控制的能力。
舉些例子:GNM(General Navigation Model,通用導航模型) 適用于多種不同的機器人導航場景,RoboCat 可針對任務目標操作不同的機器人機體,RT-X 能通過語言操控五種不同的機器人機體。盡管這些模型確實是重要進展,但它們也存在多方面的局限:它們的輸入觀察通常是預定義的且通常很有限(比如單相機輸入視頻流);它們難以有效地微調(diào)至新領(lǐng)域;這些模型中最大型的版本都沒有提供人們使用(這一點很重要)。
近日,加州大學伯克利分校、斯坦福大學、卡內(nèi)基梅隆大學和谷歌 DeepMind 的 18 位研究者組成的 Octo Model Team 發(fā)布了他們的開創(chuàng)性研究成果:Octo 模型。該項目有效地克服了上述局限。
- 論文標題:Octo: An Open-Source Generalist Robot Policy
- 論文地址:https://arxiv.org/pdf/2405.12213
- 開源項目:https://octo-models.github.io/
他們設計了一個系統(tǒng),能讓 GRP 更輕松地應對下游機器人應用的接口多樣化問題。
該模型的核心是 Transformer 架構(gòu),其可將任意輸入 token(根據(jù)觀察和任務創(chuàng)建)映射成輸出 token(然后編碼成動作),而且該架構(gòu)可使用多樣化的機器人和任務數(shù)據(jù)集進行訓練。該策略無需額外訓練就能接受不同的相機配置,也能控制不同的機器人,還能通過語言命令或目標圖像進行引導 —— 所有這些只需通過改變輸入模型的 token 即可實現(xiàn)。
最重要的是,該模型還能適應傳感器輸入、動作空間或機器人形態(tài)不同的新機器人配置,所需的只是采用適當?shù)倪m配器(adapter)并使用一個小的目標領(lǐng)域數(shù)據(jù)集和少量計算預算進行微調(diào)。
不僅如此,Octo 還已經(jīng)在迄今為止最大的機器人操控數(shù)據(jù)集上完成了預訓練 —— 該數(shù)據(jù)集包含來自 Open X-Embodiment 數(shù)據(jù)集的 80 萬個機器人演示。Octo 不僅是首個可有效微調(diào)至新觀察和動作空間的 GRP,也是首個完全開源(訓練工作流程、模型檢查點和數(shù)據(jù))的通才機器人操控策略。該團隊也在論文中強調(diào)了其組合 Octo 各組件的獨特創(chuàng)新性。
Octo 模型
下面我們來看看 Octo 這個開源的通才機器人策略是如何構(gòu)建的。總體而言,Octo 的設計目標是讓其成為一個靈活且廣泛適用的通才機器人策略,可被大量不同的下游機器人應用和研究項目使用。
架構(gòu)
Octo 的核心是基于 Transformer 的策略 π。其包含三個關(guān)鍵部分:輸入 token 化器、Transformer 骨干網(wǎng)絡和讀出頭。
如圖 2 所示,其中輸入 token 化器的作用是將語言指令、目標和觀察序列轉(zhuǎn)換成 token,Transformer 骨干會把這些 token 處理成嵌入,讀出頭則是得出所需的輸出,即動作。
任務和觀察 token 化器
為了將任務定義(比如語言指令和目標圖像)與觀察(比如相機視頻流)轉(zhuǎn)換成常用的已 token 化的格式,該團隊針對不同模態(tài)使用了不同的 token 化器:
對于語言輸入,先 token 化,然后通過一個預訓練的 Transformer 將其處理成一個語言嵌入 token 序列。具體而言,他們使用的模型是 t5-base (111M)。
對于圖像觀察和目標,則是通過一個較淺的卷積堆棧來處理,然后再拆分成平展后圖塊構(gòu)成的序列。
最后,通過向任務和觀察 token 添加可學習的位置嵌入并按一定順序排列它們來構(gòu)建 Transformer 的輸入序列。
Transformer 骨干和讀出頭
將輸入處理成一種統(tǒng)一化的 token 序列之后,就能交給 Transformer 處理了。這與之前的研究工作類似:基于觀察和動作系列來訓練基于 Transformer 的策略。
Octo 的注意力模式是逐塊掩碼式:觀察 token 只能按照因果關(guān)系關(guān)注來自同一或之前時間步驟的 token 以及任務 token。對應于不存在觀察的 token 會被完全掩蔽掉(比如沒有語言指令的數(shù)據(jù)集)。這種模塊化設計很方便,可在微調(diào)階段添加或移除觀察或任務。
除了這些輸入 token 模塊,該團隊還插入了已學習完成的讀出 token。讀出 token 會關(guān)注其之前的觀察和任務 token,但不會被任何觀察或任務 token 關(guān)注。因此,讀出 token 只能讀取和處理內(nèi)部嵌入,而無法影響內(nèi)部嵌入。讀出 token 的作用類似于 BERT 中的 [CLS] token,充當截至目前的觀察序列的緊湊向量嵌入。針對讀出 token 的嵌入,會使用一個輕量的實現(xiàn)擴散過程的「動作頭」。這個動作頭會預測多個連續(xù)動作構(gòu)成的一個「塊(chunk)」。
這樣的設計可讓用戶在下游微調(diào)時向模型靈活地添加新的任務和觀察輸入或動作輸出頭。當在下游添加新的任務、觀察或損失函數(shù)時,可以在整體上保留 Transformer 的預訓練權(quán)重,僅添加新的位置嵌入、一個新的輕量編碼器、或由于規(guī)范變化而必需的新頭的參數(shù)。這不同于之前的架構(gòu) —— 對于之前的架構(gòu),如果添加或移除圖像輸入或改變?nèi)蝿找?guī)范,就需要重新初始化或重新訓練預訓練模型的大量組件。
要讓 Octo 成為真正的「通才」模型,這種靈活性至關(guān)重要:由于我們不可能在預訓練階段覆蓋所有可能的機器人傳感器和動作配置,因此,如果能在微調(diào)階段調(diào)整 Octo 的輸入和輸出,便能讓其成為機器人社區(qū)的一種多功能工具。另外,之前使用標準 Transformer 骨干或融合使用視覺編碼器與 MLP 輸出頭的模型設計固定了模型輸入的類型和順序。相較之下,切換 Octo 的觀察或任務并不需要對大部分模型進行重新初始化。
訓練數(shù)據(jù)
該團隊從 Open X-Embodiment 中取用了包含 25 個數(shù)據(jù)集的混合數(shù)據(jù)集。圖 3 給出了數(shù)據(jù)集的組成。
有關(guān)訓練目標和訓練硬件配置等更多細節(jié)請參閱原論文。
模型檢查點和代碼
重點來了!該團隊不僅發(fā)了 Octo 的論文,還完全開源了所有資源,其中包括:
- 預訓練完成的 Octo 檢查點,包括 2700 萬參數(shù)的 Octo-Small 和 9300 萬參數(shù)的 Octo-Base。
- 用于 Octo 模型的微調(diào)腳本,基于 JAX。
- 用于在 Open X-Embodiment 數(shù)據(jù)集上預訓練 Octo 的模型預訓練工作流程,基于 JAX。用于 Open X-Embodiment 數(shù)據(jù)的數(shù)據(jù)加載器,兼容 JAX 和 PyTorch。
實驗
該團隊也通過實驗對 Octo 進行了實證分析,在多個維度上評估了其作為機器人基礎模型的性能:
- 能否直接使用 Octo 控制多臺機器人機體并解決語言和目標任務?
- Octo 權(quán)重可否作為優(yōu)質(zhì)的初始化基礎,支持面向新任務和機器人的數(shù)據(jù)高效型微調(diào),以及它是否優(yōu)于從頭開始訓練的方法和常用的預訓練表征?
- 在構(gòu)建通才機器人策略方面,Octo 中的哪種設計決策最重要?
圖 4 展示了評估 Octo 的 9 種任務。
直接使用 Octo 控制多臺機器人
該團隊比較了 Octo、RT-1-X、RT-2-X 的零樣本操控能力,結(jié)果見圖 5。
可以看到,Octo 的成功率比 RT-1-X(3500 萬參數(shù))高 29%。而在 WidowX 和 RT-1 Robot 評估上,Octo 與 550 億參數(shù)的 RT-2-X 性能相當。
此外,RT-1-X 和 RT-2-X 僅支持語言指令,而 Octo 還支持以目標圖像為條件。該團隊還發(fā)現(xiàn),在 WidowX 任務上,如果使用目標圖像為條件,成功率比使用語言為條件高 25%。這可能是因為目標圖像能提供更多有關(guān)任務完成的信息。
Octo 能高效地使用數(shù)據(jù)來適應新領(lǐng)域
表 1 給出了數(shù)據(jù)高效型微調(diào)的實驗結(jié)果。
可以看到,相比于從頭開始訓練或使用預訓練的 VC-1 權(quán)重進行預訓練,微調(diào) Octo 得到的結(jié)果更好。在 6 種評估設置上,Octo 相較于第二名基準的平均優(yōu)勢為 52%!
并且不得不提的是:針對所有這些評估任務,微調(diào) Octo 時使用的配方和超參數(shù)全都一樣,由此可見該團隊找到了一個非常好的默認配置。
通才機器人策略訓練的設計決策
上面的結(jié)果表明 Octo 確實能作為零樣本多機器人控制器,也能作為策略微調(diào)的初始化基礎。接下來,該團隊分析了不同設計決策對 Octo 策略性能的影響。具體而言,他們關(guān)注的重點是以下方面:模型架構(gòu)、訓練數(shù)據(jù)、訓練目標、模型規(guī)模。為此,他們進行了消融研究。
表 2 給出了在模型架構(gòu)、訓練數(shù)據(jù)和訓練目標的消融研究結(jié)果。
圖 6 則展現(xiàn)了模型規(guī)模對零樣本成功率的影響,可以看出來更大的模型有更好的視覺場景感知能力。
整體而言,Octo 各組件的有效性得到了證明。