RLax（發音為“ relax”）是建立在JAX之上的庫，它公開了用于實施強化學習智能體的有用構建塊。。
報道：深度強化學習實驗室
作者：DeepRL
來源：Github/DeepMind

RLax及背景

強化學習主要用于研究學習系統（智能體）的問題，該學習系統必須學習與環境進行交互的信息。智能體和環境在不連續的步驟上進行交互。在每個步驟中，智能體都會選擇一個動作，并會提供一個返回的環境狀態（觀察）狀態（部分）和標量反饋信號（獎勵）。智能體的行為以行為的概率分布為特征，該分布取決于對環境（策略）的過去觀察。智能體尋求一種策略，該策略將從任何給定步驟中最大化從該點開始（返回）將收集的折扣累積獎勵。智能體策略或環境動態本身通常是隨機的。在這種情況下，回報是一個隨機變量，并且通常將更精確的智能體策略指定為在智能體和環境的隨機性下最大化回報期望（值）的策略。

RLax（發音為“ relax”）是建立在JAX之上的庫，它公開了用于實施強化學習智能體的有用構建塊。

安裝與使用

可以使用以下命令從github直接使用pip安裝RLax：

pip install git+git://github.com/deepmind/rlax.git.

然后可以使用JAX的jax.jit函數為不同的硬件（例如CPU，GPU，TPU）及時編譯所有RLax代碼。

強化學習算法

增強學習算法包含三個原型系列：

• 那些估計狀態和動作的價值，并通過檢查推斷策略的行為（例如，通過選擇估計值最高的動作）
• 學習環境模型（能夠預測觀察結果和獎勵），并通過計劃推斷策略。
• 那些參數化可以直接執行的策略的參數，

無論如何，策略，價值或模型只是功能。在深度強化學習中，此類功能由神經網絡表示。在這種情況下，通常將強化學習更新公式化為可區分的損失函數（類似于（非）監督學習）。在自動區分下，將恢復原始更新規則。

但是請注意，尤其是只有以正確的方式對輸入數據進行采樣時，更新才有效。例如，僅當輸入軌跡是當前策略的無偏樣本時，策略梯度損失才有效。即數據是符合政策的。該庫無法檢查或強制執行此類約束。但是，在函數的文檔字符串中提供了指向描述如何使用每種操作的論文的鏈接。

命名約定和開發人員指南

我們為與單一經驗流交互的代理定義功能和操作。JAX構造vmap可用于將這些相同的功能應用于批處理（例如，支持重放和并行數據生成）。

許多功能在連續的時間步中考慮策略，行動，獎勵，價值，以便計算其輸出。在這種情況下，qm_t和tm1通常用于說明每個輸入是在哪個步驟上生成的，例如：

• q_tm1：轉換的源狀態中的操作值。
• a_tm1：在源狀態下選擇的操作。
• r_t：在目標狀態下收集的結果獎勵。
• discount_t：與轉換關聯的折扣。
• q_t：目標狀態下的操作值。

為每個功能提供了廣泛的測試。當使用jax.jit編譯為XLA以及使用jax.vmap執行批處理操作時，所有測試還應驗證rlax函數的輸出。

引用

@software{rlax2020github,
  author = {David Budden and Matteo Hessel and John Quan and Steven Kapturowski},
  title = {{RL}ax: {R}einforcement {L}earning in {JAX}},
  url = {http://github.com/deepmind/rlax},
  version = {0.0.1a0},
  year = {2020},
}

專注深度強化學習前沿技術干貨，論文，框架，學習路線等，歡迎關注微信公眾號。

審核編輯 黃昊宇