Overview#

In this tutorial, we use guide you step by step to show you how the most basic modules in Tianshou work and how they collaborate with each other to conduct a classic DRL experiment.

Run the code#

Before we get started, we must first install Tianshou’s library and Gym environment by running the commands below. Here I choose a specific version of Tianshou(0.4.8) which is the latest as of the time writing this tutorial. APIs in different versions may vary a little bit but most are the same. Feel free to use other versions in your own project.

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

Hide code cell content
import gymnasium as gym
import torch

from tianshou.data import Collector, 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
net = Net(env.observation_space.shape, hidden_sizes=[64, 64], device=device)
actor = Actor(net, env.action_space.n, device=device).to(device)
critic = Critic(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(
    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
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,
)
print(result)
<tianshou.trainer.base.OnpolicyTrainer object at 0x7f00f8b48b50>
# Let's watch its performance!
policy.eval()
result = test_collector.collect(n_episode=1, render=False)
print("Final reward: {}, length: {}".format(result["rews"].mean(), result["lens"].mean()))
Final reward: 16.0, length: 16.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

Check this webpage to find jupyter-notebook-style tutorials that will guide you through all these modules one by one. You can also read the documentation of Tianshou for more detailed explanation and advanced usages.

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.