[Paper Review] Tianshou: a Highly Modularized Deep Reinforcement Learning Library

Jiayi Weng, Huayu Chen+7 more|arXiv (Cornell University)|Jul 29, 2021

Reinforcement Learning in RoboticsComputer Science23 references114 citations

TL;DR

Tianshou is a highly modular, PyTorch-based DRL library that emphasizes building-block flexibility for online and offline learning, with a MuJoCo benchmark showing 15% higher median performance on average vs reference implementations.

ABSTRACT

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/.

Motivation & Objective

• Address the need for a modular, reliable DRL library that supports both online and offline learning for research use.
• Provide building blocks rather than monolithic training scripts to facilitate fast prototyping and flexible experimentation.
• Demonstrate reliability and usability through a MuJoCo benchmark and open-source release.

Proposed method

• Modular architecture with four layers: encapsulation, core algorithms, interaction APIs, and application scripts.
• Standardization of training paradigms into on-policy, off-policy, and offline learning with replay buffers and data collectors.
• Parallel data sampling with asynchronous options to balance simulation and inference loads, and compatibility with EnvPool for speed.
• Comprehensive utilities including normalization, GAE, automatic action scaling, partial-episode bootstrapping, and versatile loggers (TensorBoard, Weights & Biases).
• Open-source with unit tests (94% coverage) and a MuJoCo benchmark benchmark comparing 8 algorithms across 9 environments.

Experimental results

Research questions

• RQ1 How can DRL research be facilitated by a highly modular library that separates building blocks from training scripts?
• RQ2 Can a unified PyTorch-based interface efficiently support online, offline, and imitation learning techniques within one framework?
• RQ3 What is the reliability and performance of Tianshou’s MuJoCo benchmark compared with reference implementations?

Key findings

• Tianshou achieves a modular architecture that factors out shared DRL infrastructure for easy prototyping.
• The library attains 94% code coverage and demonstrates reliability via a MuJoCo benchmark.
• In the MuJoCo benchmark, Tianshou scores 15% higher on average in terms of median performance across 9 environments compared with reference implementations.
• Tianshou supports a wide range of DRL techniques including offline learning, GAIL, and ICM, through a unified interface.
• Tianshou provides extensive usability features, including simple installation, thorough documentation, and strict coding standards (PEP8).