訓練Rainbow需要1425個GPU Day?谷歌說強化學習可以降低計算成本
DeepMind 提出的 Rainbow 算法,可以讓 AI 玩 Atari 游戲的水平提升一大截,但該算法計算成本非常高,一個主要原因是學術研究發布的標準通常是需要在大型基準測試上評估新算法。來自谷歌的研究者通過添加和移除不同組件,在有限的計算預算、中小型環境下,以小規模實驗得到與 Rainbow 算法一致的結果。
人們普遍認為,將傳統強化學習與深度神經網絡結合的深度強化學習,始于 DQN 算法的開創性發布。DQN 的論文展示了這種組合的巨大潛力,表明它可以產生玩 Atari 2600 游戲的有效智能體。之后有多種方法改進了原始 DQN,而 Rainbow 算法結合了許多最新進展,在 ALE 基準測試上實現了 SOTA 的性能。然而這一進展帶來了非常高的計算成本,擁有充足計算資源的和沒有計算資源之間的差距被進一步拉大。
在 ICML 2021 的一篇論文《Revisiting Rainbow: Promoting more Insightful and Inclusive Deep Reinforcement Learning Research》中,研究者首先討論了與 Rainbow 算法相關的計算成本。研究者探討了通過結合多種算法組件,以小規模實驗得到與 Rainbow 算法一致的結果,并將該想法進一步推廣到在較小的計算預算上進行的研究如何提供有價值的科學見解。

論文地址:https://arxiv.org/abs/2011.14826
Rainbow 計算成本高的一個主要原因是學術研究發布的標準通常是需要在大型基準測試(例如 ALE,其中包含 57 款強化學習智能體能夠學會玩 Atari 2600 游戲)上評估新算法。通常使用 Tesla P100 GPU 訓練模型學會玩一個游戲大約需要五天時間。此外,如果想要建立有意義的置信邊界,通常至少執行 5 次運行。
因此,在全套 57 款游戲上訓練 Rainbow 需要大約 34,200 個 GPU hour(約 1425 天)才能提供令人信服的性能實驗數據。這樣的實驗只有能夠在多個 GPU 上并行訓練時才可行,這使得較小的研究小組望而卻步。
Rainbow 算法
與原始 Rainbow 算法的論文一樣,在 ICML 2021 的這篇論文中,研究者評估了在原始 DQN 算法中添加以下組件的效果:雙 Q 學習(double Q-learning)、優先經驗回放(prioritized experience replay,PER)、競爭網絡、多步學習、分布式強化學習和嘈雜網絡。
該研究在四個經典控制環境中進行評估。需要注意的是,相比于 ALE 游戲需要 5 天,這些環境在 10-20 分鐘內就可以完成完全訓練:

左上:在 CartPole 中,游戲任務是智能體通過左右移動平衡推車上的一根桿子;右上:在 Acrobot 中,有兩個杠桿和兩個連接點,智能體需要向兩個杠桿之間的連接點施加力以抬高下面的杠桿使其高于某個高度要求。左下:在 LunarLander 中,智能體的任務是將飛船降落在兩個旗幟之間;右下:在 MountainCar 中,智能體需要在兩座山丘之間借助一定的動力將車開到右邊的山頂。
研究者探究了將每個組件單獨添加到 DQN 以及從完整 Rainbow 算法中刪除每個組件的效果,并發現總的來說每一個算法組件的添加都確實改進了基礎 DQN 的學習效果。然而,該研究也發現了一些重要的差異,例如通常被認為能起到改進作用的分布式 RL 自身并不總是能夠產生改進。實際上,與 Rainbow 論文中的 ALE 結果相反,在經典控制環境中,分布式 RL 僅在與其他組件結合時才會產生改進。

上圖顯示了在 4 個經典控制環境中,向 DQN 添加不同組件時的訓練進度。x 軸為訓練 step,y 軸為性能(越高越好)。

上圖顯示了在 4 個經典控制環境中,從 Rainbow 中移除各種組件時的訓練進度。x 軸為訓練 step,y 軸為性能(越高越好)。
研究者還在 MinAtar 環境中重新運行了 Rainbow 實驗,MinAtar 環境由一組五個小型化的 Atari 游戲組成,實驗結果與原 Rainbow 論文類似。MinAtar 游戲的訓練速度大約是常規 Atari 2600 游戲的 10 倍,其中后者的訓練速度是在最初的 Rainbow 算法上評估的。此外,該研究的實驗結果還有一些有趣的方面,例如游戲動態和給智能體添加基于像素的輸入。因此,該研究提供了一個具有挑戰性的中級環境,介于經典控制和完整的 Atari 2600 游戲之間。
綜合來看,研究者發現現在的結果與原始 Rainbow 論文的結果一致——每個算法組件產生的影響可能因環境而異。研究者建議使用單一智能體來平衡不同算法組件之間的權衡,該研究的 Rainbow 版本可能與原始版本高度一致,這是因為將所有組件組合在一起會產生整體性能更好的智能體。然而,在不同算法組件之間,有一些重要的細節變化值得進行更徹底的探究。
「優化器 - 損失函數」不同組合實驗
DQN 被提出時,同時采用了 Huber 損失和 RMSProp 優化器。對于研究者而言,在構建 DQN 時使用相同的選擇是一種常見的做法,因為研究者將大部分時間用在了其他算法設計上。
而該研究重新討論了 DQN 在低成本、小規模經典控制和 MinAtar 環境中使用的損失函數和優化器。研究者使用 Adam 優化器進行了一些初始實驗,目前 Adam 優化器是最流行的優化器,并在實驗中結合使用了一個更簡單的損失函數,即均方誤差損失 (MSE)。由于在開發新算法時,優化器和損失函數的選擇往往被忽略,而該研究發現在所有的經典控制和 MinAtar 環境中,這二者的改變都能讓實驗結果有顯著的改進。
因此,研究者將兩個優化器(RMSProp、Adam 優化器)與兩個損失函數(Huber、MSE 損失)進行了不同的組合,并在整個 ALE 平臺(包含 60 款 Atari 2600 游戲)上進行了評估。結果發現 Adam+MSE 組合優于 RMSProp+Huber 組合。

在默認 DQN 設置下(RMSProp + Huber),評估 Adam+MSE 組合帶來的改進(越高越好)。
此外,在比較各種「優化器 - 損失函數」組合的過程中,研究者發現當使用 RMSProp 時,Huber 損失往往比 MSE 表現得更好(實線和橙色虛線之間的間隙可以說明這一點)。

對 60 款 Atari 2600 游戲的標準化得分進行匯總,比較不同的「優化器 - 損失函數」組合。
在有限的計算預算下,該研究研究者能夠在高層次上復現論文《Rainbow: Combining Improvements in Deep Reinforcement Learning》的研究,并且發現新的、有趣的現象。顯然,重新審視某事物比首次發現更容易。然而,研究者開展這項工作的目的是為了論證中小型環境實證研究的相關性和重要性。研究者相信,這些計算強度較低的環境能夠很好地對新算法的性能、行為和復雜性進行更關鍵和徹底的分析。該研究希望 AI 研究人員能夠把小規模環境作為一種有價值的工具,評審人員也要避免忽視那些專注于小規模環境的實驗工作。