Overview#

Before we get started, we must first install Tianshou’s library and Gym environment by running the commands below. This tutorials will always keep up with the latest version of Tianshou since they also serve as a test for the latest version. If you are using an older version of Tianshou, please refer to the documentation of your version.

# !pip install tianshou gym

Run the code#

Below is a short script that use a certain DRL algorithm (PPO) to solve the classic CartPole-v1 problem in Gym. Simply run it and don’t worry if you can’t understand the code very well. That is exactly what this tutorial is for.

If the script ends normally, you will see the evaluation result printed out before the first epoch is finished.

Hide code cell content 
%%capture

import gymnasium as gym
import torch

from tianshou.data import Collector, VectorReplayBuffer
from tianshou.env import DummyVectorEnv
from tianshou.policy import BasePolicy, PPOPolicy
from tianshou.trainer import OnpolicyTrainer
from tianshou.utils.net.common import ActorCritic, Net
from tianshou.utils.net.discrete import Actor, Critic

device = "cuda" if torch.cuda.is_available() else "cpu"

# environments
env = gym.make("CartPole-v1")
train_envs = DummyVectorEnv([lambda: gym.make("CartPole-v1") for _ in range(20)])
test_envs = DummyVectorEnv([lambda: gym.make("CartPole-v1") for _ in range(10)])

# model & optimizer
assert env.observation_space.shape is not None  # for mypy
net = Net(state_shape=env.observation_space.shape, hidden_sizes=[64, 64], device=device)

assert isinstance(env.action_space, gym.spaces.Discrete)  # for mypy
actor = Actor(preprocess_net=net, action_shape=env.action_space.n, device=device).to(device)
critic = Critic(preprocess_net=net, device=device).to(device)
actor_critic = ActorCritic(actor, critic)
optim = torch.optim.Adam(actor_critic.parameters(), lr=0.0003)

# PPO policy
dist = torch.distributions.Categorical
policy: BasePolicy
policy = PPOPolicy(
    actor=actor,
    critic=critic,
    optim=optim,
    dist_fn=dist,
    action_space=env.action_space,
    action_scaling=False,
)

# collector
train_collector = Collector(policy, train_envs, VectorReplayBuffer(20000, len(train_envs)))
test_collector = Collector(policy, test_envs)

# trainer
train_result = OnpolicyTrainer(
    policy=policy,
    batch_size=256,
    train_collector=train_collector,
    test_collector=test_collector,
    max_epoch=10,
    step_per_epoch=50000,
    repeat_per_collect=10,
    episode_per_test=10,
    step_per_collect=2000,
    stop_fn=lambda mean_reward: mean_reward >= 195,
).run()
Hide code cell output 
Epoch #1:   0%|          | 0/50000 [00:00<?, ?it/s]

Epoch #1:   4%|4         | 2000/50000 [00:00<00:06, 6939.36it/s]

Epoch #1:   4%|4         | 2000/50000 [00:00<00:06, 6939.36it/s, env_step=2000, gradient_step=8, len=20, n/ep=85, n/st=2000, rew=20.06]

Epoch #1:   8%|8         | 4000/50000 [00:00<00:09, 4690.40it/s, env_step=2000, gradient_step=8, len=20, n/ep=85, n/st=2000, rew=20.06]

Epoch #1:   8%|8         | 4000/50000 [00:01<00:09, 4690.40it/s, env_step=4000, gradient_step=16, len=23, n/ep=89, n/st=2000, rew=23.09]

Epoch #1:  12%|#2        | 6000/50000 [00:01<00:10, 4312.51it/s, env_step=4000, gradient_step=16, len=23, n/ep=89, n/st=2000, rew=23.09]

Epoch #1:  12%|#2        | 6000/50000 [00:01<00:10, 4312.51it/s, env_step=6000, gradient_step=24, len=24, n/ep=82, n/st=2000, rew=24.46]

Epoch #1:  16%|#6        | 8000/50000 [00:01<00:10, 4177.16it/s, env_step=6000, gradient_step=24, len=24, n/ep=82, n/st=2000, rew=24.46]

Epoch #1:  16%|#6        | 8000/50000 [00:02<00:10, 4177.16it/s, env_step=8000, gradient_step=32, len=26, n/ep=73, n/st=2000, rew=26.45]

Epoch #1:  20%|##        | 10000/50000 [00:02<00:09, 4089.16it/s, env_step=8000, gradient_step=32, len=26, n/ep=73, n/st=2000, rew=26.45]

Epoch #1:  20%|##        | 10000/50000 [00:02<00:09, 4089.16it/s, env_step=10000, gradient_step=40, len=23, n/ep=80, n/st=2000, rew=23.60]

Epoch #1:  24%|##4       | 12000/50000 [00:02<00:09, 4058.70it/s, env_step=10000, gradient_step=40, len=23, n/ep=80, n/st=2000, rew=23.60]

Epoch #1:  24%|##4       | 12000/50000 [00:03<00:09, 4058.70it/s, env_step=12000, gradient_step=48, len=28, n/ep=71, n/st=2000, rew=28.41]

Epoch #1:  28%|##8       | 14000/50000 [00:03<00:08, 4037.63it/s, env_step=12000, gradient_step=48, len=28, n/ep=71, n/st=2000, rew=28.41]

Epoch #1:  28%|##8       | 14000/50000 [00:03<00:08, 4037.63it/s, env_step=14000, gradient_step=56, len=31, n/ep=69, n/st=2000, rew=31.07]

Epoch #1:  32%|###2      | 16000/50000 [00:03<00:08, 4026.13it/s, env_step=14000, gradient_step=56, len=31, n/ep=69, n/st=2000, rew=31.07]

Epoch #1:  32%|###2      | 16000/50000 [00:04<00:08, 4026.13it/s, env_step=16000, gradient_step=64, len=29, n/ep=60, n/st=2000, rew=29.50]

Epoch #1:  36%|###6      | 18000/50000 [00:04<00:08, 3996.59it/s, env_step=16000, gradient_step=64, len=29, n/ep=60, n/st=2000, rew=29.50]

Epoch #1:  36%|###6      | 18000/50000 [00:04<00:08, 3996.59it/s, env_step=18000, gradient_step=72, len=32, n/ep=64, n/st=2000, rew=32.86]

Epoch #1:  40%|####      | 20000/50000 [00:04<00:07, 3993.51it/s, env_step=18000, gradient_step=72, len=32, n/ep=64, n/st=2000, rew=32.86]

Epoch #1:  40%|####      | 20000/50000 [00:05<00:07, 3993.51it/s, env_step=20000, gradient_step=80, len=33, n/ep=52, n/st=2000, rew=33.69]

Epoch #1:  44%|####4     | 22000/50000 [00:05<00:06, 4008.02it/s, env_step=20000, gradient_step=80, len=33, n/ep=52, n/st=2000, rew=33.69]

Epoch #1:  44%|####4     | 22000/50000 [00:05<00:06, 4008.02it/s, env_step=22000, gradient_step=88, len=54, n/ep=29, n/st=2000, rew=54.14]

Epoch #1:  48%|####8     | 24000/50000 [00:05<00:06, 4012.90it/s, env_step=22000, gradient_step=88, len=54, n/ep=29, n/st=2000, rew=54.14]

Epoch #1:  48%|####8     | 24000/50000 [00:06<00:06, 4012.90it/s, env_step=24000, gradient_step=96, len=69, n/ep=30, n/st=2000, rew=69.90]

Epoch #1:  52%|#####2    | 26000/50000 [00:06<00:05, 4001.29it/s, env_step=24000, gradient_step=96, len=69, n/ep=30, n/st=2000, rew=69.90]

Epoch #1:  52%|#####2    | 26000/50000 [00:06<00:05, 4001.29it/s, env_step=26000, gradient_step=104, len=86, n/ep=19, n/st=2000, rew=86.37]

Epoch #1:  56%|#####6    | 28000/50000 [00:06<00:05, 3995.13it/s, env_step=26000, gradient_step=104, len=86, n/ep=19, n/st=2000, rew=86.37]

Epoch #1:  56%|#####6    | 28000/50000 [00:07<00:05, 3995.13it/s, env_step=28000, gradient_step=112, len=88, n/ep=24, n/st=2000, rew=88.71]

Epoch #1:  60%|######    | 30000/50000 [00:07<00:04, 4012.01it/s, env_step=28000, gradient_step=112, len=88, n/ep=24, n/st=2000, rew=88.71]

Epoch #1:  60%|######    | 30000/50000 [00:07<00:04, 4012.01it/s, env_step=30000, gradient_step=120, len=98, n/ep=11, n/st=2000, rew=98.09]

Epoch #1:  64%|######4   | 32000/50000 [00:07<00:04, 4021.85it/s, env_step=30000, gradient_step=120, len=98, n/ep=11, n/st=2000, rew=98.09]

Epoch #1:  64%|######4   | 32000/50000 [00:08<00:04, 4021.85it/s, env_step=32000, gradient_step=128, len=128, n/ep=18, n/st=2000, rew=128.83]

Epoch #1:  68%|######8   | 34000/50000 [00:08<00:03, 4017.39it/s, env_step=32000, gradient_step=128, len=128, n/ep=18, n/st=2000, rew=128.83]

Epoch #1:  68%|######8   | 34000/50000 [00:08<00:03, 4017.39it/s, env_step=34000, gradient_step=136, len=128, n/ep=13, n/st=2000, rew=128.31]

Epoch #1:  72%|#######2  | 36000/50000 [00:08<00:03, 4024.29it/s, env_step=34000, gradient_step=136, len=128, n/ep=13, n/st=2000, rew=128.31]

Epoch #1:  72%|#######2  | 36000/50000 [00:09<00:03, 4024.29it/s, env_step=36000, gradient_step=144, len=123, n/ep=24, n/st=2000, rew=123.12]

Epoch #1:  76%|#######6  | 38000/50000 [00:09<00:02, 4035.79it/s, env_step=36000, gradient_step=144, len=123, n/ep=24, n/st=2000, rew=123.12]

Epoch #1:  76%|#######6  | 38000/50000 [00:09<00:02, 4035.79it/s, env_step=38000, gradient_step=152, len=112, n/ep=12, n/st=2000, rew=112.92]

Epoch #1:  80%|########  | 40000/50000 [00:09<00:02, 4038.09it/s, env_step=38000, gradient_step=152, len=112, n/ep=12, n/st=2000, rew=112.92]

Epoch #1:  80%|########  | 40000/50000 [00:10<00:02, 4038.09it/s, env_step=40000, gradient_step=160, len=90, n/ep=20, n/st=2000, rew=90.80]  

Epoch #1:  84%|########4 | 42000/50000 [00:10<00:01, 4024.23it/s, env_step=40000, gradient_step=160, len=90, n/ep=20, n/st=2000, rew=90.80]

Epoch #1:  84%|########4 | 42000/50000 [00:10<00:01, 4024.23it/s, env_step=42000, gradient_step=168, len=107, n/ep=22, n/st=2000, rew=107.36]

Epoch #1:  88%|########8 | 44000/50000 [00:10<00:01, 4033.19it/s, env_step=42000, gradient_step=168, len=107, n/ep=22, n/st=2000, rew=107.36]

Epoch #1:  88%|########8 | 44000/50000 [00:11<00:01, 4033.19it/s, env_step=44000, gradient_step=176, len=158, n/ep=10, n/st=2000, rew=158.00]

Epoch #1:  92%|#########2| 46000/50000 [00:11<00:00, 4043.00it/s, env_step=44000, gradient_step=176, len=158, n/ep=10, n/st=2000, rew=158.00]

Epoch #1:  92%|#########2| 46000/50000 [00:11<00:00, 4043.00it/s, env_step=46000, gradient_step=184, len=163, n/ep=10, n/st=2000, rew=163.40]

Epoch #1:  96%|#########6| 48000/50000 [00:11<00:00, 4043.60it/s, env_step=46000, gradient_step=184, len=163, n/ep=10, n/st=2000, rew=163.40]

Epoch #1:  96%|#########6| 48000/50000 [00:12<00:00, 4043.60it/s, env_step=48000, gradient_step=192, len=163, n/ep=8, n/st=2000, rew=163.25] 

Epoch #1: 100%|##########| 50000/50000 [00:12<00:00, 3239.97it/s, env_step=48000, gradient_step=192, len=163, n/ep=8, n/st=2000, rew=163.25]

Epoch #1: 100%|##########| 50000/50000 [00:12<00:00, 3239.97it/s, env_step=50000, gradient_step=200, len=219, n/ep=14, n/st=2000, rew=219.00]

Epoch #1: 50001it [00:12, 3862.34it/s, env_step=50000, gradient_step=200, len=219, n/ep=14, n/st=2000, rew=219.00]                           



Epoch #1: test_reward: 112.300000 ± 77.087029, best_reward: 112.300000 ± 77.087029 in #1

Epoch #2:   0%|          | 0/50000 [00:00<?, ?it/s]

Epoch #2:   4%|4         | 2000/50000 [00:00<00:05, 8247.06it/s]

Epoch #2:   4%|4         | 2000/50000 [00:00<00:05, 8247.06it/s, env_step=52000, gradient_step=208, len=154, n/ep=13, n/st=2000, rew=154.08]

Epoch #2:   8%|8         | 4000/50000 [00:00<00:09, 5107.36it/s, env_step=52000, gradient_step=208, len=154, n/ep=13, n/st=2000, rew=154.08]

Epoch #2:   8%|8         | 4000/50000 [00:00<00:09, 5107.36it/s, env_step=54000, gradient_step=216, len=108, n/ep=16, n/st=2000, rew=108.38]

Epoch #2:  12%|#2        | 6000/50000 [00:01<00:09, 4526.41it/s, env_step=54000, gradient_step=216, len=108, n/ep=16, n/st=2000, rew=108.38]

Epoch #2:  12%|#2        | 6000/50000 [00:01<00:09, 4526.41it/s, env_step=56000, gradient_step=224, len=176, n/ep=10, n/st=2000, rew=176.70]

Epoch #2:  16%|#6        | 8000/50000 [00:02<00:14, 2985.19it/s, env_step=56000, gradient_step=224, len=176, n/ep=10, n/st=2000, rew=176.70]

Epoch #2:  16%|#6        | 8000/50000 [00:02<00:14, 2985.19it/s, env_step=58000, gradient_step=232, len=272, n/ep=6, n/st=2000, rew=272.67] 

Epoch #2:  20%|##        | 10000/50000 [00:02<00:12, 3294.96it/s, env_step=58000, gradient_step=232, len=272, n/ep=6, n/st=2000, rew=272.67]

Epoch #2:  20%|##        | 10000/50000 [00:03<00:12, 3294.96it/s, env_step=60000, gradient_step=240, len=138, n/ep=10, n/st=2000, rew=138.50]

Epoch #2:  24%|##4       | 12000/50000 [00:03<00:12, 3077.07it/s, env_step=60000, gradient_step=240, len=138, n/ep=10, n/st=2000, rew=138.50]

Epoch #2:  24%|##4       | 12000/50000 [00:03<00:12, 3077.07it/s, env_step=62000, gradient_step=248, len=209, n/ep=19, n/st=2000, rew=209.21]

Epoch #2:  28%|##8       | 14000/50000 [00:03<00:10, 3338.91it/s, env_step=62000, gradient_step=248, len=209, n/ep=19, n/st=2000, rew=209.21]

Epoch #2:  28%|##8       | 14000/50000 [00:04<00:10, 3338.91it/s, env_step=64000, gradient_step=256, len=146, n/ep=13, n/st=2000, rew=146.77]

Epoch #2:  32%|###2      | 16000/50000 [00:04<00:09, 3528.48it/s, env_step=64000, gradient_step=256, len=146, n/ep=13, n/st=2000, rew=146.77]

Epoch #2:  32%|###2      | 16000/50000 [00:04<00:09, 3528.48it/s, env_step=66000, gradient_step=264, len=83, n/ep=34, n/st=2000, rew=83.32]  

Epoch #2:  36%|###6      | 18000/50000 [00:04<00:08, 3659.84it/s, env_step=66000, gradient_step=264, len=83, n/ep=34, n/st=2000, rew=83.32]

Epoch #2:  36%|###6      | 18000/50000 [00:05<00:08, 3659.84it/s, env_step=68000, gradient_step=272, len=59, n/ep=30, n/st=2000, rew=59.63]

Epoch #2:  40%|####      | 20000/50000 [00:05<00:07, 3762.52it/s, env_step=68000, gradient_step=272, len=59, n/ep=30, n/st=2000, rew=59.63]

Epoch #2:  40%|####      | 20000/50000 [00:05<00:07, 3762.52it/s, env_step=70000, gradient_step=280, len=51, n/ep=27, n/st=2000, rew=51.67]

Epoch #2:  44%|####4     | 22000/50000 [00:05<00:07, 3846.97it/s, env_step=70000, gradient_step=280, len=51, n/ep=27, n/st=2000, rew=51.67]

Epoch #2:  44%|####4     | 22000/50000 [00:06<00:07, 3846.97it/s, env_step=72000, gradient_step=288, len=135, n/ep=12, n/st=2000, rew=135.00]

Epoch #2:  48%|####8     | 24000/50000 [00:07<00:09, 2881.04it/s, env_step=72000, gradient_step=288, len=135, n/ep=12, n/st=2000, rew=135.00]

Epoch #2:  48%|####8     | 24000/50000 [00:07<00:09, 2881.04it/s, env_step=74000, len=202, n/ep=6, n/st=2000, rew=202.33]                    

Epoch #2:  48%|####8     | 24000/50000 [00:07<00:07, 3395.08it/s, env_step=74000, len=202, n/ep=6, n/st=2000, rew=202.33]



train_result.pprint_asdict()

InfoStats
----------------------------------------
{'best_reward': 246.8,
 'best_reward_std': 152.69564499356227,
 'gradient_step': 288,
 'test_episode': 60,
 'test_step': 7450,
 'timing': {'test_time': 2.204011917114258,
            'total_time': 20.47813868522644,
            'train_time': 18.274126768112183,
            'train_time_collect': 9.055797815322876,
            'train_time_update': 9.106553077697754,
            'update_speed': 4049.4410999231895},
 'train_episode': 1185,
 'train_step': 74000}

# Let's watch its performance!
policy.eval()
eval_result = test_collector.collect(n_episode=1, render=False)
print(f"Final reward: {eval_result.returns.mean()}, length: {eval_result.lens.mean()}")

Final reward: 330.0, length: 330.0

Tutorial Introduction#

A common DRL experiment as is shown above may require many components to work together. The agent, the environment (possibly parallelized ones), the replay buffer and the trainer all work together to complete a training task.

In Tianshou, all of these main components are factored out as different building blocks, which you can use to create your own algorithm and finish your own experiment.

Building blocks may include:

  • Batch

  • Replay Buffer

  • Vectorized Environment Wrapper

  • Policy (the agent and the training algorithm)

  • Data Collector

  • Trainer

  • Logger

These notebooks tutorials will guide you through all the modules one by one.

Further reading#

What if I am not familiar with the PPO algorithm itself?#

As for the DRL algorithms themselves, we will refer you to the Spinning up documentation, where they provide plenty of resources and guides if you want to study the DRL algorithms. In Tianshou’s tutorials, we will focus on the usages of different modules, but not the algorithms themselves.