谷歌推出分佈式強化學習框架 SEED RL，計算速度比 IMPALA 提高了 80 倍！

近年來，Go 和 Dota2 等博弈類遊戲的火爆從側面反映了強化學習（RL）在過去幾年裏取得的成就。模型或代理通過（例如在遊戲中）探索環境來學習，同時針對特定的目標進行優化。然而，目前哪怕要成功學習非常簡單的遊戲，所需要的強化學習訓練體量也非常龐大，這使得研究和產品創意的迭代計算既昂貴又耗時，因此我們亟需一個能同時解決這些問題的“超級”強化學習框架。

在論文“SEED RL：通過加速集中推理實現高效可擴展的 Deep-RL”中，我們提出了一個強化學習代理，它可以規模部署到上千臺機器上，並能以每秒數百萬幀的速度進行訓練，顯著提高了計算效率。這是通過一種新的架構實現的，它通過集中式模型推理和引入快速通信層來大規模利用加速器（GPU 和 TPU）。我們驗證了 SEED RL 在流行的強化學習基準測試上的性能，如谷歌研究足球、Arcade 學習環境和 DeepMind Lab，結果表明，使用更大的模型可以提高數據效率。這些代碼已經在 Github 上開源，並提供了使用 GPU 在谷歌雲上運行的示例。

項目地址： https://github.com/google-research/seed_rl

當前的分佈式架構

IMPALA 等上一代分佈式強化學習代理，使用了專門用於數值計算的加速器，多年來，就像（非）監督學習一樣，在速度和效率方面都有了提升。強化學習代理的架構通常分爲行動者（actors）和學習者（learners）。行動者通常在 CPU 上運行，迭代在環境中執行動作以及在模型上運行推理來預測下一個動作的過程。行動者會經常更新推理模型的參數，在收集到足夠多的觀察數據後，將觀察和動作的軌跡發送給學習者，學習者再對模型進行優化。在這個架構中，學習者使用來自數百臺機器的分佈式推理輸入在 GPU 上訓練模型。

早期 RL 代理 IMPALA 的架構示例。推理在行動者上進行，通常使用低效的 CPU。更新後的模型參數頻繁地從學習者發送給行動者，增加了帶寬需求。

強化學習代理（如 IMPALA）的架構有如下幾個缺點：

與使用加速器相比，使用 CPU 進行神經網絡推理的效率和速度要低得多，而且隨着模型變大、變多，這會帶來問題。
在行動者和學習者之間發送參數和中間模型狀態所需的帶寬可能成爲瓶頸。
在一臺機器上處理兩個完全不同的任務（即環境渲染和推理）不太可能最優地利用機器資源。

SEED RL 架構

SEED RL 架構就是爲了解決這些缺陷而設計的。使用這種方法，神經網絡推理由學習者在專門的硬件（GPU 或 TPU）上集中完成，通過在本地保存模型參數和狀態，實現了推理加速，避免了數據傳輸瓶頸。雖然在每個環境步驟中都將觀察結果發送給學習者，但是得益基於 gRPC 框架和異步流 RPC 非常高效的網絡庫，延遲很低。這使得在一臺機器上可以實現每秒一百萬次查詢。學習者可以擴展到數千個內核（例如，在 Cloud TPUs 上可以擴展到 2048 個），行動者的數量可以擴展到數千臺機器來充分利用學習者，這使得以每秒數百萬幀的速度訓練成爲可能。SEED RL 基於 TensorFlow 2 API，在我們的實驗中，TPUs 加速了 SEED RL。

爲了使該框架成功運行，兩個最先進的算法被集成到 SEED RL 中。第一個是 V-trace，這是一種基於策略梯度的方法，首先由 IMPALA 引入。通常，基於策略梯度的方法可以預測一個動作的分佈，從中可以對動作進行抽樣。然而，由於行動者和學習者在 SEED RL 中是異步執行的，所以行動者的策略稍微落後於學習者的策略，即變成了離線策略（off-policy）。通常的基於策略梯度的方法是在線策略（on-policy），這意味着，行動者和學習者有相同的策略，在離線策略設置中會遇到收斂性和數值問題。V-trace 是一種離線策略方法，因此，在異步 SEED RL 架構中可以工作得很好。

第二種算法是 R2D2，這是一種 Q-learning 方法，它使用遞歸分佈式重播根據預測的動作未來值來選擇動作。這種方法使得 Q-learning 算法可以大規模運行，同時仍然允許使用遞歸神經網絡，該網絡可以根據一個事件中所有過去幀的信息預測未來的值。