Source-linked AI summary
Tianshou: a Highly Modularized Deep Reinforcement Learning Library
Jiayi Weng, Huayu Chen, Dong Yan, Kaichao You, Alexis Duburcq, Minghao Zhang, Yi Su, Hang Su, Jun Zhu
TL;DR
DRL research needs infrastructure that combines broad algorithm coverage with flexibility and usable training workflows. Tianshou addresses this gap with a modular PyTorch library and unified interfaces for online, offline, and related DRL techniques. Its MuJoCo benchmark reports 15% higher average median performance than reference implementations.
Problem
Existing DRL libraries may hide configurable options, complicate debugging, or cover only online or offline algorithms, limiting their fit for research workflows.
Method
Tianshou factors shared DRL infrastructure into modular building blocks with unified interfaces for training, data collection, replay, and policy updates.
Results
15% higher average median performance than reference implementations was reported in Tianshou’s MuJoCo benchmark covering 8 algorithms.
Takeaways & Limitations
Tianshou provides a flexible and reliable framework for DRL research with broad algorithmic coverage and reusable infrastructure.
Abstract
from arXiv · showhide
In this paper, we present Tianshou, a highly modularized Python library for deep reinforcement learning (DRL) that uses PyTorch as its backend. Tianshou intends to be research-friendly by providing a flexible and reliable infrastructure of DRL algorithms. It supports online and offline training with more than 20 classic algorithms through a unified interface. To facilitate related research and prove Tianshou's reliability, we have released Tianshou's benchmark of MuJoCo environments, covering eight classic algorithms with state-of-the-art performance. We open-sourced Tianshou at https://github.com/thu-ml/tianshou/.
1. Introduction
Tianshou addresses flexibility, usability, and coverage gaps in existing DRL libraries with a highly modularized, unified Python infrastructure. Its benchmark and engineering practices support the paper’s reliability claim.
- Motivation: Existing DRL libraries can hide algorithmic options, complicate debugging, or support only online or offline algorithms.These trade-offs motivate a research-oriented alternative that preserves flexibility while retaining broad coverage.
- Proposed library: Tianshou is a highly modularized PyTorch-based Python library designed to address these issues.The paper presents modularity as the foundation for its research-friendly infrastructure.
- Modularity: Tianshou factors shared DRL infrastructure into building blocks, enabling fast prototyping with only a few variable changes.The approach includes commonly used techniques such as parallel data sampling.
- Reliability: 94% code coverage and a MuJoCo benchmark covering 8 algorithms support Tianshou’s reliability, with median performance 15% higher than reference implementations on average.The reliability process also includes unit tests on multiple platforms.
- Coverage: Tianshou provides a unified Python interface spanning model-free algorithms, offline learning, data collection, and agent training.Its supported techniques include GAIL and ICM, alongside synchronous and asynchronous environment execution.
2. Architecture of Tianshou
Tianshou standardizes DRL training around replay buffers, collectors, and policy updates while organizing the library into modular layers and parallel sampling infrastructure. Its released benchmarks demonstrate performance across MuJoCo and Atari settings.
- Training process: The standardized training process uses replay buffers for transitions, collectors for data collection, and policy update functions for parameter updates.Training paradigms are classified as on-policy, off-policy, or offline learning.
- Parallel computing: Tianshou focuses its parallel computing infrastructure on environment sampling for small- to medium-scale research.It supports asynchronous sampling to ease straggler effects and can integrate the C++-based vectorized environment EnvPool.
- Architecture: Tianshou’s architecture has four layers: encapsulation, core algorithms, user-oriented interaction APIs, and application scripts.Instances of the same building block share almost the same Pythonic APIs.
- Utilities: The library incorporates utilities including partial-episode bootstrapping, observation/value normalization, automatic action scaling, and GAE.These utilities are intended to relieve users from performance-critical implementation details.
- Reproduction and performance: The released benchmark covers 8 classic algorithms and 9 MuJoCo environments, with median performance 15% higher than multiple reference implementations on average.Additional examples cover 7 algorithms in 7 Atari environments, using 10 random seeds for all experiments.
3. Usability
Tianshou emphasizes accessible use through simple installation, documentation, tutorials, and concise experiment setup across major operating systems.
- Usability: Tianshou can be installed through Pip or Conda on Windows, macOS, and Linux, with full API documentation and tutorials provided.Only a few lines of code are required to start a simple experiment.
- Usability: The project follows PEP8 and includes comments, data type annotations, contributing guidelines, and extensive GitHub Actions unit tests.These practices support maintainable and tested usage across platforms.
- Usability: Tianshou is presented as lightweight and easy to install for practical experimentation.
4. Comparison to Related Works
Tianshou is compared with PyTorch-based DRL libraries that emphasize different combinations of algorithm coverage, parallel sampling, and optimization scale. The comparison is limited by the paper’s page constraints.
- Comparison scope: The comparison is restricted to selected PyTorch-backend DRL libraries because of page limits.TensorFlow-based libraries are acknowledged but not included in the comparison.
- Related works: RLlib and rlpyt target high-throughput software with multi-CPU parallel sampling and multi-GPU optimization.
- Related works: Stable-Baselines3, PFRL, and Tianshou are included among the compared PyTorch-based libraries.
5. Conclusion
Tianshou provides a flexible and reliable modular DRL framework by factoring shared infrastructure into reusable building blocks. Its MuJoCo benchmark demonstrates reliability across classic algorithms.
- Tianshou factors shared DRL infrastructure into building blocks for a modular research framework.
- The library is presented as flexible and reliable for DRL research.
- Its released MuJoCo benchmark covers many classic algorithms and demonstrates Tianshou’s reliability.