Tianshou:一个优雅的PyTorch深度强化学习库

Tianshou:一个优雅的PyTorch深度强化学习库

Tianshou(天授)是一个基于纯PyTorch的强化学习(RL)库。与其他主要基于TensorFlow、具有复杂嵌套类结构、API不友好或运行速度慢的强化学习库不同,Tianshou提供了一个高性能、模块化的框架和易用的API,用于构建深度强化学习智能体。Tianshou的主要特点包括:

1. 为算法开发人员(RL研究人员)提供灵活、可定制且类型安全的底层接口。

2. 为RL应用(在自定义环境上训练已实现的算法)提供便捷的高级接口。

3. 广泛的涵盖范围:支持在线(on-policy和off-policy)和离线RL,实验性支持多智能体RL(MARL),实验性支持基于模型的RL等。

4. 以数据为中心的设计理念,使用统一的数据流结构实现各种算法。

Tianshou旨在实现简洁的实现,既适用于研究人员,也适用于实践者,同时不牺牲灵活性。

支持的算法

Tianshou支持多种主流强化学习算法,包括:

• DQN及其变体(Double DQN、Dueling DQN等)
• Policy Gradient方法(REINFORCE、A2C、PPO等)
• Actor-Critic方法(DDPG、TD3、SAC等)
• 离线RL方法(BCQ、CQL等)
• 多智能体RL方法(MADDPG、QMIX等)
• 模仿学习方法(BC、GAIL等)

此外,Tianshou还支持优先级经验回放、GAE、好奇心驱动探索等强化学习中的常用技巧。

主要特性

除了支持多种算法外,Tianshou还具有以下显著特性:

• 优雅的框架设计,包含双重API:
  • 高级API最大限度地简化应用开发的使用,同时保留高度灵活性
  • 底层过程API为算法开发提供最大的灵活性,同时不会过于冗长
• 在MuJoCo基准测试中实现了多种算法的SOTA结果
• 支持所有算法的矢量化环境(同步或异步)
• 支持基于EnvPool的超快速矢量化环境
• 支持actor网络和critic网络中的循环状态表示(用于POMDP的RNN式训练)
• 支持任何类型的环境状态/动作(如字典、自定义类等)
• 支持自定义训练过程
• 支持多智能体强化学习
• 支持基于TensorBoard和W&B的日志记录
• 支持多GPU训练
• 全面的文档、PEP8代码风格检查、类型检查和彻底的测试

安装

Tianshou目前托管在PyPI和conda-forge上,要求Python >= 3.11。可以通过以下方式安装:

# 使用pip安装
pip install tianshou

# 或使用conda安装  
conda install tianshou -c conda-forge

# 从GitHub安装最新版本
pip install git+https://github.com/thu-ml/tianshou.git@master

快速入门

Tianshou提供两种级别的API:

1. 高级接口:为最终用户提供易用性,方便运行深度强化学习应用
2. 过程接口:为高级用户和强化学习算法开发人员提供最大控制

以下是使用高级API在CartPole环境上应用DQN算法的示例:

from tianshou.highlevel.config import SamplingConfig
from tianshou.highlevel.env import EnvFactoryRegistered, VectorEnvType
from tianshou.highlevel.experiment import DQNExperimentBuilder, ExperimentConfig
from tianshou.highlevel.params.policy_params import DQNParams
from tianshou.highlevel.trainer import (
    EpochTestCallbackDQNSetEps,
    EpochTrainCallbackDQNSetEps,
    EpochStopCallbackRewardThreshold
)

experiment = (
    DQNExperimentBuilder(
        EnvFactoryRegistered(task="CartPole-v1", seed=0, venv_type=VectorEnvType.DUMMY),
        ExperimentConfig(
            persistence_enabled=False,
            watch=True,
            watch_render=1 / 35,
            watch_num_episodes=100,
        ),
        SamplingConfig(
            num_epochs=10,
            step_per_epoch=10000,
            batch_size=64,
            num_train_envs=10,
            num_test_envs=100,
            buffer_size=20000,
            step_per_collect=10,
            update_per_step=1 / 10,
        ),
    )
    .with_dqn_params(
        DQNParams(
            lr=1e-3,
            discount_factor=0.9,
            estimation_step=3,
            target_update_freq=320,
        ),
    )
    .with_model_factory_default(hidden_sizes=(64, 64))
    .with_epoch_train_callback(EpochTrainCallbackDQNSetEps(0.3))
    .with_epoch_test_callback(EpochTestCallbackDQNSetEps(0.0))
    .with_epoch_stop_callback(EpochStopCallbackRewardThreshold(195))
    .build()
)
experiment.run()

这个例子展示了如何使用Tianshou的高级API来配置和运行一个DQN实验。通过简单的声明式配置,我们可以轻松设置环境、算法参数、训练策略等,而无需关心底层实现细节。

为什么选择Tianshou?

1. 全面的功能:支持广泛的RL算法和范式,包括在线/离线RL、单/多智能体RL等。

2. 高软件工程标准:拥有全面的文档、代码覆盖率测试和类型提示。

3. 可复现的高质量结果:所有算法都经过严格测试,包括完整的智能体训练过程。在Atari和MuJoCo基准测试中达到或超过了大多数现有基准的性能水平。

4. 灵活的策略接口:所有算法都实现了统一的高度通用API,便于新算法的开发和实验。

5. 高性能:通过使用PyTorch和优化的实现,Tianshou提供了卓越的训练速度。

6. 易用性:提供高级API以简化应用开发,同时保留底层API以实现最大灵活性。

总之,Tianshou作为一个全面、高效、灵活的强化学习库,为研究人员和实践者提供了强大的工具来开发和应用最先进的强化学习算法。无论是进行学术研究还是工业应用,Tianshou都是一个值得考虑的优秀选择。