大模型時代，模型壓縮和加速顯得尤為重要。傳統監督學習可通過稀疏神經網絡實現模型壓縮和加速，那么同樣需要大量計算開銷的強化學習任務可以基于稀疏網絡進行訓練嗎？本文提出了一種強化學習專用稀疏訓練框架，可以節省至多 95% 的訓練開銷。

深度強化學習模型的訓練通常需要很高的計算成本，因此對深度強化學習模型進行稀疏化處理具有加快訓練速度和拓展模型部署的巨大潛力。然而現有的生成小型模型的方法主要基于知識蒸餾，即通過迭代訓練稠密網絡，訓練過程仍需要大量的計算資源。另外，由于強化學習自舉訓練的復雜性，訓練過程中全程進行稀疏訓練在深度強化學習領域尚未得到充分的研究。 清華大學黃隆波團隊提出了一種強化學習專用的動態稀疏訓練框架，“Rigged Reinforcement Learning Lottery”（RLx2），可適用于多種離策略強化學習算法。它采用基于梯度的拓撲演化原則，能夠完全基于稀疏網絡訓練稀疏深度強化學習模型。RLx2 引入了一種延遲多步差分目標機制，配合動態容量的回放緩沖區，實現了在稀疏模型中的穩健值學習和高效拓撲探索。在多個 MuJoCo 基準任務中，RLx2 達到了最先進的稀疏訓練性能，顯示出 7.5 倍至 20 倍的模型壓縮，而僅有不到 3% 的性能降低，并且在訓練和推理中分別減少了高達 20 倍和 50 倍的浮點運算數。

• 論文主頁：https://arxiv.org/abs/2205.15043
• 論文代碼：https://github.com/tyq1024/RLx2

背景 在游戲、機器人技術等領域，深度強化學習（DRL）已經取得了重要的應用。然而，深度強化學習模型的訓練需要巨大的計算資源。例如，DeepMind 開發的 AlphaGo-Zero 在圍棋游戲中擊敗了已有的圍棋 AI 和人類專家，但需要在四個 TPU 上進行 40 多天的訓練。OpenAI-Five 是 OpenAI 開發的 Dota2 AI，同樣擊敗了人類半職業 Dota 高手，但是需要高達 256 個 GPU 進行 180 天的訓練。實際上，即使是簡單的 Rainbow DQN [Hessel et al. 2018] 算法，也需要在單個 GPU 上訓練約一周時間才能達到較好的性能。

圖：基于強化學習的 AlphaGo-Zero 在圍棋游戲中擊敗了已有的圍棋 AI 和人類專家 高昂的資源消耗限制了深度強化學習在資源受限設備上的訓練和部署。為了解決這一問題，作者引入了稀疏神經網絡。稀疏神經網絡最初在深度監督學習中提出，展示出了對深度強化學習模型壓縮和訓練加速的巨大潛力。在深度監督學習中，SET [Mocanu et al. 2018] 和 RigL [Evci et al. 2020] 等常用的基于網絡結構演化的動態稀疏訓練（Dynamic sparse training - DST）框架可以從頭開始訓練一個 90% 稀疏的神經網絡，而不會出現性能下降。

圖：SET 和 RigL 等常用的稀疏訓練框架會在訓練的過程中周期性地調整神經網絡結構 在深度強化學習領域，已有的工作已經成功生成了極度稀疏的深度強化學習網絡。然而，他們的方法仍然需要迭代地訓練稠密網絡，往往需要預訓練的稠密模型作為基礎，導致深度強化學習的訓練成本仍然過高，無法直接應用于資源有限設備。

圖：迭代剪枝通過迭代地訓練稠密網絡得到稀疏的深度強化學習網絡 從頭開始訓練一個稀疏深度強化學習模型，如果能夠完美實現，將極大地減少計算開銷，并實現在資源受限設備上的高效部署，具備優秀的模型適應性。然而，在深度強化學習中從頭開始訓練一個超稀疏網絡（例如 90% 的稀疏度）具有挑戰性，原因在于自舉訓練（Bootstrap training）的非穩定性。在深度強化學習中，學習目標不是固定的，而是以自舉方式給出，訓練數據的分布也可能是非穩定的。此外，使用稀疏網絡結構意味著在一個較小的假設空間中搜索，這進一步降低了學習目標的置信度。因此，不當的稀疏化可能對學習路徑造成不可逆的傷害，導致性能較差。最近的研究 [Sokar et al. 2021] 表明，在深度強化學習中直接采用動態稀疏訓練框架仍然無法在不同環境中實現模型的良好壓縮。因此，這一重要的開放問題仍然懸而未決：

能否通過全程使用超稀疏網絡從頭訓練出高效的深度強化學習智能體？

方法 清華大學黃隆波團隊對這一問題給出了肯定的答案，并提出了一種強化學習專用的動態稀疏訓練框架，“Rigged Reinforcement Learning Lottery”（RLx2），用于離策略強化學習（Off-policy RL）。這是第一個在深度強化學習領域以 90% 以上稀疏度進行全程稀疏訓練，并且僅有微小性能損失的算法框架。RLx2 受到了在監督學習中基于梯度的拓撲演化的動態稀疏訓練方法 RigL [Evci et al. 2020] 的啟發。然而，直接應用 RigL 無法實現高稀疏度，因為稀疏的深度強化學習模型由于假設空間有限而導致價值估計不可靠，進而干擾了網絡結構的拓撲演化。 因此，RLx2 引入了延遲多步差分目標（Delayed multi-step TD target）機制和動態容量回放緩沖區（Dynamic capacity buffer），以實現穩健的價值學習（Value learning）。這兩個新組件解決了稀疏拓撲下的價值估計問題，并與基于 RigL 的拓撲演化準則一起實現了出色的稀疏訓練性能。為了闡明設計 RLx2 的動機，作者以一個簡單的 MuJoCo 控制任務 InvertedPendulum-v2 為例，對四種使用不同價值學習和網絡拓撲更新方案的稀疏訓練方法進行了比較。

圖：不同網絡結構更新方案的性能比較。其中，SS 表示采用靜態稀疏網絡，RigL 表示使用基于梯度的網絡拓撲演化的方法，RigL+Q * 表示使用 RigL 的拓撲演化且采用真實值函數引導自舉訓練的方法（真實值函數在實際算法中并不可知），RLx2 表示使用 RigL 網絡拓撲演化且采用作者所提值估引導自舉訓練的方法。可以發現，RLx2 的性能已經非常逼近 RigL+Q * 的方法。 下圖展示了 RLx2 算法的主要部分，包括基于梯度的拓撲演化、延遲多步差分目標和動態容量回放緩沖區。

圖：RLx2 算法的概覽 基于梯度的拓撲演化 在 RLx2 中，作者采用了與 RigL [Evci et al. 2020] 相同的方法來進行拓撲結構的演化。作者計算了損失函數對網絡權重的梯度值。然后，周期性地增加稀疏網絡中具有較大梯度的連接，并移除權重絕對值最小的現有連接。通過周期性的結構演化，獲得了一個結構合理的稀疏神經網絡。 延遲多步差分目標 RLx2 框架還引入了多步差分目標：

這一概念在現有研究中 [Munos et al. 2016] 已被證實能夠改善差分學習（TD learning）。作者還發現，引入多步目標可以通過折扣因子減少稀疏網絡的擬合誤差，如下式所示：

然而，訓練伊始立即采用多步差分目標可能會導致更大的策略不一致誤差。因此，作者額外采用了延遲方案來抑制策略不一致性并進一步提高值函數的學習效果。 動態容量回放緩沖區 離策略（Off-policy）算法使用回放緩沖區（Replay buffer）來存儲收集到的數據，并使用從緩沖區中抽樣的批次數據來訓練網絡。研究表明 [Fedus et al. 2020]，當使用更大的回放容量時，算法的性能通常會提高。然而，無限大小的回放緩沖區會因為多步目標的不一致性和訓練數據的不匹配導致策略不一致性。動態容量回放緩沖區是一種通過調整緩沖區容量控制緩沖區中數據的不一致性，以實現穩健值函數學習的方法。作者引入了以下策略距離度量來評估緩沖區中數據的不一致性：

隨著訓練的進行，當回放緩存中的策略距離度量大于閾值時，則停止增加緩沖區容量，使得策略距離度量始終小于設定的閾值。 實驗 作者在四個 MuJoCo 環境（HalfCheetah-v3、Hopper-v3、Walker2d-v3 和 Ant-v3），和兩個常見的深度強化學習算法 TD3 和 SAC 中進行了實驗。作者定義了一個終極壓縮比率，即在該比率下，RLx2 的性能下降在原始稠密模型的 ±3% 之內。這也可以理解為具有與原始稠密模型完全相同性能的稀疏模型的最小大小。根據終極壓縮比率，作者在下表中呈現了不同算法在不同環境采用相同參數量的神經網絡的性能。

性能 在所有四個環境中，RLx2 的性能在很大程度上優于所有基準算法（除了 Hopper 環境中與 RigL 和 SAC 的性能相近）。此外，小型稠密網絡（Tiny）和隨機靜態稀疏網絡（SS）的性能平均最差。SET 和 RigL 的性能較好，但在 Walker2d-v3 和 Ant-v3 環境中無法保持性能，這意味著在稀疏訓練下穩健的價值學習是必要的。 模型壓縮 RLx2 實現了優秀的壓縮比，并且僅有輕微的性能下降（不到 3%）。具體而言，使用 TD3 算法的 RLx2 實現了 7.5 倍至 25 倍的模型壓縮，在 Hopper-v3 環境中獲得了最佳的 25 倍壓縮比。在每個環境中，演員網絡（Actor network）可以壓縮超過 96% 的參數，評論家網絡（Critic network）可以壓縮 85% 至 95% 的參數。SAC 算法的結果類似。另外，使用 SAC 算法的 RLx2 實現了 5 倍至 20 倍的模型壓縮。 節省訓練開銷 與基于知識蒸餾或行為克隆的方法 [Vischer et al. 2021] 不同，RLx2 在整個訓練過程中使用了稀疏網絡。因此，它具有加速訓練并節省計算資源的額外潛力。四個環境的平均結果表明，表格中 RLx2-TD3 分別減少了 12 倍和 20 倍的訓練和推理浮點運算數，RLx2-SAC 分別減少了 7 倍和 12 倍的訓練和推理浮點運算數。 總結 作者提出了一種用于離策略強化學習的稀疏訓練框架 RLx2，能夠適用于各種離策略強化學習算法。這一框架利用基于梯度的結構演化方法實現了高效的拓撲探索，并通過延遲多步差分目標和動態容量回放緩沖區建立了穩健的值函數學習。RLx2 不需要像傳統剪枝方法一樣預訓練稠密網絡，卻能夠在訓練過程中使用超稀疏網絡來訓練高效的深度強化學習智能體，并且幾乎沒有性能損失。作者在使用 TD3 和 SAC 的 RLx2 上進行了實驗，結果表明其稀疏訓練性能非常出色：模型壓縮比例為 7.5 倍至 20 倍，性能下降不到 3%，訓練和推理的浮點運算數分別減少高達 20 倍和 50 倍。作者認為未來有趣的工作包括將 RLx2 框架擴展到更復雜的 RL 場景，這些場景對計算資源的需求更高，例如多智能體、離線強化學習等場景，也包括真實世界的復雜決策問題而非標準的 MuJoCo 環境。 參考文獻1.Hessel, Matteo, et al. "Rainbow: Combining improvements in deep reinforcement learning." Proceedings of the AAAI conference on artificial intelligence. Vol. 32. No. 1. 2018.2.Mocanu, Decebal Constantin, et al. "Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science." Nature communications 9.1 (2018): 2383.3.Evci, Utku, et al. "Rigging the lottery: Making all tickets winners." International Conference on Machine Learning. PMLR, 2020.4.Sokar, Ghada, et al. "Dynamic sparse training for deep reinforcement learning." arXiv preprint arXiv:2106.04217 (2021).5.Munos, Rémi, et al. "Safe and efficient off-policy reinforcement learning." Advances in neural information processing systems 29 (2016).6.Fedus, William, et al. "Revisiting fundamentals of experience replay." International Conference on Machine Learning. PMLR, 2020.7.Vischer, Marc Aurel, Robert Tjarko Lange, and Henning Sprekeler. "On lottery tickets and minimal task representations in deep reinforcement learning." arXiv preprint arXiv:2105.01648 (2021).