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

RLax及背景

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

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

安裝與使用

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

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

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

強化學習算法

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

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

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

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

命名約定和開發(fā)人員指南

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

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

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

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

引用

@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},
}

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

審核編輯 黃昊宇