RLax

CI 状态 pypi

RLax（发音为"relax"）是一个基于JAX构建的库，提供了实现强化学习代理所需的有用构建模块。完整文档可在 rlax.readthedocs.io 查看。

安装

您可以通过以下方式从PyPI安装RLax的最新发布版本：

pip install rlax

或者您可以从GitHub安装最新开发版本：

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

所有RLax代码随后可以使用jax.jit即时编译以适用于不同硬件（如CPU、GPU、TPU）。

为了运行examples/，您还需要克隆仓库并安装额外的依赖： optax、 haiku和 bsuite。

内容

提供的操作和函数并非完整算法，而是实现了构建能够学习的功能完备代理所需的强化学习特定数学操作：

值函数，包括状态值和动作值；
非线性广义Bellman方程的值函数；
回报分布，即分布式值函数；
通用值函数，适用于主要奖励以外的累积量；
策略，通过连续和离散动作空间中的策略梯度。

该库支持在策略和离策略学习（即从与代理策略不同的策略采样的数据中学习）。

有关这些函数的文档以及引入和/或使用它们的论文参考，请参阅文件级和函数级的文档字符串。

使用

查看examples/目录，了解如何使用RLax中的一些函数来实现简单的强化学习代理，并演示在BSuite版本的Catch环境（强化学习文献中常见的代理开发单元测试）上的学习过程。

使用rlax的其他JAX强化学习代理示例可以在bsuite中找到。

背景

强化学习研究学习系统（代理）如何学习与其所处环境（环境）交互的问题。

代理和环境在离散的步骤中交互。在每一步，代理选择一个动作，作为回报，它会获得环境状态的（部分）快照（观察）和一个标量反馈信号（奖励）。

代理的行为由一个基于过去环境观察条件的动作概率分布（策略）来表征。代理寻求一个策略，使得从任何给定步骤开始，都能最大化从该点开始的折扣累积奖励（回报）。

通常，代理策略或环境动态本身是随机的。在这种情况下，回报是一个随机变量，最优代理策略更准确地被指定为在代理和环境的随机性下最大化回报期望（值）的策略。

强化学习算法

有三类典型的强化学习算法：

估计状态和动作的值，并通过检查推断策略（例如，选择估计值最高的动作）
学习环境模型（能够预测观察和奖励），并通过规划推断策略
参数化可以直接执行的策略

无论如何，策略、值或模型都只是函数。在深度强化学习中，这些函数由神经网络表示。在这种设置下，通常将强化学习更新公式化为可微分的伪损失函数（类似于（无）监督学习）。在自动微分下，就可以恢复原始的更新规则。

然而请注意，只有当输入数据以正确方式采样时，这些更新才有效。例如，策略梯度损失仅在输入轨迹是当前策略的无偏样本时才有效；即数据是在策略的。该库无法检查或强制执行此类约束。但是，在函数的文档字符串中提供了描述每个操作使用方法的论文链接。

命名约定和开发者指南

我们定义了代理与单个经验流交互的函数和操作。JAX构造vmap可用于将这些函数应用于批次（例如，支持回放和并行数据生成）。

许多函数考虑连续时间步的策略、动作、奖励和值，以计算其输出。在这种情况下，后缀_t和tm1通常用于澄清每个输入在哪一步生成，例如：

q_tm1：转换源状态中的动作值。
a_tm1：在源状态中选择的动作。
r_t：在目标状态中收集的结果奖励。
discount_t：与转换相关的折扣。
q_t：目标状态中的动作值。

每个函数都提供了广泛的测试。所有测试还应验证使用jax.jit编译为XLA和使用jax.vmap执行批处理操作时rlax函数的输出。

引用RLax

本仓库是DeepMind JAX生态系统的一部分，要引用Rlax请使用以下引用：

@software{deepmind2020jax,
  title = {The {D}eep{M}ind {JAX} {E}cosystem},
  author = {DeepMind and Babuschkin, Igor and Baumli, Kate and Bell, Alison and Bhupatiraju, Surya and Bruce, Jake and Buchlovsky, Peter and Budden, David and Cai, Trevor and Clark, Aidan and Danihelka, Ivo and Dedieu, Antoine and Fantacci, Claudio and Godwin, Jonathan and Jones, Chris and Hemsley, Ross and Hennigan, Tom and Hessel, Matteo and Hou, Shaobo and Kapturowski, Steven and Keck, Thomas and Kemaev, Iurii and King, Michael and Kunesch, Markus and Martens, Lena and Merzic, Hamza and Mikulik, Vladimir and Norman, Tamara and Papamakarios, George and Quan, John and Ring, Roman and Ruiz, Francisco and Sanchez, Alvaro and Sartran, Laurent and Schneider, Rosalia and Sezener, Eren and Spencer, Stephen and Srinivasan, Srivatsan and Stanojevi\'{c}, Milo\v{s} and Stokowiec, Wojciech and Wang, Luyu and Zhou, Guangyao and Viola, Fabio},
  url = {http://github.com/deepmind},
  year = {2020},
}