機器人策略學習的Game Changer?伯克利提出Body Transformer

機器之心報道

編輯:Panda

我們知道,物理智能體在執行動作的校正和穩定時,往往會根據其感受到的外部刺激的位置給出空間上的響應。比如人類對這些刺激的響應迴路位於脊髓神經迴路層面,它們專門負責單個執行器的響應。起校正作用的局部執行是高效運動的主要因素,這對機器人來說也尤爲重要。

但之前的學習架構通常都沒有建立傳感器和執行器之間的空間關聯。鑑於機器人策略使用的架構基本是爲自然語言和計算機視覺開發的架構,它們常常無法有效地利用機器人機體的結構。

不過,Transformer 在這方面還是頗具潛力的,已有研究表明,Transformer 可以有效地處理長序列依賴關係,還能輕鬆地吸收大量數據。Transformer 架構原本是爲非結構化自然語言處理(NLP)任務開發的。在這些任務中(比如語言翻譯),輸入序列通常會被映射到一個輸出序列。

基於這一觀察,加州大學伯克利分校 Pieter Abbeel 教授領導的團隊提出了 Body Transformer(BoT),增加了對機器人機體上的傳感器和執行器的空間位置的關注。

具體來說,BoT 是將機器人機體建模成圖(graph),其中的節點即爲其傳感器和執行器。然後,其在注意力層上使用高度稀疏的掩碼,以防止每個節點關注其直接近鄰之外的部分。將多個結構相同的 BoT 層連接起來,就能彙集整個圖的信息,這樣便不會損害該架構的表徵能力。BoT 在模仿學習和強化學習方面都表現不俗,甚至被一些人認爲是策略學習的「Game Changer」。

Body Transformer

如果機器人學習策略使用原始 Transformer 架構爲骨幹,則通常會忽視機器人機體結構所提供的有用信息。但實際上,這些結構信息能爲 Transformer 提供更強的歸納偏置。該團隊在利用這些信息的同時還保留了原始架構的表徵能力。

Body Transformer(BoT)架構基於掩碼式注意力。在這個架構的每一層中,一個節點都只能看到其自身和其直接近鄰節點的信息。如此一來,信息就會依照圖的結構而流動,其中上游層會根據局部信息執行推理,下游層則能彙集更多來自更遠節點的全局信息。

如圖 1 所示,BoT 架構包含以下組件:

1.tokenizer:將傳感器輸入投射成對應的節點嵌入;

2.Transformer 編碼器:處理輸入嵌入並生成同樣維度的輸出特徵;

3.detokenizer:解除 token 化,即將特徵解碼成動作(或用於強化學習批評訓練的價值)。

tokenizer

該團隊選擇將觀察向量映射成局部觀察構成的圖。

在實踐中,他們將全局量分配給機器人機體的根元素,將局部量分配給表示對應肢體的節點。這種分配方式與之前的 GNN 方法類似。

然後,使用一個線性層將局部狀態向量投射成嵌入向量。每個節點的狀態都會被饋送給其節點特定的可學習的線性投射,從而得到一個包含 n 個嵌入的序列,其中 n 表示節點的數量(或序列長度)。這不同於之前的研究成果,它們通常僅使用單個共享的可學習的線性投射來處理多任務強化學習中不同數量的節點。

BoT 編碼器

該團隊使用的骨幹網絡是一個標準的多層 Transformer 編碼器,並且該架構有兩種變體版本:

detokenizer

Transformer 編碼器輸出的特徵會被饋送給線性層,然後被投射成與該節點的肢體關聯的動作;這些動作是根據相應執行器與肢體的接近程度來分配的。同樣,每個節點的這些可學習的線性投射層是分開的。如果將 BoT 用作強化學習設置中的批評架構,則 detokenizer 輸出的就不再是動作,而是價值,然後在機體部位上取平均值。

實驗

團隊在模仿學習和強化學習設置中評估了 BoT 的性能。他們維持了與圖 1 相同的結構,只用各種基線架構替換 BoT 編碼器,以確定編碼器的效果。

這些實驗的目標是解答以下問題:

模仿學習實驗

團隊在機體跟蹤任務上評估了 BoT 架構的模仿學習性能,該任務是通過 MoCapAct 數據集定義的。

結果如圖 3a 所示,可以看到 BoT 的表現總是優於 MLP 和 Transformer 基線。值得注意的是,在未曾見過的驗證視頻片段上,BoT 相對於這些架構的優勢還會進一步增大,這證明機體感知型歸納偏置能帶來泛化能力的提升。

而圖 3b 則表明 BoT-Hard 的規模擴展性很不錯,相較於 Transformer 基線,其在訓練和驗證視頻片段上的性能都會隨着可訓練參數量的增長而增長這進一步表明 BoT-Hard 傾向於不過擬合訓練數據,而這種過擬合是由具身偏置引起的。下面展示了更多實驗示例,詳見原論文。

強化學習實驗

該團隊在 Isaac Gym 中的 4 個機器人控制任務上評估了 BoT 與使用 PPO 的基線的強化學習性能。這 4 個任務分別是:Humanoid-Mod、Humanoid-Board、Humanoid-Hill 和 A1-Walk。

圖 5 展示了 MLP、Transformer 和 BoT(Hard 和 Mix)在訓練期間的評估 rollout 的平均情節回報。其中,實線對應於平均值,陰影區域對應於五個種子的標準誤差。

結果表明,BoT-Mix 的性能在樣本效率和漸近性能方面始終優於 MLP 和原始 Transformer 基線。這說明將來自機器人機體的偏置整合進策略網絡架構是有用的。

同時,BoT-Hard 在較簡單的任務(A1-Walk 和 Humanoid-Mod)上的表現優於原始 Transformer,但在更困難的探索任務(Humanoid-Board 和 Humanoid-Hill)上表現卻更差。考慮到掩碼式注意力會妨礙來自遠處機體部分的信息傳播,BoT-Hard 在信息通信方面的強大限制可能會妨礙強化學習探索的效率。

真實世界實驗

Isaac Gym 模擬的運動環境常被用於將強化學習策略從虛擬遷移到真實環境,並且還不需要在真實世界中進行調整。爲了驗證新提出的架構是否適用於真實世界應用,該團隊將上述訓練得到的一個 BoT 策略部署到了一臺 Unitree A1 機器人中。從如下視頻可以看出,新架構可以可靠地用於真實世界部署。

計算分析

該團隊也分析了新架構的計算成本,如圖 6 所示。這裡給出了新提出的掩碼式注意力與常規注意力在不同序列長度(節點數量)上的規模擴展結果。

可以看到,當有 128 個節點時(相當於擁有靈巧雙臂的類人機器人),新注意力能將速度提升 206%。

總體而言,這表明 BoT 架構中的源自機體的偏置不僅能提高物理智能體的整體性能,而且還可受益於架構那自然稀疏的掩碼。該方法可通過充分的並行化來大幅減少學習算法的訓練時間。