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

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

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

Menger設(shè)計思路

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

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

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

高效的模型檢索

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

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

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

高通量輸入管道

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

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

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

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

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

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

實例驗證：芯片放置

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

與芯片放置的基準(zhǔn)相比，在不同TPU核心數(shù)量時使用Menger的訓(xùn)練時間的提升。

我們認(rèn)為，Menger架構(gòu)及其在芯片放置任務(wù)中優(yōu)異的表現(xiàn)為進一步縮短芯片設(shè)計周期提供了方向。同時，我們還可以應(yīng)用該架構(gòu)實現(xiàn)其他具有挑戰(zhàn)性的現(xiàn)實問題。

本文轉(zhuǎn)自雷鋒網(wǎng)，如需轉(zhuǎn)載請至雷鋒網(wǎng)官網(wǎng)申請授權(quán)。