天天看點

Speeding up DQN on PyTorch: how to solve Pong in 30 minutes

Speeding up DQN on PyTorch: how to solve Pong in 30 minutes

Intro

  前段時間我實作了文章​​Rainbow: Combining Improvements in Deep Reinforcement Learning using PyTorch​​和我的名為PTAN的小型強化學習庫中的所有模型。如果你好奇,​​這裡​​有八個系統的代碼。

  為了調試和測試它,我使用了Atari套件中的Pong遊戲,主要是因為它的簡單性、快速收斂性和超參數魯棒性:你可以将回放緩存尺寸縮小10到100倍,它仍然會很好地收斂。這對于無法通路谷歌員工擁有的計算資源的深度強化學習愛好者來說非常有幫助。在代碼的實作和調試過程中,我需要運作大約100-200次優化,是以,一次運作需要多長時間并不重要:2-3天或一個小時。

  盡管如此,你始終應該在這裡保持平衡:嘗試盡可能多地壓縮性能,您可能會引入錯誤,這将使已經很複雜的調試和實作過程變得非常複雜。是以,在Rainbow論文上的所有系統都實作之後,我問自己一個問題:是否有可能使我的實作更快,不僅能夠在Pong上訓練,而且能夠挑戰其餘的遊戲,這至少需要50M幀進行訓練,如SeaQuest、River Raid、Breakout等。

  由于我的計算資源非常有限,隻有兩個1080Ti + 一個1080,(在現在已經非常有限了),唯一的辦法就是讓代碼更快。

Initial numbers

  作為起點,我采用了具有以下超參數的經典DQN版本:

  • 使用了gym 0.9.3的環境PongNoFrameskip-v4;
  • 前100k幀的epsilon從1.0衰減到0.02,然後epsilon保持0.02;
  • 每1k幀同步目标網絡;
  • 大小為100k的簡單回放緩存最初在訓練前預取10k次轉換;
  • Gamma=0.99;
  • Adam學習率為1e-4;
  • 每個訓練步驟,從環境中添加一個轉換到回放緩存,并在從回放緩存均勻采樣的32個轉換中執行訓練;
  • 當最近100次遊戲的平均得分大于18時,認為Pong已解決。

  應用于環境的包裝器對于速度和收斂都非常重要(前段時間,我浪費了兩天時間試圖在工作代碼中找到一個錯誤,該錯誤僅僅因為缺少"Fire at reset"包裝器而拒絕收斂。是以,我前段時間從​​OpenAI基準項目​​中​​使用的包裝器清單​​:

  • EpisodicLifeEnv:在每一次失去生命時結束回合,這有助于更快地收斂;
  • NoopResetEnv:在重置時執行随機數量的NOOP動作;
  • MaxAndSkipEnv:為4個Atari環境幀重複選擇的動作以加速訓練;
  • FireResetEnv:在開始時按下開火。一些環境需要這個才能開始遊戲。
  • ProcessFrame84:幀轉換為灰階并縮小到84*84像素;
  • FrameStack:傳遞最後4幀作為觀察;
  • ClippedRewardWrapper:裁剪獎勵到-1..+1範圍。

  在GTX 1080Ti上運作的初始代碼版本顯示訓練期間每秒154次觀察的速度,并且根據初始随機種子可以在60到90分鐘解決Pong。這就是我們的出發點。從這個角度來看,RL論文通常使用的100M幀需要我們耐心等待7.5天。

Speeding up DQN on PyTorch: how to solve Pong in 30 minutes

Change 1: larger batch size + several steps

  我們通常用于加速深度學習訓練的第一個想法是更大的批量。它适用于深度強化學習領域,但這裡需要小心。在正常的監督學習案例中,一個簡單的規則"大批量更好"通常是正确的:你隻需增加你的批量,直到你的GPU記憶體允許極限,并且更大的批量通常意味着将在機關時間内處理更多的樣本,這要歸功于巨大的GPU并行性。

  強化學習案例略有不同。在訓練期間,兩件事同時發生:

  • 你的網絡經過訓練可以更好地預測目前資料;
  • 你的智能體正在探索環境。

  當智能體探索環境并了解其動作的結果時,訓練資料正在發生變化。例如,在射擊遊戲中,你的智能體可以随機運作一段時間,被怪物擊中,在訓練緩存中隻有"死亡無處不在"的悲慘經曆。 但過了一會兒,智能體會發現他有一個可以使用的武器。這種新經驗可以極大地改變我們用于訓練的資料。

  RL收斂通常建立在訓練和探索之間的脆弱平衡上。如果我們隻是在不調整其他選項的情況下增加批量大小,我們很容易過拟合目前資料(對于上面的射擊遊戲示例,你的智能體可能會開始認為"早點去死"是最大程度減少痛苦的唯一選擇,并且永遠無法發現它有的槍)。

  是以,在​​02_play_steps.py​​中,我們在每個訓練循環中執行多個步驟,并使用批次大小乘以該步驟數。但是我們需要小心這個步數參數。更多的步驟意味着更大的批量大小,這應該會導緻更快的訓練,但同時在訓練之間做很多步驟可以用從舊網絡獲得的樣本填充我們的緩存。

  為了找到最佳位置,我使用随機種子(你需要同時通過numpy和pytorch)修複了訓練過程,并針對各個步驟對其進行了訓練。

  • steps=1: speed 154 f/s (顯然,它與原始版本相同)
  • steps=2: speed 200 f/s (+30%)
  • steps=3: speed 212 f/s (+37%)
  • steps=4: speed 227 f/s (+47%)
  • steps=5: speed 228 f/s (+48%)
  • steps=6: speed 232 f/s (+50%)

  收斂動态幾乎相同(見下圖),但速度增加在4步左右飽和,是以,我決定堅持這個數字進行後續實驗。

  好的,我們的性能提升了47%。

Speeding up DQN on PyTorch: how to solve Pong in 30 minutes

Change 2: play and train in separate processes

  在這一步中,我們将檢查我們的訓練循環,它基本上包含以下步驟的重複:

  1. 使用目前網絡在環境中玩N個步驟以選擇動作;
  2. 将這些步驟中的觀察值放入回放緩存;
  3. 從回放緩存中随機采樣批次;
  4. 在這批上訓練。

  前兩步的目的是用來自環境的樣本(觀察、動作、獎勵和下一次觀察)填充回放緩存。最後兩步是訓練我們的網絡。

  上述步驟及其與環境的通信、GPU上的DQN和回放緩存的說明如下圖所示。

Speeding up DQN on PyTorch: how to solve Pong in 30 minutes

  正如我們所看到的,環境僅被第一步使用,我們訓練的上半部分和下半部分之間的唯一連接配接是我們的回放緩存。由于這種資料獨立性,我們可以并行運作兩個程序:

  • 第一個程序将與環境通信,向回放緩存提供新資料;
  • 第二個程序将從回放緩存中采樣訓練批次并執行訓練。

  這兩個活動應該同步運作,以保持我們在上一節中讨論的訓練/探索平衡。

  這個想法是在​​03_parallel.py​​中實作的,并且正在使用torch.multiprocessing子產品來并行玩遊戲和訓練仍然能夠與GPU同時工作。為了盡量減少其他類中的修改,隻有第一步(環境通信)放在單獨的過程中。使Queue類将獲得的觀察結果傳輸到訓練循環。

Speeding up DQN on PyTorch: how to solve Pong in 30 minutes

  這個新版本的基準測試顯示令人印象深刻的395幀/秒,與之前的版本相比增加了74%,與代碼的原始版本相比增加了156%。

Change 3: async cuda transfers

  下一步很簡單:每次我們調用Tensor的cuda()方法時,我們都會傳遞async=True參數,這将禁止等待傳輸完成。它不會給你帶來非常令人印象深刻的加速,但有時會給你一些東西并且很容易實作。

  此版本位于檔案​​04_cuda_async.py​​中,唯一的差別是将cuda_async=True傳遞給calc_loss函數。

  基準測試後,我獲得了406幀/秒的訓練速度,比上一步提高了3.5%,與原始DQN相比提高了165%。

Change 4: latest Atari wrappers

  正如我之前所說,DQN的原始版本使用了一些來自OpenAI基準項目的舊Atari包裝器。幾天前,這些包裝器被更改為名為​​"更改atari預處理以使用更快的opencv"​​的送出,這絕對值得一試。

  這是​​基準存儲庫中包裝器的新代碼​​。DQN的下一個版本在​​05_new_wrapper.py​​中。由于我還沒有将新的包裝器拉入ptan庫中,是以它們在示例中位于單獨的庫中。

  基準測試結果是484幀/秒,比上一步增加了18%,最終比原始版本增加了214%。

Summary

繼續閱讀