從Transformer到擴散模型，一文了解基於序列建模的強化學習方法

大規模生成模型在近兩年為自然語言處理甚至電腦視覺帶來的巨大的突破。最近這股風潮也影響了強化學習，尤其是離線強化學習（offline RL），諸如Decision Transformer (DT)[1], Trajectory Transformer（TT）[2], Gato[3], Diffuser[4]這樣的方法，將強化學習的數據（包括狀態，動作，獎勵和return-to-go）當成了一串去結構化的序列數據，並將建模這些序列數據作為學習的核心任務。這些模型都可以用監督或自我監督學習的方法來訓練，避免了傳統強化學習中比較不穩定的梯度訊號。即便使用複雜的策略提升 (policy improvement) 和估值 (value estimation) 方法，它們在離線強化學習中都展現了非常好的表現。

本篇將簡單談談這些基於序列建模的強化學習方法，下篇作者將介紹我們新提出的，Trajectory Autoencoding Planner(TAP)，一種用Vector Quantised Variational AutoEncoder (VQ-VAE)進行序列建模並進行高效率的隱動作空間（latent action space）內規劃的方法。

Transformer 與強化學習

Transformer 架構[5] 於2017 年提出之後慢慢引發了自然語言處理的革命，後續的BERT 和GPT-3 逐漸將自監督Transformer 這個組合不斷推向新的高度，在自然語言處理領域不斷湧現出少樣本(few-shot) 學習等性質的同時，也開始向諸如計算機視覺的領域不斷擴散[6][7] 。

然而對於強化學習來說，這個過程似乎在 2021 年之前都不是特別明顯。在 2018 年，多頭注意力機制也被引入強化學習 [8]，這類工作基本上都是應用在類似半符號化(sub-symbolic) 的領域嘗試解決強化學習泛化的問題。之後這類嘗試就一直處於不溫不火的狀態。根據筆者個人的經驗，實際上 Transformer 在強化學習上也並沒有展現出穩定的壓倒性的優勢，而且還很難訓練。在20 年我們的一個用Relational GCN 做強化學習的工作中[9]，我們其實也在背後試過Transformer，但是基本上比傳統結構（類似CNN）差得多，很難穩定訓練得到一個能用的policy。為什麼 Transformer 和傳統線上強化學習（online RL）的相性比較差還是個開放問題，例如 Melo[10] 解釋說是因為傳統的 Transformer 的參數初始化不適合強化學習，在此我就不多做討論了。

2021 年中，Decision Transformer (DT)和 Trajectory Transformer（TT）的發表掀起了 Transformer 在 RL 上應用的新大潮。這兩個工作的想法其實很直接：如果 Transformer 和線上強化學習的演算法不是很搭，那要不乾脆把強化學習當成一個自監督學習的任務？趁著離線強化學習這個概念也很火熱，這兩個工作都將自己的主要目標任務鎖定為建模離線資料集（offline dataset），然後再將這個序列模型用作控制和決策。

對於強化學習來說，所謂序列就是由狀態(state) s ，動作(action) ，獎勵(reward) r 和價值( value) v 構成的軌跡(trajectory) 。其中價值目前一般是被用 return-to-go 來替代，可以被看成是一種蒙特卡洛估計（Monte Carlo estimation）。離線資料集就由這條條軌跡構成。軌跡的產生和環境的動力學模型（dynamics）以及行為策略（behaviour policy）有關。而所謂序列建模，就是要建模產生產生這個序列的機率分佈（distribution），或者嚴格上說是其中的一些條件機率。

#

Decision Transformer

DT 的做法是建模一個從過往資料和價值到動作的映射(return-conditioned policy)，也就是建模了一個動作的條件機率的數學期望 。這個想法很類似 Upside Down RL[11]，不過很有可能背後的直接動機是模仿 GPT2/3 那種根據提示詞(prompt) 完成下游任務的做法。這種做法的一個問題是要決定什麼是最好的目標價值沒有一個比較系統化的方法。然而 DT 的作者發現即使將目標價值設為整個資料集中的最高 return，最後 DT 的表現也可以很不錯。

Decision Transformer, Figure 1

對於有強化學習背景的人來說，DT 這樣的方法能取得強烈的表現是非常反直覺的。如果說 DQN，策略梯度（Policy Gradient）這類方法還可以只把神經網路當成一個能做插值泛化的擬合函數，強化學習中的策略提升、估值仍然是構造策略的核心的話。 DT 就完全可以說是以神經網路為核心的了，背後它如何把一個可能不切實際的高目標價值聯繫到一個合適的動作的整個過程都完全是黑箱。 DT 的成功可以說從強化學習的角度來看有些沒有道理，不過我覺得這也正是這種實證研究的魅力所在。筆者認為神經網絡，或者說 Transformer 的泛化能力可能超出整個 RL 社群先前的預期。

DT 在所有序列建模方法中也是非常簡單的，幾乎所有強化學習的核心問題都在 Transformer 內部被解決了。這種簡單性是它目前最受青睞的原因之一。不過它黑盒的性質也導致我們在演算法設計層面失去了許多抓手，傳統的強化學習中的一些成果很難被融入其中。而這些成果的有效性已經在一些超大規模的實驗（如 AlphaGo, AlphaStar, VPT）中被反覆證實了。

Trajectory Transformer

TT 的做法則更類似傳統的模型為基礎的強化學習 (model-based RL) 的規劃（planning）方法。在建模方面，它將整個序列中的元素都離散化，然後用了像 GPT-2 那樣的離散的自回歸（auto-regressive）方式來建模整個離線資料集。這使得它能夠建模任意給定除去 return-to-go 的序列的後續 。因為建模了後續序列的分佈，TT 其實就成為了一個序列生成模型。透過在產生的序列中尋找擁有更好的估值（value estimation）的序列，TT 就可以輸出一個「最優規劃」。至於尋找最優序列的方法，TT 用了自然語言常用的方法：beam search 的變種。基本上就是永遠保留已經展開的序列中最優的一部分序列，然後在它們的基礎上尋找下一步的最優序列集 。

從強化學習的角度來說，TT 沒有 DT 那麼離經叛道。它的有趣之處在於（和 DT 一樣）完全拋棄了原本強化學習中馬可夫決策過程（Markov Decision Process）的因果圖結構。先前的模型為基礎的方法例如，PETS, world model, dramerv2 等，都會遵循馬可夫過程（或隱式馬可夫）中策略函數、轉移函數、獎勵函數等的定義，也就是狀態分佈的條件是上一步的狀態，而動作、獎勵、價值都由當下的狀態決定。整個強化學習社群一般相信這樣能提高樣本效率，不過這樣的圖結構其實也可能是一種限制。自然語言領域從 RNN 到 Transformer 以及電腦視覺領域 CNN 到 Transformer 的轉變其實都體現了：隨著資料增加，讓網路自己學習圖結構更有利於獲得表現更好的模型。

DreamerV2, Figure 3由於TT 基本上把所有序列預測的任務交給了Transformer，Transformer 就能更有彈性地從資料中學習出更好的圖結構。如下圖，TT 建模出的行為策略根據不同的任務和資料集展現出不同的圖結構。圖左對應了傳統的馬可夫策略，圖右對應了一種動作滑動平均的策略。

Trajectory Transformer, Figure 4

Transformer 強大的序列建模能力帶來了更高的長序列建模精度，下圖展示了TT 在100 步以上的預測仍然保持了高精度，而遵循馬可夫性質的單步預測模型很快就因為預測誤差疊加的問題崩潰了。

Trajectory Transformer, Figure 2TT 雖然在具體建模和預測方面和傳統方法有所不同，它提供的預測能力還是給未來融入強化學習的其它成果留出了很好的抓手。然而TT 在預測速度上有一個重要問題：因為需要建模整個序列的分佈，它將序列中所有的元素按照維度進行離散化，這也就是說一個100 維的狀態就需要佔用序列中的100 個位置，這使得被建模的序列的實際長度很容易變得特別長。而對於 Transformer，它關於序列長度 N 的運算複雜度是 ，這使得從 TT 中取樣一個對未來的預測變得非常昂貴。即使 100 維以下的任務 TT 也需要數秒甚至數十秒來進行一步決策，這樣的模型很難被投入實時的機器人控製或在線學習之中。

Gato

Gato 是 Deepmind 發表的“通才模型”，其實就是跨模態多任務生成模型。用同一個 Transformer 它可以完成從自然語言問答，圖片描述，玩電子遊戲到機器人控制等各類工作。在針對連續控制（continous control）的建模方面 Gato 的做法基本上和 TT 類似。只不過 Gato 嚴格意義不是在做強化學習，它只是建模了專家策略產生的序列數據，然後在行動時它只需要採樣下一個動作，其實是對專家策略的一種模仿。

#Gato Blog

其它序列產生模型：擴散模型

最近在圖片生成領域擴散模型（Diffusion Model）可以說是大紅大紫，DALLE-2 和Stable Diffusion 都是基於它進行圖片生成的。 Diffuser 就將這個方法也運用到了離線強化學習當中，其思路和 TT 類似，先建模序列的條件分佈，然後根據當前狀態採樣未來可能的序列。

Diffuser 相比TT 又擁有了更強的靈活性：它可以在設定起點和終點的情形下讓模型填充出中間的路徑，這樣就能實現目標驅動（而非最大化獎勵函數）的控制。它還可以將多個目標和先驗的達成目標的條件混合起來幫助模型完成任務。

Diffuser Figure 1Diffuser 相對於傳統的強化學習模式也是比較顛覆的，它產生的計劃不是在時間軸上逐步展開，而是從整個序列意義上的模糊變得逐漸精確。擴散模型本身的進一步研究也是電腦視覺中的一個火熱的話題，在其模型本身上很可能未來幾年也會有突破。

不過擴散模型本身目前相比於其它生成模型有一個特別的缺陷，那就是它的生成速度相比於其它生成模型會更慢。很多相關領域的專家認為這一點可能在未來幾年內會被緩解。不過數秒的生成時間目前對於強化學習需要即時控制的情境來說是很難接受的。 Diffuser 提出了能夠提升生成速度的方法：從上一個步驟的計劃開始增加少量噪音來重新產生下一步的計劃，不過這樣做會一定程度上降低模型的表現。

以上是從Transformer到擴散模型，一文了解基於序列建模的強化學習方法的詳細內容。更多資訊請關注PHP中文網其他相關文章！