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.
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.
Show code cell content
%%capture
import gymnasium as gym
import torch
from tianshou.data import Collector, CollectStats, VectorReplayBuffer
from tianshou.env import DummyVectorEnv
from tianshou.policy import 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: PPOPolicy = PPOPolicy(
actor=actor,
critic=critic,
optim=optim,
dist_fn=dist,
action_space=env.action_space,
action_scaling=False,
)
# collector
train_collector = Collector[CollectStats](
policy,
train_envs,
VectorReplayBuffer(20000, len(train_envs)),
)
test_collector = Collector[CollectStats](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()
Show code cell output
Epoch #1: 0%| | 0/50000 [00:00<?, ?it/s]
Epoch #1: 4%|4 | 2000/50000 [00:00<00:18, 2588.06it/s]
Epoch #1: 4%|4 | 2000/50000 [00:00<00:18, 2588.06it/s, env_episode=83, env_step=2000, gradient_step=8, len=21, n/ep=83, n/st=2000, rew=21.69]
Epoch #1: 8%|8 | 4000/50000 [00:01<00:17, 2633.44it/s, env_episode=83, env_step=2000, gradient_step=8, len=21, n/ep=83, n/st=2000, rew=21.69]
Epoch #1: 8%|8 | 4000/50000 [00:01<00:17, 2633.44it/s, env_episode=174, env_step=4000, gradient_step=16, len=18, n/ep=91, n/st=2000, rew=20.86]
Epoch #1: 12%|#2 | 6000/50000 [00:02<00:16, 2631.30it/s, env_episode=174, env_step=4000, gradient_step=16, len=18, n/ep=91, n/st=2000, rew=20.86]
Epoch #1: 12%|#2 | 6000/50000 [00:02<00:16, 2631.30it/s, env_episode=273, env_step=6000, gradient_step=24, len=18, n/ep=99, n/st=2000, rew=21.20]
Epoch #1: 16%|#6 | 8000/50000 [00:03<00:15, 2653.77it/s, env_episode=273, env_step=6000, gradient_step=24, len=18, n/ep=99, n/st=2000, rew=21.20]
Epoch #1: 16%|#6 | 8000/50000 [00:03<00:15, 2653.77it/s, env_episode=361, env_step=8000, gradient_step=32, len=20, n/ep=88, n/st=2000, rew=22.42]
Epoch #1: 20%|## | 10000/50000 [00:03<00:15, 2660.61it/s, env_episode=361, env_step=8000, gradient_step=32, len=20, n/ep=88, n/st=2000, rew=22.42]
Epoch #1: 20%|## | 10000/50000 [00:03<00:15, 2660.61it/s, env_episode=447, env_step=10000, gradient_step=40, len=20, n/ep=86, n/st=2000, rew=23.06]
Epoch #1: 24%|##4 | 12000/50000 [00:04<00:14, 2652.63it/s, env_episode=447, env_step=10000, gradient_step=40, len=20, n/ep=86, n/st=2000, rew=23.06]
Epoch #1: 24%|##4 | 12000/50000 [00:04<00:14, 2652.63it/s, env_episode=532, env_step=12000, gradient_step=48, len=20, n/ep=85, n/st=2000, rew=23.41]
Epoch #1: 28%|##8 | 14000/50000 [00:05<00:13, 2665.46it/s, env_episode=532, env_step=12000, gradient_step=48, len=20, n/ep=85, n/st=2000, rew=23.41]
Epoch #1: 28%|##8 | 14000/50000 [00:05<00:13, 2665.46it/s, env_episode=615, env_step=14000, gradient_step=56, len=21, n/ep=83, n/st=2000, rew=24.34]
Epoch #1: 32%|###2 | 16000/50000 [00:06<00:12, 2668.05it/s, env_episode=615, env_step=14000, gradient_step=56, len=21, n/ep=83, n/st=2000, rew=24.34]
Epoch #1: 32%|###2 | 16000/50000 [00:06<00:12, 2668.05it/s, env_episode=690, env_step=16000, gradient_step=64, len=21, n/ep=75, n/st=2000, rew=24.65]
Epoch #1: 36%|###6 | 18000/50000 [00:06<00:11, 2681.70it/s, env_episode=690, env_step=16000, gradient_step=64, len=21, n/ep=75, n/st=2000, rew=24.65]
Epoch #1: 36%|###6 | 18000/50000 [00:06<00:11, 2681.70it/s, env_episode=745, env_step=18000, gradient_step=72, len=26, n/ep=55, n/st=2000, rew=33.11]
Epoch #1: 40%|#### | 20000/50000 [00:07<00:11, 2693.65it/s, env_episode=745, env_step=18000, gradient_step=72, len=26, n/ep=55, n/st=2000, rew=33.11]
Epoch #1: 40%|#### | 20000/50000 [00:07<00:11, 2693.65it/s, env_episode=778, env_step=20000, gradient_step=80, len=32, n/ep=33, n/st=2000, rew=46.70]
Epoch #1: 44%|####4 | 22000/50000 [00:08<00:10, 2700.16it/s, env_episode=778, env_step=20000, gradient_step=80, len=32, n/ep=33, n/st=2000, rew=46.70]
Epoch #1: 44%|####4 | 22000/50000 [00:08<00:10, 2700.16it/s, env_episode=813, env_step=22000, gradient_step=88, len=34, n/ep=35, n/st=2000, rew=62.80]
Epoch #1: 48%|####8 | 24000/50000 [00:08<00:09, 2714.56it/s, env_episode=813, env_step=22000, gradient_step=88, len=34, n/ep=35, n/st=2000, rew=62.80]
Epoch #1: 48%|####8 | 24000/50000 [00:08<00:09, 2714.56it/s, env_episode=840, env_step=24000, gradient_step=96, len=41, n/ep=27, n/st=2000, rew=72.59]
Epoch #1: 52%|#####2 | 26000/50000 [00:09<00:08, 2719.24it/s, env_episode=840, env_step=24000, gradient_step=96, len=41, n/ep=27, n/st=2000, rew=72.59]
Epoch #1: 52%|#####2 | 26000/50000 [00:09<00:08, 2719.24it/s, env_episode=862, env_step=26000, gradient_step=104, len=46, n/ep=22, n/st=2000, rew=79.91]
Epoch #1: 56%|#####6 | 28000/50000 [00:10<00:08, 2721.06it/s, env_episode=862, env_step=26000, gradient_step=104, len=46, n/ep=22, n/st=2000, rew=79.91]
Epoch #1: 56%|#####6 | 28000/50000 [00:10<00:08, 2721.06it/s, env_episode=877, env_step=28000, gradient_step=112, len=47, n/ep=15, n/st=2000, rew=101.07]
Epoch #1: 60%|###### | 30000/50000 [00:11<00:07, 2726.79it/s, env_episode=877, env_step=28000, gradient_step=112, len=47, n/ep=15, n/st=2000, rew=101.07]
Epoch #1: 60%|###### | 30000/50000 [00:11<00:07, 2726.79it/s, env_episode=897, env_step=30000, gradient_step=120, len=41, n/ep=20, n/st=2000, rew=113.00]
Epoch #1: 64%|######4 | 32000/50000 [00:11<00:06, 2725.05it/s, env_episode=897, env_step=30000, gradient_step=120, len=41, n/ep=20, n/st=2000, rew=113.00]
Epoch #1: 64%|######4 | 32000/50000 [00:11<00:06, 2725.05it/s, env_episode=918, env_step=32000, gradient_step=128, len=55, n/ep=21, n/st=2000, rew=113.76]
Epoch #1: 68%|######8 | 34000/50000 [00:12<00:05, 2730.79it/s, env_episode=918, env_step=32000, gradient_step=128, len=55, n/ep=21, n/st=2000, rew=113.76]
Epoch #1: 68%|######8 | 34000/50000 [00:12<00:05, 2730.79it/s, env_episode=932, env_step=34000, gradient_step=136, len=42, n/ep=14, n/st=2000, rew=98.79]
Epoch #1: 72%|#######2 | 36000/50000 [00:13<00:05, 2734.56it/s, env_episode=932, env_step=34000, gradient_step=136, len=42, n/ep=14, n/st=2000, rew=98.79]
Epoch #1: 72%|#######2 | 36000/50000 [00:13<00:05, 2734.56it/s, env_episode=945, env_step=36000, gradient_step=144, len=48, n/ep=13, n/st=2000, rew=114.31]
Epoch #1: 76%|#######6 | 38000/50000 [00:14<00:04, 2731.29it/s, env_episode=945, env_step=36000, gradient_step=144, len=48, n/ep=13, n/st=2000, rew=114.31]
Epoch #1: 76%|#######6 | 38000/50000 [00:14<00:04, 2731.29it/s, env_episode=965, env_step=38000, gradient_step=152, len=44, n/ep=20, n/st=2000, rew=122.15]
Epoch #1: 80%|######## | 40000/50000 [00:14<00:03, 2738.36it/s, env_episode=965, env_step=38000, gradient_step=152, len=44, n/ep=20, n/st=2000, rew=122.15]
Epoch #1: 80%|######## | 40000/50000 [00:14<00:03, 2738.36it/s, env_episode=974, env_step=40000, gradient_step=160, len=57, n/ep=9, n/st=2000, rew=163.11]
Epoch #1: 84%|########4 | 42000/50000 [00:15<00:02, 2716.21it/s, env_episode=974, env_step=40000, gradient_step=160, len=57, n/ep=9, n/st=2000, rew=163.11]
Epoch #1: 84%|########4 | 42000/50000 [00:15<00:02, 2716.21it/s, env_episode=991, env_step=42000, gradient_step=168, len=60, n/ep=17, n/st=2000, rew=173.06]
Epoch #1: 88%|########8 | 44000/50000 [00:16<00:02, 2722.62it/s, env_episode=991, env_step=42000, gradient_step=168, len=60, n/ep=17, n/st=2000, rew=173.06]
Epoch #1: 88%|########8 | 44000/50000 [00:16<00:02, 2722.62it/s, env_episode=1001, env_step=44000, gradient_step=176, len=47, n/ep=10, n/st=2000, rew=121.70]
Epoch #1: 92%|#########2| 46000/50000 [00:17<00:01, 2733.16it/s, env_episode=1001, env_step=44000, gradient_step=176, len=47, n/ep=10, n/st=2000, rew=121.70]
Epoch #1: 92%|#########2| 46000/50000 [00:17<00:01, 2733.16it/s, env_episode=1016, env_step=46000, gradient_step=184, len=38, n/ep=15, n/st=2000, rew=134.07]
Epoch #1: 96%|#########6| 48000/50000 [00:18<00:00, 2046.52it/s, env_episode=1016, env_step=46000, gradient_step=184, len=38, n/ep=15, n/st=2000, rew=134.07]
Epoch #1: 96%|#########6| 48000/50000 [00:18<00:00, 2046.52it/s, env_episode=1024, env_step=48000, len=58, n/ep=8, n/st=2000, rew=221.25]
Epoch #1: 96%|#########6| 48000/50000 [00:18<00:00, 2046.52it/s, env_episode=1024, env_step=48000, gradient_step=184, len=58, n/ep=8, n/st=2000, rew=221.25]
Epoch #1: 96%|#########6| 48000/50000 [00:18<00:00, 2584.09it/s, env_episode=1024, env_step=48000, gradient_step=184, len=58, n/ep=8, n/st=2000, rew=221.25]
train_result.pprint_asdict()
InfoStats
----------------------------------------
{ 'best_reward': 230.3,
'best_reward_std': 107.27912192034385,
'best_score': 230.3,
'gradient_step': 184,
'test_episode': 20,
'test_step': 2504,
'timing': { 'test_time': 0.0,
'total_time': 18.691221237182617,
'train_time': 18.691221237182617,
'train_time_collect': 0.0,
'train_time_update': 6.0345141887664795,
'update_speed': 2568.050497658931},
'train_episode': 1024,
'train_step': 48000}
# Let's watch its performance!
policy.eval()
eval_result = test_collector.collect(n_episode=3, render=False)
print(f"Final reward: {eval_result.returns.mean()}, length: {eval_result.lens.mean()}")
/home/docs/checkouts/readthedocs.org/user_builds/tianshou/checkouts/latest/tianshou/data/collector.py:514: UserWarning: n_episode=3 should be larger than self.env_num=10 to collect at least one trajectory in each environment.
warnings.warn(
Number of episodes (3) is smaller than the number of environments (10). This means that 7 environments (or, equivalently, parallel workers) will not be used!
Final reward: 197.33333333333334, length: 197.33333333333334
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.