過去幾年間，Transformer 架構已經取得了巨大的成功，同時其也衍生出了大量變體，比如擅長處理視覺任務的 Vision Transformer（ViT）。本文要介紹的 Body Transformer（BoT） 則是非常適合機器人策略學習的 Transformer 變體。
我們知道，物理智慧體在執行動作的校正和穩定時，往往會根據其感受到的外部刺激的位置給出空間上的響應。比如人類對這些刺激的響應迴路位於脊髓神經迴路層面，它們專門負責單個執行器的響應。起校正作用的區域性執行是高效運動的主要因素，這對機器人來說也尤為重要。
但之前的學習架構通常都沒有建立感測器和執行器之間的空間關聯。鑑於機器人策略使用的架構基本是為自然語言和計算機視覺開發的架構，它們常常無法有效地利用機器人機體的結構。
不過，Transformer 在這方面還是頗具潛力的，已有研究表明，Transformer 可以有效地處理長序列依賴關係，還能輕鬆地吸收大量資料。Transformer 架構原本是為非結構化自然語言處理（NLP）任務開發的。在這些任務中（比如語言翻譯），輸入序列通常會被對映到一個輸出序列。
基於這一觀察，加州大學伯克利分校 Pieter Abbeel 教授領導的團隊提出了 Body Transformer（BoT），增加了對機器人機體上的感測器和執行器的空間位置的關注。
論文標題：Body Transformer: Leveraging Robot Embodiment for Policy Learning
論文地址：https://arxiv.org/pdf/2408.06316v1
專案網站：https://sferrazza.cc/bot_site
程式碼地址：https://github.com/carlosferrazza/BodyTransformer
具體來說，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 編碼器，並且該架構有兩種變體版本：
BoT-Hard：使用一個反映該圖結構的二元掩碼來掩蔽每一層。具體來說，他們構建掩碼的方式是 M = I_n + A，其中 I_n 是 n 維單位矩陣，A 是對應於該圖的鄰接矩陣。圖 2 展示了一個示例。這讓每個節點僅能看到其自身和其直接近鄰，並且能為該問題引入相當可觀的稀疏性 —— 從計算成本角度看，這特別有吸引力。
BoT-Mix：將帶有掩碼式注意力的層（如 BoT-Hard 一樣）與帶有無掩碼式注意力的層交織在一起。
detokenizer
Transformer 編碼器輸出的特徵會被饋送給線性層，然後被投射成與該節點的肢體關聯的動作；這些動作是根據相應執行器與肢體的接近程度來分配的。同樣，每個節點的這些可學習的線性投射層是分開的。如果將 BoT 用作強化學習設定中的批評架構，則 detokenizer 輸出的就不再是動作，而是價值，然後在機體部位上取平均值。
實驗
團隊在模仿學習和強化學習設定中評估了 BoT 的效能。他們維持了與圖 1 相同的結構，只用各種基線架構替換 BoT 編碼器，以確定編碼器的效果。
這些實驗的目標是解答以下問題：
掩碼式注意力是否能提升模仿學習的效能和泛化能力？
相比於原始的 Transformer 架構，BoT 是否能表現出正面的規模擴充套件趨勢？
BoT 是否與強化學習框架相容，有哪些合理設計選擇可以儘可能地提升效能？
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 架構中的源自機體的偏置不僅能提高物理智慧體的整體效能，而且還可受益於架構那自然稀疏的掩碼。該方法可透過充分的並行化來大幅減少學習演算法的訓練時間。