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

机器之心报道

项目作者：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 目前是清华大学软件学院的本科二年级学生，同时还修了数学专业。

作为本科生，该项目的两位作者已经有了非常丰富的研究经验，Jiayi Weng 去年夏季就作为访问学生到访 MILA 实验室，并与 Yoshua Bengio 开展了关于意识先验相关的研究。在 Jiayi Weng 的主页中，我们可以看到在本科期间已经发了 IJCAI 的 Oral 论文。

Minghao Zhang 也有丰富的研究经验，之前他在软件学院 iMoon Lab 做关于 3D 视觉相关的研究，而后目前在清华交叉信息学院做研究助理，从事强化学习方面的研究。尽管离毕业还有不短的时间，Minghao Zhang 已经做出了自己的研究成果。

所以综合来看，因为在本科已经有了丰富的科研经验，并且做过多个项目，那么在这个阶段能做一个非常不错的强化学习开源项目也就理所当然了。

接下来的工作

天授目前还处于初期开发阶段，尚有一些未实现的功能或有待完善的地方。项目作者表示今后主要在以下几个方面来完善该 RL 框架：

• Prioritized replay buffer
• RNN support
• Imitation Learning
• Multi-agent
• Distributed training

它们分别是提供更多 RL 环境的 benchmark、优先经验回放、循环神经网络支持、模仿学习、多智能体学习以及分布式训练。