從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 提出了能夠提升生成速度的方法:從上一步的計劃開始增加少量噪音來重新生成下一步的計劃,不過這樣做會一定程度上降低模型的表現。