DeepMind发布强化学习库RLax

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

专注深度强化学习前沿技术干货，论文，框架，学习路线等，欢迎关注微信公众号。

审核编辑 黄昊宇