在過去的十年中，強化學習(RL)成為機器學習中最受矚目的研究領域之一，應用RL能夠很好地解決芯片放置和資源管理等復雜的問題，以及Go/Dota 2/hide-and-seek等有挑戰性的游戲。簡單來說，RL基礎架構就是數據采集和訓練的循環，Actor根據環境收集樣本數據，然后將其傳輸給Learner來訓練和更新模型。當前大多數RL實現都需要對環境中成千上萬個樣本進行多次迭代，以學習目標任務，如Dota 2每2秒要學習成千上萬幀樣本。這樣，RL架構不僅要有很強的數據處理能力，例如增加Actor實現大量樣本的采集，而且還應能夠在訓練過程中快速迭代這些樣本。

Actor與Learner交互的RL架構。Learner使用采樣數據訓練模型，并將更新后的模型傳遞給Actor(例如TF-Agents，IMPALA)。

今天，我們要介紹的是Menger——一種具有本地化推理能力的大規模分布式RL架構，可通過多個處理集群(如Borg單元)擴展數千個Actors，從而減少了芯片放置任務的訓練時間。在接下來的章節，我們介紹了如何使用Google TPU配置Menger，從而提高訓練速度，進一步我們通過芯片放置任務驗證框架的性能和可擴展性。可以發現，與基準模型相比，Menger將培訓時間減少了8.6倍。

Menger設計思路

當前有各種各樣的分布式RL系統，如Acme和SEED RL，然而，這些系統往往只從一個特定角度對分布式強化學習系統進行優化。例如，Acme從頻繁的Learner獲取模型，使每個Actor都進行本地推理，而SEED RL則通過分配一部分TPU內核執行批量調用，進行集中推理。對通信成本和推理成本的衡量是不同優化系統的區別，具體包括：(1)向/從集中式推理服務器發送/接收觀察和動作的通信成本，或從Larner獲取模型的通信成本;(2)相比加速器(TPU/GPU)成本，Actor的推理成本大小。考慮到觀察值、動作和模型大小等目標程序要求，Menger使用類似Acme的局部推理，但同時盡可能的增加Actor的可擴展性。要實現良好擴展性和訓練速度，主要挑戰包括以下兩點：

Actor向Learner進行大量讀取請求以進行模型檢索，這就造成Learner的負擔，隨著Actor數量的增加模型表現明顯受限(如收斂時間的顯著增長)。

在將訓練數據輸送給TPU計算核心時，TPU性能通常受到輸入管道效率的限制。隨著TPU計算核心數量的增加(如TPU Pod)，輸入管道的性能對于訓練時間的影響更加明顯。

高效的模型檢索

為應對第一個挑戰，在TensorFlow代碼中，我們在Learner和Actor之間引入了透明的分布式緩存組件，并通過Reverb進行優化(類似于Dota中使用的方法)。緩存組件的主要職責是對Actor的大量請求和Learner的處理能力進行平衡。通過添加這些緩存組件，不僅顯著減輕了過多請求對Learner的壓力，而且以少量的通信成本將Actor分配給多個Borg單元。我們的研究表明，對有512個Actors、大小為16MB的模型，引入緩存組件可以將平均讀取延遲降低約4.0倍，從而實現更快的訓練迭代，在PPO等策略算法中效果更加明顯。

若干Actors放置在不同Borg單元的分布式RL系統。 不同Borg單元大量Actors的頻繁的模型更新請求限制了Learner性能以及Learner于Actor之間的通信網絡，從而導致總體收斂時間顯著增加。虛線表示不同機器之間的gRPC通信。

引入透明分布式緩存服務的分布式RL系統。多個Actor放置在不同的Borg單元中，Learner僅將更新的模型發送給分布式緩存組件。每個緩存組件應對鄰近Actor和緩存區的模型請求更新。緩存區不僅減輕了Learner對模型更新請求提供服務的負擔，而且減少了Actor的平均讀取延遲。

高通量輸入管道

為提高輸入數據管道的吞吐量，Menger使用了Reverb——一種專為機器學習應用設計的、新型開源數據存儲系統。在Reverb中，可以以在線或離線算法進行經驗回放。但是，單個Reverb當前無法擴展到有成千上萬Actors的分布式RL系統，并且Actor的寫入吞吐量效率很低。

具有單個回放緩沖區的分布式RL系統。在Actor提出大量思維寫請求后，回放緩沖區會受到限制并降低總體吞吐量。此外，一旦我們將Learner擴展到具有多個計算引擎(TPU Pod)，單個回放緩沖區給這些引擎提供數據的效率就變得很低，嚴重影響總體收斂時間。

為了更好地了解回放緩沖器在分布式RL系統中的效率，我們評估了在不同負載大小(16 MB-512 MB)和不同Actor(16-2048)情況下的平均寫入延遲。我們將回放緩沖區和Actor放置在同一個Borg單元中。，可以發現，隨著Actor數量的增加，平均寫入延遲顯著增加。將Actor的數量從16擴展2048，16MB和512MB大小的有效負載的平均寫入延遲分別增加了約6.2倍和約18.9倍。這樣的寫入等待時間的增加影響了數據收集時間，導致訓練效率低下。

將不同大小的有效負載(16 MB-512 MB)和不同數量的Actor(16至2048)放置在同一個Borg單元上時，單個Reverb重回放緩沖區的平均寫入延遲。

為緩解這種情況，我們使用Reverb的分片功能來增加Actor、Learner和和回放緩沖區之間的吞吐量。分片可在多個回放緩沖服務器之間平衡大量Actor的寫入負載，而不是僅僅作用于單個回放緩沖服務器，同時由于少數Actor共享同一服務器，可以最小化每個回放緩沖服務器的平均寫入延遲。這樣Menger就可以在多個Borg單元中擴展數千個Actor。

具有分片回放緩沖區的分布式RL系統。每個回放緩沖區用于存儲位于同一Borg單元上的特定Actor。此外，分片回放緩沖區為加速器內核提供了具有更高吞吐量的輸入管道。

實例驗證：芯片放置

我們面向大型網表在芯片放置任務中測試了Menger。與基準相比，Menger使用512個TPU內核，在訓練時間上有了顯著改善(最高提升約8.6倍，即在最優配置下，可以將訓練時間從約8.6小時減少到1小時)。盡管Menger針對TPU進行了優化，但該框架才是性能提升的關鍵因素，我們預計在GPU上實驗也會看到類似的提升。

與芯片放置的基準相比，在不同TPU核心數量時使用Menger的訓練時間的提升。

我們認為，Menger架構及其在芯片放置任務中優異的表現為進一步縮短芯片設計周期提供了方向。同時，我們還可以應用該架構實現其他具有挑戰性的現實問題。

本文轉自雷鋒網，如需轉載請至雷鋒網官網申請授權。