單GPU實現20Hz線上決策，最新基於序列生成模型的高效軌跡規劃方法解讀

之前我們介紹了基於 Transformer 和擴散模型（Diffussion Model）的序列建模（sequence modelling）方法在強化學習，特別是離線連續控制領域的應用。這其中 Trajectory Transformer（TT）和 Diffusser 屬於基於模型的規劃型演算法，它們展現出了非常高精度的軌跡預測已經很好的靈活性，但是相對來說決策延遲也比較高。尤其是 TT 將每個維度獨立離散化作為序列中的符號，這使得整個序列變得非常長，而且序列生成的耗時會隨著狀態和動作的維度提升而快速升高。

為了讓軌跡生成模型能被達到實用級別的決策速度，我們在和Diffusser 平行（有重疊但是應該稍晚）的時候開始了高效軌跡生成與決策的項目。我們首先想到的是用連續空間內的 Transformer Mixture of Gaussian 而非離散分佈來擬合整個軌跡分佈。雖然不排除實作上的問題，但這種思路下我們沒能得到一個比較穩定的生成模型。隨後我們嘗試了 Variational Autoencoder(VAE)，並且取得了一定的突破。不過 VAE 的重建（reconstruction）精度並不是特別理想，使得下游的控製表現和 TT 相差比較大。在幾輪迭代之後，我們最終選定了VQ-VAE 作為軌跡生成的基礎模型，最終得到了一個能高效採樣和規劃，並且在高維度控制任務上表現遠超其它基於模型方法的新演算法，我們稱為Trajectory Autoencoding Planner（TAP）。

• 計畫首頁：https://sites.google .com/view/latentplan
• 論文首頁：https://arxiv.org/abs/2208.10291

#規劃效率與高維度下的表現

在單一GPU 下，TAP 能輕鬆以20Hz 的決策效率進行線上決策，在低維度的D4RL 任務中下決策延遲只有TT 的1% 左右。更重要的是隨著任務狀態和動作維度D 的增加，TT 的理論決策延遲會以三次方增長，Diffusser 理論上會線性增長 ，而TAP 的決策速度則不受維度影響 。而在智能體的決策表現方面，隨著動作維度增高，TAP 相對於其它方法的表現出現了提升，相對於基於模型方法（如 TT）的提升尤為明顯。

決策延遲對決策和控制任務的重要性是非常明顯的，像MuZero 這樣的演算法雖然在模擬環境中表現優異，但是面對現實世界中需要即時快速回應的任務，過高的決策延遲就會成為它部署的一大困難。此外，在擁有模擬環境的前提下，決策速度慢也會導致類似的演算法的測試成本偏高，同時運用在線上強化學習的成本也會比較高。

此外，我們認為讓序列產生建模方法能順利擴展到維度較高的任務上也是 TAP 一個很重要的貢獻。現實世界中我們希望強化學習能最終解決的問題其實大都有較高的狀態和動作維度。例如對於自動駕駛來說，各路感測器的輸入即使經過各種感知層面的預處理也不太可能小於 100。複雜的機器人控制往往也有很高的動作空間，人類的所有關節自由度大概是 240 左右，也就對應了至少 240 維的動作空間，一個和人一樣靈活的機器人也需要同樣高維的動作空間。

#四組維度逐漸升高的任務

#

決策延遲和相對模型表現隨著任務維度增長的變化

##方法概述

#首先，訓練VQ-VAE 的自我編碼器(autoencoders) 部分，這裡和原本VQ-VAE 有兩個不同。第一個不同是編碼器和解碼器都是基於 Causal Transformer，而不是 CNN。第二個不同則是我們學習了一個條件機率分佈，被建模的可能的軌跡都必須從當前狀態 出發。自編碼器學習一個從目前狀態 開始的軌跡和隱編碼（latent codes）之間的雙向映射。這些隱編碼和原本軌跡一樣按時間順序排列，每個隱編碼會被對應到實際 步驟軌跡。因為我們使用了Causal Transformer，時間排位靠後的隱編碼(如 ) 不會將訊息傳到排位靠前的序列(如 # )，這使得TAP 可以透過前 N 個隱編碼部分解碼出長度為 NL 的軌跡，這在後續用它進行規劃時是非常有用的。

接著我們會用另一個GPT-2 式的Transformer 來建模這些隱編碼的條件機率分佈 ：

#在決策時，我們可以透過在隱變數空間內進行最佳化尋找最好的未來軌跡，而不是在原始動作空間進行優化。一個非常簡單但也能起效的方法就是直接從隱編碼的分佈中進行取樣，然後選取表現最好的軌跡，如下圖：

在選擇最優軌跡的時候參考的目標分數(objective score) 會同時考慮軌跡的預期收益（獎勵加最後一步的估值）和這條軌跡本身的可行性或者說機率。如以下這個公式，其中 是一個遠大於最高return 的數，當軌跡的機率高於一個閾值 ，評判這條軌跡的標準會是它的預期收益（紅色高亮），否則這條軌跡的機率本身將會是主導部分（藍色高亮）。也就是說 TAP 會在大於閾值的軌跡中選擇一條預期收益最高的。

#

雖然採樣數量足夠大了以後直接採樣在預測序列較短的時候效果也可以很好，在限制採樣數量和規劃所需的總時間的前提下，用更好的優化器還是會帶來更好的表現。以下兩個動圖顯示了當要預測未來 144 步驟以後用直接取樣和 beam search 產生的軌跡的差異。這些軌跡被依照最後的目標分數排序，在最上層最前面的軌跡分數越高，在被疊在後面的軌跡分數越低。另外分數低的軌跡同時透明度也會越低。

在圖中我們可以看到直接採樣產生的軌跡很多動態都不太穩定也不太符合物理規律，尤其是背景裡比較淡的軌跡幾乎都是飄著走的。這些都是機率比較低的軌跡，在最終選擇方案的時候會被剔除。在最前排的軌跡看起來動態要更真實一些，但是相對應的表現就比較差，似乎要跌倒了。而相較之下，beam search 在展開下一個隱變數的時候就會動態考慮軌跡的機率，使得機率很低的分支會被提前終結，這樣產生的候選軌跡就都集中在表現較好而且可能性比較大的軌跡周圍了。

#直接取樣

##Beam search

實驗結果

在沒有更高級的估值以及策略提升的情況下，僅依靠預測精度的優勢，在低維度的任務上TAP 就取得了和其它離線強化學習相當的表現：

gym locomotion control

在高維度的任務上，TAP 取得了遠超其它基於模型的方法的表現，同時也勝過了常見的無模型方法。這裡其實有兩個還未有解答的開放性問題。首先是為什麼先前基於模型的方法在這些高維度的離線強化學習任務中表現較差，其次是為什麼 TAP 在這些任務上表現又能反超很多無模型方法。我們的一個假設是因為在高維度問題上進行策略最佳化又要考慮防止策略偏離行為策略太多是非常困難的。當學習了一個模型，模型本身的誤差可能還會放大這種困難。而 TAP 將最佳化空間搬到了一個很小的離散隱變數空間，這使得整個最佳化過程的穩健性更強了。

adroit robotic hand control

一些切片研究######## ##對於TAP 裡面的諸多設計，我們也在gym locomotion control 的任務上做了一系列切片研究。首先是每個隱編碼實際對應的軌跡的步數#########（黃色柱狀圖），事實證明讓一個隱變數對應多步驟狀態轉移不光有計算上的優勢，在最後模型表現上也有提升。透過調節搜尋的目標函數中觸發低機率軌跡懲罰的閾值 #########（紅色長條圖），我們也確認了目標函數中兩個部分確實都對模型最後表現是有幫助的。另外一點就是向未來規劃的步數（planning horizo​​n，藍色長條圖）對模型表現的影響反而不大，在部署後的搜尋中哪怕只展開一個隱變量最後智能體的表現也只會降低10%左右。 ######

最後我們嘗試了在直接採樣的情況下 TAP 的表現（綠色長條圖）。注意這裡的採樣的樣本數量是 2048 而上面的動圖裡只有 256，而且上面的動圖是生成了未來 144 步的規劃，但是實際我們的基礎模型指揮規劃 15 步。結論是直接採樣在樣本數量足夠的情況下，且規劃路徑不長，那麼直接採樣也能獲得和 beam search 相近的表現。但這是從學到的隱變量條件分佈中採樣的情況，如果直接從隱編碼中直接等機率採樣，那最後還是會比完整的 TAP 模型差很多。

#切片研究的結果

以上是單GPU實現20Hz線上決策，最新基於序列生成模型的高效軌跡規劃方法解讀的詳細內容。更多資訊請關注PHP中文網其他相關文章！