再掀強化學習變革!DeepMind提出「算法蒸餾」:可探索的預訓練強化學習Transformer
在當下的序列建模任務上,Transformer可謂是最強大的神經(jīng)網(wǎng)絡架構,并且經(jīng)過預訓練的Transformer模型可以將prompt作為條件或上下文學習(in-context learning)適應不同的下游任務。
大型預訓練Transformer模型的泛化能力已經(jīng)在多個領域得到驗證,如文本補全、語言理解、圖像生成等等。
從去年開始,已經(jīng)有相關工作證明,通過將離線強化學習(offline RL)視為一個序列預測問題,那么模型就可以從離線數(shù)據(jù)中學習策略。
但目前的方法要么是從不包含學習的數(shù)據(jù)中學習策略(如通過蒸餾固定的專家策略),要么是從包含學習的數(shù)據(jù)(如智能體的重放緩沖區(qū))中學習,但由于其context太小,以至于無法捕捉到策略提升。
DeepMind的研究人員通過觀察發(fā)現(xiàn),原則上強化學習算法訓練中學習的順序性(sequential nature)可以將強化學習過程本身建模為一個「因果序列預測問題」。
具體來說,如果一個Transformer的上下文足夠長到可以包含由于學習更新而產(chǎn)生的策略改進,那它應該不僅能夠表示一個固定的策略,而且能夠通過關注之前episodes的狀態(tài)、行動和獎勵表示為一個策略提升算子(policy improvement operator)。
這也提供了一種技術上的可行性,即任何RL算法都可以通過模仿學習蒸餾成一個足夠強大的序列模型,并將其轉化為一個in-context RL算法。
基于此,DeepMind提出了算法蒸餾(Algorithm Distillation, AD) ,通過建立因果序列模型將強化學習算法提取到神經(jīng)網(wǎng)絡中。
論文鏈接:?https://arxiv.org/pdf/2210.14215.pdf?
算法蒸餾將學習強化學習視為一個跨episode的序列預測問題,通過源RL算法生成一個學習歷史數(shù)據(jù)集,然后根據(jù)學習歷史作為上下文,通過自回歸預測行為來訓練因果Transformer。
與蒸餾后學習(post-learning)或專家序列的序列策略預測結構不同,AD能夠在不更新其網(wǎng)絡參數(shù)的情況下完全在上下文中改進其策略。
- Transfomer收集自己的數(shù)據(jù),并在新任務上最大化獎勵;
- 無需prompting或微調;
- 在權重凍結的情況下,Transformer可探索、利用和最大化上下文的返回(return)!諸如Gato類的專家蒸餾(Expert Distillation)方法無法探索,也無法最大化返回。
實驗結果證明了AD可以在稀疏獎勵、組合任務結構和基于像素觀察的各種環(huán)境中進行強化學習,并且AD學習的數(shù)據(jù)效率(data-efficient)比生成源數(shù)據(jù)的RL算法更高。
AD也是第一個通過對具有模仿?lián)p失(imitation loss)的離線數(shù)據(jù)進行序列建模來展示in-context強化學習的方法。
算法蒸餾
2021年,有研究人員首先發(fā)現(xiàn)Transformer可以通過模仿學習從離線RL數(shù)據(jù)中學習單任務策略,隨后又被擴展為可以在同域和跨域設置中提取多任務策略。
這些工作為提取通用的多任務策略提出了一個很有前景的范式:首先收集大量不同的環(huán)境互動數(shù)據(jù)集,然后通過序列建模從數(shù)據(jù)中提取一個策略。
把通過模仿學習從離線RL數(shù)據(jù)中學習策略的方法也稱之為離線策略蒸餾,或者簡稱為策略蒸餾(Policy Distillation, PD)。
盡管PD的思路非常簡單,并且十分易于擴展,但PD有一個重大的缺陷:生成的策略并沒有從與環(huán)境的額外互動中得到提升。
例如,MultiGame Decision Transformer(MGDT)學習了一個可以玩大量Atari游戲的返回條件策略,而Gato通過上下文推斷任務,學習了一個在不同環(huán)境中解決任務的策略,但這兩種方法都不能通過試錯來改進其策略。
MGDT通過微調模型的權重使變壓器適應新的任務,而Gato則需要專家的示范提示才能適應新的任務。
簡而言之,Policy Distillation方法學習政策而非強化學習算法。
研究人員假設Policy Distillation不能通過試錯來改進的原因是,它在沒有顯示學習進展的數(shù)據(jù)上進行訓練。
算法蒸餾(AD)通過優(yōu)化一個RL算法的學習歷史上的因果序列預測損失來學習內涵式策略改進算子的方法。
AD包括兩個組成部分:
1、通過保存一個RL算法在許多單獨任務上的訓練歷史,生成一個大型的多任務數(shù)據(jù)集;
2、將Transformer使用前面的學習歷史作為其背景對行動進行因果建模。
由于策略在源RL算法的整個訓練過程中不斷改進,AD必須得學習如何改進算子,才能準確模擬訓練歷史中任何給定點的行動。
最重要的是,Transformer的上下文大小必須足夠大(即跨周期),以捕捉訓練數(shù)據(jù)的改進。
在實驗部分,為了探索AD在in-context RL能力上的優(yōu)勢,研究人員把重點放在預訓練后不能通過zero-shot 泛化解決的環(huán)境上,即要求每個環(huán)境支持多種任務,且模型無法輕易地從觀察中推斷出任務的解決方案。同時episodes需要足夠短以便可以訓練跨episode的因果Transformer。
在四個環(huán)境Adversarial Bandit、Dark Room、Dark Key-to-Door、DMLab Watermaze的實驗結果中可以看到,通過模仿基于梯度的RL算法,使用具有足夠大上下文的因果Transformer,AD可以完全在上下文中強化學習新任務。
AD能夠進行in-context中的探索、時間上的信用分配和泛化,AD學習的算法比產(chǎn)生Transformer訓練的源數(shù)據(jù)的算法更有數(shù)據(jù)效率。
PPT講解
為了方便論文理解,論文的一作Michael Laskin在推特上發(fā)表了一份ppt講解。
算法蒸餾的實驗表明,Transformer可以通過試錯自主改善模型,并且不用更新權重,無需提示、也無需微調。單個Transformer可以收集自己的數(shù)據(jù),并在新任務上將獎勵最大化。
盡管目前已經(jīng)有很多成功的模型展示了Transformer如何在上下文中學習,但Transformer還沒有被證明可以在上下文中強化學習。
為了適應新的任務,開發(fā)者要么需要手動指定一個提示,要么需要調整模型。
如果Transformer可以適應強化學習,做到開箱即用豈不美哉?
但Decision Transformers或者Gato只能從離線數(shù)據(jù)中學習策略,無法通過反復實驗自動改進。
使用算法蒸餾(AD)的預訓練方法生成的Transformer可以在上下文中強化學習。
首先訓練一個強化學習算法的多個副本來解決不同的任務和保存學習歷史。
一旦收集完學習歷史的數(shù)據(jù)集,就可以訓練一個Transformer來預測之前的學習歷史的行動。
由于策略在歷史上有所改進,因此準確地預測行動將會迫使Transformer對策略提升進行建模。
整個過程就是這么簡單,Transformer只是通過模仿動作來訓練,沒有像常見的強化學習模型所用的Q值,沒有長的操作-動作-獎勵序列,也沒有像 DTs 那樣的返回條件。
在上下文中,強化學習沒有額外開銷,然后通過觀察 AD 是否能最大化新任務的獎勵來評估模型。
Transformer探索、利用、并最大化返回在上下文時,它的權重是凍結的!
另一方面,專家蒸餾(最類似于Gato)不能探索,也不能最大化回報。
AD 可以提取任何 RL 算法,研究人員嘗試了 UCB、DQNA2C,一個有趣的發(fā)現(xiàn)是,在上下文 RL 算法學習中,AD更有數(shù)據(jù)效率。
用戶還可以輸入prompt和次優(yōu)的demo,模型會自動進行策略提升,直到獲得最優(yōu)解!
而專家蒸餾ED只能維持次優(yōu)的demo表現(xiàn)。
只有當Transformer的上下文足夠長,跨越多個episode時,上下文RL才會出現(xiàn)。
AD需要一個足夠長的歷史,以進行有效的模型改進和identify任務。
通過實驗,研究人員得出以下結論:
- Transformer可以在上下文中進行 RL
- 帶 AD 的上下文 RL 算法比基于梯度的源 RL 算法更有效
- AD提升了次優(yōu)策略
- in-context強化學習產(chǎn)生于長上下文的模仿學習