強化學習中的Transformer發展到哪一步了?清華、北大等聯合發布TransformRL綜述
強化學習(RL)為順序決策提供了一種數學形式,深度強化學習(DRL)近年來也取得巨大進展。然而,樣本效率問題阻礙了在現實世界中廣泛應用深度強化學習方法。為了解決這個問題,一種有效的機制是在 DRL 框架中引入歸納偏置。
在深度強化學習中,函數逼近器是非常重要的。然而,與監督學習(SL)中的架構設計相比,DRL 中的架構設計問題仍然很少被研究。大多數關于 RL 架構的現有工作都是由監督學習 / 半監督學習社區推動的。例如,在 DRL 中處理基于高維圖像的輸入,常見的做法是引入卷積神經網絡(CNN)[LeCun et al., 1998; Mnih et al., 2015];處理部分可觀測性(partial observability)圖像的常見做法則是引入遞歸神經網絡(RNN) [Hochreiter and Schmidhuber, 1997; Hausknecht and Stone, 2015]。
近年來,Transformer 架構 [Vaswani et al., 2017] 展現出優于 CNN 和 RNN 的性能,成為越來越多 SL 任務中的學習范式 [Devlin et al., 2018; Dosovitskiy et al., 2020; Dong et al., 2018]。Transformer 架構支持對長程(long-range)依賴關系進行建模,并具有優異的可擴展性 [Khan et al., 2022]。受 SL 成功的啟發,人們對將 Transformer 應用于強化學習產生了濃厚的興趣,希望將 Transformer 的優勢應用于 RL 領域。
Transformer 在 RL 中的使用可以追溯到 Zambaldi 等人 2018 年的一項研究,其中自注意力(self-attention)機制被用于結構化狀態表征的關系推理。隨后,許多研究人員尋求將自注意力應用于表征學習,以提取實體之間的關系,從而更好地進行策略學習 [Vinyals et al., 2019; Baker et al., 2019]。
除了利用 Transformer 進行表征學習,之前的工作還使用 Transformer 捕獲多時序依賴,以處理部分可觀測性問題 [Parisotto et al., 2020; Parisotto and Salakhutdinov, 2021]。離線 RL [Levine et al., 2020] 因其使用離線大規模數據集的能力而受到關注。受離線 RL 的啟發,最近的研究表明,Transformer 結構可以直接作為順序決策的模型 [Chen et al., 2021; Janner et al., 2021] ,并推廣到多個任務和領域 [Lee et al., 2022; Carroll et al., 2022]。
實際上,在強化學習中使用 Transformer 做函數逼近器面臨一些特殊的挑戰,包括:
- 強化學習智能體(agent)的訓練數據通常是當前策略的函數,這在學習 Transformer 的時候會導致不平穩性(non-stationarity);
- 現有的 RL 算法通常對訓練過程中的設計選擇高度敏感,包括模型架構和模型容量 [Henderson et al., 2018];
- 基于 Transformer 的架構經常受制于高性能計算和內存成本,這使得 RL 學習過程中的訓練和推理都很昂貴。
例如,在用于視頻游戲的 AI 中,樣本生成的效率(在很大程度上影響訓練性能)取決于 RL 策略網絡和估值網絡(value network)的計算成本 [Ye et al., 2020a; Berner et al., 2019]。
為了更好地推動強化學習領域發展,來自清華大學、北京大學、智源人工智能研究院和騰訊公司的研究者聯合發表了一篇關于強化學習中 Transformer(即 TransformRL)的綜述論文,歸納總結了當前的已有方法和面臨的挑戰,并討論了未來的發展方向,作者認為 TransformRL 將在激發強化學習潛力方面發揮重要作用。

論文地址:https://arxiv.org/pdf/2301.03044.pdf

論文的總體結構如下:
- 第 2 章介紹了 RL 和 Transformer 的背景知識,然后簡要介紹了這兩者是如何結合在一起的;
- 第 3 章描述了 RL 中網絡架構的演變,以及長期以來 RL 中阻礙廣泛探索 Transformer 架構的挑戰;
- 第 4 章論文作者對 RL 中的 Transformer 進行了分類,并討論了目前具有代表性的方法;
- 第 5 章總結并指出了未來潛在的研究方向。
核心內容從論文第 3 章開始,下面我們來看一下論文的主要內容。
RL 中的網絡架構
在介紹 TransformRL 的分類方法之前,論文回顧了 RL 中網絡架構設計的早期進展,并總結了其存在的挑戰。作者認為 Transformer 是一種先進的神經網絡架構,將有助于深度強化學習(DRL)的發展。
函數逼近器的架構
自 Deep Q-Network [Mnih et al., 2015] 的開創性工作以來,人們為 DRL 智能體的網絡架構做了許多努力。強化學習中網絡架構的改進主要分為兩類。
一類是設計新的結構,結合 RL 歸納偏置來降低訓練策略或價值函數的難度。例如 [Wang et al. 2016] 提出決斗(dueling)網絡架構,其中一個網絡用于狀態價值函數,另一個用于狀態相關的行動優勢函數(action advantage function),這種架構設計結合了歸納偏置。
另一類是研究常用的神經網絡技術(如正則化、殘差連接(skip connection)、批歸一化)是否可以應用于 RL。例如,[Ota et al. 2020] 發現在使用在線特征提取器增強狀態表征的同時增加輸入維度,會有助于提高 DRL 算法的性能和樣本效率。[Sinha et al. 2020] 為 DRL 智能體提出了一種深度密集架構,使用殘差連接進行有效學習,并使用歸納偏置來緩解數據處理不平等問題。[Ota et al. 2021] 使用 DenseNet [Huang et al., 2017] 和解耦表征學習來改善大型網絡的信息流和梯度。最近,由于 Transformer 的優越性能,研究人員嘗試將 Transformer 架構應用于策略優化算法,但發現普通的 Transformer 設計無法在 RL 任務中達到理想的性能 [Parisotto et al., 2020]。
面臨的挑戰
雖然過去幾年基于 Transformer 的架構在 SL 領域取得了諸多進展,但將 Transformer 應用于 RL 并不簡單。實際上,這存在多個特有的挑戰。
從 RL 的角度看,許多研究指出現有的 RL 算法對深度神經網絡的架構非常敏感 [Henderson et al., 2018; Engstrom et al., 2019; Andrychowicz et al., 2020]。首先,RL 中數據收集和策略優化之間的范式交替導致訓練的不平穩。其次,RL 算法通常對訓練過程中的設計選擇高度敏感。[Emmons et al. 2021] 證明仔細選擇模型架構和正則化對于 DRL 智能體的性能至關重要。
從 Transformer 的角度看,基于 Transformer 的架構存在內存占用大、延遲高的問題,這阻礙了它們的高效部署和推理。最近,許多研究圍繞原始 Transformer 架構對計算和內存效率進行改進,但其中大部分工作都集中在 SL 領域。
在 RL 領域,Parisotto 和 Salakhutdinov 提出把基于大容量 Transformer 的學習器模型轉變為小容量的 actor 模型,以避免 Transformer 的高推理延遲。然而,這種方法在內存和計算方面仍然很昂貴。目前,RL 社區還未充分探索高效或輕量級的 Transformer。
強化學習中的 Transformer
盡管 Transformer 已成為大多數監督學習研究的基礎模型,但由于前述挑戰,它在 RL 社區長期未得到廣泛應用。實際上,TransformRL 的大多數早期嘗試都將 Transformer 用于狀態表征學習或提供記憶信息(memory information),同時仍然將標準 RL 算法用于智能體學習,例如時序差分學習和策略優化。
因此,盡管引入 Transformer 作為函數逼近器,這些方法仍然受到傳統 RL 框架的挑戰。直到最近,離線 RL 使得從大規模離線數據中學習最優策略成為可能。受離線 RL 的啟發,最近的工作進一步將 RL 問題視作固定經驗的條件序列建模問題。這樣做有助于繞過傳統 RL 中的 bootstrapping error 挑戰,從而使 Transformer 架構釋放其強大的順序建模能力。
論文回顧了 TransformRL 的進展,并按分類展示現有方法。作者將現有方法分為四類:表征學習、模型學習、順序決策和通用智能體。圖 2 顯示相關分類的概覽。

用于表征學習的 Transformer
考慮到 RL 任務的順序性,使用 Transformer 編碼器模塊是合理的。事實上,RL 任務中的各種序列都需要處理,例如局部 per-timestep 序列(多實體序列 [Vinyals et al., 2019; Baker et al., 2019]、多智能體序列 [Wen et al., 2022])、時序序列([Parisotto et al., 2020; Banino et al., 2021])等。
局部 per-timestep 序列的編碼器
這種方法早期顯著的成功體現在使用 Transformer 處理智能體觀察到的可變數量實體的復雜信息。[Zambaldi et al. 2018a] 首先提出用多頭點積注意力捕獲結構化觀察的關系推理,隨后 AlphaStar [Vinyals et al., 2019] 實現了在具有挑戰性的多智能體環境(星際爭霸 II)中處理多實體觀察。在這種稱為實體 Transformer 的機制中,觀察結果以如下形式編碼:

其中 e_i 代表智能體對實體 i 的觀察,要么直接從整個觀察中切片,要么由實體 tokenizer 給出。
一些后續工作豐富了實體 Transformer 機制。[Hu et al. 2020] 提出了一種兼容的解耦策略,以明確地將行動與各種實體相關聯,并利用注意力機制進行策略解釋。為了實現具有挑戰性的 one-shot 視覺模仿,Dasari 和 Gupta [2021] 使用 Transformer 來學習專注于特定任務元素的表征。
類似于分散在觀察中的實體,一些研究利用 Transformer 來處理其他局部的 per-timestep 序列。Tang 和 Ha [2021] 利用 Transformer 的注意力機制來處理感知序列并構建一個置換不變性輸入策略。在不兼容的多任務 RL 中,[Kurin et al., 2020] 提出使用 Transformer 來提取形態域知識 。
時序編碼器
同時,用 Transformer 處理時序序列也是合理的。時序編碼器被用作存儲架構,

其中 o_t 表示智能體在時間 t 的觀察,Emb_0:t 表示從初始觀察到當前觀察的歷史觀察的嵌入。
在早期的工作中,[Mishra et al. 2018] 無法使用 vanilla Transformer 處理時序序列,并且發現它在某些特定任務中甚至比隨機策略表現更差。Gated Transformer-XL (GTrXL) [Parisotto et al., 2020] 是第一個使用 Transformer 作為存儲架構來處理軌跡的有效方案。GTrXL 通過 Identity Map Reordering 修改 Transformer-XL 架構 [Dai et al., 2019],以提供從時序輸入到 Transformer 輸出的「skip」路徑,這可能有助于形成從一開始就穩定的訓練過程。[Loynd et al. 2020] 提出了一種用于長期依賴的記憶向量快捷機制,[Irie et al. 2021] 將線性 Transformer 與快速加權編程器(Fast Weight Programmer)相結合以獲得更好的性能。[Melo 2022] 提出使用自注意力機制來模擬基于存儲的元 RL 的存儲恢復。
雖然隨著存儲的增長和參數規模的擴大,Transformer 的性能優于 LSTM/RNN,但它在 RL 上的數據效率不佳。后續工作利用一些輔助自監督任務來促進學習 [Banino et al., 2021] 或使用預訓練的 Transformer 架構作為時序編碼器 [Li et al., 2022; Fan et al.,2022]。
用于模型學習的 Transformer
除了使用 Transformer 作為序列嵌入的編碼器,Transformer 架構還在基于模型的算法中作為環境模型的 backbone。與以單步觀察和行動為條件的預測不同,Transformer 使環境模型能夠預測以一定長度的歷史信息為條件的變換(transition)。
實際上,Dreamer 及其后續算法的成功 [Hafner et al., 2020, 2021; Seo et al., 2022] 已經在一些部分可觀察的環境或需要記憶機制的任務中證明了基于歷史信息的世界模型的優點。以歷史信息為條件的世界模型由一個捕獲抽象信息的觀察編碼器和一個學習潛在空間中變換的變換模型組成。
已有研究使用 Transformer 架構而不是 RNN 來構建基于歷史的世界模型。[Chen et al. 2022] 用基于 Transformer 的模型 TSSM(Transformer State-Space Model)替換 Dreamer 中基于 RNN 的循環狀態空間模型(RSSM)。IRIS(Imagination with autoRegression over an Inner Speech)[Micheli et al., 2022] 通過對 rollout 經驗的自回歸學習來學習基于 Transformer 的世界模型,而沒有像 Dreamer 那樣的 KL 平衡,并在 Atari [Bellemare et al., 2013] 上取得了很好的結果。
此外,還有研究嘗試用基于 Transformer 的世界模型做規劃。[Ozair et al. 2021] 驗證了使用 Transformer 變換模型進行規劃來完成隨機任務的有效性。[Sun et al. 2022] 提出了一種以目標為條件的 Transformer 變換模型,該模型在程序任務的視覺基礎規劃中是很有效的。
RNN 和 Transformer 都適合學習基于歷史信息的世界模型。然而,[Micheli et al. 2022] 發現與 Dreamer 相比,Transformer 架構是數據效率更高的世界模型。TSSM 的實驗結果表明,Transformer 架構在需要長期記憶的任務中表現出眾。
用于順序決策的 Transformer
除了融入到傳統 RL 算法中作為高性能架構以外,Transformer 還可以直接用作順序決策模型。這是因為可以把 RL 看作一個條件序列建模問題:生成可以產生高回報的行動序列。

鑒于 Transformer 在序列預測方面的優異準確性,Bootstrapped Transformer (BooT) [Wang et al., 2022] 提出通過 bootstrap Transformer 來生成數據,同時優化數據以進行順序決策。Bootstrapping Transformer 用于數據增強可以擴大離線數據集的數量和覆蓋范圍,從而提升性能。具體地說,BooT 比較了不同的數據生成方案和 bootstraping 方案,以分析 BooT 如何助力策略學習。結果表明,它可以生成與底層 MDP 一致的數據,而無需額外的約束。
用于通用智能體的 Transformer
Decision Transformer 已經在離線數據的各種任務中發揮巨大作用,有研究者開始考慮 Transformer 是否可以像 CV 和 NLP 領域那樣讓通用智能體解決多個不同任務或問題。
泛化到多個任務
一些研究者借鑒了 CV 和 NLP 中對大規模數據集進行預訓練的思想,并嘗試從大規模多任務數據集中抽象出通用策略。Multi-Game Decision Transformer (MGDT) [Lee et al., 2022] 是 DT 的一個變體,該模型在由專家和非專家數據組成的多樣化數據集上學習 DT,并使用一組參數在多個 Atari 游戲上實現接近人類的水平。為了在非專家級數據集上獲得專家級的表現,MGDT 設計了專家行動推理機制,從 return-to-go 的先驗分布計算專家級的 return-to-go 后驗分布并根據貝葉斯公式預設專家級 return-to-go 的概率。
同樣,Switch Trajectory Transformer (SwitchTT) [Lin et al., 2022] 是 TT 的多任務擴展,利用稀疏激活模型,將 FFN 層替換為混合專家層,以實現高效的多任務離線學習。此外,SwitchTT 還采用分布式 trajectory 值估計器對值估計的不確定性進行建模。依靠這兩個增強功能,SwitchTT 在性能和訓練速度方面都比 TT 提升了很多。MGDT 和 SwitchTT 利用從多個任務和各種性能級策略中收集的經驗來學習通用策略。然而,構建大規模的多任務數據集并非易事。
與 CV 和 NLP 中的大規模數據集通常使用來自互聯網的海量數據和簡單的人工標記不同,RL 中的順序決策數據總是缺少行動信息,并且不易標記。因此,[Baker et al. 2022] 提出了一種半監督方案,利用沒有行動信息的大規模在線數據,學習基于 Transformer 的逆向動態模型(IDM)。該模型利用對過去和未來的觀察來預測行動信息,能夠標記大量在線視頻數據。IDM 是在包含手動標記行動的小型數據集上學習的,并且足夠準確。
NLP 的許多已有工作證明了 prompt 在適應新任務方面的有效性,一些工作利用基于 DT 方法的 prompt 技術來實現快速適應。Prompt-based Decision Transformer (Prompt-DT) [Xu et al., 2022] 從少樣本(few-shot)演示數據集中采樣一系列變換作為 prompt,并將少樣本策略泛化到離線元 RL 任務上。[Reed et al. 2022] 進一步利用基于 prompt 的架構,通過在涵蓋自然語言、圖像、時間決策和多模態數據的超大規模數據集上進行自回歸序列建模來學習通用智能體(Gato)。Gato 能夠執行來自不同領域的一系列任務,包括文本生成和決策。
[Laskin et al. 2022] 提出了算法蒸餾 (AD),以在單任務 RL 算法學習過程的 across-episode 序列上訓練 Transformer。因此,即使在新任務中,Transformer 也可以學會在自回歸生成過程中逐步改進其策略。
泛化到更廣泛領域
除了泛化到多個任務,Transformer 還是一個強大的「通用」模型,可以用于與順序決策相關的一系列領域。受 NLP 中掩碼語言建模(masked language modeling)[Devlin et al., 2018] 技術的啟發,[Carroll et al. 2022] 提出了 Uni [MASK],它將各種常用研究領域統一為 mask 推理問題,包括行為克隆、離線 RL、GCRL、過去 / 未來推理和動態預測。Uni [MASK] 比較了不同的 mask 方案,包括任務特定 mask、隨機 mask 和微調變體。結果表明,用隨機 mask 訓練的單個 Transformer 可以解決任意推理任務。
此外,[Reid et al. 2022] 發現,使用在語言數據集或包含語言模態的多模態數據集上預訓練的 Transformer 對 DT 進行微調是有益的。這表明,即使是來自非 RL 領域的知識也可以通過 transformer 進行 RL 訓練。
感興趣的讀者可以閱讀論文原文,了解更多研究細節。



































