How to use tianshou: install and train a reinforcement learning agent

What is tianshou and who is it for?

Tianshou is a reinforcement learning (RL) library based on pure PyTorch and Gymnasium. It offers modular low-level interfaces for algorithm developers and higher-level interfaces for applying an implemented algorithm to a custom environment.

It is best suited to researchers who want to implement or modify RL algorithms, practitioners training a supported algorithm on their own Gymnasium-style environment, students studying how common algorithms are built, and developers who want vectorized or parallel environment sampling.

• Requires Python 3.11 or newer and PyTorch.
• Covers online (on- and off-policy) and offline RL, with experimental support for multi-agent and model-based RL.
• The README is a library reference, not a beginner course.

Install tianshou with pip

Tianshou is published on PyPI and conda-forge and requires Python 3.11 or newer. Install it with pip from PyPI.

For the most recent version, the project's README suggests cloning the repository and installing from source, using poetry.

Choose an algorithm for your problem

Tianshou supports a long list of algorithms. Common starting points include DQN for discrete actions, and PPO, SAC and TD3 for continuous control. Its family also includes Double DQN, Rainbow, A2C, TRPO, DDPG, REDQ and imitation learning methods such as GAIL, plus offline methods such as CQL and TD3+BC.

Pick the algorithm whose assumptions match your action space and whether you can collect new experience during training.

• DQN, Double DQN, Dueling DQN, Rainbow: Q-learning family.
• PPO, A2C, TRPO: policy gradient methods.
• DDPG, TD3, SAC: off-policy methods for continuous actions.
• CQL, BCQ, TD3+BC: offline RL from fixed datasets.

Use the high-level interface or the procedural API

The high-level interface maximizes ease of use for applying existing algorithms to your own environment. The procedural API gives more flexibility for algorithm development without being overly verbose.

Version 2 separates learning algorithms from policies through the Algorithm and Policy abstractions, and it is not backwards compatible with earlier versions. Check the change log before migrating older code.

• Use the high-level API to train a supported algorithm quickly.
• Use the procedural API when you need to change how an algorithm works.
• Consult the CHANGELOG.md file in the repository when migrating from an older release.

Speed up sampling and extend your setup

Tianshou supports synchronous or asynchronous vectorized environments for all algorithms, so you can sample experience in parallel. It also supports EnvPool-based vectorized environments, multi-GPU training, and logging to TensorBoard or W&B.

Other documented options include recurrent networks for partially observable problems, custom state and action types, and customized training processes. The cheatsheet in the documentation shows how to use each one.

• Parallel sampling with vectorized environments.
• EnvPool integration for faster environments.
• RNN-style training for POMDPs.
• Logging with TensorBoard or Weights & Biases.
• Multi-GPU training and multi-agent RL (experimental areas noted in the README).