機器人策略學習的Game Changer？伯克利提出Body Transformer
我們知道，物理智能體在執(zhí)行動作的校正和穩(wěn)定時，往往會根據(jù)其感受到的外部刺激的位置給出空間上的響應。比如人類對這些刺激的響應回路位于脊髓神經(jīng)回路層面，它們專門負責單個執(zhí)行器的響應。起校正作用的局部執(zhí)行是高效運動的主要因素，這對機器人來說也尤為重要。
但之前的學習架構通常都沒有建立傳感器和執(zhí)行器之間的空間關聯(lián)。鑒于機器人策略使用的架構基本是為自然語言和計算機視覺開發(fā)的架構，它們常常無法有效地利用機器人機體的結構。
不過，Transformer 在這方面還是頗具潛力的，已有研究表明，Transformer 可以有效地處理長序列依賴關系，還能輕松地吸收大量數(shù)據(jù)。Transformer 架構原本是為非結構化自然語言處理（NLP）任務開發(fā)的。在這些任務中（比如語言翻譯），輸入序列通常會被映射到一個輸出序列。
基于這一觀察，加州大學伯克利分校 Pieter Abbeel 教授領導的團隊提出了 Body Transformer（BoT），增加了對機器人機體上的傳感器和執(zhí)行器的空間位置的關注。
- 論文標題：Body Transformer: Leveraging Robot Embodiment for Policy Learning
- 論文地址：https://arxiv.org/pdf/2408.06316v1
- 項目網(wǎng)站：https://sferrazza.cc/bot_site
- 代碼地址：https://github.com/carlosferrazza/BodyTransformer
具體來說，BoT 是將機器人機體建模成圖（graph），其中的節(jié)點即為其傳感器和執(zhí)行器。然后，其在注意力層上使用高度稀疏的掩碼，以防止每個節(jié)點關注其直接近鄰之外的部分。將多個結構相同的 BoT 層連接起來，就能匯集整個圖的信息，這樣便不會損害該架構的表征能力。BoT 在模仿學習和強化學習方面都表現(xiàn)不俗，甚至被一些人認為是策略學習的「Game Changer」。
Body Transformer
如果機器人學習策略使用原始 Transformer 架構為骨干，則通常會忽視機器人機體結構所提供的有用信息。但實際上，這些結構信息能為 Transformer 提供更強的歸納偏置。該團隊在利用這些信息的同時還保留了原始架構的表征能力。
Body Transformer（BoT）架構基于掩碼式注意力。在這個架構的每一層中，一個節(jié)點都只能看到其自身和其直接近鄰節(jié)點的信息。如此一來，信息就會依照圖的結構而流動，其中上游層會根據(jù)局部信息執(zhí)行推理，下游層則能匯集更多來自更遠節(jié)點的全局信息。
如圖 1 所示，BoT 架構包含以下組件：
1.tokenizer：將傳感器輸入投射成對應的節(jié)點嵌入；
2.Transformer 編碼器：處理輸入嵌入并生成同樣維度的輸出特征；
3.detokenizer：解除 token 化，即將特征解碼成動作（或用于強化學習批評訓練的價值）。
tokenizer
該團隊選擇將觀察向量映射成局部觀察構成的圖。
在實踐中，他們將全局量分配給機器人機體的根元素，將局部量分配給表示對應肢體的節(jié)點。這種分配方式與之前的 GNN 方法類似。
然后，使用一個線性層將局部狀態(tài)向量投射成嵌入向量。每個節(jié)點的狀態(tài)都會被饋送給其節(jié)點特定的可學習的線性投射，從而得到一個包含 n 個嵌入的序列，其中 n 表示節(jié)點的數(shù)量（或序列長度）。這不同于之前的研究成果，它們通常僅使用單個共享的可學習的線性投射來處理多任務強化學習中不同數(shù)量的節(jié)點。
BoT 編碼器
該團隊使用的骨干網(wǎng)絡是一個標準的多層 Transformer 編碼器，并且該架構有兩種變體版本：
- BoT-Hard：使用一個反映該圖結構的二元掩碼來掩蔽每一層。具體來說，他們構建掩碼的方式是 M = I_n + A，其中 I_n 是 n 維單位矩陣，A 是對應于該圖的鄰接矩陣。圖 2 展示了一個示例。這讓每個節(jié)點僅能看到其自身和其直接近鄰，并且能為該問題引入相當可觀的稀疏性 —— 從計算成本角度看，這特別有吸引力。
- BoT-Mix：將帶有掩碼式注意力的層（如 BoT-Hard 一樣）與帶有無掩碼式注意力的層交織在一起。
detokenizer
Transformer 編碼器輸出的特征會被饋送給線性層，然后被投射成與該節(jié)點的肢體關聯(lián)的動作；這些動作是根據(jù)相應執(zhí)行器與肢體的接近程度來分配的。同樣，每個節(jié)點的這些可學習的線性投射層是分開的。如果將 BoT 用作強化學習設置中的批評架構，則 detokenizer 輸出的就不再是動作，而是價值，然后在機體部位上取平均值。
實驗
團隊在模仿學習和強化學習設置中評估了 BoT 的性能。他們維持了與圖 1 相同的結構，只用各種基線架構替換 BoT 編碼器，以確定編碼器的效果。
這些實驗的目標是解答以下問題：
- 掩碼式注意力是否能提升模仿學習的性能和泛化能力？
- 相比于原始的 Transformer 架構，BoT 是否能表現(xiàn)出正面的規(guī)模擴展趨勢？
- BoT 是否與強化學習框架兼容，有哪些合理設計選擇可以盡可能地提升性能？
- BoT 策略是否可以應用于真實世界機器人任務？
- 掩碼式注意力在計算方面有哪些優(yōu)勢？
模仿學習實驗
團隊在機體跟蹤任務上評估了 BoT 架構的模仿學習性能，該任務是通過 MoCapAct 數(shù)據(jù)集定義的。
結果如圖 3a 所示，可以看到 BoT 的表現(xiàn)總是優(yōu)于 MLP 和 Transformer 基線。值得注意的是，在未曾見過的驗證視頻片段上，BoT 相對于這些架構的優(yōu)勢還會進一步增大，這證明機體感知型歸納偏置能帶來泛化能力的提升。
而圖 3b 則表明 BoT-Hard 的規(guī)模擴展性很不錯，相較于 Transformer 基線，其在訓練和驗證視頻片段上的性能都會隨著可訓練參數(shù)量的增長而增長這進一步表明 BoT-Hard 傾向于不過擬合訓練數(shù)據(jù)，而這種過擬合是由具身偏置引起的。下面展示了更多實驗示例，詳見原論文。
強化學習實驗
該團隊在 Isaac Gym 中的 4 個機器人控制任務上評估了 BoT 與使用 PPO 的基線的強化學習性能。這 4 個任務分別是：Humanoid-Mod、Humanoid-Board、Humanoid-Hill 和 A1-Walk。
圖 5 展示了 MLP、Transformer 和 BoT（Hard 和 Mix）在訓練期間的評估 rollout 的平均情節(jié)回報。其中，實線對應于平均值，陰影區(qū)域對應于五個種子的標準誤差。
結果表明，BoT-Mix 的性能在樣本效率和漸近性能方面始終優(yōu)于 MLP 和原始 Transformer 基線。這說明將來自機器人機體的偏置整合進策略網(wǎng)絡架構是有用的。
同時，BoT-Hard 在較簡單的任務（A1-Walk 和 Humanoid-Mod）上的表現(xiàn)優(yōu)于原始 Transformer，但在更困難的探索任務（Humanoid-Board 和 Humanoid-Hill）上表現(xiàn)卻更差?？紤]到掩碼式注意力會妨礙來自遠處機體部分的信息傳播，BoT-Hard 在信息通信方面的強大限制可能會妨礙強化學習探索的效率。
真實世界實驗
Isaac Gym 模擬的運動環(huán)境常被用于將強化學習策略從虛擬遷移到真實環(huán)境，并且還不需要在真實世界中進行調整。為了驗證新提出的架構是否適用于真實世界應用，該團隊將上述訓練得到的一個 BoT 策略部署到了一臺 Unitree A1 機器人中。從如下視頻可以看出，新架構可以可靠地用于真實世界部署。
計算分析
該團隊也分析了新架構的計算成本，如圖 6 所示。這里給出了新提出的掩碼式注意力與常規(guī)注意力在不同序列長度（節(jié)點數(shù)量）上的規(guī)模擴展結果。
可以看到，當有 128 個節(jié)點時（相當于擁有靈巧雙臂的類人機器人），新注意力能將速度提升 206%。
總體而言，這表明 BoT 架構中的源自機體的偏置不僅能提高物理智能體的整體性能，而且還可受益于架構那自然稀疏的掩碼。該方法可通過充分的并行化來大幅減少學習算法的訓練時間。
