具有Keras和Tensorflow Eager的功能性RL

作者 | 梁启超

来源 | Medium

编辑 | 代码医生团队

在此博客文章中，探索了用于实现强化学习（RL）算法的功能范例。范例是开发人员将其算法的数值写为独立的纯函数，然后使用库将其编译为可以大规模训练的策略。分享了如何在RLlib的策略构建器API中实现这些想法，消除了数千行“胶水”代码，并为Keras和TensorFlow 2.0提供支持。

为什么要进行函数式编程？

函数式编程的主要思想之一是程序可以主要由纯函数组成，即，其输出完全由其输入决定的函数。少得多的是：通过对功能可以执行的限制，获得了更容易地推理和操纵其执行的能力。

在TensorFlow中，可以使用占位符输入象征性地执行张量的此类功能，也可以使用实际的张量值急切地执行这些功能。由于此类函数没有副作用，因此无论是符号调用还是多次调用它们，它们对输入都具有相同的效果。

功能强化学习

考虑代理状态数据的以下损失函数，其中包括当前状态s，操作a，返回r和策略π：

L（s，a，r）=-[log π（s，a）] * r

如果不熟悉RL，那么所有这些功能就是说，应该尝试提高采取良好行动（即增加未来收益的行动）的可能性。这种损失是策略梯度算法的核心。正如将看到的，定义损失几乎是开始在RLlib中训练RL策略所需要的全部。

给定一系列部署，策略梯度损失将设法提高采取良好行动的可能性（即，在上面的此Pong示例中导致成功的行动）。

到Python的直接翻译如下。在这里，损失函数取（π，s，a，r），计算π（s，a）作为离散动作分布，并返回动作的对数概率乘以收益：

def loss(model, s: Tensor, a: Tensor, r: Tensor) -> Tensor:
    logits = model.forward(s)
    action_dist = Categorical(logits)
    return -tf.reduce_mean(action_dist.logp(a) * r)

此功能定义有很多好处。首先请注意，损失是很自然的理解- 在RL实现中通常没有占位符，控制循环，外部变量访问或类成员。其次，由于它不会改变外部状态，因此它与TF图和渴望模式执行兼容。

与基于类的API（其中类方法可以访问类状态的任意部分）相反，功能性API从松散耦合的纯函数中构建策略。

在此博客中，探索将RL算法定义为此类纯函数的集合。范例是开发人员将算法的数字编写为独立的纯函数，然后使用RLlib帮助器函数将其编译为可以大规模训练的策略。该建议在RLlib库中具体实现。

带有RLlib的功能性RL

RLlib是一个用于强化学习的开源库，它为各种应用程序提供高可伸缩性和统一的API。它提供了多种可扩展的RL算法。

RLlib如何扩展算法的示例，在这种情况下为分布式同步采样。

鉴于PyTorch（即命令执行）的日益普及和TensorFlow 2.0的发布，看到了通过功能性地重写RLlib算法来改善RLlib开发人员体验的机会。主要目标是：

改善RL调试经验

• 只需要一个eager标志，就可以将eager执行用于任何算法，从而实现简单的print（）调试。

简化新算法的开发

• 通过用从纯函数（例如TRFL提供的原语）集合构建的策略替换单片“ Agent”类，使算法更易于自定义和理解。
• 无需手动声明TF的张量占位符。
• 统一定义TF和PyTorch策略的方式。

Policy Builder API

用于功能性RL的RLlib策略构建器API（在RLlib 0.7.4中稳定）仅涉及两个关键功能：

• build_tf_policy（）
• build_torch_policy（）

从较高的角度来看，这些构建器将许多函数对象作为输入，包括与之前看到的相似的loss_fn，给定算法配置以返回神经网络模型的model_fn以及给定模型输出以生成动作样本的action_fn。实际的API需要更多的参数，但这是主要的参数。构建器将这些功能编译为一个策略，可以查询操作并在给定经验的情况下随着时间的推移进行改进：

这些策略可用于RLlib中的单代理，矢量和多代理训练，并要求它们确定如何与环境交互：

发现策略构建器模式足够通用，可以移植几乎所有RLlib参考算法，包括TensorFlow中的A2C，APPO，DDPG，DQN，PG，PPO，SAC和IMPALA，以及PyTorch的PG / A2C。尽管代码的可读性在一定程度上是主观的，但用户报告说，构建器模式使自定义算法更加容易，尤其是在Jupyter笔记本电脑等环境中。此外，这些重构已经高达几百行代码的减少了算法的大小每个。

香草政策梯度示例

RLlib中香草策略梯度损失函数的可视化。

看一下如何使用构建器模式来具体实现前面的损失示例。定义policy_gradient_loss，需要进行一些调整以实现一般性：（1）RLlib提供适当的distribution_class，以便算法可以处理任何类型的操作空间（例如，连续或分类），以及（2）将经验数据保存在其中一个train_batch字典，其中包含状态，动作等张量：

def policy_gradient_loss(
        policy, model, distribution_cls, train_batch):
    logits, _ = model.from_batch(train_batch)
    action_dist = distribution_cls(logits, model)
    return -tf.reduce_mean(
        action_dist.logp(train_batch[“actions”]) *
        train_batch[“returns”])

要将“ returns”数组添加到批处理中，需要定义一个后处理函数，将其计算为轨迹上的时间折扣奖励：

在代码中计算以下R（T）时，将γ= 0.99设置为：

from ray.rllib.evaluation.postprocessing import discount
# Run for each trajectory collected from the environment
def calculate_returns(policy,
                      batch,
                      other_agent_batches=None,
                      episode=None):
   batch[“returns”] = discount(batch[“rewards”], 0.99)
   return batch

有了这些功能，就可以构建RLlib策略和训练师（协调整个训练工作流程）。如果未指定，则模型和操作分布由RLlib自动提供：

MyTFPolicy = build_tf_policy(
   name="MyTFPolicy",
   loss_fn=policy_gradient_loss,
   postprocess_fn=calculate_returns)
MyTrainer = build_trainer(
   name="MyCustomTrainer", default_policy=MyTFPolicy)

现在，可以使用Tune在所需的规模上运行此示例，在此示例中显示了在群集中使用128个CPU和1个GPU的配置：

tune.run(MyTrainer,
    config={“env”: “CartPole-v0”,
            “num_workers”: 128,
            “num_gpus”: 1})

尽管此示例（可运行代码）只是一种基本算法，但它演示了功能API如何变得简洁，可读和可高度扩展。与以前使用TF占位符在RLlib中定义策略的方法相比，该功能性API使用的代码行减少了大约3倍（23行对81行），并且还非常有用：

将旧的基于类的API与新的功能策略构建器API进行比较。两种策略都实现相同的行为，但是功能定义要短得多。

策略生成器的工作方式

在底层，build_tf_policy接受提供的构建块（model_fn，action_fn，loss_fn等），并将其编译为DynamicTFPolicy或EagerTFPolicy，具体取决于是否启用了TF急切执行。前者实现图形模式执行（动态地自动定义占位符），后者渴望执行。

DynamicTFPolicy和EagerTFPolicy之间的主要区别是它们调用传入的函数的次数。在两种情况下，一次调用一次model_fn来创建Model类。但是，涉及张量运算的函数要么在图模式下调用一次以构建符号计算图，要么在实际张量下以急切模式多次调用。在下图中，以蓝色和橙色显示这些操作如何一起工作：

生成的EagerTFPolicy概述。该策略通过model.forward（）传递环境状态，该状态发出输出logit。模型输出参数化了动作的概率分布（“ ActionDistribution”），可在对动作或训练进行采样时使用。损失函数是在大量经验中运行的。该模型可以根据损失函数的需要提供其他方法，例如值函数（浅橙色）或其他用于计算Q值的方法等（未显示）。

RLlib启动和扩展RL训练所需的所有政策对象。直观地讲，这是因为它封装了如何计算操作和改进策略的方法。外部状态（例如环境状态和RNN隐藏状态）由RLlib从外部进行管理，并且不需要成为策略定义的一部分。根据是在计算部署还是在给定大量部署数据的情况下尝试改进策略，以两种方式之一使用策略对象：

推论：正向传递以计算单个动作。这仅涉及查询模型，生成动作分布以及从该分布中采样动作。在急切模式下，这涉及到调用action_fn（动作采样器的DQN示例），该函数创建一个相关的动作分配/动作采样器，然后从中进行采样。

训练：前进和后退，以学习一系列经验。在这种模式下，调用损失函数以生成标量输出，该标量输出可用于通过SGD优化模型变量。在紧急模式下，将同时调用action_fn和loss_fn来分别生成操作分配和策略丢失。请注意这里没有显示通过action_fn进行的区分，但这确实发生在DQN之类的算法中。

松散的结局：国家管理

RL训练固有地涉及很多状态。如果使用纯函数定义算法，那么状态将保持在哪里？在大多数情况下，它可以由框架自动管理。RLlib中需要管理三种状态：

1. 环境状态：这包括环境的当前状态以及在策略步骤之间传递的任何重复状态。RLlib在其推出工作程序实现中内部进行管理。
2. 模型状态：这些是我们试图通过RL损失学习的策略参数。对于图形和急切模式，必须以相同的方式访问和优化这些变量。幸运的是，Keras模型可以在任何一种模式下使用。RLlib 基于面向对象的Keras样式提供了可定制的模型类（TFModelV2），用于保存策略参数。
3. 训练工作流状态：用于管理训练的状态，例如，各种超参数的退火时间表，自上次更新以来的步骤等。RLlib允许算法作者将混合类添加到可以容纳任何此类额外变量的策略。

松散的结局：渴望开销

接下来，通过打开或关闭快速跟踪来研究RLlib的快速模式性能。如下图所示，跟踪大大提高了性能。但是，要权衡的是可能不会每次都调用诸如print之类的Python操作。因此，默认情况下，RLlib中的跟踪处于关闭状态，但可以使用“ eager_tracing”启用：True。另外，还可以设置“ no_eager_on_workers”，使其仅对学习启用渴望，而对推理禁用：

在笔记本电脑处理器上使用“ rllib train -run = PG -env = <env> [-eager [-trace]]”测量的急切推断和梯度开销。随着时间的推移，热切地为小批量操作增加了可观的开销。但是，启用跟踪时，它通常比图形模式快或快。

结论

回顾一下，在这篇博客文章中，建议使用函数式编程的思想来简化RL算法的开发。在RLlib中实现并验证了这些想法。除了使支持新功能（如渴望执行）变得容易之外，还发现功能范式导致代码更加简洁和易于理解。使用“ pip install ray [rllib]”或通过检查文档和源代码自己尝试一下。

https://ray.readthedocs.io/en/latest/rllib.html

https://github.com/ray-project/ray/tree/master/rllib

推荐阅读

终版API已定型，TensorFlow 2.0 Beta蜕变归来