之前我們介紹了基于 Transformer 和擴(kuò)散模型（Diffussion Model）的序列建模（sequence modelling）方法在強(qiáng)化學(xué)習(xí)，特別是離線連續(xù)控制領(lǐng)域的應(yīng)用。這其中 Trajectory Transformer（TT）和 Diffusser 屬于基于模型的規(guī)劃型算法，它們展現(xiàn)出了非常高精度的軌跡預(yù)測(cè)已經(jīng)很好的靈活性，但是相對(duì)來(lái)說(shuō)決策延遲也比較高。尤其是 TT 將每個(gè)維度獨(dú)立離散化作為序列中的符號(hào)，這使得整個(gè)序列變得非常長(zhǎng)，而且序列生成的耗時(shí)會(huì)隨著狀態(tài)和動(dòng)作的維度提升快速升高。

為了讓軌跡生成模型能被達(dá)到實(shí)用級(jí)別的決策速度，我們?cè)诤?Diffusser 平行（有重疊但是應(yīng)該稍晚）的時(shí)候開(kāi)始了高效軌跡生成與決策的項(xiàng)目。我們首先想到的是用連續(xù)空間內(nèi)的 Transformer+Mixture of Gaussian 而非離散分布來(lái)擬合整個(gè)軌跡分布。雖然不排除實(shí)現(xiàn)上的問(wèn)題，但是這種思路下我們沒(méi)能獲得一個(gè)比較穩(wěn)定的生成模型。隨后我們嘗試了 Variational Autoencoder(VAE)，并且取得了一定的突破。不過(guò) VAE 的重建（reconstruction）精度不是特別理想，使得下游的控制表現(xiàn)和 TT 相差比較大。在幾輪迭代之后，我們最終選定了 VQ-VAE 作為軌跡生成的基礎(chǔ)模型，最終得到了一個(gè)能高效采樣和規(guī)劃，并且在高維度控制任務(wù)上表現(xiàn)遠(yuǎn)超其它基于模型方法的新算法，我們稱為 Trajectory Autoencoding Planner（TAP）。

?

• 項(xiàng)目主頁(yè)：https://sites.google.com/view/latentplan
• 論文主頁(yè)：https://arxiv.org/abs/2208.10291

規(guī)劃效率與高維下的表現(xiàn)

在單個(gè) GPU 下，TAP 能輕松以 20Hz 的決策效率進(jìn)行在線決策，在低維度的 D4RL 任務(wù)中下決策延遲只有 TT 的 1% 左右。更重要的是隨著任務(wù)狀態(tài)和動(dòng)作維度 D 的增加，TT 的理論決策延遲會(huì)以三次方增長(zhǎng)，Diffusser 理論上會(huì)線性增長(zhǎng)  ，而 TAP 的決策速度則不受維度影響 。而在智能體的決策表現(xiàn)方面，隨著動(dòng)作維度增高，TAP 相對(duì)于其它方法的表現(xiàn)出現(xiàn)了提升，相對(duì)于基于模型方法（如 TT）的提升尤為明顯。

決策延遲對(duì)決策和控制任務(wù)的重要性是非常明顯的，像 MuZero 這樣的算法雖然在模擬環(huán)境中表現(xiàn)優(yōu)異，但是面對(duì)現(xiàn)實(shí)世界中需要實(shí)時(shí)快速響應(yīng)的任務(wù)，過(guò)高的決策延遲就會(huì)成為它部署的一大困難。此外，在擁有模擬環(huán)境的前提下，決策速度慢也會(huì)導(dǎo)致類似的算法的測(cè)試成本偏高，同時(shí)被運(yùn)用在在線強(qiáng)化學(xué)習(xí)中的成本也會(huì)比較高。

此外，我們認(rèn)為讓序列生成建模方法能順利擴(kuò)展到維度較高的任務(wù)上也是 TAP 一個(gè)很重要的貢獻(xiàn)。現(xiàn)實(shí)世界中我們希望強(qiáng)化學(xué)習(xí)能最終解決的問(wèn)題其實(shí)大都有較高的狀態(tài)和動(dòng)作維度。比如對(duì)于自動(dòng)駕駛來(lái)說(shuō)，各路傳感器的輸入哪怕經(jīng)過(guò)各種感知層面的預(yù)處理也不太可能小于 100。復(fù)雜的機(jī)器人控制往往也有很高的動(dòng)作空間，人類的所有關(guān)節(jié)自由度大概是 240 左右，也就對(duì)應(yīng)了至少 240 維的動(dòng)作空間，一個(gè)和人一樣靈活的機(jī)器人也需要同樣高維的動(dòng)作空間。

四組維度逐漸升高的任務(wù)

決策延遲和相對(duì)模型表現(xiàn)隨著任務(wù)維度增長(zhǎng)的變化

方法概述

首先，訓(xùn)練 VQ-VAE 的自編碼器 (autoencoders) 部分，這里和原本 VQ-VAE 有兩個(gè)不同。第一個(gè)不同是編碼器和解碼器都是基于 Causal Transformer，而不是 CNN。第二個(gè)不同則是我們學(xué)習(xí)了一個(gè)條件概率分布，被建模的可能的軌跡都必須從當(dāng)前狀態(tài) 出發(fā)。自編碼器學(xué)習(xí)一個(gè)從當(dāng)前狀態(tài) 開(kāi)始的軌跡和隱編碼（latent codes）之間的雙向映射。這些隱編碼和原本軌跡一樣按時(shí)間順序排列，每個(gè)隱編碼會(huì)被對(duì)應(yīng)到實(shí)際 步軌跡。因?yàn)槲覀兪褂昧?Causal Transformer，時(shí)間排位靠后的隱編碼 (如 ) 不會(huì)將信息傳到排位靠前的序列(如 )，這使得 TAP 可以通過(guò)前 N 個(gè)隱編碼部分解碼出長(zhǎng)度為 NL 的軌跡，這在后續(xù)用它進(jìn)行規(guī)劃時(shí)是非常有用的。

隨后我們會(huì)用另一個(gè) GPT-2 式的 Transformer 來(lái)建模這些隱編碼的條件概率分布 ：

在決策時(shí)，我們可以通過(guò)在隱變量空間內(nèi)進(jìn)行優(yōu)化尋找最好的未來(lái)軌跡，而不是在原動(dòng)作空間進(jìn)行優(yōu)化。一個(gè)非常簡(jiǎn)單但也能起效的方法就是直接從隱編碼的分布中進(jìn)行采樣，然后選取表現(xiàn)最好的軌跡，如下圖：

在選擇最優(yōu)軌跡的時(shí)候參考的目標(biāo)分?jǐn)?shù) (objective score) 會(huì)同時(shí)考慮軌跡的預(yù)期收益（獎(jiǎng)勵(lì)加最后一步的估值）和這條軌跡本身的可行性或者說(shuō)概率。如下面這個(gè)公式，其中 是一個(gè)遠(yuǎn)大于最高 return 的數(shù)，當(dāng)軌跡的概率高于一個(gè)閾值 ，評(píng)判這條軌跡的標(biāo)準(zhǔn)會(huì)是它的預(yù)期收益（紅色高亮），否則這條軌跡的概率本身將會(huì)是主導(dǎo)部分（藍(lán)色高亮）。也就是說(shuō) TAP 會(huì)在大于閾值的軌跡中選擇一條預(yù)期收益最高的。

雖然采樣數(shù)量足夠大了以后直接采樣在預(yù)測(cè)序列較短的時(shí)候效果也可以很好，在限制采樣數(shù)量和規(guī)劃所需的總時(shí)間的前提下，用更好的優(yōu)化器還是會(huì)帶來(lái)更好的表現(xiàn)。以下兩個(gè)動(dòng)圖展示了當(dāng)要預(yù)測(cè)未來(lái) 144 步以后用直接采樣和 beam search 產(chǎn)生的軌跡的區(qū)別。這些軌跡被按最后的目標(biāo)分?jǐn)?shù)排序，在最上層最前面的軌跡分?jǐn)?shù)越高，在被疊在后面的軌跡分?jǐn)?shù)越低。另外分?jǐn)?shù)低的軌跡同時(shí)透明度也會(huì)越低。

在圖中我們可以看到直接采樣產(chǎn)生的軌跡很多動(dòng)態(tài)都不太穩(wěn)定也不太符合物理規(guī)律，尤其是背景里比較淡的軌跡幾乎都是飄著走的。這些都是概率比較低的軌跡，在最終選擇方案的時(shí)候會(huì)被剔除。在最前排的軌跡看起來(lái)動(dòng)態(tài)要更真實(shí)一些，但是相對(duì)應(yīng)的表現(xiàn)就比較差，似乎要摔倒了。而相比之下，beam search 在展開(kāi)下一個(gè)隱變量的時(shí)候就會(huì)動(dòng)態(tài)考慮軌跡的概率，使得概率很低的分支會(huì)被提前終結(jié)，這樣產(chǎn)生的候選軌跡就都集中在表現(xiàn)較好而且可能性比較大的軌跡周圍了。

直接采樣

Beam search

實(shí)驗(yàn)結(jié)果

在沒(méi)有更高級(jí)的估值以及策略提升的情況下，僅僅依靠預(yù)測(cè)精度的優(yōu)勢(shì)，在低維度的任務(wù)上 TAP 就取得了和其它離線強(qiáng)化學(xué)習(xí)相當(dāng)?shù)谋憩F(xiàn)：

gym locomotion control

在高維的任務(wù)上，TAP 取得了遠(yuǎn)超其它基于模型的方法的表現(xiàn)，同時(shí)也勝過(guò)了常見(jiàn)的無(wú)模型方法。這里其實(shí)有兩個(gè)還未有解答的開(kāi)放性問(wèn)題。首先是為什么此前基于模型的方法在這些高維度的離線強(qiáng)化學(xué)習(xí)任務(wù)中表現(xiàn)較差，其次是為什么 TAP 在這些任務(wù)上表現(xiàn)又能反超很多無(wú)模型方法。我們的一個(gè)假設(shè)是因?yàn)樵诟呔S問(wèn)題上進(jìn)行策略優(yōu)化又要考慮防止策略偏離行為策略太多是非常困難的。當(dāng)學(xué)習(xí)了一個(gè)模型，模型本身的誤差可能還會(huì)放大這種困難。而 TAP 將優(yōu)化空間搬到了一個(gè)很小的離散隱變量空間，這使得整個(gè)優(yōu)化過(guò)程的魯棒性更強(qiáng)了。

adroit robotic hand control

一些切片研究

對(duì)于 TAP 里面的諸多設(shè)計(jì)，我們也在 gym locomotion control 的任務(wù)上做了一系列切片研究。首先是每個(gè)隱編碼實(shí)際對(duì)應(yīng)的軌跡的步數(shù)（黃色柱狀圖），事實(shí)證明讓一個(gè)隱變量對(duì)應(yīng)多步狀態(tài)轉(zhuǎn)移不光有計(jì)算上的優(yōu)勢(shì)，在最后模型表現(xiàn)上也有提升。通過(guò)調(diào)節(jié)搜索的目標(biāo)函數(shù)中觸發(fā)低概率軌跡懲罰的閾值 （紅色柱狀圖），我們也確認(rèn)了目標(biāo)函數(shù)中兩個(gè)部分確實(shí)都對(duì)模型最后表現(xiàn)是有幫助的。另外一點(diǎn)就是向未來(lái)規(guī)劃的步數(shù)（planning horizon，藍(lán)色柱狀圖）對(duì)模型表現(xiàn)的影響反而不大，在部署后的搜索中哪怕只展開(kāi)一個(gè)隱變量最后智能體的表現(xiàn)也只會(huì)降低 10% 左右。

最后我們嘗試了直接采樣的情況下 TAP 的表現(xiàn)（綠色柱狀圖）。注意這里的采樣的樣本數(shù)量是 2048 而上面的動(dòng)圖里只有 256，而且上面的動(dòng)圖是生成了未來(lái) 144 步的規(guī)劃，但是實(shí)際我們的基礎(chǔ)模型指揮規(guī)劃 15 步。結(jié)論是直接采樣在樣本數(shù)量足夠的情況下，且規(guī)劃路徑不長(zhǎng)，那么直接采樣也能獲得和 beam search 相近的表現(xiàn)。但是這是從學(xué)到的隱變量條件分布中采樣的情況，如果直接從隱編碼中直接等概率采樣，那最后還是會(huì)比完整的 TAP 模型差很多。

切片研究的結(jié)果