来自本科生的暴击：清华开源「天授」强化学习平台，纯PyTorch实现

机器之心

发布于 2020-04-02 11:07:13

5920

发布于 2020-04-02 11:07:13

文章被收录于专栏：机器之心机器之心

机器之心报道

项目作者：thu-ml

参与：思、肖清

训练模型的极速，与 1500 行源代码的精简，清华大学新开源强化学习平台「天授」。值得注意的是，该项目的两位主要作者目前都是清华大学的本科生。

是否你也有这样的感觉，成熟 ML 工具的源码很难懂，各种继承与处理关系需要花很多时间一点点理清。在清华大学开源的「天授」项目中，它以极简的代码实现了很多极速的强化学习算法。重点是，天授框架的源码很容易懂，不会有太复杂的逻辑关系。

项目地址：https://github.com/thu-ml/tianshou

天授（Tianshou）是纯基于 PyTorch 代码的强化学习框架，与目前现有基于 TensorFlow 的强化学习库不同，天授的类继承并不复杂，API 也不是很繁琐。最重要的是，天授的训练速度非常快，我们试用 Pythonic 的 API 就能快速构建与训练 RL 智能体。

目前天授支持的 RL 算法有如下几种：

Policy Gradient (PG)
Deep Q-Network (DQN)
Double DQN (DDQN) with n-step returns
Advantage Actor-Critic (A2C)
Deep Deterministic Policy Gradient (DDPG)
Proximal Policy Optimization (PPO)
Twin Delayed DDPG (TD3)
Soft Actor-Critic (SAC)

另外，对于以上代码天授还支持并行收集样本，并且所有算法均统一改写为基于 replay-buffer 的形式。

速度与轻量：「天授」的灵魂

天授旨在提供一个高速、轻量化的 RL 开源平台。下图为天授与各大知名 RL 开源平台在 CartPole 与 Pendulum 环境下的速度对比。所有代码均在配置为 i7-8750H + GTX1060 的同一台笔记本电脑上进行测试。值得注意的是，天授实现的 VPG（vanilla policy gradient）算法在 CartPole-v0 任务中，训练用时仅为 3 秒。

以上测试使用了 10 个不同的 seed。CartPole 和 Pendulum 任务中的累积奖赏阈值分别设置为 195.0 与-250.0。可能会有读者感觉这两个任务比较简单，不太能突出框架的优势。该项目也表示，在这几天内，他们会更新天授在 Atari Pong / Mujoco 任务上的性能。

天授，只需 1500 行代码

非常令人惊讶的是，天授平台整体代码量不到 1500 行，其实现的 RL 算法大多数都少于百行代码。单从数量上来说，这样的代码量已经非常精简了，各种类与函数之间的关系应该也容易把握住。

项目表示，天授虽然代码量少，但可读性并不会有损失。我们可以快速浏览整个框架，并理解运行的流程与策略到底是什么样的。该项目提供了很多灵活的 API，例如可以便捷地使用如下代码令策略与环境交互 n 步：

result = collector.collect(n_step=n)

或者，如果你想通过采样的批量数据训练给定的策略，可以这样写：

result = policy.learn(collector.sample(batch_size))

正是通过大量精简的 API 构造 RL 模型，天授才能保持在 1500 行代码内。例如我们可以看看 DQN 的模型代码，它是非常流行的一种强化学习模型，在天授内部，DQN 模型真的只用了 99 行代码就完成了。当然，这 99 行代码是不包含其它公用代码块的。

如下为 DQN 的主要代码结构，我们省略了部分具体代码，各个 RL 策略都会继承基本类的结构，然后重写就够了。可以发现，在常规地定义好模型后，传入这个类就能创建策略。DQN 策略的各种操作都会写在一起，后续配置 Collector 后就能直接训练。

项目作者把所有策略算法都模块化为 4 部分：

__init__：初始化策略
process_fn：从 replay buffer 中处理数据
__call__：给定环境观察结果计算对应行动
learn：给定批量数据学习策略

实际体验

天授很容易安装，直接运行「pip install tianshou」就可以。下面我们将该项目克隆到本地，实际测试一下。

!git clone https://github.com/thu-ml/tianshou
!pip3 install tianshou
import os
os.chdir('tianshou')

该项目在 test 文件夹下提供了诸多算法的测试示例，下面我们在 CartPole 任务下逐个测试一番。

!python test/discrete/test_pg.py

!python test/discrete/test_ppo.py

!python test/discrete/test_a2c.py

!python test/discrete/test_dqn.py

以上分别为 VPG、PPO、A2C 与 DQN 在 P100 GPU 上的训练结果。可以看到，我们的测试结果与项目提供的结果出入不大。

由于 CartPole 任务在强化学习中相对简单，相当于图像识别中的 MNIST。为更进一步测试该 RL 框架的性能，我们也在 MinitaurBulletEnv-v0 任务中对其进行了测试。

Minitaur 是 PyBullet 环境中一个四足机器人运动控制任务，其观测值为该机器人的位置、姿态等 28 个状态信息，控制输入为电机的转矩（每条腿 2 个电机，总共 8 个电机），策略优化的目标为最大化机器人移动速度的同时最小化能量消耗。也就是说，agent 需要根据奖赏值自主地学习到由 28 个状态信息到 8 个控制输入的映射关系。

使用 SAC 算法在 Minitaur 任务中的训练结果如下图所示：

需要注意的是，天授的 SAC 实现在 Minitaur 任务中仅训练了不到 200k 步即能获得以上控制策略，效果可以说是很不错的。

项目作者，清华本科生

在 GitHub 中，其展示了该项目的主要作者是 Jiayi Weng 与 Minghao Zhang，他们都是清华的本科生。其中 Jiayi Weng 今年 6 月份本科毕业，在此之前作为本科研究者与清华大学苏航、朱军等老师开展强化学习领域的相关研究。Minghao Zhang 目前是清华大学软件学院的本科二年级学生，同时还修了数学专业。