替代離線RL？Transformer進(jìn)軍決策領(lǐng)域，「序列建模」成關(guān)鍵

Transformer 開(kāi)始進(jìn)軍決策領(lǐng)域了，它能否替代離線強(qiáng)化學(xué)習(xí)呢？近日，UC 伯克利、FAIR 和谷歌大腦的研究者提出了一種通過(guò)序列建模進(jìn)行強(qiáng)化學(xué)習(xí)的 Transformer 架構(gòu)，并在 Atari、OpenAI Gym 等強(qiáng)化學(xué)習(xí)實(shí)驗(yàn)平臺(tái)上媲美甚至超越 SOTA 離線 RL 基線方法。

自 2016 年 AlphaGo 擊敗李世石開(kāi)始，強(qiáng)化學(xué)習(xí)（Reinforcement Learning）在優(yōu)化決策領(lǐng)域可謂是風(fēng)頭無(wú)兩。同年，基于強(qiáng)化學(xué)習(xí)算法的 AlphaGo 升級(jí)版 AlphaGo Zero 更是采用「從零開(kāi)始」、「無(wú)師自通」的學(xué)習(xí)模式，以 100:0 的比分輕而易舉打敗了之前的 AlphaGo。

與此同時(shí)，BERT、GPT 等語(yǔ)言模型在 NLP 領(lǐng)域掀起狂潮。去年 DETR 和 ViT 出來(lái)之后，Transformer 在計(jì)算機(jī)視覺(jué)領(lǐng)域也是大殺四方。而多虧了深度學(xué)習(xí)，我們今天能做的比幾年前要多出許多。處理序列數(shù)據(jù)的能力，如音樂(lè)歌詞、句子翻譯、理解評(píng)論或構(gòu)建聊天機(jī)器人，所有這些都要?dú)w功于序列建模（Sequential Modeling）。

那么能否結(jié)合強(qiáng)化學(xué)習(xí)與序列建模，并構(gòu)建優(yōu)化決策的 Transformer 架構(gòu)呢？

近日，UC 伯克利、FAIR 和谷歌大腦的研究者發(fā)布了一篇論文，就此問(wèn)題展開(kāi)了研究并提出了 Decision Transformer——一種通過(guò)序列建模進(jìn)行強(qiáng)化學(xué)習(xí)的架構(gòu)。

首先，為了使用 Transformer 架構(gòu)的簡(jiǎn)單性和可擴(kuò)展性，以及 GPT-x 和 BERT 等語(yǔ)言建模的優(yōu)勢(shì)，研究者引入了一個(gè)框架，將強(qiáng)化學(xué)習(xí)抽象為序列建模問(wèn)題。然后將強(qiáng)化學(xué)習(xí)問(wèn)題轉(zhuǎn)化為條件序列建模的架構(gòu)，提出了 Decision Transformer。與先前擬合值函數(shù)或計(jì)算策略梯度的方法不同，Decision Transformer 通過(guò)利用因果掩蔽的 Transformer 來(lái)輸出最佳行動(dòng)。

根據(jù)期望獎(jiǎng)勵(lì)、過(guò)去的狀態(tài)和行動(dòng)來(lái)調(diào)節(jié)自回歸模型，Decision Transformer 模型能夠生成實(shí)現(xiàn)期望回報(bào)的未來(lái)行動(dòng)。盡管很簡(jiǎn)單，Decision Transformer 在 Atari、OpenAI Gym 和 Key-to-Door 任務(wù)上達(dá)到甚至超過(guò)了最先進(jìn)的無(wú)模型離線強(qiáng)化學(xué)習(xí)基線性能。

論文鏈接：https://arxiv.org/pdf/2106.01345.pdf

GitHub 鏈接：https://github.com/kzl/decision-transformer

當(dāng)前，該研究已經(jīng)引起業(yè)內(nèi)極大的關(guān)注，官方代碼庫(kù)已經(jīng)有 750 + 的 star 量。

作為序列建模問(wèn)題的離線強(qiáng)化學(xué)習(xí)

研究者首先研究了如何通過(guò)在語(yǔ)言建?？蚣苤刑岢鲰樞驔Q策問(wèn)題來(lái)改變對(duì)強(qiáng)化學(xué)習(xí)的觀點(diǎn)。雖然強(qiáng)化學(xué)習(xí)中的傳統(tǒng)工作使用了依賴 Bellman 備份的特定框架，但用序列建模來(lái)作為替代對(duì)軌跡進(jìn)行建模使其能夠使用強(qiáng)大且經(jīng)過(guò)充分研究的架構(gòu)（如 Transformer）來(lái)生成行為。

為了說(shuō)明這一點(diǎn)，研究者研究了離線強(qiáng)化學(xué)習(xí)，從固定的數(shù)據(jù)集上訓(xùn)練模型，只需最少的更改就能夠使用與語(yǔ)言建?？蚣芟嗤拇a來(lái)訓(xùn)練強(qiáng)化學(xué)習(xí)策略。

為何選用 Transformer

最近的各項(xiàng)成果表明，Transformer 可以大規(guī)模地對(duì)語(yǔ)義概念的高維分布進(jìn)行建模，包括語(yǔ)言中的有效零樣本泛化和分布外圖像生成。鑒于此類模型成功應(yīng)用的多樣性，研究者想要驗(yàn)證它們能否用于形式化為強(qiáng)化學(xué)習(xí)的序列決策問(wèn)題。與以往使用 Transformer 作為傳統(tǒng)強(qiáng)化學(xué)習(xí)算法中組件的架構(gòu)選擇相比，他們?cè)噲D研究生成軌跡建模，即對(duì)狀態(tài)、動(dòng)作和獎(jiǎng)勵(lì)的序列聯(lián)合分布進(jìn)行建模，以替代傳統(tǒng)強(qiáng)化學(xué)習(xí)算法。

此外，研究者還考慮了以下范式轉(zhuǎn)變：使用序列建模目標(biāo)，根據(jù)采集的經(jīng)驗(yàn)來(lái)訓(xùn)練 Transformer 模型，而不是通過(guò)傳統(tǒng)的強(qiáng)化學(xué)習(xí)算法（如時(shí)序差分學(xué)習(xí)）來(lái)訓(xùn)練策略。這將使研究者繞過(guò)對(duì)長(zhǎng)期信用分配進(jìn)行自舉的需要，從而避免已知會(huì)破壞強(qiáng)化學(xué)習(xí)穩(wěn)定的「deadly triad」之一。它還避免了時(shí)序差分學(xué)習(xí)（temporal difference，TD）中可能會(huì)導(dǎo)致不受歡迎的短視行為，減少未來(lái)獎(jiǎng)勵(lì)的需求。此外，利用在語(yǔ)言和視覺(jué)領(lǐng)域廣泛應(yīng)用且易于擴(kuò)展的 Transformer 框架，可以進(jìn)行大量穩(wěn)定的訓(xùn)練。

除了長(zhǎng)序列建模能力之外，Transformer 還具有其他優(yōu)勢(shì)，比如可以通過(guò)自注意力（self-attention）直接執(zhí)行信用分配。這與緩慢傳播獎(jiǎng)勵(lì)并容易產(chǎn)生干擾信號(hào)的 Bellman 備份相反，可以使 Transformer 在獎(jiǎng)勵(lì)稀少或分散注意力的情況下仍然有效地工作。Transformer 還可以對(duì)廣泛的行為分布進(jìn)行建模，從而實(shí)現(xiàn)更好的泛化和轉(zhuǎn)移。

離線強(qiáng)化學(xué)習(xí)是從次優(yōu)數(shù)據(jù)中學(xué)習(xí)策略來(lái)分配代理，即從固定、有限的經(jīng)驗(yàn)中產(chǎn)生最大有效的行為。由于錯(cuò)誤傳播和價(jià)值高估，探索非常具有挑戰(zhàn)性。但是，在使用序列建模目標(biāo)進(jìn)行訓(xùn)練時(shí)，這是一項(xiàng)自然的任務(wù)。通過(guò)在狀態(tài)、動(dòng)作和返回序列上訓(xùn)練自回歸模型，研究者將策略抽樣減少到自回歸生成建模，選擇作為生成的提示的返回 token 來(lái)指定策略的專業(yè)知識(shí)。

Decision Transformer：強(qiáng)化學(xué)習(xí)的自回歸序列建模

研究者采用了一種簡(jiǎn)單的方法：每個(gè)模態(tài)（返回、狀態(tài)或動(dòng)作）都被傳遞到一個(gè)嵌入網(wǎng)絡(luò)（圖像的卷積編碼器和連續(xù)狀態(tài)的線性層），然后嵌入通過(guò)自回歸 Transformer 模型處理，在給定先前 token 的情況下，使用線性輸出層進(jìn)行訓(xùn)練以預(yù)測(cè)下一個(gè)動(dòng)作。

評(píng)估也很容易：通過(guò)期望的目標(biāo)返回值（例如成功或失敗的 1 或 0）和環(huán)境中的起始狀態(tài)進(jìn)行初始化，展開(kāi)序列（類似于語(yǔ)言模型中的標(biāo)準(zhǔn)自回歸生成）以產(chǎn)生一系列要在環(huán)境中執(zhí)行的動(dòng)作。

Decision Transformer 架構(gòu)。

拼接子序列以產(chǎn)生最佳軌跡

研究者考慮了固定圖上找到最短路徑的任務(wù)的強(qiáng)化學(xué)習(xí)問(wèn)題（累積獎(jiǎng)勵(lì) = 邊權(quán)重之和）。在由隨機(jī)游走組成的訓(xùn)練數(shù)據(jù)集中，他們觀察到了許多次優(yōu)軌跡。如果在這些序列上訓(xùn)練 Decision Transformer，可以要求模型通過(guò)調(diào)節(jié)高回報(bào)來(lái)生成最佳路徑。如果僅對(duì)隨機(jī)游走進(jìn)行訓(xùn)練，Decision Transformer 可以學(xué)習(xí)將來(lái)自不同訓(xùn)練軌跡的子序列拼接在一起，以便在測(cè)試時(shí)產(chǎn)生最佳軌跡。

事實(shí)上，這與離線強(qiáng)化學(xué)習(xí)框架中常用的離策略 Q-learning 算法所期望的行為相同。然而，無(wú)需引入 TD 學(xué)習(xí)算法、價(jià)值悲觀主義或行為正則化，就可以使用序列建?？蚣軐?shí)現(xiàn)相同的行為。

如下圖所示，圖左為強(qiáng)化學(xué)習(xí)為固定圖尋找最短路徑，圖中顯示由隨機(jī)游走軌跡和每個(gè)節(jié)點(diǎn)的返回組成的訓(xùn)練數(shù)據(jù)集，圖右顯示了以起始狀態(tài)和每個(gè)節(jié)點(diǎn)產(chǎn)生的最大可能回報(bào)為條件，Decision Transformer 對(duì)最佳路徑進(jìn)行了排序。

離線強(qiáng)化學(xué)習(xí)基準(zhǔn)的比較

研究者擴(kuò)展到了離線強(qiáng)化學(xué)習(xí)文獻(xiàn)中常用的基準(zhǔn)，即 Atari 學(xué)習(xí)環(huán)境、OpenAI Gym、Minigrid Key-To-Door 任務(wù)。在離散和連續(xù)控制以及狀態(tài)和圖像觀察的多樣化任務(wù)中，他們發(fā)現(xiàn) Decision Transformer 的性能可以媲美經(jīng)過(guò)充分研究的專業(yè) TD 學(xué)習(xí)算法的性能。

主要比較點(diǎn)是基于 TD 學(xué)習(xí)的無(wú)模型離線強(qiáng)化學(xué)習(xí)算法，因?yàn)?Decision Transformer 架構(gòu)本質(zhì)上也是無(wú)模型的。此外，TD 學(xué)習(xí)是強(qiáng)化學(xué)習(xí)中提高樣本效率的主要范式，并且作為一個(gè)子程序在許多基于模型的強(qiáng)化學(xué)習(xí)算法中也很突出。研究者還與行為克隆和變體進(jìn)行了比較，因?yàn)檫@些也涉及了基于似然的策略學(xué)習(xí)公式。確切的算法取決于環(huán)境，但研究者的動(dòng)機(jī)如下：

TD 學(xué)習(xí)：這些方法中的大多數(shù)使用動(dòng)作空間約束或價(jià)值悲觀主義，并且將是與 Decision Transformer 最忠實(shí)的比較，代表標(biāo)準(zhǔn)的強(qiáng)化學(xué)習(xí)方法。最先進(jìn)的無(wú)模型方法是 Conservative Q-Learning (CQL)，它作為主要的比較方法。此外，研究者還與其他的無(wú)模型強(qiáng)化學(xué)習(xí)算法（如 BEAR 和 BRAC ）進(jìn)行了比較；

模仿學(xué)習(xí)：這種機(jī)制類似地使用監(jiān)督損失進(jìn)行訓(xùn)練（而不是 Bellman 備份），并在這里使用行為克隆。

關(guān)于評(píng)估離散（Atari）和連續(xù)（OpenAI Gym）控制任務(wù)，前者涉及高維觀察空間，需要長(zhǎng)期的信用分配，而后者需要細(xì)粒度的連續(xù)控制，代表不同的任務(wù)集。如下圖所示，主要結(jié)果總結(jié)了每個(gè)域的平均歸一化性能。

序列建模和多任務(wù)學(xué)習(xí)

此類型建模的一個(gè)效果是執(zhí)行條件生成：通過(guò)輸入想要的回報(bào)來(lái)初始化一個(gè)軌跡。Decision Transformer 不產(chǎn)生單個(gè)策略，相反它模擬了廣泛的政策分布。如果繪制訓(xùn)練后 Decision Transformer 的目標(biāo)回報(bào)與平均獲得的回報(bào)之間的關(guān)系圖，就會(huì)發(fā)現(xiàn)可以合理地匹配目標(biāo)，并且僅使用監(jiān)督學(xué)習(xí)進(jìn)行訓(xùn)練。此外，在某些任務(wù)（例如 Qbert 和 Seaquest）上，研究者發(fā)現(xiàn) Decision Transformer 實(shí)際上可以在數(shù)據(jù)集和模型策略之外進(jìn)行推理，從而獲得更高的回報(bào)。

研究者通過(guò)在很大的范圍內(nèi)改變所需的目標(biāo)回報(bào)來(lái)評(píng)估 Decision Transformer 理解返回 token 的能力，即評(píng)估 Transformer 的多任務(wù)分布建模能力。下圖顯示了當(dāng)以指定的目標(biāo)（期望）回報(bào)為條件時(shí)，Decision Transformer 累積的平均采樣（評(píng)估）回報(bào)，上部為 Atari，底部為 D4RL 中重放數(shù)據(jù)集。在每項(xiàng)任務(wù)中，期望的目標(biāo)回報(bào)和真實(shí)觀察到的回報(bào)是高度相關(guān)的。

在 Pong、HalfCheetah 和 Walker 等一些任務(wù)上，Decision Transformer 生成的軌跡幾乎完美匹配所需的回報(bào)（如圖中與 oracle 線重疊所示）。此外，在諸如 Seaquest 之類的一些 Atari 任務(wù)中，Decision Transformer 有時(shí)能夠進(jìn)行外推。