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

tianshou

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

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

为算法开发人员(RL研究人员)提供灵活、可定制且类型安全的底层接口。
为RL应用(在自定义环境上训练已实现的算法)提供便捷的高级接口。
广泛的涵盖范围:支持在线(on-policy和off-policy)和离线RL,实验性支持多智能体RL(MARL),实验性支持基于模型的RL等。
以数据为中心的设计理念,使用统一的数据流结构实现各种算法。

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:

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

以下是使用高级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实验。通过简单的声明式配置,我们可以轻松设置环境、算法参数、训练策略等,而无需关心底层实现细节。