RLax

RLax（发音为“relax”）是一个基于JAX构建的库，它公开了实现强化学习代理的有用构建块。完整文档可以在rlax.readthedocs.io找到。

安装

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

pip install rlax

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

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

然后，可以使用jax.jit将RLax代码即时编译为针对不同硬件（例如CPU、GPU、TPU）的版本。

为了运行examples/，您还需要克隆存储库并安装以下附加需求：optax、haiku和bsuite。

内容

提供的操作和函数并不是完整的算法，而是强化学习特定数学操作的实现，这些操作是构建能够学习的完整代理时所需的

• 值，包括状态和动作值；
• Bellman方程的非线性推广值。
• 回报分布，也称为分布值函数；
• 一般值函数，用于累积量（除了主要奖励以外的累积量）；
• 策略，通过策略梯度在连续和离散动作空间中进行。

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

请参阅函数级别的文档字符串以了解这些函数的文档以及引用介绍了这些函数的论文。

用法

请参阅examples/，了解使用RLax的一些函数实现几个简单的强化学习代理的示例，并在BSuite的Catch环境中演示学习（强化学习文献中代理开发的常见单元测试）

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

背景

强化学习研究学习系统（即代理）的问题，该系统必须学习与环境（即环境）相互作用。

代理和环境在离散步骤上交互。在每一步，代理选择一个动作，并得到环境状态（即观察）的（部分）快照和一个标量反馈信号（即奖励）。

代理的行为可以用一个概率分布来描述，该概率分布在给定过去的环境观察值下选择动作（即策略）。代理寻求一个策略，从任何给定的步骤开始，最大化从该点开始收集的折扣累积奖励（即回报）。

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

强化学习算法

强化学习算法有三个典型系列

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

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

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

命名约定和开发者指南

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

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

• q_tm1：转换中 source 状态的动作值。
• a_tm1：在 source 状态中选定的动作。
• r_t：在 destination 状态中收集的奖励。
• discount_t：与转换相关的 discount。
• q_t：在 destination 状态中的动作值。

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

引用 RLax

RLax 是 DeepMind JAX 生态系统的一部分，要引用 RLax，请使用 DeepMind JAX 生态系统引用。